a[mask] = f(a[mask]) on NEON. faster than the obvious blend
## Problem
Apply an operation to elements that satisfy a condition:
for (size_t i = 0; i < n; ++i)
if (mask(a[i])) a[i] = f(a[i]);
### Notes
- `a[i] ∈ (0, 1)`, `thd ∈ (0, 1)`, `mask = a[i] < thd`; uniform distribution (except at the end of the article)
- `f` is one of `sqrt`, `frfrexp` (mantissa), `sin` 3.5 ULP, `sin` 1 ULP, `pow` 1 ULP (from SLEEF)
- `f` and `mask` are passed as runtime values, so they are wrapped in a lambda with `always_inline`, otherwise they may
not be inlined
- The array size `n` is a multiple of every unroll, tile etc. The tail is trivial to handle(BSL/scalar)
- In tables `*` = best, units = GiB/s
- Don't compare numbers across tables. Different conditions, values fluctuate
- All benchmarks: Apple M5; clang++ -O3 -std=c++23 -march=native; GiB/s = (n * 4 bytes) / time, min of 720 runs
(During bench, functions run in a changing order, data is restored ofc); n=1e7 + 2432;
## BSL blend
If the problem is memory bound (cheap function or high density), the standard algorithm is optimal:
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]));
}
}
It computes `f` on every element, but stores only the selected ones.
Skip helps on sparse masks, but otherwise mispredictions will kill performance. We'll need it later.
But for expensive `f` this algo does too much extra work
## Detour
To avoid unnecessary work, we compress selected elements, apply only to them, and expand back.
avx512 does this in two instructions. NEON doesn't, so we'll emulate and optimize.
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; // idx, D and B come in later
constexpr double D = 0.845;
constexpr double B = 0.3;
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) {
Post #25678
6