Оптимизация RIPEMD-160 используя ARM Neon и не только
У меня есть хобби-проект — ecloop — "калькулятор" Bitcoin-ключей, предназначенный для поиска Bitcoin Puzzles, проверки brain wallets и тому подобного. Математические шансы нахождения приватного ключа от использованного адреса стремятся к нулю. Эта программа мне интересна как набор трюков над эллиптической кривой (secp256k1) и как способ попрактиковаться в программировании, близком к процессору (быстрая 256-битная арифметика), поэтому я периодически продолжаю её развивать.
Для вычисления Bitcoin-адреса из private key нужно выполнить несколько операций:
- вычислить точку на эллиптической кривой (Public key) —
P = G * PrivKey - посчитать SHA256 от PubKey:
(P.y % 2 == 0 ? 0x02 : 0x03) + P.x - посчитать RMD160 от полученного SHA256
В результате получится так называемый hash160, который затем кодируется в Bitcoin-адрес с использованием base58 или bech32.
Как бы ни были медленны операции на эллиптической кривой, самая медленная часть генерации адреса — это вычисление RMD160. Фактически, оно занимает примерно половину времени работы (SHA256 на современных процессорах имеет аппаратное ускорение).
Практически все современные процессоры поддерживают SIMD: AVX2 на amd64 и Neon на arm64. Поэтому я решил, что было бы неплохо ускорить RMD160, реализовав его в виде параллельных вычислений. Тем более, я раньше никогда не писал SIMD-код, и мне было интересно это попробовать.
Мой основной компьютер — MacBook на Apple Silicon (M-чипах). Изначально я хотел реализовать RMD160 SIMD на SVE (256 bit / 8 lane), но оказалось, что чипы Apple не поддерживают SVE 🤦 (следует пояснить, что M2-чип реализует стандарт ARMv8.6, а SVE был добавлен в ARMv8.2 с пометкой optional), поэтому пришлось использовать Neon (128 bit / 4 lane). Если это не так, я где-то ошибся, и есть способ запускать SVE-инструкции на M-чипах — буду рад комментариям.
Что такое RMD160?
RIPEMD-160 (RMD160) — это криптографическая хеш-функция, создающая 160-битный хеш из произвольных данных. Она была разработана как безопасная альтернатива более ранним алгоритмам, таким как MD5 и SHA-1. RMD160 широко применяется в блокчейн-технологиях, особенно в Bitcoin, где используется для создания адресов кошельков: публичный ключ сначала хешируется с помощью SHA-256, а затем — RMD160, для повышения безопасности и сокращения длины.
Алгоритм RMD160 состоит из 5 раундов, каждый из которых включает в себя базовые логические функции, циклические сдвиги (ROTL) и сложение. Особенностью RMD160 является то, что каждый раунд выполняется в двух параллельных ветках: основной (левая) и параллельной (правая). Эти две ветки используют разные константы, порядок обработки слов и логические функции, после чего их результаты объединяются.
Есть классическая C-реализация RMD160 с использованием кучи макросов для объявления раундов, логических функций и т.п., но такой код сложно читать, поэтому мне больше нравится реализация в Golang, которую я ранее уже портировал в ecloop. Дальнейшее портирование на Neon я планирую делать на этой основе.
Если отойти немного в сторону — SIMD-инструкции не выглядят чем-то экстремально сложным, но у них нет синтаксического сахара, поэтому вместо a + b нужно писать что-то в духе vaddq_u32(a, b). Такие специальные функции есть для каждой стандартной операции × количество числовых типов (u/i 8/16/32/64, f16/32/64).
RMD160 (как и другие хеш-функции) должно быть не слишком сложно портировать на SIMD, потому что в их алгоритмах нет ветвлений. По сути, алгоритм остаётся таким же — только все операции нужно заменить на SIMD-специфические инструкции.
Простая Neon программа
Чтобы понять, как писать с использованием Neon, следует начать с максимально простой программы — например, умножить 42 × 2. Так как SIMD — это параллельные вычисления, его операции применяются на весь вектор сразу, и результат в отдельных частях вектора должен быть одинаковым. Чтобы убедиться в этом, результат можно вывести в консоль.
#include <arm_neon.h>
#include <stdint.h>
#include <stdio.h>
void print_check(uint32x4_t *a) {
uint32_t arr[4];
vst1q_u32(arr, *a); // store 4x32-bit vector into a regular array
for (int i = 0; i < 4; i++) {
printf("%x%c", arr[i], i == 3 ? '\n' : ' ');
}
}
int main() {
uint32_t a = 42;
uint32x4_t b = vdupq_n_u32(42); // load u32 to all 4 lanes (42, 42, 42, 42)
printf("%x = ", a);
print_check(&b); // out: 2a = 2a 2a 2a 2a
a = a * 2;
b = vmulq_n_u32(b, 2); // multiply each lane by 2
printf("%x = ", a);
print_check(&b); // out: 54 = 54 54 54 54
return 0;
}
В общем, идея того, как работают SIMD-вычисления, думаю, понятна. Далее print_check будет использоваться часто для проверки корректности алгоритма.
Базовые функции и ROTL
В вычислении RMD160 используются 5 базовых функций и ROTL — всё остальное это перемешивание данных в определённом порядке. GPT переписал эти макросы, а я проверил их корректность:
// original functions
#define OLD_F1(x, y, z) ((x) ^ (y) ^ (z))
#define OLD_F2(x, y, z) (((x) & (y)) | (~(x) & (z)))
#define OLD_F3(x, y, z) (((x) | ~(y)) ^ (z))
#define OLD_F4(x, y, z) (((x) & (z)) | ((y) & ~(z)))
#define OLD_F5(x, y, z) ((x) ^ ((y) | ~(z)))
#define OLD_ROTL(x, n) (((x) << (n)) | ((x) >> (32 - (n))))
// simd functions
#define F1(x, y, z) veorq_u32(veorq_u32(x, y), z)
#define F2(x, y, z) vorrq_u32(vandq_u32(x, y), vandq_u32(vmvnq_u32(x), z))
#define F3(x, y, z) veorq_u32(vorrq_u32(x, vmvnq_u32(y)), z)
#define F4(x, y, z) vorrq_u32(vandq_u32(x, z), vandq_u32(y, vmvnq_u32(z)))
#define F5(x, y, z) veorq_u32(x, vorrq_u32(y, vmvnq_u32(z)))
#define ROTL(x, n) vorrq_u32(vshlq_n_u32(x, n), vshrq_n_u32(x, 32 - (n)))
void print_check(char *l, uint32_t c, uint32x4_t a) {
printf("%s: %08x = ", l, c);
uint32_t arr[4];
vst1q_u32(arr, a); // store 4x32-bit vector into a regular array
for (int i = 0; i < 4; i++) {
printf("%08x%c", arr[i], i == 3 ? '\n' : ' ');
// assert(arr[i] == c);
}
}
uint32_t a1, b1, c1;
uint32x4_t a2, b2, c2;
a1 = 0x67452301, b1 = 0xefcdab89, c1 = 0x98badcfe;
a2 = vdupq_n_u32(a1), b2 = vdupq_n_u32(b1), c2 = vdupq_n_u32(c1); // loading vectors
// // compare original and simd functions
print_check("F1", OLD_F1(a1, b1, c1), F1(a2, b2, c2));
print_check("F2", OLD_F2(a1, b1, c1), F2(a2, b2, c2));
print_check("F3", OLD_F3(a1, b1, c1), F3(a2, b2, c2));
print_check("F4", OLD_F4(a1, b1, c1), F4(a2, b2, c2));
print_check("F5", OLD_F5(a1, b1, c1), F5(a2, b2, c2));
print_check("RL", OLD_ROTL(a1, 12), ROTL(a2, 12));
// output:
// F1: 10325476 = 10325476 10325476 10325476 10325476
// F2: ffffffff = ffffffff ffffffff ffffffff ffffffff
// F3: efcdab89 = efcdab89 efcdab89 efcdab89 efcdab89
// F4: 67452301 = 67452301 67452301 67452301 67452301
// F5: 88888888 = 88888888 88888888 88888888 88888888
// RL: 52301674 = 52301674 52301674 52301674 52301674
Проблема с портированием Golang-реализации
Golang-реализации состоит из 5 больших циклов, которые выполняют по 16 левых и правых раундов; внутри каждой операции миксуются входные данные по индексу _n[i] и происходит ROTL по индексу _r[i]. В reference C-implementation используется куча подряд идущих макросов, из-за чего, на мой взгляд, читать такое сложнее.
// Golang implementation
static const u8 _n[80] = { /* ... */ }; // Left DATA indexes
static const u8 _r[80] = { /* ... */ }; // Left ROTL indexes
// round 1
for (; i < 16; ++i) {
// left branch
alpha = a1 + F1(b1, c1, d1) + x[_n[i]];
alpha = rotl32(alpha, _r[i]) + e1;
beta = rotl32(c1, 10);
a1 = e1, c1 = b1, e1 = d1, b1 = alpha, d1 = beta;
// right branch
// ...
}
// Reference C-implementation
#define F(x, y, z) ((x) ^ (y) ^ (z))
#define FF(a, b, c, d, e, x, s) {\
(a) += F((b), (c), (d)) + (x);\
(a) = ROL((a), (s)) + (e);\
(c) = ROL((c), 10);\
}
// round 1 - left branch
FF(aa, bb, cc, dd, ee, X[ 0], 11);
FF(ee, aa, bb, cc, dd, X[ 1], 14);
// ...
FF(bb, cc, dd, ee, aa, X[14], 9);
FF(aa, bb, cc, dd, ee, X[15], 8);
В общем, если переписать код напрямую с использованием Neon-инструкций, то компиляция падает с ошибкой: vshlq_n_u32 и vshrq_n_u32 требуют, чтобы значение поворота (второй аргумент) было известно на момент компиляции. Пример для сравнения:
// Golang implementation (original)
#define F1(x, y, z) ((x) ^ (y) ^ (z))
#define rotl32(x, n) (((x) << (n)) | ((x) >> (32 - (n))))
alpha = a1 + F1(b1, c1, d1) + x[_n[i]];
alpha = rotl32(alpha, _r[i]) + e1;
// Golang implementation (SIMD)
#define F1(x, y, z) veorq_u32(veorq_u32(x, y), z)
#define ROTL(x, n) vorrq_u32(vshlq_n_u32(x, n), vshrq_n_u32(x, 32 - (n)))
alpha = vaddq_u32(a1, F1(b1, c1, d1));
alpha = vaddq_u32(alpha, X[_n[i]]);
alpha = vaddq_u32(ROTL(alpha, _r[i]), e1);
// err: argument to '__builtin_neon_vshlq_n_v' must be a constant integer
// err: argument to '__builtin_neon_vshrq_n_v' must be a constant integer
Так что придётся использовать версию на макросах, так как там индексы передаются напрямую (последний аргумент в макросе FF) и раскрываются в константные значения во время компиляции. Возможно, это изменение и к лучшему (позже увидим почему).
Обобщённый макрос раунда
Если посмотреть на раунды RMD160, то там происходят одни и те же действия, но меняются: базовая функция, константа, индекс данных и поворот. В общем, макрос раунда был выше (я взял немного другой код с GitHub). Моя цель — портировать макрос раунда на SIMD.
В раунде мы складываем 4 переменные, делаем ROTL + ещё одно сложение и отдельный ROTL для другой переменной. Так как операции сложения "+" нет в SIMD, нужно использовать специальные инструкции.
Я добавил несколько макросов для сложения векторов и описал сам раунд:
#define ADD2(a, b) vaddq_u32(a, b)
#define ADD3(a, b, c) vaddq_u32(vaddq_u32(a, b), c)
#define ADD4(a, b, c, d) vaddq_u32(vaddq_u32(vaddq_u32(a, b), c), d)
#define RN(a, b, c, d, e, f, x, k, r) \
u = ADD4(a, f, x, vdupq_n_u32(k)); \
a = ADD2(ROTL(u, r), e); \
c = ROTL(c, 10);
В макросе a, b, c, d, e — это переменные состояния, f — значение после вычисления базовой функции, x — это uint32 данных по индексу для текущей итерации, k — константа и r — значение поворота для ROTL.
vdupq_n_u32(k) загружает константу в вектор (одинаковое значение во все 4 lanes). Ранее мы писали код, чтобы умножить вектор на число; для этого используется vmulq_n_u32. Логично предположить, что инструкция для добавления числа к вектору должна быть vaddq_n_u32, но её нет. Вместо этого нужно писать в стиле vaddq_u32(vec1, vdupq_n_u32(2)) (если кто знает, почему так — пишите в комментарии).
Далее, на основании этого макроса, можно определить левые и правые раунды. Тут код аналогичен любой другой реализации на макросах (разве что названия раундов я сделал как Li/Ri).
#define L1(a, b, c, d, e, x, r) RN(a, b, c, d, e, F1(b, c, d), x, 0, r)
#define L2(a, b, c, d, e, x, r) RN(a, b, c, d, e, F2(b, c, d), x, 0x5A827999ul, r)
#define L3(a, b, c, d, e, x, r) RN(a, b, c, d, e, F3(b, c, d), x, 0x6ED9EBA1ul, r)
#define L4(a, b, c, d, e, x, r) RN(a, b, c, d, e, F4(b, c, d), x, 0x8F1BBCDCul, r)
#define L5(a, b, c, d, e, x, r) RN(a, b, c, d, e, F5(b, c, d), x, 0xA953FD4Eul, r)
#define R1(a, b, c, d, e, x, r) RN(a, b, c, d, e, F5(b, c, d), x, 0x50A28BE6ul, r)
#define R2(a, b, c, d, e, x, r) RN(a, b, c, d, e, F4(b, c, d), x, 0x5C4DD124ul, r)
#define R3(a, b, c, d, e, x, r) RN(a, b, c, d, e, F3(b, c, d), x, 0x6D703EF3ul, r)
#define R4(a, b, c, d, e, x, r) RN(a, b, c, d, e, F2(b, c, d), x, 0x7A6D76E9ul, r)
#define R5(a, b, c, d, e, x, r) RN(a, b, c, d, e, F1(b, c, d), x, 0, r)
Теперь, используя эти макросы, можно написать первую итерацию первого раунда, сравнить её с работающей реализацией. Если всё ок, тогда можно скопировать весь раунд целиком, проверить его, а затем и все оставшиеся раунды. Результат я сравнивал функцией print_check, которую делал ранее.
Первый левый раунд, первая итерация:
#define K1 0x67452301
#define K2 0xEFCDAB89
#define K3 0x98BADCFE
#define K4 0x10325476
#define K5 0xC3D2E1F0
void rmd160_block(uint32x4_t *s, const uint32_t x[4][16]) {
// a1-e1 left rounds state, a2-e2 right rounds state, u - temp varible used in RD macro
uint32x4_t a1, b1, c1, d1, e1, a2, b2, c2, d2, e2, u;
// Load initial constants
a1 = a2 = vdupq_n_u32(K1);
b1 = b2 = vdupq_n_u32(K2);
c1 = c2 = vdupq_n_u32(K3);
d1 = d2 = vdupq_n_u32(K4);
e1 = e2 = vdupq_n_u32(K5);
uint32x4_t w[16]; // Load data to vector
for (int i = 0; i < 16; i++) {
// Load 4x32-bit integers from x[0][i], x[1][i], x[2][i], x[3][i]
// w[i] = vsetq_lane_u32(x[0][i], w[i], 0);
// w[i] = vsetq_lane_u32(x[1][i], w[i], 1);
// w[i] = vsetq_lane_u32(x[2][i], w[i], 2);
// w[i] = vsetq_lane_u32(x[3][i], w[i], 3);
w[i] = vld1q_u32(((uint32_t[4]){x[0][i], x[1][i], x[2][i], x[3][i]})); // A bit faster
}
L1(a1, b1, c1, d1, e1, w[0], 11);
print_check("a1", 0, a1);
print_check("b1", 0, b1);
print_check("c1", 0, c1);
print_check("d1", 0, d1);
print_check("e1", 0, e1);
}
uint32x4_t s[5] = {0}; // initial state
s[0] = vdupq_n_u32(K1);
s[1] = vdupq_n_u32(K2);
s[2] = vdupq_n_u32(K3);
s[3] = vdupq_n_u32(K4);
s[4] = vdupq_n_u32(K5);
uint32_t x[4][16] = {0}; // data block, filled with zeros
rmd160_block((uint32x4_t *)s, x);
Следует заметить, что обычно хеш-функции тестируют на нулевых данных (для простоты работы). Данные в хеш-функциях обрабатываются по блокам. Блок в RMD160 — это 32×16 = 512 бит. rmd160_block можно вызывать несколько раз с тем же состоянием (он меняется) и новыми данными — для случаев, когда нужно посчитать хеш сообщения, большего чем один раунд. В моей задаче (генерация адресов) все сообщения помещаются в один блок. Результат первого раунда vs текущая реализация:
// a1: 1602f864 1602f864 1602f864 1602f864 vs c3d2e1f0
// a1: efcdab89 efcdab89 efcdab89 efcdab89 vs 1602f864
// a1: eb73fa62 eb73fa62 eb73fa62 eb73fa62 vs efcdab89
// a1: 10325476 10325476 10325476 10325476 vs eb73fa62
// a1: c3d2e1f0 c3d2e1f0 c3d2e1f0 c3d2e1f0 vs 10325476
В общем, эти значения ± похожи на значения из текущей версии, отличаются на одно смещение. Это не проблема, так как переменных 5, и к концу смещения они выравняются. Просто разница в реализациях.
Все раунды я писать не буду, их там по 80 с каждой стороны (всего 160) — выйдет слишком длинно. Левые и правые раунды не зависят друг от друга, и вычислять их можно в любом порядке: либо сначала все левые / все правые, либо чередовать левые / правые, либо вообще чередовать итерации внутри раунда. На финальный результат это не повлияет.
Финализация RMD160
В конце блока RMD160 нужно объединить старое состояние с локальным состоянием — это тоже три сложения со смещением индексов.
void rmd160_block(uint32x4_t *s, const uint32_t x[4][16]) {
// ... 160 rounds
uint32x4_t t = s[0];
s[0] = ADD3(s[1], c1, d2);
s[1] = ADD3(s[2], d1, e2);
s[2] = ADD3(s[3], e1, a2);
s[3] = ADD3(s[4], a1, b2);
s[4] = ADD3(t, b1, c2);
}
Финально остаётся изменить endianness у значений (RMD160 использует не тот endianness) и выгрузить значения из вектора в результирующий массив.
// ... init & rmd160_block
for (int i = 0; i < 5; ++i) {
// swap32 for uint32x4_t (can it be shorter?)
s[i] = vreinterpretq_u32_u8(vrev32q_u8(vreinterpretq_u8_u32(s[i])));
}
uint32_t r[4][5] = {0}; // result stored as 4x5 uint32_t
for (int i = 0; i < 5; i++) { // load it from uint32x4_t
r[0][i] = vgetq_lane_u32(s[i], 0);
r[1][i] = vgetq_lane_u32(s[i], 1);
r[2][i] = vgetq_lane_u32(s[i], 2);
r[3][i] = vgetq_lane_u32(s[i], 3);
}
На этом, в общем, и всё — r может использоваться дальше, где нужно (r[0], r[1], r[2], r[3] — это посчитанные хеши, 4 сразу).
Подводя итог по этой секции, полный алгоритм параллельного RMD160 выглядит так:
- Инициализировать стейт длиной 160 бит × 4 lanes (
uint32x4_t s[5]). - Разбить сообщение (данные) на блоки по 512 бит × 4 lanes (
uint32_t x[4][16]). - Прокрутить RMD160 раунды, пока данные не закончатся (
rmd160_blockсам считывает данные в вектор). - Изменить endianness в финальном стейте.
- Выгрузить вектор финального стейта в массив хешей (
int32_t r[4][5]).
Производительность RMD160 SIMD
Теперь остаётся замерить то, зачем это затевалось — сравнить производительность оригинального и SIMD-кода. Для этого я сделал небольшой benchmark:
size_t tsnow() {
struct timespec ts;
clock_gettime(CLOCK_REALTIME, &ts);
return ts.tv_sec * 1000 + ts.tv_nsec / 1e6;
}
void rmd160_simd() {
uint32_t r[4][5] = {0};
uint32_t x[4][16] = {0};
size_t stime = tsnow();
size_t iters = 1000 * 1000 * 32;
for (size_t i = 0; i < iters; ++i) rmd160_4w(r, x);
double dt = (tsnow() - stime) / 1000.0;
double ir = iters / dt / 1000000;
double hr = ir * 4; // 4 hash per iter
printf("%.2fM it/s ~ %.2fM h/s ~ %.2fs\n", ir, hr, dt);
printf("s[0]: %08x\n", r[0][0]);
printf("s[1]: %08x\n", r[0][1]);
printf("s[2]: %08x\n", r[0][2]);
printf("s[3]: %08x\n", r[0][3]);
printf("s[4]: %08x\n", r[0][4]);
}
void rmd160_naive() {
uint32_t s[5] = {0};
uint32_t x[16] = {0};
size_t stime = tsnow();
size_t iters = 1000 * 1000 * 32;
for (size_t i = 0; i < iters; ++i) rmd160_1w(s, x);
double dt = (tsnow() - stime) / 1000.0;
double ir = iters / dt / 1000000;
double hr = ir * 1; // 1 hash per iter
printf("%.2fM it/s ~ %.2fM h/s ~ %.2fs\n", ir, hr, dt);
printf("s[0]: %08x\n", s[0]);
printf("s[1]: %08x\n", s[1]);
printf("s[2]: %08x\n", s[2]);
printf("s[3]: %08x\n", s[3]);
printf("s[4]: %08x\n", s[4]);
}
Скомпилировал обе программы с -O3 и запустил (на базовом Apple M2):
❯ clang -O3 -march=native ./lib/rmd160.c && ./a.out # original
5.50M it/s ~ 5.50M h/s ~ 5.81s
❯ clang -O3 -march=native ./lib/rmd160s.c && ./a.out # neon
2.14M it/s ~ 8.55M h/s ~ 14.98s
Вышло, что версия на Neon (128 бит / 4 lanes) работает на 55% быстрее. Что, конечно, крутой результат, но печально, что на M-chips нет SVE для 256/512 бит (8/16 lanes) — так было бы ещё лучше!
One more thing
Пока я записывал результаты работы выше, мне стало любопытно, что будет, если поиграться с порядком раундов в RMD160. Изначальный порядок раундов был такой: сначала все левые раунды, затем все правые. Мне казалось, что это хорошо для процессора, потому что, на первый взгляд, нужно меньше "переключений контекста".
Я поменял порядок раундов на чередование левых и правых (L1 / R1, L2 / R2), и скорость работы значительно возросла. Я изначально подумал, что ошибка в данных, но print_check (в качестве тестов) говорит, что всё в порядке.
Я решил попробовать чередовать итерации (L1_1 R1_1 L1_2 R1_2 и т.д.). Честно говоря, переставлять 160 строчек — не самое веселое занятие, но результат удивил меня ещё больше.
Сравнение разных размещений (три теста самой RMD160 функции и полный цикл работы логики ecloop):
# L1_1 L2_2 .. L5_16 R1_1 R2_2 .. R5_16 (+56%)
2.25M it/s ~ 9.02M h/s ~ 14.19s
2.23M it/s ~ 8.93M h/s ~ 14.33s
2.23M it/s ~ 8.94M h/s ~ 14.32s
ecloop (addr33 x 8 core) ~ 19.53M it/s (+22%)
# L1_1-L1_16 R1_1-R1_16 L2_1-L2_16 .. (+165%)
3.70M it/s ~ 14.80M h/s ~ 8.65s
3.85M it/s ~ 15.42M h/s ~ 8.30s
3.87M it/s ~ 15.46M h/s ~ 8.28s
ecloop (addr33 x 8 core) ~ 22.46M it/s (+40%)
# L1_1 R1_1 L1_2 R1_2 .. L5_16 R5_16 (+175%)
3.96M it/s ~ 15.82M h/s ~ 8.09s
3.94M it/s ~ 15.78M h/s ~ 8.11s
3.94M it/s ~ 15.76M h/s ~ 8.12s
ecloop (addr33 x 8 core) ~ 24.83M it/s (+55%)
Для меня загадка, почему это так работает, и, возможно, есть ещё более эффективная расстановка? Кто знает — пишите в комментарии.
Поддержка AVX2 (AMD64)
Изначально у меня не было этого в планах, но алгоритм RMD160 на макросах вышел довольно абстрактным, и дальнейшее портирование на AVX2 выглядело довольно простым. Основное отличие AVX2 от Neon (кроме другого набора инструкций) — это размер вектора: 256 бит против 128 бит. То есть, можно посчитать 8 хешей параллельно (против 4-х у Neon).
Дальнейшее портирование состоит из таких этапов:
- вынести весь Neon-специфический код в макросы
- переписать все макросы под AVX2
- проверить корректность работы на amd64
Сейчас напрямую в коде алгоритма используются такие Neon-инструкции: тип вектора (uint32x4_t), инициализация состояния через vdupq_n_u32, endian-swap и load / dump данных в вектор.
Переношу эти вещи в макросы (на всякий случай я добавил префикс RMD_, чтобы не было конфликтов с другими файлами):
#define RMD_LEN 4 // vector length
#define RMD_VEC uint32x4_t // vector type
#define RMD_LD_NUM(x) vdupq_n_u32(x) // load same number into all lanes
#define RMD_SWAP(x) vreinterpretq_u32_u8(vrev32q_u8(vreinterpretq_u8_u32(x)))
#define RMD_LOAD(x, i) vld1q_u32(((uint32_t[4]){x[0][i], x[1][i], x[2][i], x[3][i]}))
#define RMD_DUMP(r, s, i) \
do { \
r[0][i] = vgetq_lane_u32(s[i], 0); \
r[1][i] = vgetq_lane_u32(s[i], 1); \
r[2][i] = vgetq_lane_u32(s[i], 2); \
r[3][i] = vgetq_lane_u32(s[i], 3); \
} while (0);
И обновить текущий код на что-то такого рода:
void rmd160_block(RMD_VEC *s, const uint32_t x[RMD_LEN][16]) {
RMD_VEC a1, b1, c1, d1, e1, a2, b2, c2, d2, e2, u;
a1 = a2 = RMD_LD_NUM(RMD_K1);
b1 = b2 = RMD_LD_NUM(RMD_K2);
c1 = c2 = RMD_LD_NUM(RMD_K3);
d1 = d2 = RMD_LD_NUM(RMD_K4);
e1 = e2 = RMD_LD_NUM(RMD_K5);
RMD_VEC w[16];
for (int i = 0; i < 16; i++) w[i] = RMD_LOAD(x, i);
// ... rounds and iterations
}
// new function to process full single block
void rmd160_batch(uint32_t r[RMD_LEN][5], const uint32_t x[RMD_LEN][16]) {
RMD_VEC s[5] = {0}; // load initial state
s[0] = RMD_LD_NUM(RMD_K1);
s[1] = RMD_LD_NUM(RMD_K2);
s[2] = RMD_LD_NUM(RMD_K3);
s[3] = RMD_LD_NUM(RMD_K4);
s[4] = RMD_LD_NUM(RMD_K5);
rmd160_block((RMD_VEC *)s, x); // round
for (int i = 0; i < 5; ++i) s[i] = RMD_SWAP(s[i]); // change endian
for (int i = 0; i < 5; ++i) RMD_DUMP(r, s, i); // dump data to array
}
Вышло, конечно, уже достаточно магически, но самих изменений реально не так много. Теперь следует добавить переопределённые макросы для AVX2. Также я дополнительно обернул серии макросов, специфичных для конкретной архитектуры, в #ifdef. Таким образом, по сути, у меня есть один код алгоритма RMD160, и нужные макросы подключаются в зависимости от того, на какой архитектуре компилируется программа.
#if defined(__aarch64__) && defined(__ARM_NEON)
#include <arm_neon.h>
#define RMD_LEN 4 // vector length
#define RMD_VEC uint32x4_t // vector type
// ... move all current Neon related macros here
#elif defined(__x86_64__) && defined(__AVX2__)
#include <immintrin.h>
#define RMD_LEN 8
#define RMD_VEC __m256i
#define RMD_LD_NUM(x) _mm256_set1_epi32(x)
#define RMD_SWAP(x) \
_mm256_shuffle_epi8((x), _mm256_setr_epi8(3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8, 15, 14, 13, \
12, 19, 18, 17, 16, 23, 22, 21, 20, 27, 26, 25, 24, \
31, 30, 29, 28))
#define RMD_LOAD(x, i) \
_mm256_set_epi32(x[0][i], x[1][i], x[2][i], x[3][i], x[4][i], x[5][i], x[6][i], x[7][i])
#define RMD_DUMP(r, s, i) \
do { \
r[0][i] = _mm256_extract_epi32(s[i], 0); \
r[1][i] = _mm256_extract_epi32(s[i], 1); \
r[2][i] = _mm256_extract_epi32(s[i], 2); \
r[3][i] = _mm256_extract_epi32(s[i], 3); \
r[4][i] = _mm256_extract_epi32(s[i], 4); \
r[5][i] = _mm256_extract_epi32(s[i], 5); \
r[6][i] = _mm256_extract_epi32(s[i], 6); \
r[7][i] = _mm256_extract_epi32(s[i], 7); \
} while (0);
#define _mm256_not_si256(x) _mm256_xor_si256((x), _mm256_set1_epi32(0xffffffff))
#define RMD_F1(x, y, z) _mm256_xor_si256(x, _mm256_xor_si256(y, z))
#define RMD_F2(x, y, z) _mm256_or_si256(_mm256_and_si256(x, y), _mm256_andnot_si256(x, z))
#define RMD_F3(x, y, z) _mm256_xor_si256(_mm256_or_si256(x, _mm256_not_si256(y)), z)
#define RMD_F4(x, y, z) _mm256_or_si256(_mm256_and_si256(x, z), _mm256_andnot_si256(z, y))
#define RMD_F5(x, y, z) _mm256_xor_si256(x, _mm256_or_si256(y, _mm256_not_si256(z)))
#define RMD_ROTL(x, n) _mm256_or_si256(_mm256_slli_epi32(x, n), _mm256_srli_epi32(x, 32 - (n)))
#define RMD_ADD2(a, b) _mm256_add_epi32(a, b)
#define RMD_ADD3(a, b, c) _mm256_add_epi32(_mm256_add_epi32(a, b), c)
#define RMD_ADD4(a, b, c, d) _mm256_add_epi32(_mm256_add_epi32(a, b), _mm256_add_epi32(c, d))
#else
#error "Unsupported arch for RIPEMD-160 (AVX2 or NEON required)"
#endif
Основные отличия следующие:
- другой заголовочный файл
- другой тип вектора (8 lanes вместо 4 у Neon) и названия intrinsics
- у AVX2 нет Bitwise NOT, поэтому пришлось добавить его отдельно как
_mm256_not_si256. - нет отдельной функции для endian-swap, но есть более обобщённая функция, чтобы переставлять биты в заданном порядке —
_mm256_shuffle_epi8(первый аргумент — где переставить биты, второй аргумент — как переставить). _mm256_set_epi32позволяет удобнее загружать данные в разные lanes, в Neon пришлось использовать временный массив (вариант с установкой каждой lane по отдельности (vsetq_lane_u32) оказался более медленным).
С переписыванием макросов почти справился GPT, я лишь в очередной раз проверил их корректность.
Производительность AVX2 версии
У меня есть небольшой Fanless PC с Linux на Intel N100, который я использую для нативного тестирования приложений. На нем я запустил бенчмарк, написанный ранее, и получил такие результаты:
❯ clang -O3 -march=native ./lib/rmd160.c && ./a.out # original (on Intel N100)
4.26M it/s ~ 4.26M h/s ~ 7.51s
❯ clang -O3 -march=native ./lib/rmd160s.c && ./a.out # avx2 (on Intel N100)
2.25M it/s ~ 17.96M h/s ~ 14.25s
8 lanes AVX2 и правильная расстановка раундов в алгоритме (итерации "зеброй") дают прирост 320% в количестве хешей в секунду по сравнению с оригинальным кодом. Что интересно, AVX2 на Intel N100 работает на 20% быстрее, чем Neon на Apple M2 (в основном из-за размера вектора). Ускорение работы программы в целом составило: 5.45M it/s vs 7.73M it/s (+42%).
Fallback реализация для старых процессоров / VMs
В #ifdef выше я оставил секцию #else с #error, чтобы компиляция не происходила на неподдерживаемых системах. В общем, это не совсем хорошее решение, и хотелось бы, чтобы программа работала везде (в основном это касается потенциального запуска в виртуальных машинах). Так как весь алгоритм уже написан на макросах, добавление новой реализации не составит труда. Просто переопределяю все макросы на размер вектора 1 и uint32_t в качестве "векторного" типа. В реальности программа будет работать с единичным массивом, что с точки зрения памяти в C примерно то же самое, что и просто использование uint32_t.
#else
#warning "Fallback RIPEMD-160 implementation used. AVX2 or NEON required for SIMD"
#define RMD_LEN 1
#define RMD_VEC uint32_t
#define RMD_LD_NUM(x) x
#define RMD_SWAP(x) __builtin_bswap32(x)
#define RMD_LOAD(x, i) x[0][i]
#define RMD_DUMP(r, s, i) r[0][i] = s[i]
#define RMD_F1(x, y, z) ((x) ^ (y) ^ (z))
#define RMD_F2(x, y, z) (((x) & (y)) | (~(x) & (z)))
#define RMD_F3(x, y, z) (((x) | ~(y)) ^ (z))
#define RMD_F4(x, y, z) (((x) & (z)) | ((y) & ~(z)))
#define RMD_F5(x, y, z) ((x) ^ ((y) | ~(z)))
#define RMD_ROTL(x, n) (((x) << (n)) | ((x) >> (32 - (n))))
#define RMD_ADD2(a, b) (a + b)
#define RMD_ADD3(a, b, c) (a + b + c)
#define RMD_ADD4(a, b, c, d) (a + b + c + d)
#endif
Запустил, проверил — работает корректно, разве что скорость немного возросла из-за нового порядка раундов. Точно таким же образом несложно добавить реализацию для AVX512 (16 lanes), но у меня нет подходящего процессора для проверки, поэтому я этого не сделал. Да и без этого статья уже получилась достаточно длинной.
Выводы
SIMD-программирование оказалось проще, чем я ожидал. Нужные intrinsics можно легко найти с помощью GPT, что значительно упрощает задачу. Самостоятельно их искать сложно, потому что существует множество возможных комбинаций инструкций. Алгоритм RMD160 в основном используется в криптовалюте (по крайней мере, я не знаю других популярных сценариев использования), поэтому практическая ценность полученного кода вне обучающего контекста может быть сомнительной.
SIMD-вычисления дают отличный прирост скорости выполнения, но, конечно, важно учитывать специфику задачи: они работают эффективно, когда нужно обработать большое количество данных одинакового размера. Нет смысла использовать SIMD, если один блок данных для хеширования имеет размер 1, а другой — 100 (например, при обработке различных файлов). Основная программа, в свою очередь, должна быть способна обрабатывать данные батчами.
Финальный код одним файлом на Github.