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

时间:2018-06-29 14:55:42

标签: c++ intel sse 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());
}

Wojciech Muła, Nathan Kurz, Daniel Lemire, Faster Population Counts Using AVX2 Instructions, Computer Journal 61 (1), 2018

问题在于它返回的是8 short到long的总和,而不是4 short到int的总和。

正在发生的事情:
我有__m256i x,其中包含这8个32位int:

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

__ m256i res = count(x);

res包含:

  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 = count(x);

res包含:

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

结果是8 int 32位。

希望我很清楚,不要犹豫,请我提供更多精度。

谢谢。

1 个答案:

答案 0 :(得分:2)

您引用的原始代码依赖于_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);
}

所以我们用乘法和移位代替_mm256_sad_epu8。那应该是合理的。在我的测试中,it is slightly slower than the original 64-bit version, but the difference is relatively small

您可以使用更多的内在函数来获得更好的性能:

__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));
}