Planar NEON kernels for the A2 fast path on Apple Silicon: ~2.4x and ~2.0x, bit-identical

[rikkus]

Draft -- opened for discussion and measurement on other people's machines, not as a finished proposal. Happy to split, rework or drop any part of it.

Note: I spotted the revert of an earlier optimisation attempt and have made sure this is gated on APPLE && aarch64 -- though I haven't done any checks on what compiler was used. I'm on MacOS 27 dev beta and beta Xcode so this could definitely use some testing elsewhere.

What this is

Two NEON kernels for the A2 fast path on Apple Silicon, for the 3-channel and 8-channel submodels. They are about 2.4× (A2 standard_ and 2.0× (A2 nano) the speed of - with their output bit-identical to - a2_fast.

Where it is active, and where it is not

This is the first thing worth checking, so it is at the top rather than the bottom.

NAM_A2_PLANAR is defined only when NAM_ENABLE_A2_FAST is on and the target is Apple Silicon (__APPLE__ && __aarch64__). On every other target — x86, and also every other AArch64 target — the header declares nothing, a2_planar.cpp compiles to an object with no symbols in it, the call site in a2_fast.cpp is preprocessed away, and the A2 path is byte for byte the code that is there today.

Off Apple Silicon, the only thing that changes anywhere is that two lines of A2FastConfig::create now live in a named function (create_a2_fast_reference_model) so a test can call them directly. Same construction, same order, same object.

Checked rather than asserted, by cross-building this branch for x86_64:

$ nm build-x86/.../NAM/wavenet/a2_planar.cpp.o
build-x86/.../NAM/wavenet/a2_planar.cpp.o: no symbols

$ ./build-x86/tools/run_tests
Running tests...
Success!

Every target builds, and bench_a2_planar builds too — it prints that this build has no planar kernel to measure and exits 0.

The gate stops at __APPLE__ rather than plain __aarch64__ deliberately. These kernels are very likely correct and probably faster on any AArch64 part, but they have only been built and measured on Apple Silicon, and two of the things they lean on are toolchain properties rather than architectural ones: the tile widths are M2 measurements, and bit-identity relies on the compiler contracting a*b+c into an FMA inside a2_fast's own 3-channel branch — which clang and gcc do by default and MSVC at /fp:precise does not. Opening it up to Linux arm64 or Windows on ARM is a one-line change for whoever can measure there; I would rather not claim a target nobody has run.

-DNAM_DISABLE_A2_PLANAR also turns it off on Apple Silicon, which is what makes an A/B against the reference a one-flag change. The test suite passes in that configuration too.

The numbers

Every figure below comes from tools/bench_a2_planar, which is in this PR, run on this branch. Apple M2 (4P + 4E), macOS 27, Release, 48 kHz. One pass is a 10.9 s render of a three-tone test signal the tool generates itself; each figure is the mean of the fastest 70% of 20 passes.

64-frame blocks

	a2_fast	planar NEON
A2 standard (8 ch)	418.3 ms	172.3 ms	2.43×
A2 nano (3 ch)	57.2 ms	28.4 ms	2.01×

32-frame blocks — closer to what a plugin actually runs

	a2_fast	planar NEON
A2 standard (8 ch)	470.2 ms	177.6 ms	2.65×
A2 nano (3 ch)	61.4 ms	30.3 ms	2.03×

The 8-channel advantage is larger at the smaller block, because a2_fast slows from 418 ms to 470 ms as the block shrinks while the planar kernel barely moves. So this is not a win that only exists at unrealistically large buffers.

Run-to-run spread on this machine is around ±2%, so please read the third digit as decoration.

Why it is faster

a2_fast stores history column-major — the channels of one frame adjacent — and a SIMD register naturally spans channels. These kernels give each channel its own plane, so a register holds four consecutive frames of one channel.

That has two consequences. Each lane independently executes a2_fast's own per-frame scalar chain, so nothing is reassociated (see below). And the conv accumulators stay in registers across all K taps instead of making a round trip to memory per tap per frame, which is where most of the time was going.

On top of the layout there are three tuning choices — a frame tile of 32 at C=3 and 8 at C=8, residuals written straight into the next layer's ring rather than through a scratch buffer the next layer copies in, and a linear ring rewound by memmove instead of the power-of-two ring with an eagerly mirrored tail that a2_fast ships with.

Provenance, since it differs from the rest of this PR: those three constants were chosen by sweeping them in a separate benchmark harness of mine, one variable at a time, which is not part of this branch and whose numbers are therefore not reproducible from it. What is reproducible here is the end-to-end result in the table above. I have deliberately not quoted the sweep's own figures. The constants are three named constexpr values in one file, so re-tuning them on other hardware is easy; they affect speed only, never output.

One finding that is reproducible from this branch, and that needs none of the above: building a2_fast itself with -DNAM_A2_RING_MODE=0 rather than its shipped default measures 51.3 ms against 57.2 ms on the 3-channel submodel — about 11% faster, from a one-line change. On the 8-channel submodel it makes no difference worth reporting (417.6 ms against 418.3 ms). That is entirely independent of this PR and might be worth a look on its own.

Why it is bit-identical, and how that was established

The two channel counts reproduce two different orders of arithmetic, because a2_fast itself branches.

Channels == 3 is a hand-written scalar 3×3 GEMV, so reproducing it is a matter of writing the same chain: bias first, then per tap the three input channels in increasing order, everything contracted into FMAs.

Channels == 8 is Eigen, and its reduction order had to be established rather than assumed. The order that matches is:

per tap k:  t_i = 0;  for j = 0..7:  t_i = fma(W_k(i,j), x_k(j), t_i)
then:       z_i = 0;  for k:         z_i = z_i + t_i^(k)
            z_i = z_i + bias_i
            z_i = z_i + (mixin_i * cond)     <- product then add, two roundings
            z_i = LeakyReLU(z_i)
            head_sum_i = head_sum_i + z_i
            u_i = 0;  for j:  u_i = fma(L(i,j), z_j, u_i)
            lin_i = (lin_i + u_i) + l1x1_b_i

The per-tap partial t, summed into the running total only at the end of the tap, is the part that is easy to get wrong — folding the taps into one long chain is the obvious thing to write, and it is a different association that does move the bits. So is treating the mixin as an FMA when Eigen rounds the product and the sum separately.

That order was found by replicating those expressions on those exact shapes and comparing candidate orderings bit-for-bit against Eigen's own output; the one above matched every time and four plausible alternatives matched zero times. That probe also lives in the separate harness rather than here — but unlike a timing, its conclusion is checkable directly from this branch, because if the order were wrong the parity tests below would fail.

Verification in this PR

tools/test/test_a2_planar.cpp, wired into run_tests. For each channel count it renders 20,000 frames through both engines and asserts memcmp equality — not a tolerance — at 14 block sizes: 1, 3, 4, 7, 8, 15, 16, 31, 32, 33, 64, 65, 128, 512. The small and awkward ones matter, because they are what exercise the partial-tile and single-frame tail paths. A third test asserts the dispatcher actually routes to the planar kernel, so the other two are not checking something nothing uses.

That test earned its keep: it failed immediately at block size 1, on a mixin in the 8-channel scalar tail written as a = a + m * cf, which the compiler contracts into an FMA where Eigen rounds the product and the sum separately. Every block size that is a multiple of the tile width had been passing.

tools/bench_a2_planar renders a whole signal through both engines, compares bit for bit, and reports speed only if they matched. It unwraps a SlimmableContainer and picks a submodel by width rather than by index:

bench_a2_planar --submodel widest    model.nam
bench_a2_planar --submodel narrowest model.nam

It is built at -O3 rather than -Ofast deliberately: -ffast-math lets the compiler contract a multiply and an add into an FMA across statement boundaries, which is exactly the freedom the parity result is checking has not been taken.

Beyond the synthetic weights in the unit test, both kernels were checked bit-identical on four real captures — both submodels of each, at 64- and 100-frame blocks, over 523,200 frames each.

What it touches

NAM/wavenet/a2_planar.{h,cpp}	new; the two kernels. Empty translation unit off Apple Silicon
NAM/wavenet/a2_fast.cpp	A2FastConfig::create prefers the planar model where one exists
NAM/wavenet/a2_fast.h	adds create_a2_fast_reference_model, so a test can reach the portable implementation directly now that the dispatcher may not return it
tools/	the test and the benchmark

Notes and caveats

• The tuning constants are M2 measurements. They should hold across the M-series; if they turn out to be wrong on some other Apple part, they are three named constants in one file.
• fp16, int8, AMX/SME, Accelerate and threading were all considered and deliberately not used. The point here is a drop-in replacement that needs no listening test, and a format change forfeits that by construction.
• Bit-identity is a property of a normal -O3 build. Under -ffast-math the compiler is free to contract across statements in either engine and the guarantee no longer applies — to a2_fast either.