Stream compaction on NEON: vectorizing copy_if by hand (30x)
## 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.
```cpp
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:
```text
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:
1. \> 0 (always true) and > 1 (always false): branch predictor never misses → 33 GB/s. The lack of vectorization costs
3x.
2. \> 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:
1. The size of `out` must be >= `n`.
2. 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](https://lemire.me/blog/2017/01/20/how-quickly-can-you-remove-spaces-from-a-string/) 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:
1. Accept a register from `a` and a mask register that says which elements to keep.
2. Return the number of elements we selected (to move the `out` pointer).
3. 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:
1. `table` - the bytes to select from.
2. `index` - the positions of the bytes to take.
In other
Post #25586
13