Можно ли подсчитать __m256i и сохранить результат в 8 32-битных словах вместо 4 64-битных, используя алгоритм Войцеха Мула?

Question

Недавно я обнаружил, что у AVX2 нет поп-счета для __m256i, и единственный способ сделать что-то похожее - следовать алгоритму Войцеха Мула:

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

Войцех Мула, Натан Курц, Даниэль Лемир, Ускорение подсчета населения с использованием инструкций AVX2, Компьютерный журнал 61 (1), 2018

Проблема в том, что он возвращает мне сумму 8 коротких в long вместо суммы 4 коротких в 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 * ​​1048 * * * 21 тысяча сорок-девять

Результат 4 длинных 64-битных

Expectation:

У меня есть __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. 11
2. 13
3. 10
4. 11
5. 12
6. 9
7. 11
8. 10

Результат - 8 int 32-bit.

Надеюсь, я был ясен, не стесняйтесь спрашивать меня о большей точности.

Спасибо.

Answer 1

Исходный код, который вы цитируете, основан на свойстве _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 умножением и сдвигом. Это должно быть разумно. В моих тестах он немного медленнее, чем оригинальная 64-битная версия, но разница относительно мала .

Вы можете получить чуть лучшую производительность за счет использования большего количества встроенных функций:

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