Работа с битовыми векторами с AVX2 и SSE2 - PullRequest
Новая
       58

Работа с битовыми векторами с AVX2 и SSE2

3 голосов
/ 04 ноября 2019

Я новичок в наборах команд AVX2 и SSE2 и хочу узнать больше о том, как использовать такие наборы команд для ускорения операций с битовыми векторами.

До сих пор я успешно использовал их для векторизациикоды с операциями типа double / float.

В этом примере у меня есть код C ++, который проверяет условие перед тем, чтобы установить или нет бит в векторе битов (используя unsigned int) для конкретного значения: 

int process_bit_vetcor(unsigned int *bitVector, float *value, const float threshold, const unsigned int dim)
{
       int sum = 0, cond = 0;

       for (unsigned int i = 0; i < dim; i++) {
            unsigned int *word = bitVector + i / 32;
            unsigned int bitValue = ((unsigned int)0x80000000 >> (i & 0x1f));
            cond = (value[i] <= threshold);
            (*word) = (cond) ? (*word) | bitValue : (*word);
            sum += cond;
        }

        return sum;
}

Переменная sum просто возвращает количество случаев, когда условие TRUE.

Я пытался переписать эту процедуру с SSE2 и AVX2, но она не сработала...: - (

Можно ли переписать такой код C ++ с использованием AVX2 и SSE2? Стоит ли использовать векторизацию для такого типа битовых операций? Битовый вектор может содержать много тысяч битов, поэтому я надеюсь, чтобыло бы интересно использовать SSE2 и AVX2 для ускорения.

Заранее спасибо!

1 Ответ

2 голосов
/ 07 ноября 2019

Следующее должно работать, если dim кратно 8 (чтобы обработать остаток, добавьте тривиальный цикл в конце). Незначительные изменения в API:

  • Использование long вместо unsigned int для индексов цикла (это помогает лязгнуть развертывание цикла)
  • Предположим, что bitvector имеет младший порядок (какпредлагается в комментариях)

Внутри цикла, bitVector доступен побайтно. Возможно, стоит объединить 2 или 4 результата movemask и битовых или их одновременно (вероятно, зависит от целевой архитектуры).

Для вычисления sum 8 частных сумм рассчитываются непосредственно изрезультат операции cmp_ps. Поскольку битовая маска в любом случае вам нужна, возможно, стоит использовать popcnt (в идеале, после объединения 2, 4 или 8 байт - опять же, это, вероятно, зависит от вашей целевой архитектуры). 

int process_bit_vector(uint32_t *bitVector32, float *value,
                       const float threshold_float, const long dim) {
  __m256i sum = _mm256_setzero_si256();
  __m256 threshold_vector = _mm256_set1_ps(threshold_float);
  uint8_t *bitVector8 = (uint8_t *)bitVector32;

  for (long i = 0; i <= dim-8; i += 8) {
    // compare next 8 values with threshold
    // (use threshold as first operand to allow loading other operand from memory)
    __m256 cmp_mask = _mm256_cmp_ps(threshold_vector, _mm256_loadu_ps(value + i), _CMP_GE_OQ);
    // true values are `-1` when interpreted as integers, subtract those from `sum`
    sum = _mm256_sub_epi32(sum, _mm256_castps_si256(cmp_mask));
    // extract bitmask
    int mask = _mm256_movemask_ps(cmp_mask);
    // bitwise-or current mask with result bit-vector
    *bitVector8++ |= mask;
  }

  // reduce 8 partial sums to a single sum and return
  __m128i sum_reduced = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum,1));
  sum_reduced = _mm_add_epi32(sum_reduced, _mm_srli_si128(sum_reduced, 8));
  sum_reduced = _mm_add_epi32(sum_reduced, _mm_srli_si128(sum_reduced, 4));

  return _mm_cvtsi128_si32(sum_reduced);
}

Godbolt-Link: https://godbolt.org/z/ABwDPe

  • По какой-то причине GCC делает vpsubd ymm2, ymm0, ymm1; vmovdqa ymm0, ymm2; вместо просто vpsubd ymm0, ymm0, ymm1.
  • Clang не может присоединиться к load с vcmpps(и использует LE вместо GE сравнения) - если вам не важно, как обрабатываются NaN, вы можете использовать _CMP_NLT_US вместо _CMP_GE_OQ.

Пересмотренная версия с выходом с прямым порядком байтов (непроверенный): 

int process_bit_vector(uint32_t *bitVector32, float *value,
                       const float threshold_float, const long dim) {
  int sum = 0;
  __m256 threshold_vector = _mm256_set1_ps(threshold_float);

  for (long i = 0; i <= dim-32; i += 32) {
    // compare next 4x8 values with threshold
    // (use threshold as first operand to allow loading other operand from memory)
    __m256i cmp_maskA = _mm256_castps_si256(_mm256_cmp_ps(threshold_vector, _mm256_loadu_ps(value + i+ 0), _CMP_GE_OQ));
    __m256i cmp_maskB = _mm256_castps_si256(_mm256_cmp_ps(threshold_vector, _mm256_loadu_ps(value + i+ 8), _CMP_GE_OQ));
    __m256i cmp_maskC = _mm256_castps_si256(_mm256_cmp_ps(threshold_vector, _mm256_loadu_ps(value + i+16), _CMP_GE_OQ));
    __m256i cmp_maskD = _mm256_castps_si256(_mm256_cmp_ps(threshold_vector, _mm256_loadu_ps(value + i+24), _CMP_GE_OQ));

    __m256i cmp_mask = _mm256_packs_epi16(
        _mm256_packs_epi16(cmp_maskA,cmp_maskB), // b7b7b6b6'b5b5b4b4'a7a7a6a6'a5a5a4a4 b3b3b2b2'b1b1b0b0'a3a3a2a2'a1a1a0a0
        _mm256_packs_epi16(cmp_maskC,cmp_maskD)  // d7d7d6d6'd5d5d4d4'c7c7c6c6'c5c5c4c4 d3d3d2d2'd1d1d0d0'c3c3c2c2'c1c1c0c0
    );                                // cmp_mask = d7d6d5d4'c7c6c5c4'b7b6b5b4'a7a6a5a4 d3d2d1d0'c3c2c1c0'b3b2b1b0'a3a2a1a0

    cmp_mask = _mm256_permute4x64_epi64(cmp_mask, 0x8d);
                // cmp_mask = [b7b6b5b4'a7a6a5a4 b3b2b1b0'a3a2a1a0  d7d6d5d4'c7c6c5c4 d3d2d1d0'c3c2c1c0]
    __m256i shuff_idx = _mm256_broadcastsi128_si256(_mm_set_epi64x(0x00010203'08090a0b,0x04050607'0c0d0e0f));
    cmp_mask = _mm256_shuffle_epi8(cmp_mask, shuff_idx);

    // extract bitmask
    uint32_t mask = _mm256_movemask_epi8(cmp_mask);
    sum += _mm_popcnt_u32 (mask);
    // bitwise-or current mask with result bit-vector
    *bitVector32++ |= mask;
  }

  return sum;
}

Идея состоит в том, чтобы перемешать байты перед применением vpmovmskb к нему. Для этого требуется 5 операций тасования (включая 3 vpacksswb) для 32 входных значений, но вычисление суммы выполняется с использованием popcnt вместо 4 vpsubd. vpermq (_mm256_permute4x64_epi64), вероятно, можно было бы избежать путем стратегической загрузки 128-битных половин в 256-битные векторы перед их сравнением. Еще одна идея (поскольку в любом случае вам необходимо перетасовать конечный результат) состоит в том, чтобы смешать частичные результаты (это обычно требует p5 или 2*p015 на архитектурах, которые я проверял, поэтому, вероятно, не стоит).

...