a[mask] = f(a[mask]) на NEON. compress вместо blend

Задача

Применить операцию к тем элементам массива, которые подходят под условие:

for (size_t i = 0; i < n; ++i)    if (mask(a[i])) a[i] = f(a[i]);

Примечание:

В статье:

• a[i] ∈ (0, 1), thd ∈ (0, 1), mask = a[i] < thd; распределения равномерны (кроме конца статьи)

• f из sqrt, frfrexp (мантисса числа), sin 3.5 ULP, sin 1 ULP, pow 1 ULP (из SLEEF)

• Операции обернуты в лямбду с always_inline, они передаются рантайм значением и иначе могут не заинлайниться.

• Размер массива n кратен всем анроллам, тайлам итд, хвосты обрабатывать элементарно (BSL/скаляр)

• В таблицах * — лучший, единицы измерения — GiB/s

• Числа между таблицами не сравниваются, из-за разных условий значения колеблются

• Все бенчи: Apple M5; clang++ -O3 -std=c++23 -march=native; GiB/s = (n * 4 bytes) / time, минимум из 720 запусков (во время замеров все функции чередуются в меняющемся порядке, и данные восстанавливаются, конечно); n=1e7 + 2432

BSL blend

Если задача упирается в память (функция легкая или плотность высокая), то стандартный алгоритм оптимален:

    template <bool Skip>    void bsl(float* dst, const size_t n, auto f, auto mask) {        for (size_t i = 0; i < n; i += 16) {            std::array<float32x4_t, 4> v;            for (size_t j = 0; j < 4; ++j) v[j] = vld1q_f32(dst + i + 4 * j);            std::array<uint32x4_t, 4> m;            for (size_t j = 0; j < 4; ++j) m[j] = mask(v[j]);            if constexpr (Skip) {                if (vmaxvq_u32(vaddq_u32(vaddq_u32(m[0], m[1]), vaddq_u32(m[2], m[3]))) == 0) continue;            }            for (size_t j = 0; j < 4; ++j) vst1q_f32(dst + i + 4 * j, vbslq_f32(m[j], f(v[j]), v[j]));        }    }

Считает для всех элементов, но сохраняет только по маске. Skip — для разреженных масок поможет, но иначе миспредикты убьют скорость. Дальше он пригодится.

Однако bsl на тяжелых операциях будет выполнять слишком много лишней работы.

Detour

Чтобы не считать лишнее, сожмем нужные (compress), применим только к ним, разожмем обратно (expand).

На avx512 делается это в пару инструкций, но на NEON — нет. Так что придется эмулировать и оптимизировать.

    constexpr double D = 0.845; // D и B дальше нужны будут    constexpr double B = 0.3;    constexpr size_t tile = 4096;    constexpr std::array<uint32_t, 4> weights{1 + 16, 2 + 16, 4 + 16, 8 + 16};    constexpr auto cps_tbl = compress_table();    constexpr auto exp_tbl = expand_table();    std::array<float, tile + 16> tmp;    std::array<uint8_t, tile / 4 + 3> s;    std::array<uint16_t, tile / 4 + 3> idx; // в будущем пригодится    template <bool Skip>    size_t detour(float* dst, const size_t n, const auto w, auto f, auto mask) {        float* ptr = tmp.data();        for (size_t i = 0; i < n; i += 16) {            std::array<float32x4_t, 4> v;            for (size_t j = 0; j < 4; ++j) v[j] = vld1q_f32(dst + i + 4 * j);            std::array<uint32x4_t, 4> m;            for (size_t j = 0; j < 4; ++j) m[j] = mask(v[j]);                        if constexpr (Skip)                if (vmaxvq_u32(vaddq_u32(vaddq_u32(m[0], m[1]), vaddq_u32(m[2], m[3]))) == 0) {                    s[i / 4] = s[i / 4 + 1] = s[i / 4 + 2] = s[i / 4 + 3] = 0;                    continue;                }                        std::array<uint32_t, 4> sk;            for (size_t j = 0; j < 4; ++j) {                sk[j] = vaddvq_u32(vandq_u32(m[j], w));                s[i / 4 + j] = sk[j];            }            std::array<size_t, 4> off; off[0] = 0;            for (size_t j = 1; j < 4; ++j) off[j] = off[j - 1] + (sk[j - 1] >> 4);             std::array<uint8x16_t, 4> index;            for (size_t j = 0; j < 4; ++j) index[j] = vld1q_u8(cps_tbl[sk[j] & 15].data());                        for (size_t j = 0; j < 4; ++j) vst1q_f32(ptr + off[j], vreinterpretq_f32_u8(vqtbl1q_u8(vreinterpretq_u8_f32(v[j]), index[j])));            ptr += off[3] + (sk[3] >> 4);        }        const size_t size = ptr - tmp.data();        if (size == 0) return size;                ptr = tmp.data();        for (size_t i = 0; i < size; i += 16) {            std::array<float32x4_t, 4> v;            for (size_t j = 0; j < 4; ++j) v[j] = vld1q_f32(ptr + i + 4 * j);            for (size_t j = 0; j < 4; ++j) vst1q_f32(ptr + i + 4 * j, f(v[j]));        }                for (size_t i = 0; i < n; i += 16) {            std::array<uint32_t, 4> sk;            for (size_t j = 0; j < 4; ++j) sk[j] = s[i / 4 + j];            std::array<size_t, 4> off{};            for (size_t j = 1; j < 4; ++j) off[j] = off[j - 1] + (sk[j - 1] >> 4);             std::array<float32x4_t, 4> v;            for (size_t j = 0; j < 4; ++j) v[j] = vld1q_f32(ptr + off[j]);            std::array<float32x4_t, 4> a;            for (size_t j = 0; j < 4; ++j) a[j] = vld1q_f32(dst + i + 4 * j);            std::array<uint8x16_t, 4> index;            for (size_t j = 0; j < 4; ++j) index[j] = vld1q_u8(exp_tbl[sk[j] & 15].data());            std::array<uint8x16x2_t, 4> tbl;            for (size_t j = 0; j < 4; ++j) tbl[j] = {{vreinterpretq_u8_f32(v[j]), vreinterpretq_u8_f32(a[j])}};            for (size_t j = 0; j < 4; ++j) vst1q_f32(dst + i + 4 * j, vreinterpretq_f32_u8(vqtbl2q_u8(tbl[j], index[j])));            ptr += off[3] + (sk[3] >> 4);        }        return size;    }    void tiled_detour(float* dst, const size_t n, auto f, auto mask) {        const auto w = vld1q_u32(weights.data());        for (size_t i = 0; i < n; i += tile)            detour<false>(dst + i, tile, w, f, mask);    }

compress здесь тот же что и в первой статье.

expand_table — для true лейнов выбирает очередной элемент сжатого вектора (байты из [0, 15]), а для false — те же байты + 16. Потом tbl2 по {обработанные, исходные}, в общем тот же трюк, что и в compress

И еще:

• бесплатно скипается expand на пустых тайлах.

• массив разбивается на тайлы, которые помещаются в кеш

• Эмпирически tile=4096 оптимален.

• s бесплатно сохраняется, чтобы заново не считать addv

скорость BSL от плотности не зависит. sqrt 32.1, sin35 - 9.7, pow10 1.02.

Для дешевого sqrt всегда выгоден BSL (кроме thd = 0, очевидно), а для очень дорогого pow — detour (если thd != 1, конечно)

Когда окупается detour

Обозначим:

• B — оверхед BSL(vld + vbsl + vst) на один регистр.

• D — оверхед detour(сжатие + разжатие) на один регистр.

• T — цена f на один регистр (считаем цену f >> цены mask, влияет только на точность калибровки).

• d — доля элементов к которым применяется операция.

BSL применяет f ко всем регистрам, а detour только к d — экономит T * (1 - d). Значит detour выгоден когда экономия перекрывает D - B:

T * (1 - d) > D - B

detour выгоден когда d меньше d_max = 1 - (D - B) / T.

B и D — зависят только от железа, стабильны, померим заранее (у меня — B = 0.3, D = 0.845 ns/регистр) T — узнаем за первые несколько тайлов, меряя время работы BSL. А по нему узнаем и d_max:

    template<size_t tile>    double bsl_calibrate(float* dst, const size_t len, auto f, auto mask) {        double ns = 1e18;        for (size_t i = 0; i < len; i += tile) {            const auto st = std::chrono::high_resolution_clock::now();            bsl<false>(dst, tile, f, mask);            const auto ed = std::chrono::high_resolution_clock::now();            ns = std::min(ns, std::chrono::duration<double, std::nano>(ed - st).count());            dst += tile;        }        const double t = std::max(1e-9, ns / (tile / 4.0) - B);        return 1.0 - (D - B) / t;    }

Здесь:

• эмпирически 4 тайла по 2048 достаточно.

• min для замеров более устойчив к шуму чем среднее.

• вместе с калибровкой обработаем первые несколько тайлов.

• меряем bsl, ведь он всегда считает f, так что T = t_bsl - B.

• std::max тут защищает от деления на 0 и от t < 0 при очень дешевом T.

pilot v1

bsl всегда выгоднее, когда d_max < 0:

    void pilot_v1(float* dst, const size_t n, auto f, auto mask) {        const auto w = vld1q_u32(weights.data());        const float d_max = bsl_calibrate<tile / 2>(dst, 2 * tile, f, mask);        dst += 2 * tile;        for (size_t i = 2 * tile; i < n; i += tile) {            if (d_max < 0)                bsl<false>(dst, tile, f, mask);            else                detour<false>(dst, tile, w, f, mask);            dst += tile;        }    }

Алгоритм для sqrt верно определил, что BSL лучше. Однако для sin35 при больших thd выбирается detour, и теряется 30%. Проблема — v1 переключается на bsl только если он оптимален при любой плотности.

pilot v2

Плотность всего тайла считать — дорого, так что будем брать первые 256 (это чтение 6% тайла, на pow — шум, но на sqrt ощутимо). Для более-менее равномерного распределения этого достаточно:

    size_t density(float* dst, const size_t n, auto mask) {        std::array<uint32x4_t, 4> acc;        acc.fill(vdupq_n_u32(0));        for (size_t i = 0; i < n; i += 16) {            for (size_t j = 0; j < 4; ++j)                 acc[j] = vsubq_u32(acc[j], mask(vld1q_f32(dst + i + 4 * j)));        }        return vaddvq_u32(vaddq_u32(vaddq_u32(acc[0], acc[1]), vaddq_u32(acc[2], acc[3])));    }

Трюк здесь: true лейн битмаски — 0xFFFFFFFF = -1, вычитание лейна, по сути, делает прибавление.

Skip выгодно включать, когда предиктор редко промахивается, то есть 80% на то, что все 4 регистра пусты. Плотность примерно 0.014 ((1 - x)^16 = 0.8)

    void pilot_v2(float* dst, const size_t n, auto f, auto mask) {        const auto w = vld1q_u32(weights.data());        const float d_max = bsl_calibrate<tile / 2>(dst, 2 * tile, f, mask);        constexpr size_t probe = 256;        constexpr size_t xlo = 0.014 * probe;        const long long hi = d_max * probe;        dst += 2 * tile;        for (size_t i = 2 * tile; i < n; i += tile) {            const long long cnt = density(dst, probe, mask);            if (cnt > hi) {                if (cnt < xlo)                    bsl<true>(dst, tile, f, mask);                else                    bsl<false>(dst, tile, f, mask);            } else if (cnt < xlo)                detour<true>(dst, tile, w, f, mask);            else                detour<false>(dst, tile, w, f, mask);            dst += tile;        }    }

тут hi и xlo — пороги d_max и 0.014, но умноженные на длину пробы(256).

При thd = 0 благодаря скипу огромный выигрыш. Для больших thd верно выбирается bsl. Но сейчас на неплотных масках зря делается expand для пустых регистров.

pilot v3

Сделаем еще одну версию detour — во время compress будет сохранять индексы только непустых регистров (в буфер idx), и expand будет итерировать по ним:

    template <bool Skip>    size_t detour_compact(float* dst, const size_t n, const auto w, auto f, auto mask) {        float* ptr = tmp.data();        size_t k = 0;        for (size_t i = 0; i < n; i += 16) {            // начало цикла как в detour                        if constexpr (Skip)                if (vmaxvq_u32(vaddq_u32(vaddq_u32(m[0], m[1]), vaddq_u32(m[2], m[3]))) == 0) continue;            std::array<uint32_t, 4> sk;            for (size_t j = 0; j < 4; ++j) sk[j] = vaddvq_u32(vandq_u32(m[j], w));            for (size_t j = 0; j < 4; ++j) {                s[k] = sk[j];                idx[k] = i + 4 * j;                k += bool(sk[j]);            }            // конец цикла как в detour        }        // тут тоже как в detour                for (size_t j = 0; j < 3; ++j) s[k + j] = 0, idx[k + j] = 0;        // 2 фаза как в detour        k = (k + 3) & ~size_t(3);        for (size_t i = 0; i < k; i += 4) {            std::array<uint32_t, 4> sk;            for (size_t j = 0; j < 4; ++j) sk[j] = s[i + j];            std::array<float32x4_t, 4> a;            for (size_t j = 0; j < 4; ++j) a[j] = vld1q_f32(dst + idx[i + j]);            // и здесь как в detour            for (size_t j = 0; j < 4; ++j)                vst1q_f32(dst + idx[i + j], vreinterpretq_f32_u8(vqtbl2q_u8(tbl[j], index[j])));            ptr += off[3] + (sk[3] >> 4);        }            return size;    }

k (кол-во непустых регистров) до expand округляется вверх до кратного 4, чтобы после анролла не обрабатывать остаток.

Нету веток в горячих циклах — они бы убили скорость, поэтому compress придется выполнить для всех четырех регистров. Вместо веток позиция текущего элемента сдвигается на bool(sk).

Эту версию выгодно включать, когда хотя бы 50% всех регистров — пустые. Примерная плотность — 0.16 ((1 - x) ^ 4 = 0.5).

    void pilot_v3(float* dst, const size_t n, auto f, auto mask) {        // начало как в v2        constexpr size_t lo = 0.16 * probe;        for (size_t i = 2 * tile; i < n; i += tile) {            const long long cnt = density(dst, probe, mask);            if (cnt > hi) {                if (cnt < xlo)                    bsl<true>(dst, tile, f, mask);                else                    bsl<false>(dst, tile, f, mask);            } else if (cnt < xlo)                detour_compact<true>(dst, tile, w, f, mask);            else if (cnt < lo)                detour_compact<false>(dst, tile, w, f, mask);            else                detour<false>(dst, tile, w, f, mask);            dst += tile;        }    }

На маленьких плотностях работает заметно быстрее:

Теперь алгоритм довольно быстрый, но есть одна существенная проблема — мы допускаем, что данные равномерны. Поэтому легко подобрать тест, на котором v3 провалится:

for (size_t i = 0; i < n; i++)    dst[i] = i % 4096 >= 256;

Этот простой сценарий заставляет алгоритм всегда выбирать BSL, хоть на 3840 элементах тайла он не выгоден.

pilot v3.5

По самой первой таблице detour в худшем менее чем в 3 раза медленнее (на sqrt thd = 1), но на pow, thd = 0, он в 37 раз быстрее. Поэтому ошибочно выбранный BSL гораздо страшнее ошибочно выбранного detour.

Значит, для bsl будем страховаться — обрабатывать его блоками по tile / 8 и проверять плотность. Если она упала сильно ниже d_max — переключимся на detour. Стоить это будет 1 инструкцию на регистр. acc = vsubq_u32(acc, m), маска ведь 0/-1, и в отличие от probe, тут плотность точная.

    size_t bsl_verified(float* dst, const size_t n, float d_max, auto f, auto mask) {        for (size_t i0 = 0; i0 < 8; ++i0) {            std::array<uint32x4_t, 4> acc;            acc.fill(vdupq_n_u32(0));            for (size_t i = 0; i < n / 8; i += 16) {                std::array<float32x4_t, 4> v;                for (size_t j = 0; j < 4; ++j) v[j] = vld1q_f32(dst + 4 * j);                std::array<uint32x4_t, 4> m;                for (size_t j = 0; j < 4; ++j) m[j] = mask(v[j]);                for (size_t j = 0; j < 4; ++j) acc[j] = vsubq_u32(acc[j], m[j]);                for (size_t j = 0; j < 4; ++j)                    vst1q_f32(dst + 4 * j, vbslq_f32(m[j], f(v[j]), v[j]));                                dst += 16;            }            const size_t cur = vaddvq_u32(vaddq_u32(vaddq_u32(acc[0], acc[1]), vaddq_u32(acc[2], acc[3])));            if (static_cast<double>(cur) / (n / 8) < 0.75 * d_max) return (i0 + 1) * n / 8;        }        return n;    }

Тут BSL прерывается, когда плотность оказалась меньше 0.75 * d_max. Математического объяснения этому у меня нет, вариант выбран, как в среднем лучший.

И тогда pilot v3.5 на плотности > hi будет юзать bsl_verified:

    void pilot_v3_5(float* dst, const size_t n, auto f, auto mask) {        // начало как в v3            for (size_t i = 2 * tile; i < n; i += tile) {            const long long cnt = density(dst, probe, mask);            if (cnt > hi) {                if (hi < 0) {                    if (cnt < xlo)                        bsl<true>(dst, tile, f, mask);                    else                        bsl<false>(dst, tile, f, mask);                } else {                    auto done = bsl_verified(dst, tile, d_max, f, mask);                    if (done < tile)                        detour<false>(dst + done, tile - done, w, f, mask);                }            }            // остальные ветки как v3        }    }

Но и здесь легко подобрать контртест:

for (size_t i = 0; i < n; ++i)    dst[i] = i % 4096 >= 390;

Первый блок проверку пройдет (в нем > 75% нулей), а следующий — нет, на нем будет зря выполняться bsl, для pow это дорого. Проблема v3.5 — у него нету памяти. После промаха алгоритм продолжает доверять первым 256 и каждый раз промахивается.

pilot v4

Так что v4 это исправит: bsl сорвался — перестаем верить пробе следующие 16 тайлов, и вместо нее берем плотность прошлого тайла:

    void pilot_v4(float* dst, const size_t n, auto f, auto mask) {        // тут как в v3        size_t distrust = 0;        size_t prev = 0;        for (size_t i = 2 * tile; i < n; i += tile) {            if (distrust) --distrust;            const long long cnt = distrust ? prev : density(dst, probe, mask);            if (cnt > hi) {                if (hi < 0) {                    if (cnt < xlo) {                        bsl<true>(dst, tile, f, mask);                    } else {                        bsl<false>(dst, tile, f, mask);                    }                } else {                    if (distrust == 0) {                        auto done = bsl_verified(dst, tile, d_max, f, mask);                        if (done < tile) {                            distrust = 16;                            const size_t rem = detour<false>(dst + done, tile - done, w, f, mask);                            prev = rem * probe / (tile - done);                        }                    } else {                        prev = detour<false>(dst, tile, w, f, mask) * probe / tile;                    }                }            } else if (cnt < xlo) {                prev = detour_compact<true>(dst, tile, w, f, mask) * probe / tile;            } else if (cnt < lo) {                prev = detour_compact<false>(dst, tile, w, f, mask) * probe / tile;            } else {                prev = detour<false>(dst, tile, w, f, mask) * probe / tile;            }            dst += tile;        }    }

И эта версия легко проходит тот тест, pow10:

btw тут thd уже не совпадает с плотностью, при thd = 0 — плотность 0, при 1 — 9.5%.

Выигрывает у v3.5 более чем на 80%. Заодно работает быстрее чем обычный detour, ведь выбирается detour compact. Это финальная версия кода.

Итог

Сравнение против BSL:

Против BSL он проигрывает в худшем на 6%, однако гораздо чаще сильно выигрывает. Чтобы уменьшить проигрыш, можно диспатч ускорить — в density и bsl_verified брать каждый 4 регистр. Однако польза будет только если плотность равномерна.

Самый большой промах v4, который я нашел — 512 плотных, 512 пустых, дальше все плотное до конца цикла.

const size_t cycle = 17 * 4096;for (size_t i = 0; i < n; ++i) {    dst[i] = i % cycle >= 512 && i % cycle < 1024;}

Существеннее всего это при максимально дешевой функции (с d_max > 0 конечно же):

const auto a = vdupq_n_f32(0.5f);#pragma unrollfor (int i = 0; i < 13; ++i) x = vfmaq_f32(a, x, a);return x;

v3.5 в худшем упирается в BSL, а v4 в detour. Выше посчитано, что ошибочный detour менее страшен. Так что v4 не всегда лучше, но его ошибки дешевле.

Стоит сказать, что у поставленной задачи нету идеального решения. На любой алгоритм с диспатчем можно дать контртест.

Код целиком: godbolt.