C++ - Reddit

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

```cpp
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 greater
- `vst1q_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:
```cpp
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](https://godbolt.org/z/n8E6ocKoj)):

```cpp
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

godbolt.org

Compiler Explorer - C++ (armv8-a clang (trunk))

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…

8 views18:23