是否可以使用 Wojciech Mula 算法 popcount __m256i 并将结果存储在 8 个 32 位字而不是 4 个 64 位字中?

yat*_*ino 5 c++ sse intel avx avx2


我最近发现 AVX2 没有 __m256i 的 popcount,我发现做类似事情的唯一方法是遵循 Wojciech Mula 算法:

__m256i count(__m256i v) {
    __m256i lookup = _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2,
                     2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3,
                     1, 2, 2, 3, 2, 3, 3, 4);
    __m256i low_mask = _mm256_set1_epi8(0x0f);
    __m256i lo =_mm256_and_si256(v,low_mask);
    __m256i hi = _mm256_and_si256( _mm256_srli_epi32(v, 4), low_mask);
    __m256i popcnt1 = _mm256_shuffle_epi8(lookup,lo);
    __m256i popcnt2 = _mm256_shuffle_epi8(lookup,hi);
    __m256i total = _mm256_add_epi8(popcnt1,popcnt2);

    return _mm256_sad_epu8(total,_mm256_setzero_si256());
}
Run Code Online (Sandbox Code Playgroud)

Wojciech Mu?a、Nathan Kurz、Daniel Lemire,使用 AVX2 指令加快人口计数,计算机杂志 61 (1),2018

问题是它返回 8 个短到长的总和,而不是 4 个短到 int 的总和。

当前发生的事情:
我有 __m256i x ,其中包含 8 个 32 位整数:

  1. 01101011111000011100000000000000
  2. 01110101011010010111100000000000
  3. 10100100011011000101010000000000
  4. 11101010100001001111000000000000
  5. 10010011111111001001010000000000
  6. 00011110101100101000000000000000
  7. 00011101011000111011000000000000
  8. 10011011100010100000110000000000

__m256i res = 计数(x);

资源包含:

  1. 24
  2. 21
  3. 22
  4. 21

结果是4长64位

期待:

我有 __m256i x ,其中包含 8 个 32 位整数:

  1. 01101011111000011100000000000000
  2. 01110101011010010111100000000000
  3. 10100100011011000101010000000000
  4. 11101010100001001111000000000000
  5. 10010011111111001001010000000000
  6. 00011110101100101000000000000000
  7. 00011101011000111011000000000000
  8. 10011011100010100000110000000000

__m256i res = 计数(x);

资源包含:

  1. 11
  2. 13
  3. 10
  4. 11
  5. 12
  6. 9
  7. 11
  8. 10

结果是 8 int 32 位。

希望我很清楚,不要犹豫,问我更精确。

谢谢。

Dan*_*ire 3

AVX-512VPOPCNTDQ 必须_mm256_popcnt_epi32以 32 位块(也是 64 位块大小版本)进行 popcount。除了 Xeon Phi 之外,Ice Lake 中还引入了 AVX512BITALG,它的字节和字(16 位)块大小也为vpopcnt.


与AVX2

您引用的原始代码依赖于_mm256_sad_epu8内在函数,它专门用于对 64 位字内的字节进行求和。

要使用 32 位字的总和获得相同的结果,您需要做一些稍微不同的事情。以下应该有效:

__m256i popcount_pshufb32(__m256i v) {

  __m256i lookup = _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2,
                 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3,
                 1, 2, 2, 3, 2, 3, 3, 4);
  __m256i low_mask = _mm256_set1_epi8(0x0f);
  __m256i lo = _mm256_and_si256(v, low_mask);
  __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
  __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);
  __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);
  __m256i sum8 = _mm256_add_epi8(popcnt1, popcnt2);
  return _mm256_srli_epi32(
      _mm256_mullo_epi32(sum8, _mm256_set1_epi32(0x01010101)), 24);
      // vpmulld is slowish (2 uops) on most recent Intel CPUs
      // but still single-uop on AMD
}
Run Code Online (Sandbox Code Playgroud)

所以我们_mm256_sad_epu8用乘法和移位来代替。这应该是合理的。在我的测试中,它比原来的64位版本稍慢一些,但差别比较小

通过使用不同的两条指令从字节累加到 32 位块,您可以在 Intel 上获得稍微更好的性能,但代价是多一个向量常量。AMD Zen1/2/3 的上述版本至少与以下版本一样高效。

在最新的 Intel CPU 上,32 位 SIMD 整数乘法为 2 uop(均适用于 SIMD 整数乘法单元),但成对乘法累加指令(8->16 和 16->32)各为一个 uop。( https://uops.info/ ) 这需要多一个常量,但指令数量相同,以减少 uops,特别是如果编译器可以在循环中重用这些常量的话。

__m256i popcount_pshufb32(__m256i v) {

  __m256i lookup = _mm256_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2,
                 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3,
                 1, 2, 2, 3, 2, 3, 3, 4);
  __m256i low_mask = _mm256_set1_epi8(0x0f);
  __m256i lo = _mm256_and_si256(v, low_mask);
  __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
  __m256i popcnt1 = _mm256_shuffle_epi8(lookup, lo);
  __m256i popcnt2 = _mm256_shuffle_epi8(lookup, hi);
  __m256i sum8 = _mm256_add_epi8(popcnt1, popcnt2);
  return _mm256_madd_epi16(_mm256_maddubs_epi16(sum8, _mm256_set1_epi8(1)),
                       _mm256_set1_epi16(1));
}
Run Code Online (Sandbox Code Playgroud)

  • 您确定“_mm256_mul_epi32”会起作用吗?它忽略位 32-63。我期望使用 `_mm256_mullo_epi32` 而不是 `_mm256_mul_epi32`。请注意,“_mm256_mullo_epi32”在 Intel 上相当慢。使用 shift-add-shift-add-mask_0xFF 的水平求和可能同样快。 (2认同)