在 AVX 寄存器 (__m256i) 中设置单个位，需要“随机访问”运算符

Question

所以，我想设置__m256i 寄存器的单个位。

说，我的__m256i 包含：[ 1 0 1 0 | 1 0 1 0 | ... | 1 0 1 0 ]，如何设置和取消设置第 n 位？

Answer 1

这是一个可以在向量中设置单个位的函数的实现：

#include <immintrin.h>
#include <assert.h>

void SetBit(__m256i & vector, size_t position, bool value)
{
    assert(position <= 255);
    uint8_t lut[32] = { 0 };
    lut[position >> 3] = 1 << (position & 7);
    __m256i mask = _mm256_loadu_si256((__m256i*)lut);
    if (value)
        vector = _mm256_or_si256(mask, vector);
    else
        vector = _mm256_andnot_si256(mask, vector);
}

int main(int argc, char* argv[])
{
    __m256i a = _mm256_set1_epi8(-1);
    SetBit(a, 54, false);

    __m256i b = _mm256_set1_epi8(0);
    SetBit(b, 54, true);

    return 0;
}

Answer 2

如果您想避免使用 LUT，可以使用 BTS 设置单个位（或分别使用 BTR 清除它）。这条指令似乎没有内在意义（至少在 GCC 中），因此需要内联汇编（因此仅适用于 x86 架构）。

0F AB /r --- BTS r/m32, r32 --- 将选定位存储在 CF 标志中并设置。

它们对内存操作数非常慢，但这些位字符串指令允许位偏移超出​​寻址模式引用的字节或双字。手册说明：

一些汇编程序通过将立即位偏移字段与内存操作数的位移字段结合使用来支持大于 31 的立即位偏移。在这种情况下，立即位偏移量的低 3 位或 5 位（16 位操作数为 3，32 位操作数为 5）存储在立即位偏移字段中，高位位为由汇编程序在寻址模式中移位并与字节位移组合。如果高位不为零，处理器将忽略它们。

当访问内存中的一个位时，处理器可以访问从内存地址开始的 4 个字节，用于 32 位操作数大小，使用如下关系：

有效地址 + (4 ∗ (BitOffset DIV 32))

在纯汇编程序（Intel-MASM-syntax）中，它看起来像这样：

.data
  .align 16
  save db 32 dup(0)    ; 256bit = 32 byte YMM/__m256i temp variable space
  bitNumber dd 254     ; use an UINT for the bit to set (here the second to last)
.code
  mov eax, bitNumber
  ...
  lea edx, save
  movdqa xmmword ptr [edx], xmm0    ; save __m256i to to memory
  bts dword ptr [edx], eax          ; set the 255st bit
  movdqa xmm0, xmmword ptr [edx]    ; read __m256i back to register
  ...

如果变量已经在内存中，这会更容易。

---

使用内联汇编，这将产生以下函数：

static inline
void set_m256i_bit(__m256i * value, uint32_t bit)
{
    // doesn't need to be volatile: we only want to run this for its effect on *value.
    __asm__ ("btsl %[bit], %[memval]\n\t"
             : [memval] "+m" (*value) : [bit] "ri" (bit));
}

static inline
void clear_m256i_bit(__m256i * value, uint32_t bit)
{
    __asm__ ( "btrl %[bit], %[memval]\n\t"
              : [memval] "+m" (*value) : [bit] "ri" (bit));
}

这些编译为您所期望的on the Godbolt compiler explorer

还有一些类似于上面的汇编代码的测试代码：

__m256i value = _mm256_set_epi32(0,0,0,0,0,0,0,0);
set_m256i_bit(&value,254);
clear_m256i_bit(&value,254);

Answer 3

还有另一种实现方式：

#include <immintrin.h>
#include <assert.h>

template <bool value> void SetMask(const __m256i & mask, __m256i & vector);

template <> inline void SetMask<true>(const __m256i & mask, __m256i & vector)
{
    vector = _mm256_or_si256(mask, vector);
}

template <> inline void SetMask<false>(const __m256i & mask, __m256i & vector)
{
    vector = _mm256_andnot_si256(mask, vector);
}

template <int position, bool value> void SetBit(__m256i & vector)
{
    const uint8_t mask8 = 1 << (position & 7);
    const __m128i mask128 = _mm_insert_epi8(_mm_setzero_si128(), mask8, (position >> 3)&15);
    const __m256i mask256 = _mm256_inserti128_si256(_mm256_setzero_si256(), mask128, position >> 7);
    SetMask<value>(mask256, vector);
}

int main(int argc, char* argv[])
{
    __m256i a = _mm256_set1_epi8(-1);
    SetBit<50, false>(a);

    __m256i b = _mm256_set1_epi8(0);
    SetBit<50, true>(b);

    return 0;
}

Answer 4

如果您想避免 LUT 和/或存储转发停顿，您可以这样做来设置 avx-256 寄存器的第 k 位：

inline __m256i setbit_256(__m256i x,int k){
// constants that will (hopefully) be hoisted out of a loop after inlining  
  __m256i indices = _mm256_set_epi32(224,192,160,128,96,64,32,0);
  __m256i one = _mm256_set1_epi32(-1);
  one = _mm256_srli_epi32(one, 31);    // set1(0x1)


  __m256i kvec = _mm256_set1_epi32(k);  
// if 0<=k<=255 then kvec-indices has exactly one element with a value between 0 and 31
  __m256i shiftcounts = _mm256_sub_epi32(kvec, indices);
  __m256i kbit        = _mm256_sllv_epi32(one, shiftcounts);   // shift counts outside 0..31 shift the bit out of the element
                                                               // kth bit set, all 255 other bits zero.
  return _mm256_or_si256(kbit, x);                             // use _mm256_andnot_si256 to unset the k-th bit
}

---

以下是我之前的答案，它不那么直截了当，现在已经过时了。

#include <immintrin.h>

inline __m256i setbit_256(__m256i x,int k){
  __m256i c1, c2, c3;
  __m256i t, y, msk;

  // constants that will (hopefully) be hoisted out of a loop after inlining
  c1=_mm256_set_epi32(7,6,5,4,3,2,1,0);
  c2=_mm256_set1_epi32(-1);
  c3=_mm256_srli_epi32(c2,27);     // set1(0x1f) mask for the shift within elements
  c2=_mm256_srli_epi32(c2,31);     // set1(0x1)

  // create a vector with the kth bit set
  t=_mm256_set1_epi32(k);
  y=_mm256_and_si256(c3,t);        // shift count % 32: distance within each elem
  y=_mm256_sllv_epi32(c2,y);       // set1( 1<<(k%32) )

  t=_mm256_srli_epi32(t,5);        // set1( k>>5 )
  msk=_mm256_cmpeq_epi32(t,c1);    // all-ones in the selected element
  y=_mm256_and_si256(y,msk);       // kth bit set, all 255 other bits zero.

  x=_mm256_or_si256(y,x);   /* use _mm256_andnot_si256 to unset the k-th bit */
  return x;
}

我不确定这是否会比其他答案中建议的方法更快。

考虑到常量将被提升出循环，这可以使用 clang 或 gcc (Godbolt compiler explorer) 编译为非常好的 asm。像往常一样，clang 阻止了动态生成常量的尝试，并从内存中广播加载它们（这在现代 CPU 上非常有效）。