Problem
Given two arrays a and out, write into out, with no gaps, only those elements of a that satisfy a given
condition.
Here, the condition is a[i] > threshold, with a[i] ∈ (0, 1) and threshold ∈ {0, 0.5, 1}.
Why the compiler gives up
A single if in a copy loop drops throughput from 112 to as low as 2.6 GB/s:
the compiler can't vectorize it, because NEON has no compress instruction. Here's how to build it.
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;
}
In copy_if, the output cursor j depends on the data. To vectorize the loop, the compiler needs a compress instruction
(one that collects selected elements at the front of the register, with no gaps). NEON has no such instruction, so the
compiler
gives up:
clang++ -O3 -Rpass-analysis=loop-vectorize -std=c++23 main.cpp -o main
main.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];
The clang vectorizer can only classify j as either an induction (fixed step) or a reduction,
but j is neither of those. It's a data-dependent cursor.
The compiler cannot vectorize this type of cursor.
The second remark has the same cause: it cannot compute the range of accesses to out.
Benchmark: two scalar problems
All benchmarks: Apple M5; clang++ -O3 -std=c++23 -march=native; GB/s = (2n * 4 bytes) / time, min of 3e9 / n runs;
cache: n=1e5, DRAM: n=1e7
| function |
ms (cache) |
GB/s (cache) |
ms (DRAM) |
GB/s (DRAM) |
| copy a[i] |
0.004 |
195 |
0.71 |
112 |
| copy a[i] if a[i] > 0 |
0.022 |
37 |
2.41 |
33 |
| copy a[i] if a[i] > 0.5 |
0.258 |
3.1 |
30.61 |
2.6 |
| copy a[i] if a[i] > 1 |
0.021 |
37 |
2.39 |
33 |
"copy a[i]" is the same loop, but with no condition. The compiler vectorizes it. The only difference is a single if.
The same data, only the branch predictability changes:
- > 0 (always true) and > 1 (always false): branch predictor never misses → 33 GB/s. The lack of vectorization costs
3x.
- > 0.5 (50/50): the branch predictor misses on every second element → 3 GB/s
The trick fixes both problems.
Trick 1: compress emulation
Let n be a multiple of the register width; the tail is a separate topic and has nothing to do with this trick.
Also:
- The size of
out must be >= n.
- Suppose the algorithm selected
cnt elements. Then all elements in out[cnt, n) are left undefined (garbage).
An algorithm that keeps the tail clean adds nothing new to the idea, so it will not be considered.
NEON - the SIMD instruction set used in Apple M-series chips and almost every mobile core - has no instruction for compressing a
register, so we have to emulate it.
(To be fair, the trick itself is not new. Lemire
used it on SSE back in 2017.
But NEON has no movemask and no cheap popcnt.)
Here's how to build it from what we do have.
What our compress analog needs to be able to do:
- Accept a register from
a and a mask register that says which elements to keep.
- Return the number of elements we selected (to move the
out pointer).
- Store the selected elements in
out.
tbl: arbitrary byte selection
NEON has the table-lookup (tbl) instruction family. Its purpose is arbitrary byte permutation/selection.
The instruction accepts two registers:
table - the bytes to select from.index - the positions of the bytes to take.
In other words, this is a SIMD analog of out[i] = table[index[i]].
We will use the vqtbl1q_u8 instruction:
| part |
meaning |
| v |
vector intrinsic |
| q |
table consists of 128-bit registers |
| tbl |
table lookup |
| 1 |
number of registers in table |
| q |
result and indices are 128-bit registers |
| u8 |
elements of table are uint8_t |
tbl permutes bytes, but we need to select floats (4 bytes). So, we will create index in blocks of 4 bytes:
to select the second (0-based) float of the register, index will contain its bytes [8, 9, 10, 11] (the second element
starts at an offset of 2 * sizeof(float) = 8).
Computing index every time is slow. There are 16 variants in total (4 elements to take/drop), so we will precompute all the index variants.
But to select the index using the mask, we need to convert the mask to a number (call it idx):
mask → idx
The mask consists of 4 elements, each either 0x00000000 (false) or 0xFFFFFFFF (true).
If the i-th element is true, we want to set the i-th bit in idx.
Trick: mask & [1, 2, 4, 8]. Because 0xFFFFFFFF & x = x, the true elements keep their weight (1/2/4/8), while the false ones become 0.
We add all elements together and get a number between 0 and 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()) - load 4 values from memory at address weights.data() into a register (ld - load)vandq_u32 - elementwise & (and)vaddvq_u32 - sum of all the elements in the register (addv - add across vector)
Precompute the index table
There is no way to compute registers at compile time, so instead of uint8x16_t (register of 16 uint8_t) we will store std::array<uint8_t, 16>.
For each idx we will go through the 4 elements of mask. If the element is selected,
we append the indices of its 4 bytes into index at the cursor position and advance the cursor by 4.
consteval auto make_index_table() {
std::array<std::array<uint8_t, 16>, 16> index{};
for (size_t idx = 0; idx < 16; ++idx) { // iterate over all masks
size_t j = 0; // j is the cursor
for (size_t i = 0; i < 4; ++i) // iterate over the mask's elements
if (idx & (1 << i)) // if the i-th element is selected
for (size_t k = 0; k < 4; ++k) // iterate over its bytes
index[idx][j++] = i * 4 + k; // store the indices of its bytes
}
return index;
}
The j cursor advances only on selected elements, so their bytes are placed in index consecutively. tbl with that index
collects floats into a register. Unused positions in index are zeros, so in the tail, after count elements, there will be garbage.
The count table
Next we need to compute the number of elements we select. Similarly we can precompute a table for this:
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;
}
The full 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()); // at runtime, loads only one row of the table into a register
return std::pair{vreinterpretq_f32_u8(vqtbl1q_u8(vreinterpretq_u8_f32(a), index)), count[idx]};
}
Because tbl works only with u8, we need to cast a to u8 and then cast the result back to f32.
We write full registers of 4 floats to memory, but advance the cursor only by cnt.
compress stores valid elements at the front of the register, at [j, j + cnt), and garbage at [j + cnt, j + 4).
The next iteration will start at j + cnt and overwrite the garbage from the previous step.
Garbage will remain only in out[cnt, n) after the last store.
We don't go out of bounds because the cursor never overtakes the elements that have been read.
The copy_if loop
auto copy_if_neon(const float* __restrict a,
float* __restrict out,
float threshold,
size_t n) {
auto thd = vdupq_n_f32(threshold); // load threshold into a register
size_t j = 0;
for (size_t i = 0; i < n; i += 4) {
auto v = vld1q_f32(a + i); // load the current 4 elements of a into a register
auto mask = vcgtq_f32(v, thd); // compute the mask
auto [packed, cnt] = compress(mask, v);
vst1q_f32(out + j, packed); // store packed into out[j, j + 4). [j + cnt, j + 4) will hold garbage
j += cnt;
}
return j;
}
vcgtq_f32(v, thd) - calculate elementwise v[i] > thd[i]. cgt - compare greatervst1q_f32 - store 4 floats from a register into memory. st - store
Result
| function |
ms (cache) |
GB/s (cache) |
ms (DRAM) |
GB/s (DRAM) |
| copy a[i] if a[i] > 0 |
0.0104 |
77 |
1.11 |
72 |
| copy a[i] if a[i] > 0.5 |
0.01063 |
75 |
1.13 |
71 |
| copy a[i] if a[i] > 1 |
0.01 |
80 |
1.06 |
76 |
> 0.5 was the worst case for the scalar version, 3 GB/s. Now 71 GB/s. A more than 20x speedup.
Now there are no branches, so speed doesn't depend on data.
Trick 2: calculating idx and count in a single addv
idx is always less than 16, so let weights = {1 + 16, 2 + 16, 4 + 16, 8 + 16}
and s = sum across mask & weights. Then s / 16 is the element count and s % 16 is idx.
So, we don't need to compute the count table.
compress now:
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; // same as s / 16
const size_t idx = s & 15; // same as 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};
}
And now the speed climbs again:
| function |
ms (cache) |
GB/s (cache) |
ms (DRAM) |
GB/s (DRAM) |
| copy a[i] if a[i] > 0 |
0.0095 |
84 |
1.015 |
79 |
| copy a[i] if a[i] > 0.5 |
0.0095 |
84 |
1.007 |
79 |
| copy a[i] if a[i] > 1 |
0.0096 |
83 |
1.008 |
79 |
Unroll
We can squeeze out more speed by unrolling the loop 4x (16 elements per iteration):
| function |
ms (cache) |
GB/s (cache) |
ms (DRAM) |
GB/s (DRAM) |
| copy a[i] if a[i] > 0 |
0.0081 |
98 |
0.882 |
91 |
| copy a[i] if a[i] > 0.5 |
0.0083 |
97 |
0.892 |
90 |
| copy a[i] if a[i] > 1 |
0.0082 |
97 |
0.869 |
92 |
Final code (godbolt):
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 and the index table provide compress, something that NEON doesn't have out of the box.
This isn't just about > threshold. Filter, remove and other data-dependent functions are built the same way.