Stream compaction на NEON. Векторизуем copy_if

Задача

Даны массивы a и out. Нужно записать в out подряд, без пробелов, только те элементы a, которые удовлетворяют заданному условию. В статье — условие a[i] > threshold, a[i] ∈ (0, 1), threshold ∈ {0, 0.5, 1}.

Почему компилятор сдается

Один if внутри цикла копирования роняет скорость с 112 до 2.6 GB/s. Компилятор не может векторизовать — в NEON нет подходящей для этого инструкции. Статья о том, как собрать ее.

auto copy_if(const float* a, float* out, size_t n) {    size_t j = 0;    for (size_t i = 0; i < n; ++i) {        if (a[i] > 0) out[j++] = a[i];    }    return j;}

В задаче copy_if выходной курсор (в коде — j) зависит от данных, поэтому, чтобы компилятор векторизовал, нужна compress-инструкция (собирает выбранные элементы регистра подряд в его начало, без дырок), а в NEON ее нет, поэтому компилятор сдается:

clang++ -O3 -Rpass-analysis=loop-vectorize -std=c++23 main.cpp -o mainmain.cpp:5:5: remark: loop not vectorized: value that could not be identified as reduction is used outside the loop [-Rpass-analysis=loop-vectorize]5 | for (size_t i = 0; i < n; ++i) {| ^main.cpp:6:23: remark: loop not vectorized: cannot identify array bounds [-Rpass-analysis=loop-vectorize]6 | if (a[i] > 0) out[j++] = a[i];

Векторизатор clang умеет классифицировать j либо как induction (фиксированный шаг), либо как reduction (свертка). j не является ни тем, ни другим — он курсор, зависящий от данных. Компилятор не умеет векторизовать такой вид курсоров. Вторая ремарка отсюда же: он не может вычислить диапазон обращений к out.

Замер: две проблемы скаляра

все замеры: Apple M5; clang++ -O3 -std=c++23 -march=native; GB/s = (2n * 4 байта) / время, минимум из 3e9 / n прогонов; cache: n=1e5, DRAM: n=1e7

“copy a[i]” — тот же цикл, но без условия, его компилятор векторизует. Вся разница с ним в одном if. Данные те же, меняется только предсказуемость ветки:

1. > 0 (всегда true) и > 1 (всегда false): branch predictor не промахивается → 33 GB/s. Из-за отсутствия векторизации скорость падает в 3 раза

2. > 0.5 (50/50): промах на каждом втором → 3 GB/s

Трюк исправит обе эти проблемы.

Трюк 1: эмулируем compress

Дальше считаем n кратным ширине регистра; хвост — отдельная тема, к трюку отношения не имеет. Еще:

1. Размер out должен быть >= n.

2. Пусть алгоритм выбрал cnt элементов, тогда все в out[cnt, n) будет не определено (мусор). Алгоритм, сохраняющий чистоту хвоста, не добавляет к идее ничего нового, поэтому не рассматривается.

В NEON нет инструкции для сжатия регистра, поэтому нам придется эмулировать ее.

(Стоит сказать, что сам приём — не новый. Лемир использовал его еще в 2017 году на SSE link. Но на NEON нет ни movemask, ни дешевого popcnt — дальше о том, как собрать все это из того, что есть)

Что должен уметь делать наш аналог compress:

1. Принимать регистр элементов a и регистр-маску, по которой будет выбор.

2. Сказать, сколько элементов из регистра мы взяли (чтобы переместить указатель out).

3. Собственно, записать выбранные элементы в out.

tbl: произвольная выборка байт

В NEON есть семейство инструкций table-lookup (tbl), их назначение — произвольная выборка/перестановка байт. Инструкция принимает два регистра:

1. Таблица (table) — байты, откуда будем выбирать.

2. Индексы (index) — номера байтов, которые нужно взять.

Проще говоря — это SIMD-аналог out[i] = table[index[i]].

Мы будем использовать инструкцию vqtbl1q_u8:

tbl переставляет байты, а нам нужно выбрать float (4 байта). Значит, index будем делать блоками по 4 байта: чтобы взять второй (0-based) float регистра, в index пойдут его байты [8, 9, 10, 11] (второй элемент начинается со смещения 2 *sizeof(float) = 8).

Считать index каждый раз долго. Всего вариантов 16 (4 элемента взять/выкинуть), поэтому посчитаем заранее все варианты index. Но чтобы по маске выбирать index, надо превратить маску в число (назовем его idx):

Маска → idx

Маска состоит из 4 элементов, каждый элемент — либо 0x00000000 (false), либо 0xFFFFFFFF (true). Хотим: если i-й элемент true — выставить в idx i-й бит. Трюк: mask & [1, 2, 4, 8]. Так как 0xFFFFFFFF & x = x, в истинных элементах останется вес (1/2/4/8), а в ложных 0. Складываем все элементы между собой — получаем число 0…15.

std::array<uint32_t, 4> weights{1, 2, 4, 8};size_t idx = vaddvq_u32(vandq_u32(mask, vld1q_u32(weights.data())));

• vld1q_u32(weights.data()) — загрузить 4 значения из памяти по адресу weights.data() в регистры (ld — load)

• vandq_u32 — поэлементное & (and)

• vaddvq_u32 — сумма всех элементов регистра (addv — add across vector)

Строим таблицу index

Регистры нельзя вычислить на этапе компиляции, поэтому вместо uint8x16_t(регистр из 16 uint8_t) будем хранить std::array<uint8_t, 16>. Для каждого idx идем по 4 элементам маски. Если элемент выбран — дописываем номера его 4 байт в index на позицию курсора и сдвигаем курсор на 4.

consteval auto make_index_table() {    std::array<std::array<uint8_t, 16>, 16> index{};    for (size_t idx = 0; idx < 16; ++idx) {      // перебираем все маски        size_t j = 0;                               // j - курсор        for (size_t i = 0; i < 4; ++i)              // перебираем элементы маски            if (idx & (1 << i))                    // если i элемент выбран                for (size_t k = 0; k < 4; ++k)      // перебираем байты элемента                    index[idx][j++] = i * 4 + k;     // пишем номера его байт    }    return index;}

Курсор j двигается только на выбранных, поэтому байты выбранных элементов ложатся в index подряд, а tbl по index соберет float в регистр. Незанятые позиции в index — нули, поэтому в хвосте, после count элементов, будет мусор.

Таблица count

Дальше нужно научиться считать количество элементов, которые берем. Таким же образом посчитаем таблицу:

consteval auto make_count_table() {    std::array<uint8_t, 16> count{};    for (size_t idx = 0; idx < 16; ++idx)        for (size_t i = 0; i < 4; ++i)            if (idx & 1 << i)                ++count[idx];    return count;}

Compress целиком

auto compress(uint32x4_t mask, float32x4_t a) {    static constexpr std::array<uint32_t, 4> weights{1, 2, 4, 8};    const size_t idx = vaddvq_u32(vandq_u32(mask, vld1q_u32(weights.data())));    static constexpr auto count = make_count_table();    static constexpr auto index_table = make_index_table();    const auto index = vld1q_u8(index_table[idx].data()); // в рантайме одну строку таблицы грузим в регистр    return std::pair{vreinterpretq_f32_u8(vqtbl1q_u8(vreinterpretq_u8_f32(a), index)), count[idx]};}

Так как tbl работает только с u8, надо сначала a перевести в u8, а потом результат обратно в f32.

Записывать в память будем полными регистрами — по 4 float, а курсор двигать только на cnt. Compress кладет валидные элементы в начало регистра, то есть [j, j + cnt), а в [j + cnt, j + 4) — мусор. А следующая итерация начнется с j + count, то есть затрет мусор предыдущего шага. Мусор останется только после последнего store в out[cnt, n), а за границы out[0, n) не выходим — курсор никогда не обгоняет прочитанные элементы.

Цикл copy_if

auto copy_if_neon(const float* __restrict a,                  float* __restrict out,                  float threshold,                  size_t n) {    auto thd = vdupq_n_f32(threshold);          // загружаем threshold в регистр    size_t j = 0;     for (size_t i = 0; i < n; i += 4) {        auto v = vld1q_f32(a + i);              // загружаем текущие 4 элемента a в регистр        auto mask = vcgtq_f32(v, thd);          // вычисляем маску в регистрах        auto [packed, cnt] = compress(mask, v);        vst1q_f32(out + j, packed);             // сохраняем packed в out[j, j + 4), в [j + cnt, j + 4) будет мусор         j += cnt;     }    return j;}

— vcgtq_f32(v, thd) — вычисление поэлементного v[i] > thd[i]. cgt — compare greater — vst1q_f32 — сохранение в память 4 float из регистра. st — store

Результат

> 0.5 для скаляра был худшим случаем. 3 GB/s. Стало 71 GB/s. Прирост более чем в 20 раз. Теперь нет веток, поэтому скорость больше не зависит от данных.

Трюк 2: Считаем idx и count за один addv

idx всегда меньше 16, поэтому сделаем weights = {1 + 16, 2 + 16, 4 + 16, 8 + 16} и тогда s = (сумма внутри mask & weights) / 16 даст количество элементов, а s % 16 — сам idx. Поэтому можно не считать таблицу count, и теперь compress:

auto compress(uint32x4_t mask, float32x4_t a) {    static constexpr std::array<uint32_t, 4> weights{1 + 16, 2 + 16, 4 + 16, 8 + 16};    const size_t s = vaddvq_u32(vandq_u32(mask, vld1q_u32(weights.data())));    const size_t count = s >> 4; // то же самое, что и s / 16    const size_t idx = s & 15;   // то же самое, что и s % 16        static constexpr auto index_table = make_index_table();        const auto index = vld1q_u8(index_table[idx].data());    return std::pair{vreinterpretq_f32_u8(vqtbl1q_u8(vreinterpretq_u8_f32(a), index)), count};}

И теперь скорость поднимается еще:

Unroll

Еще можно скорость выжать разворотом цикла по 4 итерации (16 элементов), числа:

Итоговый код:

consteval auto make_index_table() {    std::array<std::array<uint8_t, 16>, 16> index{};    for (size_t idx = 0; idx < 16; ++idx) {        size_t j = 0;        for (size_t i = 0; i < 4; ++i)            if (idx & (1 << i))                for (size_t k = 0; k < 4; ++k)                    index[idx][j++] = i * 4 + k;    }    return index;}auto compress(uint32x4_t mask, float32x4_t a) {    static constexpr std::array<uint32_t, 4> weights{1 + 16, 2 + 16, 4 + 16, 8 + 16};    const size_t s = vaddvq_u32(vandq_u32(mask, vld1q_u32(weights.data())));    const size_t count = s >> 4;    const size_t idx = s & 15;        static constexpr auto index_table = make_index_table();        const auto index = vld1q_u8(index_table[idx].data());    return std::pair{vreinterpretq_f32_u8(vqtbl1q_u8(vreinterpretq_u8_f32(a), index)), count};}auto copy_if_neon_unroll(const float* __restrict a,                          float* __restrict out,                          float threshold,                          size_t n) {    auto thd = vdupq_n_f32(threshold);    size_t j = 0;    size_t i = 0;    for (; i + 16 <= n; i += 16) {#pragma unroll        for (size_t i0 = 0; i0 < 16; i0 += 4) {            auto v = vld1q_f32(a + i + i0);            auto mask = vcgtq_f32(v, thd);            auto [packed, cnt] = compress(mask, v);            vst1q_f32(out + j, packed);            j += cnt;        }    }    for (; i + 4 <= n; i += 4) {        auto v = vld1q_f32(a + i);        auto mask = vcgtq_f32(v, thd);        auto [packed, cnt] = compress(mask, v);        vst1q_f32(out + j, packed);        j += cnt;    }    return j;}

tbl и таблица индексов дают то, чего в NEON нет из коробки — compress. И это не про >threshold: так же делаются filter, remove и прочие функции, где выход зависит от данных.