AVX(2)/SIMD 方式获取/设置（至 1）256 位寄存器中的单个位

Question

当前但hacky的方法是这样的：

__m256i bitset(__m256i source, uint8_t index) {
    uint8_t pos_in_64 = index % 64;
    uint8_t location = index / 64;
    uint64_t bitmask = 1ULL << pos_in_64;

    __m256i mask = _mm256_setzero_si256();
    switch (location) {
    case 0: mask = _mm256_set_epi64x(0, 0, 0, bitmask); break;
    case 1: mask = _mm256_set_epi64x(0, 0, bitmask, 0); break;
    case 2: mask = _mm256_set_epi64x(0, bitmask, 0, 0); break;
    case 3: mask = _mm256_set_epi64x(bitmask, 0, 0, 0); break;
    }

    return _mm256_or_si256(source, mask);
}

bool bitget(__m256i source, uint8_t index) {
    uint8_t pos_in_64 = index % 64;
    uint8_t location = index / 64;
    uint64_t bitmask = 1ULL << pos_in_64;

    uint64_t extracted = 0;
    switch (location) {
    case 0: extracted = _mm256_extract_epi64(source, 0); break;
    case 1: extracted = _mm256_extract_epi64(source, 1); break;
    case 2: extracted = _mm256_extract_epi64(source, 2); break;
    case 3: extracted = _mm256_extract_epi64(source, 3); break;
    }
    return (extracted & bitmask) != 0;
}

但我确信有一种更明智的方法可以实现此目的，而无需使用 switch 语句和大量标量代码

Answer 1

这是一个不同的策略。

通过将索引广播到每个位置并与 0..7 进行比较，制作索引指向哪个双字的掩码。通过将索引左移 1（索引模 32）来制作索引指向双字中的哪一位的掩码。这两个掩码的按位 AND 给出了一个仅设置 1 位的掩码，即既位于正确的双字中又在该双字中具有正确位置的位。

__m256i bitset(__m256i source, uint8_t index) {
    __m256i iota = _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0);
    __m256i dwmask = _mm256_cmpeq_epi32(_mm256_set1_epi32(index >> 5), iota);
    __m256i bitmask = _mm256_sllv_epi32(_mm256_set1_epi32(1), _mm256_set1_epi32(index & 31));
    return _mm256_or_si256(source, _mm256_and_si256(dwmask, bitmask));
}

相同的技术可用于

bitget

bool bitget(__m256i source, uint8_t index) {
    __m256i iota = _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0);
    __m256i dwmask = _mm256_cmpeq_epi32(_mm256_set1_epi32(index >> 5), iota);
    __m256i bitmask = _mm256_sllv_epi32(_mm256_set1_epi32(1), _mm256_set1_epi32(index & 31));
    __m256i mask = _mm256_and_si256(dwmask, bitmask);
    return _mm256_testz_si256(source, mask) == 0;
}