Два универсальных SIMD алгоритма / Хабр

Большинство SIMD инструкций узконаправленны, например применяют бинарную операцию параллельно для нескольких чисел, упакованных в длинный регистр. Применение таких операций прямолинейно и в большинстве случаев компилятор сам оптимизирует код с использованием таких инструкций. Например компилятор легко соптимизирует таким образом проверку несложного предиката на массиве или например суммирование элементов массива. Есть однако и более универсальные инструкции, в частности довольно много всякого рода манипуляций с битами внутри регистра. В этой статье хочу рассказать о двух таких инструкциях: уже давно присутствующей PSHUFB и довольно новой GF2P8AFFINEQB, расскажу как с их помощью делать побайтовую обработку общего вида и приведу пару примеров с известными операциями такими как popcount, подсчет четности, разворот битов числа.

Параллельный просмотр таблицы с помощью PSHUFB

Давайте для начала познакомимся с самой инструкцией, нам нужна её 128-битная версия без маскирования, согласно документации интел (зеркало на гитхаб)

Synopsis

__m128i _mm_shuffle_epi8 (__m128i a, __m128i b)
#include <tmmintrin.h>
Instruction: pshufb xmm, xmm
CPUID Flags: SSSE3

Description

Shuffle packed 8-bit integers in a according to shuffle control mask in the corresponding 8-bit element of b, and store the results in dst.

Operation

FOR j := 0 to 15
    i := j*8
    IF b[i+7] == 1
        dst[i+7:i] := 0
    ELSE
        index[3:0] := b[i+3:i]
        dst[i+7:i] := a[index*8+7:index*8]
    FI
ENDFOR

Latency and Throughput

Architecture	Latency	Throughput (CPI)
Alderlake	1	0.5
Icelake Xeon	1	0.5
Sapphire Rapids	1	0.5
Skylake	1	1

Итак, доступна начиная с SSSE3, легковесна и делает "перемешивание 8-битных чисел из a в соответствии с маской b". Как-то не очень понятно, но если присмотреться к строкам 6-7 описании операции, то становится понятным, что это просмотр таблицы. Итак, индекс iпробегает все значения кратные 8, в index записывается первые 4 бита b в позиции i, после чего в результат записывается 8 бит из a начиная с 8*index. Вот неоптимизированный эквивалент инструкции

std::array<uint8_t, 16> _mm_shuffle_epi8(const std::array<uint8_t, 16>& a,
                                         const std::array<uint8_t, 16>& b) {
  std::array<uint8_t, 16> result;
  for (size_t i = 0; i < 16; ++i) {
    if (b[i] & (1 << 7)) {
      result[i] = 0;
    } else {
      result[i] = a[b[i] & 15];
    }
  }
  return result;
}

По сути это параллельный просмотр таблицы размера 16. Интересным тут является то, что в качестве индексов берутся только 4 бита. Есть всего 16 возможных значений, которые могут принимать 4 бита и, соответственно, все эти значения умещаются в 128-регистр -- так по сути эта инструкция и работает. Интересно то, что таким образом мы можем представить любую операцию над 4-мя битами. В некоторых случаях этого достаточно чтобы сделать и более сложное вычисление например над байтами. Вот пример как можно сделать подсчет побайтового попкаунта

// Таблица popcount для 4 бит
const __m128i lookup_popcount_4 = _mm_setr_epi8(
    0, 1, 1, 2,  // 0000, 0001, 0010, 0011
    1, 2, 2, 3,  // 0100, 0101, 0110, 0111
    1, 2, 2, 3,  // 1000, 1001, 1010, 1011
    2, 3, 3, 4   // 1100, 1101, 1110, 1111
);

// Побайтовый popcnt
void popcount_16x8(const uint8_t* x, uint8_t* result) {
  __m128i data = _mm_loadu_si128((__m128i const*)x);

  const __m128i low_bits_mask = _mm_set1_epi8(0x0F);

  // Popcount 4 нижних битов
  __m128i low_bits = _mm_and_si128(data, low_bits_mask);
  __m128i low_count = _mm_shuffle_epi8(lookup_popcount_4, low_bits);

  // Popcount 4 верхних битов
  __m128i high_bits = _mm_srli_epi16(data, 4);
  high_bits = _mm_and_si128(high_bits, low_bits_mask);
  __m128i high_count = _mm_shuffle_epi8(lookup_popcount_4, high_bits);

  __m128i result_vec = _mm_add_epi8(low_count, high_count);
  _mm_storeu_si128((__m128i*)result, result_vec);
}

Отмечу, что для ARM-процессоров аналогом являются инструкции vqtbl1q_u8 и смежные.

А вот тут можно найти более сложный пример использования этого метода для ускорения векторных операций над полем $\mathbb{F}_{2^8}$ и, насколько мне известно, используется для кодов Рида-Соломона в RAID6. А раз уж начали про $\mathbb{F}_{2^8}$ , то ….

Произвольные аффинные операции надс помощью GF2P8AFFINEQB

$\mathbb{F}_2$ -- это поле размера 2, умножение в котором соответствует логическому "и", а сложение производится по модулю 2 (aka XOR). Относительно недавно Intel (а следом и AMD) пополнил линейку AVX-512 набором Galois Field New Instructions, основное назначение которого побайтовая обработка в поле $\mathbb{F}_{2^8}$ для чего-то в духе кодов Рида-Соломона или Rijndael энкрипции (хотя для этого выпустили отдельное расширение VAES). Две из трёх инструкций GFNI завязаны непосредственно на $\mathbb{F}_{2^8}$ , но последняя инструкция универсальна и работает над $\mathbb{F}_2$ , снова обращаемся к гайду от Intel.

Synopsis

__m512i _mm512_gf2p8affine_epi64_epi8 (__m512i x, __m512i A, int b)
#include <immintrin.h>
Instruction: vgf2p8affineqb zmm, zmm, zmm, imm8
CPUID Flags: GFNI + AVX512F

Description

Compute an affine transformation in the Galois Field 2^8. An affine transformation is defined by A * x + b, where A represents an 8 by 8 bit matrix, x represents an 8-bit vector, and b is a constant immediate byte. Store the packed 8-bit results in dst.

Operation

DEFINE parity(x) {
    t := 0
    FOR i := 0 to 7
        t := t XOR x.bit[i]
    ENDFOR
    RETURN t
}
DEFINE affine_byte(tsrc2qw, src1byte, imm8) {
    FOR i := 0 to 7
        retbyte.bit[i] := parity(tsrc2qw.byte[7-i] AND src1byte) XOR imm8.bit[i]
    ENDFOR
    RETURN retbyte
}
FOR j := 0 TO 7
    FOR i := 0 to 7
        dst.qword[j].byte[i] := affine_byte(A.qword[j], x.qword[j].byte[i], b)
    ENDFOR
ENDFOR
dst[MAX:512] := 0

Latency and Throughput

Architecture	Latency	Throughput (CPI)
Icelake Intel Core	-	1
Icelake Xeon	3	1
Sapphire Rapids	-	1

Здесь описание посложнее, да еще и немного неточное, все-таки операции проводятся не в $\mathbb{F}_{2^8}$ , а над матрицами в $\mathbb{F}_2$ , давайте разбираться: вот пример из гайда от Интел

Базовая операция GF2P8AFFINEQB: перемножение двух бинарных матриц 8х8, на самом деле там еще есть bias, смотрите документацию инструкции

Здесь приводится простой пример как аффинным преобразованием можно развернуть биты в каждом из байтов числа. GF2P8AFFINEQBумеет делать 2/4/8 таких операций при использовании 128/256/512-битных регистров соответственно, вот например как это выглядит для 256-битного регистра.

256-битная версия GF2P8AFFINEQB делает 4 аффинных преобразования за раз.

Фактически идея с разворотом была взята на вооружение в последних версиях Clang, посмотрим во что компилируется __builtin_bitreverse64:

#include <cstdint>
uint64_t reverse_bits(uint64_t x) {
    return __builtin_bitreverse64(x);
}

.LCPI0_1:
        .byte   1
        .byte   2
        .byte   4
        .byte   8
        .byte   16
        .byte   32
        .byte   64
        .byte   128
reverse_bits(unsigned long):
        vmovq   xmm0, rdi
        vgf2p8affineqb  xmm0, xmm0, qword ptr [rip + .LCPI0_1]{1to2}, 0
        vmovq   rax, xmm0
        bswap   rax
        ret

Константа LCPI0_1соответствует левой матрицы с единичками на диагонали, применении аффинного преобразования в строчке 12 меняет порядок битов в каждом байте, а в строчке 14 разворачиваются байты, что в итоге даёт полный разворот в 64-битном числе. Вот краткий список похожих операций, которые можно реализовать через умножение на матрицу:

Побайтовая четность
Побайтовые сдвиги влево вправо
Побайтовые циклические сдвиги

В качестве более сложного примера приведу пример как сделать так, чтобы GF2P8AFFINEQB работала как другая инструкция GFNI GF2P8MULB, предназначенная именно для $\mathbb{F}_{2^8}$ . Напомню вкратце, что элементы $\mathbb{F}_{2^8}$ - это многочлены степени меньше 8 над $\mathbb{F}_2$ c обычным сложением многочленов и умножением по модулю некоторого неприводимого многочлена степени 8. Эти многочлены можно трактовать и как элементы векторного пространства $\mathbb{F}_2^8$ , если мы зафиксируем элемент $a\in \mathbb{F}_{2^8}$ , то умножение нав $\mathbb{F}_{2^8}$ является линейной операцией в $\mathbb{F}_2^8$ , а значит должна существовать матрица $\Psi_a\in \mathbb{F}_2^{8\times 8}$ такая, что

$a\times x = \Psi_ax$

Слева подразумевается умножение в $\mathbb{F}_{2^8}$ , справа -- умножение матрицы на вектор в $\mathbb{F}_2$ . Собственно идея в том, чтобы для каждого элемента поля посчитать соответствующую матрицу. После этого мы можем сгруппировать умножение навплоть до 64 других элементов, что позволит например делать $x[0..n] = x[0..n]+a \times y[0 .. n]$ , т.е. прибавить к одному вектору другой, помноженный на скаляр - типичная операция например для перемножения матриц или метода Гаусса решения линейных систем. Для подсчета матриц $\Psi_a$ воспользуемся простым фактом из линейной алгебры: столбцы этой матрицы являются результатом применения преобразования на стандартном базисе

/* Бинарные матрицы 8x8 соответствующие умножению  */
static uint64_t gfni_matrix[256];
const uint8_t irreducible_poly = 0x1b; // x^4+x^3+x+1

void InitGFNI() {
  for (int16_t y = 0; y < 256; ++y) {
    uint64_t mt = 0;
    uint8_t row = y;
    // На каждой итерации цикла вычисляется произведение y
    // на очердной базисный вектор, результат записывается в mt
    for (size_t i = 0, shift = 0; i < 8; ++i, shift += 8) {
      mt |= ((uint64_t)row << shift);
      row = (row << 1) ^ ((row >> 7) * irreducible_poly);
    }

    // Фактически в mt уже лежит нужная матрица, но порядок элементов
    // в GF2P8AFFINEQB немного странный, нужно немного попереставлять биты
    gfni_matrix[y] = 0;
    for (size_t i = 0; i < 8; ++i) {
      for (size_t j = 0; j < 8; ++j) {
        gfni_matrix[y] |= ((mt >> (8 * i + j)) & 1) << (8 * j + (7 - i));
      }
    }
  }
}

// Производит x[0..63] += y[0..63] * c
void FMA_64(uint8_t* x, const uint8_t* y, uint8_t c) {
  __m512i c_matrix = _mm512_set1_epi64(gfni_matrix[c]);
  auto x_reg = _mm512_loadu_epi8(x);
  auto y_reg = _mm512_loadu_epi8(y);
  x_reg = _mm512_xor_epi64(x_reg,
                           _mm512_gf2p8affine_epi64_epi8(y_reg, c_matrix, 0));
  _mm512_storeu_epi8(x, x_reg);
}

Подробнее на этот пример можно посмотреть в репозитории. Отдельно для тех, кто работает с вычислением в бинарных полях, отмечу, что указанный выше метод может быть использован для работы с нормальным базисом или базисом Кантора.

Ну и под конец отмечу, что умножение матриц может быть использовано само по себе, наличие GF2P8AFFINEQB скорее всего делает бесполезным метод четырех русских.

Ссылки

Intel Intrinsics Guide.
Документация по SIMD-инструкциям Intel (зеркало на github).
URL: https://intrinsics.idma88.ru/
Intel Corporation. Galois Field New Instructions (GFNI) Technology Guide.
Документация по GFNI
URL: https://builders.intel.com/docs/networkbuilders/galois-field-new-instructions-gfni-technology-guide-1-1639042826.pdf
ARM Architecture Reference.
Документация по ARM NEON Intrinsics (vqtbl1qu8 и др.).
URL: https://developer.arm.com/architectures/instruction-sets/intrinsics/vqtbl1qu8
catid / gf256.
Реализация операций над полем GF(256) на C++.
URL: https://github.com/catid/gf256/blob/master/gf256.cpp#L475
H. Peter Anvin. The mathematics of RAID-6.
Коды Рида–Соломона для RAID 6 с ускорением векторных операций через PSHUFB, использующихся в том числе в catid/gf256
URL: http://kernel.org/pub/linux/kernel/people/hpa/raid6.pdf
Intel ISA-L (Intel Intelligent Storage Acceleration Library).
Коды Рида-Соломона от Intel на С/ASM с использованием GFNI.
URL: https://github.com/intel/isa-l/blob/master/erasurecode/gfvect_gfni.inc#L60
Malkovsky / galois.
Мой репозиторий со сравнением производительности векторных операций над GF(256).
URL: https://github.com/Malkovsky/galois/tree/main

Друзья и коллеги! С удовольствием хотел бы прорекламировать CS Space — открытый научный клуб по CS-related темам; идейных последователей питерского Computer Science Club (и CS Center), расформировавшегося после известных событий. Ребята организуют крутые лекции и курсы по CS от профессионалов своего дела, да еще и помогают мне с написанием научно-популярных статей!

Сайт сообщества: csspace.io
Telegram-канал: t.me/csspace

Если вам понравилась статья — поставьте плюс, автору всегда приятно когда его работу ценят. Возможно вас также заинтересует мой канал А зачем это нужно? где я рассказываю о том, что математику и алгоритмы придумали не только для собеседований в бигтехи.