屏蔽 AVX2 中单个位的最佳方法？

Question

例如，使用输入 ymm 向量 x 和位索引 i 我想要一个仅保留第 i 位且其他所有内容都归零的输出向量.

使用 AVX512 k 寄存器，我可以编写以下内容，但 AVX2 及以下版本没有 k 寄存器，那么您认为最好的方法是什么？

__m512i m512i_maskBit(__m512i x, unsigned i) {
    __mmask8 m = _cvtu32_mask8(1u << i / 64);
    __m512i vm = _mm512_maskz_set1_epi64(m, 1ull << i % 64);
    return _mm512_and_si512(x, vm);
}

Answer 1

最简单的方法如何：

__m256i m256i_create_mask(unsigned i) {
    // Get the required bit in every byte of the vector
    __m256i vm = _mm256_broadcastb_epi8(_mm_cvtsi32_si128(1u << (i & 7u)));
    // Mask off the bytes that are outside the index
    __m256i vi = _mm256_broadcastb_epi8(_mm_cvtsi32_si128(i >> 3u));
    __m256i vm1 = _mm256_cmpeq_epi8(vi,
        _mm256_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
            16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31));
    return _mm256_and_si256(vm, vm1);
}

Answer 2

这是一种使用变量移位的方法（仅创建掩码）：

__m256i create_mask(unsigned i) {
    __m256i ii = _mm256_set1_epi32(i);
    ii = _mm256_sub_epi32(ii,_mm256_setr_epi32(0,32,64,96,128,160,192,224));
    __m256i mask = _mm256_sllv_epi32(_mm256_set1_epi32(1), ii);
    return mask;
}

_mm256_sllv_epi32 (vpsllvd) 由 AVX2 引入，它将每个 32 位元素移动可变的位数。如果 (unsigned) shift-amount 大于 31（即，也适用于有符号负数），则相应的结果为 0.

Godbolt link 带有小测试代码：https://godbolt.org/z/a5xfqTcGs

Answer 3

这是另一种方法。不确定它是否一定更好，这取决于 CPU 模型和周围的代码，但它可能是。

// A buffer to load vectors with a single bit set in one lane
alignas( 64 ) static const std::array<int, 16> s_oneBuffer =
{
    0, 0, 0, 0, 0, 0, 0, 0,
    1, 0, 0, 0, 0, 0, 0, 0
};

__m256i maskSingleBit( __m256i x, uint32_t bitIndex )
{
    // Load `1` into a single 32-bit lane of the vector
    // The buffer aligned by 64 bytes, contained in a single cache line, no unaligned load penalty.
    __m256i one = _mm256_loadu_si256( ( const __m256i* )( ( s_oneBuffer.data() + 8 ) - ( bitIndex / 32 ) ) );

    // Left shift to move the `1` into the correct location
    __m128i shift = _mm_cvtsi32_si128( bitIndex % 32 );
    __m256i bit = _mm256_sll_epi32( one, shift );

    // Bitwise AND with the value
    return _mm256_and_si256( x, bit );
}

屏蔽 AVX2 中单个位的最佳方法？

Best way to mask a single bit in AVX2?

c

x86

simd

avx

avx2