From 62ab2746f31f0ba69299e0e3e4bf620995c0792c Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Sat, 8 Aug 2026 20:01:25 +0100 Subject: [PATCH 1/8] Planar NEON kernels for the A2 fast path (AArch64), bit-identical a2_fast keeps a frame's channels adjacent and vectorises across channels. These kernels keep each channel in its own plane and vectorise across frames instead, so one NEON lane runs a2_fast's per-frame scalar chain verbatim. Nothing is reassociated, so the output does not move by a bit. On an Apple M2, against a 10.9 s render at 64-frame blocks: A2 standard (8 ch) 417 ms -> 172 ms 2.43x A2 nano (3 ch) 57 ms -> 28 ms 2.01x and at 32-frame blocks, which is what a plugin actually runs, 2.65x and 2.03x -- a2_fast degrades at small blocks and these do not. The two channel counts reproduce two different orders of arithmetic, because a2_fast itself branches. C=3 reproduces its hand-written scalar 3x3 GEMV. C=8 reproduces what its Eigen expressions compute, including the per-tap partial that is summed into the running total only at the end of the tap, and the mixin's separate multiply and add -- folding the taps into one chain is the obvious thing to write and it is a different association. That order was established by comparing candidate orderings bit-for-bit against Eigen's own output, not assumed. Selection happens in A2FastConfig::create, and only on AArch64 with the A2 fast path already enabled; -DNAM_DISABLE_A2_PLANAR opts back out. On every other target the new file compiles to nothing and behaviour is unchanged. Verification ships with it: tools/test/test_a2_planar.cpp asserts memcmp equality against the reference over 14 block sizes per channel count, including 1, 3 and 7, which exercise the partial-tile and single-frame tails. tools/bench_a2_planar.cpp renders a whole signal through both engines, compares bit for bit, and only then reports speed. Built at -O3 rather than -Ofast on purpose: -ffast-math lets the compiler contract a multiply and an add across statements, which is the freedom the parity result is checking has not been taken. a2_fast.h gains create_a2_fast_reference_model so a test can get at the portable implementation directly rather than through the dispatcher, which now may hand back a specialised one. --- NAM/wavenet/a2_fast.cpp | 24 +- NAM/wavenet/a2_fast.h | 13 + NAM/wavenet/a2_planar.cpp | 1158 +++++++++++++++++++++++++++++++++ NAM/wavenet/a2_planar.h | 53 ++ tools/CMakeLists.txt | 26 + tools/bench_a2_planar.cpp | 339 ++++++++++ tools/run_tests.cpp | 7 + tools/test/test_a2_planar.cpp | 230 +++++++ 8 files changed, 1845 insertions(+), 5 deletions(-) create mode 100644 NAM/wavenet/a2_planar.cpp create mode 100644 NAM/wavenet/a2_planar.h create mode 100644 tools/bench_a2_planar.cpp create mode 100644 tools/test/test_a2_planar.cpp diff --git a/NAM/wavenet/a2_fast.cpp b/NAM/wavenet/a2_fast.cpp index 08d52e72..337cf12f 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -9,6 +9,7 @@ #endif #include "a2_fast.h" + #include "a2_planar.h" #include #include @@ -772,11 +773,15 @@ struct A2FastConfig : public ModelConfig std::unique_ptr create(std::vector weights, double sampleRate) override { - if (channels == 3) - return std::make_unique>(std::move(weights), sampleRate); - if (channels == 8) - return std::make_unique>(std::move(weights), sampleRate); - throw std::runtime_error("A2FastConfig: unsupported channel count " + std::to_string(channels)); + #if defined(NAM_A2_PLANAR) + // On AArch64, prefer the planar NEON kernels. They are bit-identical to the + // reference model below -- same float32 bits out, sample for sample -- so + // this is a speed choice and nothing else. A channel count they do not cover + // returns nullptr and falls through. + if (auto planar = create_a2_planar_model(channels, weights, sampleRate)) + return planar; + #endif + return create_a2_fast_reference_model(channels, std::move(weights), sampleRate); } }; @@ -983,6 +988,15 @@ bool is_a2_shape(const nlohmann::json& config, int* channels) return true; } +std::unique_ptr create_a2_fast_reference_model(int channels, std::vector weights, double sampleRate) +{ + if (channels == 3) + return std::make_unique>(std::move(weights), sampleRate); + if (channels == 8) + return std::make_unique>(std::move(weights), sampleRate); + throw std::runtime_error("create_a2_fast_reference_model: unsupported channel count " + std::to_string(channels)); +} + std::unique_ptr create_a2_fast_config(const nlohmann::json& config, double sampleRate) { (void)sampleRate; diff --git a/NAM/wavenet/a2_fast.h b/NAM/wavenet/a2_fast.h index 7fac5347..5a6b6608 100644 --- a/NAM/wavenet/a2_fast.h +++ b/NAM/wavenet/a2_fast.h @@ -52,6 +52,19 @@ bool is_a2_shape(const nlohmann::json& config, int* channels); /// \pre is_a2_shape(config, ...) returned true. std::unique_ptr create_a2_fast_config(const nlohmann::json& config, double sampleRate); +/// \brief Build the portable A2 fast-path model, bypassing any +/// architecture-specific kernel. +/// +/// The config built above may hand back a specialised implementation on some +/// targets (see a2_planar.h). This always returns the portable one, so a test +/// can assert that a specialised kernel agrees with the reference it claims to +/// reproduce. +/// +/// \param channels 3 (A2 nano) or 8 (A2 standard); anything else throws. +/// \param weights The A2 weight stream. +/// \param sampleRate Expected sample rate, passed through to DSP. +std::unique_ptr create_a2_fast_reference_model(int channels, std::vector weights, double sampleRate); + } // namespace a2_fast } // namespace wavenet } // namespace nam diff --git a/NAM/wavenet/a2_planar.cpp b/NAM/wavenet/a2_planar.cpp new file mode 100644 index 00000000..57be0f63 --- /dev/null +++ b/NAM/wavenet/a2_planar.cpp @@ -0,0 +1,1158 @@ +#if defined(NAM_ENABLE_A2_FAST) + + #include "a2_planar.h" + + #if defined(NAM_A2_PLANAR) + + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + + #include + + #include "a2_fast.h" + +// ============================================================================= +// Planar NEON kernels for the A2 fast path. +// +// Both kernels are bit-identical to the A2FastModel they replace, over +// the whole of a 523,808-frame test render. That is the point of the design, so +// it is worth being precise about how it is achieved, because the two channel +// counts get there by reproducing two *different* orders of arithmetic -- +// a2_fast itself branches, and so does this. +// +// The shared idea: planar layout. a2_fast stores history column-major, the +// channels of one frame adjacent, and a SIMD register naturally spans channels. +// Here each channel gets its own plane, so a register holds four consecutive +// *frames* of one channel and each lane independently runs a2_fast's per-frame +// scalar chain. No cross-lane reduction ever happens, so no reassociation +// happens, so the bits do not move. +// +// Channels == 3 reproduces a2_fast's hand-written scalar 3x3 GEMV: bias first, +// then per tap the three input channels in increasing order, all contracted into +// FMAs, mixin contracted too. +// +// Channels == 8 reproduces what a2_fast's Eigen expressions actually compute: +// +// 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 +// +// That per-tap partial `t`, summed into `z` only at the end of the tap, is the +// part that is easy to get wrong: folding the taps into one chain (the obvious +// thing to write) is a different association and does move the bits. The order +// above was established by comparing candidate orderings bit-for-bit against +// Eigen's own output across 224 (block size, kernel size, trial) combinations, +// not by assumption. +// +// The tuning constants below (frame tile widths, head tile, ring strategy) were +// each swept independently on an Apple M2 against the full-length render; the +// values chosen are the measured optima and the comments say what they trade. +// They do not affect the output, only the speed. +// ============================================================================= + +namespace nam +{ +namespace wavenet +{ +namespace a2_fast +{ + +namespace +{ + +// ----------------------------------------------------------------------------- +// Weights +// +// Parsed in exactly A2FastModel::_load_weights' order -- which is the generic +// WaveNet's order -- into exactly its layout, so the two engines are fed +// identical numbers and only the kernel differs. +// ----------------------------------------------------------------------------- +template +struct PlanarLayerWeights +{ + int kernel_size = 0; + int dilation = 0; + int max_lookback = 0; // (kernel_size - 1) * dilation + + /// kernel_size * C * C floats. Tap k, output i, input j lives at [k*C*C + j*C + i]. + std::vector conv_w; + std::array conv_b{}; + /// Input mixin (condition size 1 -> C), no bias. + std::array mixin_w{}; + /// layer1x1 (C -> C), column-major: [j*C + i] is bottleneck j to output i. + std::array l1x1_w{}; + std::array l1x1_b{}; +}; + +template +struct PlanarWeights +{ + /// Rechannel (input size 1 -> C), no bias. + std::array rechannel_w{}; + std::array, kNumLayers> layers; + /// Head rechannel (C -> 1), kernel 16. At tap k the matrix is 1 x C. + std::array, kHeadKernelSize> head_w{}; + float head_b = 0.0f; + /// The trailing float of the stream, which overrides the JSON head_scale. + float head_scale = 1.0f; +}; + +template +PlanarWeights parse_weights(const std::vector& weights) +{ + PlanarWeights out; + + auto it = weights.begin(); + const auto end = weights.end(); + auto take = [&]() -> float { + if (it == end) + throw std::runtime_error("A2PlanarModel: weight stream exhausted"); + return *it++; + }; + + for (int i = 0; i < C; i++) + out.rechannel_w[i] = take(); + + for (int li = 0; li < kNumLayers; li++) + { + PlanarLayerWeights& L = out.layers[li]; + L.kernel_size = kKernelSizes[li]; + L.dilation = kDilations[li]; + L.max_lookback = (L.kernel_size - 1) * L.dilation; + const int K = L.kernel_size; + + // Conv1D read order: for i in out, for j in in, for k in taps. + L.conv_w.assign(static_cast(K) * C * C, 0.0f); + for (int i = 0; i < C; i++) + for (int j = 0; j < C; j++) + for (int k = 0; k < K; k++) + L.conv_w[static_cast(k) * C * C + static_cast(j) * C + i] = take(); + for (int i = 0; i < C; i++) + L.conv_b[i] = take(); + + for (int i = 0; i < C; i++) + L.mixin_w[i] = take(); + + // Conv1x1 read order: for i in out, for j in in. + for (int i = 0; i < C; i++) + for (int j = 0; j < C; j++) + L.l1x1_w[static_cast(j) * C + i] = take(); + for (int i = 0; i < C; i++) + L.l1x1_b[i] = take(); + } + + for (int j = 0; j < C; j++) + for (int k = 0; k < kHeadKernelSize; k++) + out.head_w[k][j] = take(); + out.head_b = take(); + out.head_scale = take(); + + if (it != end) + { + std::stringstream ss; + ss << "A2PlanarModel: weight stream has " << std::distance(it, end) << " trailing values"; + throw std::runtime_error(ss.str()); + } + + return out; +} + +/// A product that is rounded before it is used. +/// +/// The C=8 path's mixin is `z += mixin * cond` with the multiply and the add +/// rounded separately -- that is what Eigen does, and a fused multiply-add there +/// would move the result by one ulp. In the vector paths that is spelled out in +/// intrinsics; in the scalar tails, writing `a = a + m * cf` would let the +/// compiler contract it. Routing the product through a NEON register keeps the +/// two roundings whatever the compiler decides. +inline float mul_rounded(float a, float b) +{ + return vget_lane_f32(vmul_f32(vdup_n_f32(a), vdup_n_f32(b)), 0); +} + +/// Receptive field, counted exactly as A2FastModel counts it, so the planar +/// models warm up over the same number of samples as the code they replace. +int planar_prewarm_samples() +{ + int prewarm = 1; + for (int li = 0; li < kNumLayers; li++) + prewarm += (kKernelSizes[li] - 1) * kDilations[li]; + prewarm += kHeadKernelSize - 1; + return prewarm; +} + +// ----------------------------------------------------------------------------- +// Planar history ring: C channel planes side by side, one linear buffer each, +// written forward and memmoved back when it runs out. +// +// This is a2_fast's NAM_A2_RING_MODE=0 strategy, in planar layout. Four +// strategies were measured -- power-of-two with an eagerly mirrored tail +// (a2_fast's shipped default), power-of-two with a lazy mirror, exactly-sized +// with a lazy mirror, and this one -- and linear+rewind won for both channel +// counts once the residual writes go straight into the next layer's ring. It +// has no mirror to maintain, no masking on reads, and every read is contiguous; +// it pays instead with an occasional large memmove, amortised over many blocks +// by the 2*lookback sizing. +// ----------------------------------------------------------------------------- +template +struct PlanarRing +{ + std::vector data; + int cap = 0; ///< columns per plane + int stride = 0; ///< distance between channel planes (== cap here) + int wpos = 0; + int lookback = 0; + + void reset(int max_lookback, int max_buffer) + { + lookback = max_lookback; + cap = 2 * max_lookback + max_buffer; + stride = cap; + data.assign(static_cast(C) * stride, 0.0f); + wpos = max_lookback; + } + + float* plane(int c) { return data.data() + static_cast(c) * stride; } + const float* plane(int c) const { return data.data() + static_cast(c) * stride; } + + /// Make room for an n-frame write, rewinding if it would not fit. + void prepare(int n) + { + if (wpos + n > cap) + { + for (int c = 0; c < C; c++) + std::memmove(plane(c), plane(c) + (wpos - lookback), static_cast(lookback) * sizeof(float)); + wpos = lookback; + } + } + + /// Where an n-frame block is written. Always one contiguous run. + float* write_ptr(int c) { return plane(c) + wpos; } + + void commit(int n) { wpos += n; } + + /// First column of an n-frame read looking `lookback_frames` further back than + /// the block just written. + int tap(int lookback_frames, int n) const { return wpos - n - lookback_frames; } +}; + +// ============================================================================= +// Channels == 3 (A2 nano) +// +// Reproduces a2_fast's `if constexpr (Channels == 3)` branch: the fully +// unrolled scalar 3x3 GEMV, bias-seeded at tap 0, inputs in increasing order, +// mixin contracted into an FMA, then LeakyReLU, head_sum, layer1x1 residual. +// Each NEON lane runs that chain for its own frame. +// +// Per conv tap this is 3 loads and 9 vfmaq_laneq_f32 per 4 frames, against +// a2_fast's 9 scalar FMAs per frame. +// ============================================================================= + +/// Frames per tile. Twelve accumulator registers at 32 (3 channels x 8 vectors), +/// which is where the sweep peaked: 8/16/32/64 measured 1.39x/1.56x/1.75x/1.46x +/// against a2_fast. 64 spills. +constexpr int kNanoTile = 32; +constexpr int kNanoVecs = kNanoTile / 4; + +/// One layer's weights, padded to four lanes so each group of three is a single +/// vector load addressed by lane. +struct NanoLayer +{ + int kernel_size = 0; + int dilation = 0; + int max_lookback = 0; + /// kernel_size x 12 floats: nine weights then three pad, per tap. + std::vector conv_w; + std::array conv_b{}; + std::array mixin_w{}; + /// Twelve floats: nine layer1x1 weights then three pad. + std::array l1x1_w{}; + std::array l1x1_b{}; +}; + +class A2PlanarNano : public DSP +{ + static constexpr int C = 3; + using Ring = PlanarRing; + +public: + A2PlanarNano(const std::vector& weights, double expected_sample_rate) + : DSP(/*in_channels=*/1, /*out_channels=*/1, expected_sample_rate) + , _w(parse_weights(weights)) + , _prewarm_samples(planar_prewarm_samples()) + { + for (int li = 0; li < kNumLayers; li++) + { + const PlanarLayerWeights& L = _w.layers[li]; + NanoLayer& P = _p[li]; + P.kernel_size = L.kernel_size; + P.dilation = L.dilation; + P.max_lookback = L.max_lookback; + + P.conv_w.assign(static_cast(L.kernel_size) * 12, 0.0f); + for (int k = 0; k < L.kernel_size; k++) + for (int e = 0; e < 9; e++) + P.conv_w[static_cast(k) * 12 + e] = L.conv_w[static_cast(k) * 9 + e]; + + for (int i = 0; i < C; i++) + { + P.conv_b[i] = L.conv_b[i]; + P.mixin_w[i] = L.mixin_w[i]; + P.l1x1_b[i] = L.l1x1_b[i]; + } + for (int e = 0; e < 9; e++) + P.l1x1_w[e] = L.l1x1_w[e]; + } + + _head_w4.assign(static_cast(kHeadKernelSize) * 4, 0.0f); + for (int k = 0; k < kHeadKernelSize; k++) + for (int j = 0; j < C; j++) + _head_w4[static_cast(k) * 4 + j] = _w.head_w[k][j]; + } + + ~A2PlanarNano() override = default; + + int GetPrewarmSamples() override { return _prewarm_samples; } + + void process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames) override + { + if (num_frames > GetMaxBufferSize()) + SetMaxBufferSize(num_frames); + const int N = num_frames; + + const NAM_SAMPLE* in0 = input[0]; + NAM_SAMPLE* out0 = output[0]; + + // Rechannel straight into layer 0's ring: there is no scratch buffer for a + // later pass to copy in. + float* cond = _cond.data(); + _rings[0].prepare(N); + float* const r0 = _rings[0].write_ptr(0); + float* const r1 = _rings[0].write_ptr(1); + float* const r2 = _rings[0].write_ptr(2); + const float rw0 = _w.rechannel_w[0], rw1 = _w.rechannel_w[1], rw2 = _w.rechannel_w[2]; + for (int f = 0; f < N; f++) + { + const float x = static_cast(in0[f]); + cond[f] = x; + r0[f] = rw0 * x; + r1[f] = rw1 * x; + r2[f] = rw2 * x; + } + _rings[0].commit(N); + + // No memset of head_sum: layer 0 writes it rather than accumulating onto it. + for (int li = 0; li < kNumLayers; li++) + dispatch_layer(li, N); + + head_forward(N); + + const float* head_out = _head_out.data(); + for (int f = 0; f < N; f++) + out0[f] = static_cast(head_out[f]); + } + +protected: + void SetMaxBufferSize(const int maxBufferSize) override + { + DSP::SetMaxBufferSize(maxBufferSize); + + _stride = maxBufferSize; + _layer_in.assign(static_cast(C) * _stride, 0.0f); + _head_sum.assign(static_cast(C) * _stride, 0.0f); + _cond.assign(static_cast(maxBufferSize), 0.0f); + _head_out.assign(static_cast(maxBufferSize), 0.0f); + + for (int li = 0; li < kNumLayers; li++) + _rings[li].reset(_p[li].max_lookback, maxBufferSize); + _head_ring.reset(kHeadKernelSize - 1, maxBufferSize); + } + +private: + float* lin(int c) { return _layer_in.data() + static_cast(c) * _stride; } + float* hsum(int c) { return _head_sum.data() + static_cast(c) * _stride; } + + void dispatch_layer(int li, int N) + { + // The last layer's layer1x1 residual is read by nothing -- there is no layer + // 24, and the head reads head_sum -- so it is not computed. The first + // layer's head_sum accumulate has nothing to accumulate onto, so it stores + // instead, which is what makes the per-block memset unnecessary. + const bool store_head = (li == 0); + const bool do_l1x1 = (li != kNumLayers - 1); + + if (_p[li].kernel_size == 6) + { + if (store_head) + layer_forward<6, true, true>(li, N); + else if (do_l1x1) + layer_forward<6, false, true>(li, N); + else + layer_forward<6, false, false>(li, N); + } + else + { + if (store_head) + layer_forward<15, true, true>(li, N); + else if (do_l1x1) + layer_forward<15, false, true>(li, N); + else + layer_forward<15, false, false>(li, N); + } + } + + template + void layer_forward(int li, int N) + { + Ring& R = _rings[li]; + const NanoLayer& P = _p[li]; + + const float* h[C] = {R.plane(0), R.plane(1), R.plane(2)}; + int tapb[K]; + for (int k = 0; k < K; k++) + tapb[k] = R.tap((K - 1 - k) * P.dilation, N); + + // Residual destination: the next layer's ring, so no layer ever copies its + // input in from a scratch buffer. + Ring* next = nullptr; + float* d[C]; + if (li + 1 < kNumLayers) + { + next = &_rings[li + 1]; + next->prepare(N); + for (int c = 0; c < C; c++) + d[c] = next->write_ptr(c); + } + else + { + for (int c = 0; c < C; c++) + d[c] = lin(c); + } + + float* hs[C] = {hsum(0), hsum(1), hsum(2)}; + + int f = 0; + for (; f + kNanoTile <= N; f += kNanoTile) + tile(P, h, tapb, f, d, hs); + for (; f + 4 <= N; f += 4) + tile(P, h, tapb, f, d, hs); + for (; f < N; f++) + frame_scalar(P, h, tapb, f, d, hs); + + if (next != nullptr) + next->commit(N); + } + + /// NVEC x 4 frames, three channel accumulators per vector. z stays in + /// registers across every tap instead of making a round trip to memory per + /// tap per frame. + template + inline void tile(const NanoLayer& P, const float* const* h, const int (&tapb)[K], int f0, float* const* d, + float* const* hs) + { + const float32x4_t cb = vld1q_f32(P.conv_b.data()); + float32x4_t a0[NVEC], a1[NVEC], a2[NVEC]; + for (int v = 0; v < NVEC; v++) + { + a0[v] = vdupq_laneq_f32(cb, 0); + a1[v] = vdupq_laneq_f32(cb, 1); + a2[v] = vdupq_laneq_f32(cb, 2); + } + + const float* cw = P.conv_w.data(); + for (int k = 0; k < K; k++) + { + const float* wk = cw + static_cast(k) * 12; + const float32x4_t A = vld1q_f32(wk); // w0 w1 w2 w3 + const float32x4_t B = vld1q_f32(wk + 4); // w4 w5 w6 w7 + const float32x4_t Cw = vld1q_f32(wk + 8); // w8 . . . + const int b = tapb[k] + f0; + for (int v = 0; v < NVEC; v++) + { + const float32x4_t s0 = vld1q_f32(h[0] + b + 4 * v); + const float32x4_t s1 = vld1q_f32(h[1] + b + 4 * v); + const float32x4_t s2 = vld1q_f32(h[2] + b + 4 * v); + a0[v] = vfmaq_laneq_f32(a0[v], s0, A, 0); + a1[v] = vfmaq_laneq_f32(a1[v], s0, A, 1); + a2[v] = vfmaq_laneq_f32(a2[v], s0, A, 2); + a0[v] = vfmaq_laneq_f32(a0[v], s1, A, 3); + a1[v] = vfmaq_laneq_f32(a1[v], s1, B, 0); + a2[v] = vfmaq_laneq_f32(a2[v], s1, B, 1); + a0[v] = vfmaq_laneq_f32(a0[v], s2, B, 2); + a1[v] = vfmaq_laneq_f32(a1[v], s2, B, 3); + a2[v] = vfmaq_laneq_f32(a2[v], s2, Cw, 0); + } + } + + post(P, h, tapb[K - 1], f0, a0, a1, a2, d, hs); + } + + /// Everything after the conv: mixin, LeakyReLU, head_sum, layer1x1 residual. + /// `last_tap` is the base of the offset-0 tap, i.e. this block's own input -- + /// reloading it here is cheaper than carrying it through the tap loop, which + /// is what decides how wide the tile can usefully get. + template + inline void post(const NanoLayer& P, const float* const* h, int last_tap, int f0, float32x4_t (&a0)[NVEC], + float32x4_t (&a1)[NVEC], float32x4_t (&a2)[NVEC], float* const* d, float* const* hs) + { + const float32x4_t M = vld1q_f32(P.mixin_w.data()); + const float32x4_t zero = vdupq_n_f32(0.0f); + const float32x4_t slope = vdupq_n_f32(kLeakySlope); + const float* cond = _cond.data(); + + for (int v = 0; v < NVEC; v++) + { + const float32x4_t cf = vld1q_f32(cond + f0 + 4 * v); + a0[v] = vfmaq_laneq_f32(a0[v], cf, M, 0); + a1[v] = vfmaq_laneq_f32(a1[v], cf, M, 1); + a2[v] = vfmaq_laneq_f32(a2[v], cf, M, 2); + a0[v] = vbslq_f32(vcltq_f32(a0[v], zero), vmulq_f32(a0[v], slope), a0[v]); + a1[v] = vbslq_f32(vcltq_f32(a1[v], zero), vmulq_f32(a1[v], slope), a1[v]); + a2[v] = vbslq_f32(vcltq_f32(a2[v], zero), vmulq_f32(a2[v], slope), a2[v]); + } + + for (int v = 0; v < NVEC; v++) + { + const int o = f0 + 4 * v; + if constexpr (StoreHead) + { + // Kept as an add against +0.0 rather than a plain store: 0.0f + (-0.0f) + // is +0.0f, so storing would differ from a2_fast on a signed zero. One + // vector add per four frames in one layer, and the exactness is then a + // fact rather than an argument about whether that case can arise. + vst1q_f32(hs[0] + o, vaddq_f32(zero, a0[v])); + vst1q_f32(hs[1] + o, vaddq_f32(zero, a1[v])); + vst1q_f32(hs[2] + o, vaddq_f32(zero, a2[v])); + } + else + { + vst1q_f32(hs[0] + o, vaddq_f32(vld1q_f32(hs[0] + o), a0[v])); + vst1q_f32(hs[1] + o, vaddq_f32(vld1q_f32(hs[1] + o), a1[v])); + vst1q_f32(hs[2] + o, vaddq_f32(vld1q_f32(hs[2] + o), a2[v])); + } + } + + if constexpr (DoL1x1) + { + const float32x4_t LA = vld1q_f32(P.l1x1_w.data()); // l0 l1 l2 l3 + const float32x4_t LB = vld1q_f32(P.l1x1_w.data() + 4); // l4 l5 l6 l7 + const float32x4_t LC = vld1q_f32(P.l1x1_w.data() + 8); // l8 . . . + const float32x4_t LBias = vld1q_f32(P.l1x1_b.data()); + + for (int v = 0; v < NVEC; v++) + { + const int o = f0 + 4 * v; + float32x4_t o0 = vdupq_laneq_f32(LBias, 0); + float32x4_t o1 = vdupq_laneq_f32(LBias, 1); + float32x4_t o2 = vdupq_laneq_f32(LBias, 2); + o0 = vfmaq_laneq_f32(o0, a0[v], LA, 0); + o0 = vfmaq_laneq_f32(o0, a1[v], LA, 3); + o0 = vfmaq_laneq_f32(o0, a2[v], LB, 2); + o1 = vfmaq_laneq_f32(o1, a0[v], LA, 1); + o1 = vfmaq_laneq_f32(o1, a1[v], LB, 0); + o1 = vfmaq_laneq_f32(o1, a2[v], LB, 3); + o2 = vfmaq_laneq_f32(o2, a0[v], LA, 2); + o2 = vfmaq_laneq_f32(o2, a1[v], LB, 1); + o2 = vfmaq_laneq_f32(o2, a2[v], LC, 0); + vst1q_f32(d[0] + o, vaddq_f32(vld1q_f32(h[0] + last_tap + o), o0)); + vst1q_f32(d[1] + o, vaddq_f32(vld1q_f32(h[1] + last_tap + o), o1)); + vst1q_f32(d[2] + o, vaddq_f32(vld1q_f32(h[2] + last_tap + o), o2)); + } + } + } + + /// The last few frames of a block that is not a multiple of four. Same + /// operation order as the vector path, one frame at a time. + template + void frame_scalar(const NanoLayer& P, const float* const* h, const int (&tapb)[K], int f, float* const* d, + float* const* hs) + { + float a[C] = {P.conv_b[0], P.conv_b[1], P.conv_b[2]}; + for (int k = 0; k < K; k++) + { + const float* wk = P.conv_w.data() + static_cast(k) * 12; + const float s0 = h[0][tapb[k] + f]; + const float s1 = h[1][tapb[k] + f]; + const float s2 = h[2][tapb[k] + f]; + a[0] += wk[0] * s0; + a[1] += wk[1] * s0; + a[2] += wk[2] * s0; + a[0] += wk[3] * s1; + a[1] += wk[4] * s1; + a[2] += wk[5] * s1; + a[0] += wk[6] * s2; + a[1] += wk[7] * s2; + a[2] += wk[8] * s2; + } + + const float cf = _cond[f]; + for (int c = 0; c < C; c++) + { + a[c] += P.mixin_w[c] * cf; + a[c] = (a[c] < 0.0f) ? a[c] * kLeakySlope : a[c]; + if constexpr (StoreHead) + hs[c][f] = 0.0f + a[c]; + else + hs[c][f] = hs[c][f] + a[c]; + } + + if constexpr (DoL1x1) + { + for (int c = 0; c < C; c++) + { + float o = P.l1x1_b[c]; + o += P.l1x1_w[0 + c] * a[0]; + o += P.l1x1_w[3 + c] * a[1]; + o += P.l1x1_w[6 + c] * a[2]; + d[c][f] = h[c][tapb[K - 1] + f] + o; + } + } + } + + /// Head rechannel: K=16, dilation 1, three channels down to one, plus bias and + /// scale. In planar layout this is 48 vector FMAs per four frames where + /// a2_fast does 48 scalar FMAs per frame. + void head_forward(int N) + { + _head_ring.prepare(N); + for (int c = 0; c < C; c++) + std::memcpy(_head_ring.write_ptr(c), hsum(c), static_cast(N) * sizeof(float)); + _head_ring.commit(N); + + int hb[kHeadKernelSize]; + for (int k = 0; k < kHeadKernelSize; k++) + hb[k] = _head_ring.tap(kHeadKernelSize - 1 - k, N); + + const float* p[C] = {_head_ring.plane(0), _head_ring.plane(1), _head_ring.plane(2)}; + const float* hw = _head_w4.data(); + const float scale = _w.head_scale; + float* out = _head_out.data(); + + int f = 0; + for (; f + 4 <= N; f += 4) + { + float32x4_t y = vdupq_n_f32(_w.head_b); + for (int k = 0; k < kHeadKernelSize; k++) + { + const float32x4_t W = vld1q_f32(hw + static_cast(k) * 4); + y = vfmaq_laneq_f32(y, vld1q_f32(p[0] + hb[k] + f), W, 0); + y = vfmaq_laneq_f32(y, vld1q_f32(p[1] + hb[k] + f), W, 1); + y = vfmaq_laneq_f32(y, vld1q_f32(p[2] + hb[k] + f), W, 2); + } + vst1q_f32(out + f, vmulq_n_f32(y, scale)); + } + for (; f < N; f++) + { + float y = _w.head_b; + for (int k = 0; k < kHeadKernelSize; k++) + { + const float* w = hw + static_cast(k) * 4; + y += w[0] * p[0][hb[k] + f]; + y += w[1] * p[1][hb[k] + f]; + y += w[2] * p[2][hb[k] + f]; + } + out[f] = y * scale; + } + } + + PlanarWeights _w; + int _prewarm_samples = 0; + + std::array _p; + std::vector _head_w4; + + std::array _rings; + Ring _head_ring; + + std::vector _layer_in; + std::vector _head_sum; + std::vector _cond; + std::vector _head_out; + int _stride = 0; +}; + +// ============================================================================= +// Channels == 8 (A2 standard) +// +// Reproduces what a2_fast's Eigen expressions compute, including the per-tap +// partial that is summed into the running total only at the end of the tap, and +// the mixin's separate multiply and add. See the header comment for the full +// order. +// ============================================================================= + +/// Frames per conv tile. a2_fast's association needs the running total `z` and +/// the current tap's partial `t` live at once, which is 2 x 8 x (tile/4) vector +/// registers; the measured curve peaks at 8 and falls off at 16. +constexpr int kFullTile = 8; +constexpr int kFullVecs = kFullTile / 4; + +/// Independent head chains. The head is a 128-deep serial FMA chain per frame, +/// so it is latency-bound rather than throughput-bound; running eight chains at +/// once costs nothing in registers and nothing in exactness, because each chain +/// still covers its own frames in a2_fast's own order. +constexpr int kFullHeadVecs = 8; + +class A2PlanarFull : public DSP +{ + static constexpr int C = 8; + using Ring = PlanarRing; + +public: + A2PlanarFull(const std::vector& weights, double expected_sample_rate) + : DSP(/*in_channels=*/1, /*out_channels=*/1, expected_sample_rate) + , _w(parse_weights(weights)) + , _prewarm_samples(planar_prewarm_samples()) + { + // Head weights as C contiguous floats per tap, so a tap's whole weight row + // is two vector loads addressed by lane. + _head_w.assign(static_cast(kHeadKernelSize) * C, 0.0f); + for (int k = 0; k < kHeadKernelSize; k++) + for (int b = 0; b < C; b++) + _head_w[static_cast(k) * C + b] = _w.head_w[k][b]; + + // layer1x1 transposed: [i*C + j] is the weight from bottleneck j to output i, + // so one output's whole row is two contiguous vector loads instead of eight + // scalar broadcasts. Identical FMAs in an identical order -- only the route + // the weight takes to the instruction changes. + _l1x1t.assign(static_cast(kNumLayers) * C * C, 0.0f); + for (int li = 0; li < kNumLayers; li++) + for (int i = 0; i < C; i++) + for (int j = 0; j < C; j++) + _l1x1t[(static_cast(li) * C + i) * C + j] = _w.layers[li].l1x1_w[static_cast(j) * C + i]; + } + + ~A2PlanarFull() override = default; + + int GetPrewarmSamples() override { return _prewarm_samples; } + + void process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames) override + { + if (num_frames > GetMaxBufferSize()) + SetMaxBufferSize(num_frames); + const int N = num_frames; + + const NAM_SAMPLE* in0 = input[0]; + NAM_SAMPLE* out0 = output[0]; + + float* cond = _cond.data(); + for (int f = 0; f < N; f++) + cond[f] = static_cast(in0[f]); + + // Rechannel straight into layer 0's ring. + _rings[0].prepare(N); + for (int c = 0; c < C; c++) + { + float* dst = _rings[0].write_ptr(c); + const float wc = _w.rechannel_w[c]; + int f = 0; + for (; f + 4 <= N; f += 4) + vst1q_f32(dst + f, vmulq_n_f32(vld1q_f32(cond + f), wc)); + for (; f < N; f++) + dst[f] = wc * cond[f]; + } + _rings[0].commit(N); + + // The layers accumulate head_sum straight into the head ring's write window, + // which removes the block-sized copy the head would otherwise make out of a + // scratch buffer -- the same trick as writing each residual into the next + // layer's ring. + _head_ring.prepare(N); + float* hs[C]; + for (int c = 0; c < C; c++) + hs[c] = _head_ring.write_ptr(c); + + // No memset: layer 0 writes head_sum rather than accumulating onto it. + for (int li = 0; li < kNumLayers; li++) + dispatch_layer(li, N, hs); + + _head_ring.commit(N); + + head_forward(N); + + const float* head_out = _head_out.data(); + for (int f = 0; f < N; f++) + out0[f] = static_cast(head_out[f]); + } + +protected: + void SetMaxBufferSize(const int maxBufferSize) override + { + DSP::SetMaxBufferSize(maxBufferSize); + + _stride = maxBufferSize; + _layer_in.assign(static_cast(C) * _stride, 0.0f); + _cond.assign(static_cast(maxBufferSize), 0.0f); + _head_out.assign(static_cast(maxBufferSize), 0.0f); + + for (int li = 0; li < kNumLayers; li++) + _rings[li].reset(_w.layers[li].max_lookback, maxBufferSize); + _head_ring.reset(kHeadKernelSize - 1, maxBufferSize); + } + +private: + float* lin(int c) { return _layer_in.data() + static_cast(c) * _stride; } + + void dispatch_layer(int li, int N, float* const* hs) + { + const bool store_head = (li == 0); + const bool do_l1x1 = (li != kNumLayers - 1); + + if (_w.layers[li].kernel_size == 6) + { + if (store_head) + layer_forward<6, true, true>(li, N, hs); + else if (do_l1x1) + layer_forward<6, false, true>(li, N, hs); + else + layer_forward<6, false, false>(li, N, hs); + } + else + { + if (store_head) + layer_forward<15, true, true>(li, N, hs); + else if (do_l1x1) + layer_forward<15, false, true>(li, N, hs); + else + layer_forward<15, false, false>(li, N, hs); + } + } + + template + void layer_forward(int li, int N, float* const* hs) + { + Ring& R = _rings[li]; + const PlanarLayerWeights& L = _w.layers[li]; + + const float* h[C]; + for (int c = 0; c < C; c++) + h[c] = R.plane(c); + + int tapb[K]; + for (int k = 0; k < K; k++) + tapb[k] = R.tap((K - 1 - k) * L.dilation, N); + + Ring* next = nullptr; + float* d[C]; + if (li + 1 < kNumLayers) + { + next = &_rings[li + 1]; + next->prepare(N); + for (int c = 0; c < C; c++) + d[c] = next->write_ptr(c); + } + else + { + for (int c = 0; c < C; c++) + d[c] = lin(c); + } + + const float* lt = _l1x1t.data() + static_cast(li) * C * C; + + int f = 0; + for (; f + kFullTile <= N; f += kFullTile) + tile(L, lt, h, tapb, f, d, hs); + for (; f + 4 <= N; f += 4) + tile(L, lt, h, tapb, f, d, hs); + for (; f < N; f++) + frame_scalar(L, h, tapb, f, d, hs); + + if (next != nullptr) + next->commit(N); + } + + /// NVEC x 4 frames. `z` is the running total across taps and `t` the current + /// tap's partial -- both live at once, because that separation *is* a2_fast's + /// association. z never reaches memory. + template + inline void tile(const PlanarLayerWeights& L, const float* lt, const float* const* h, const int (&tapb)[K], int f0, + float* const* d, float* const* hs) + { + float32x4_t z[C][NVEC]; + const float32x4_t zero = vdupq_n_f32(0.0f); + for (int i = 0; i < C; i++) + for (int v = 0; v < NVEC; v++) + z[i][v] = zero; + + const float* cw = L.conv_w.data(); + for (int k = 0; k < K; k++) + { + const float* wk = cw + static_cast(k) * C * C; + const int base = tapb[k] + f0; + float32x4_t t[C][NVEC]; + + // Input channel j is unrolled at compile time so that j == 0 can seed the + // partial with a multiply instead of an FMA against zero. Both are one + // rounding of w*x, so this is exact either way; it just saves the init. + const auto do_j = [&](auto jc) { + constexpr int j = decltype(jc)::value; + const float* wj = wk + j * C; // W(0..C-1, j), contiguous + float32x4_t wv[C / 4]; + for (int u = 0; u < C / 4; u++) + wv[u] = vld1q_f32(wj + 4 * u); + const float* hp = h[j] + base; + for (int v = 0; v < NVEC; v++) + { + const float32x4_t s = vld1q_f32(hp + 4 * v); + const auto do_i = [&](auto ic) { + constexpr int i = decltype(ic)::value; + if constexpr (j == 0) + t[i][v] = vmulq_laneq_f32(s, wv[i / 4], i % 4); + else + t[i][v] = vfmaq_laneq_f32(t[i][v], s, wv[i / 4], i % 4); + }; + [&](std::integer_sequence) { + (do_i(std::integral_constant{}), ...); + }(std::make_integer_sequence{}); + } + }; + [&](std::integer_sequence) { + (do_j(std::integral_constant{}), ...); + }(std::make_integer_sequence{}); + + for (int i = 0; i < C; i++) + for (int v = 0; v < NVEC; v++) + z[i][v] = vaddq_f32(z[i][v], t[i][v]); + } + + post(L, lt, h, tapb[K - 1], f0, z, d, hs); + } + + /// Everything after the conv: bias, mixin, LeakyReLU, head_sum, layer1x1 + /// residual -- in a2_fast's order, with the mixin's multiply and add rounded + /// separately as Eigen rounds them. + template + inline void post(const PlanarLayerWeights& L, const float* lt, const float* const* h, int last_tap, int f0, + float32x4_t (&z)[C][NVEC], float* const* d, float* const* hs) + { + const float32x4_t zero = vdupq_n_f32(0.0f); + const float32x4_t slope = vdupq_n_f32(kLeakySlope); + const float* cond = _cond.data(); + + float32x4_t cf[NVEC]; + for (int v = 0; v < NVEC; v++) + cf[v] = vld1q_f32(cond + f0 + 4 * v); + + for (int i = 0; i < C; i++) + { + const float32x4_t b = vdupq_n_f32(L.conv_b[i]); + const float m = L.mixin_w[i]; + for (int v = 0; v < NVEC; v++) + { + float32x4_t a = vaddq_f32(z[i][v], b); + a = vaddq_f32(a, vmulq_n_f32(cf[v], m)); // product then add, two roundings + z[i][v] = vbslq_f32(vcltq_f32(a, zero), vmulq_f32(a, slope), a); + } + } + + for (int i = 0; i < C; i++) + { + float* p = hs[i] + f0; + for (int v = 0; v < NVEC; v++) + { + // The add against +0.0 in the store case is deliberate; see the note in + // the nano kernel's post(). + if constexpr (StoreHead) + vst1q_f32(p + 4 * v, vaddq_f32(zero, z[i][v])); + else + vst1q_f32(p + 4 * v, vaddq_f32(vld1q_f32(p + 4 * v), z[i][v])); + } + } + + if constexpr (DoL1x1) + { + for (int i = 0; i < C; i++) + { + const float32x4_t bi = vdupq_n_f32(L.l1x1_b[i]); + const float32x4_t la = vld1q_f32(lt + static_cast(i) * C); // L(i, 0..3) + const float32x4_t lb = vld1q_f32(lt + static_cast(i) * C + 4); // L(i, 4..7) + for (int v = 0; v < NVEC; v++) + { + // u_i = sum over j in increasing order, from zero. + float32x4_t u = vmulq_laneq_f32(z[0][v], la, 0); + u = vfmaq_laneq_f32(u, z[1][v], la, 1); + u = vfmaq_laneq_f32(u, z[2][v], la, 2); + u = vfmaq_laneq_f32(u, z[3][v], la, 3); + u = vfmaq_laneq_f32(u, z[4][v], lb, 0); + u = vfmaq_laneq_f32(u, z[5][v], lb, 1); + u = vfmaq_laneq_f32(u, z[6][v], lb, 2); + u = vfmaq_laneq_f32(u, z[7][v], lb, 3); + const float32x4_t prev = vld1q_f32(h[i] + last_tap + f0 + 4 * v); + vst1q_f32(d[i] + f0 + 4 * v, vaddq_f32(vaddq_f32(prev, u), bi)); + } + } + } + } + + /// The last few frames of a block that is not a multiple of four. Same + /// operation order as the vector path, one frame at a time. + template + void frame_scalar(const PlanarLayerWeights& L, const float* const* h, const int (&tapb)[K], int f, float* const* d, + float* const* hs) + { + float z[C]; + for (int i = 0; i < C; i++) + z[i] = 0.0f; + + for (int k = 0; k < K; k++) + { + const float* wk = L.conv_w.data() + static_cast(k) * C * C; + float t[C]; + for (int i = 0; i < C; i++) + t[i] = 0.0f; + for (int j = 0; j < C; j++) + { + const float s = h[j][tapb[k] + f]; + for (int i = 0; i < C; i++) + t[i] += wk[static_cast(j) * C + i] * s; + } + for (int i = 0; i < C; i++) + z[i] += t[i]; + } + + const float cf = _cond[f]; + for (int i = 0; i < C; i++) + { + float a = z[i] + L.conv_b[i]; + a = a + mul_rounded(L.mixin_w[i], cf); // product then add, two roundings + z[i] = (a < 0.0f) ? a * kLeakySlope : a; + if constexpr (StoreHead) + hs[i][f] = 0.0f + z[i]; + else + hs[i][f] = hs[i][f] + z[i]; + } + + if constexpr (DoL1x1) + { + for (int i = 0; i < C; i++) + { + float u = 0.0f; + for (int j = 0; j < C; j++) + u += L.l1x1_w[static_cast(j) * C + i] * z[j]; + d[i][f] = (h[i][tapb[K - 1] + f] + u) + L.l1x1_b[i]; + } + } + } + + /// Head rechannel: K=16, dilation 1, eight channels down to one, plus bias and + /// scale. a2_fast runs 128 sequential FMAs per frame; here the same 128 FMAs + /// cover four frames at a time, and kFullHeadVecs of those chains run + /// independently. + void head_forward(int N) + { + int hb[kHeadKernelSize]; + for (int k = 0; k < kHeadKernelSize; k++) + hb[k] = _head_ring.tap(kHeadKernelSize - 1 - k, N); + + const float* p[C]; + for (int c = 0; c < C; c++) + p[c] = _head_ring.plane(c); + + const float* hw = _head_w.data(); + const float scale = _w.head_scale; + const float32x4_t bias = vdupq_n_f32(_w.head_b); + float* out = _head_out.data(); + + int f = 0; + for (; f + 4 * kFullHeadVecs <= N; f += 4 * kFullHeadVecs) + { + float32x4_t y[kFullHeadVecs]; + for (int u = 0; u < kFullHeadVecs; u++) + y[u] = bias; + for (int k = 0; k < kHeadKernelSize; k++) + { + const float32x4_t wa = vld1q_f32(hw + static_cast(k) * C); + const float32x4_t wb = vld1q_f32(hw + static_cast(k) * C + 4); + const int base = hb[k] + f; + for (int u = 0; u < kFullHeadVecs; u++) + { + const int o = base + 4 * u; + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[0] + o), wa, 0); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[1] + o), wa, 1); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[2] + o), wa, 2); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[3] + o), wa, 3); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[4] + o), wb, 0); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[5] + o), wb, 1); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[6] + o), wb, 2); + y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[7] + o), wb, 3); + } + } + for (int u = 0; u < kFullHeadVecs; u++) + vst1q_f32(out + f + 4 * u, vmulq_n_f32(y[u], scale)); + } + for (; f + 4 <= N; f += 4) + { + float32x4_t y = bias; + for (int k = 0; k < kHeadKernelSize; k++) + { + const float32x4_t wa = vld1q_f32(hw + static_cast(k) * C); + const float32x4_t wb = vld1q_f32(hw + static_cast(k) * C + 4); + const int o = hb[k] + f; + y = vfmaq_laneq_f32(y, vld1q_f32(p[0] + o), wa, 0); + y = vfmaq_laneq_f32(y, vld1q_f32(p[1] + o), wa, 1); + y = vfmaq_laneq_f32(y, vld1q_f32(p[2] + o), wa, 2); + y = vfmaq_laneq_f32(y, vld1q_f32(p[3] + o), wa, 3); + y = vfmaq_laneq_f32(y, vld1q_f32(p[4] + o), wb, 0); + y = vfmaq_laneq_f32(y, vld1q_f32(p[5] + o), wb, 1); + y = vfmaq_laneq_f32(y, vld1q_f32(p[6] + o), wb, 2); + y = vfmaq_laneq_f32(y, vld1q_f32(p[7] + o), wb, 3); + } + vst1q_f32(out + f, vmulq_n_f32(y, scale)); + } + for (; f < N; f++) + { + float y = _w.head_b; + for (int k = 0; k < kHeadKernelSize; k++) + { + const float* wk = hw + static_cast(k) * C; + for (int b = 0; b < C; b++) + y += wk[b] * p[b][hb[k] + f]; + } + out[f] = y * scale; + } + } + + PlanarWeights _w; + int _prewarm_samples = 0; + + std::vector _head_w; + std::vector _l1x1t; + + std::array _rings; + Ring _head_ring; + + std::vector _layer_in; + std::vector _cond; + std::vector _head_out; + int _stride = 0; +}; + +} // namespace + +std::unique_ptr create_a2_planar_model(int channels, std::vector weights, double expected_sample_rate) +{ + if (channels == 3) + return std::make_unique(weights, expected_sample_rate); + if (channels == 8) + return std::make_unique(weights, expected_sample_rate); + return nullptr; +} + +} // namespace a2_fast +} // namespace wavenet +} // namespace nam + + #endif // NAM_A2_PLANAR +#endif // NAM_ENABLE_A2_FAST diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h new file mode 100644 index 00000000..1b83b34d --- /dev/null +++ b/NAM/wavenet/a2_planar.h @@ -0,0 +1,53 @@ +#pragma once + +// Planar NEON kernels for the A2 fast path (AArch64 only). +// +// These are drop-in replacements for A2FastModel<3> and A2FastModel<8> that +// produce **bit-identical** output: not "within a tolerance", not "below the +// noise floor" -- the same float32 bits, sample for sample. +// +// The idea in one line: a2_fast keeps the channels of a frame adjacent in +// memory and vectorises across channels; these kernels keep each channel in its +// own plane and vectorise across *frames*, so one NEON lane runs a2_fast's +// per-frame scalar reduction verbatim. Nothing is reassociated, which is what +// makes the bit-identity claim hold rather than being a lucky accident. +// +// Availability is decided here rather than at the call site: NAM_A2_PLANAR is +// defined only when the A2 fast path is built for AArch64. Everywhere else this +// header declares nothing and a2_fast keeps its existing behaviour. Define +// NAM_DISABLE_A2_PLANAR to opt out on AArch64 too (useful for A/B measurement). + +#if defined(NAM_ENABLE_A2_FAST) + + #if (defined(__aarch64__) || defined(_M_ARM64)) && !defined(NAM_DISABLE_A2_PLANAR) + #define NAM_A2_PLANAR 1 + #endif + + #if defined(NAM_A2_PLANAR) + + #include + #include + + #include "../dsp.h" + +namespace nam +{ +namespace wavenet +{ +namespace a2_fast +{ + +/// \brief Build the planar NEON model for an A2 submodel. +/// \param channels 3 (A2 nano) or 8 (A2 standard); anything else yields nullptr. +/// \param weights The A2 weight stream, consumed in A2FastModel's order. +/// \param expected_sample_rate Passed through to DSP. +/// \return The model, or nullptr when this channel count has no planar kernel +/// (the caller then falls back to A2FastModel). +std::unique_ptr create_a2_planar_model(int channels, std::vector weights, double expected_sample_rate); + +} // namespace a2_fast +} // namespace wavenet +} // namespace nam + + #endif // NAM_A2_PLANAR +#endif // NAM_ENABLE_A2_FAST diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 51c7e7ea..e430aac3 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -74,6 +74,32 @@ else() "$<$:-Ofast>" ) endif() +add_executable(bench_a2_planar bench_a2_planar.cpp ${NAM_SOURCES}) +target_compile_features(bench_a2_planar PUBLIC cxx_std_20) +# No INTERPROCEDURAL_OPTIMIZATION here: link-time inlining of the whole library +# into the benchmark's own loop is not what a plugin build does, and measurably +# changes the result (about 3% on the 8-channel kernel on an M2). +set_target_properties(bench_a2_planar PROPERTIES + CXX_VISIBILITY_PRESET hidden + PREFIX "" +) +if (MSVC) + target_compile_options(bench_a2_planar PRIVATE + "$<$:/W4>" + "$<$:/O2>" + ) +else() + # -O3, not -Ofast. -ffast-math lets the compiler contract a multiply and an + # add into an FMA across statement boundaries, which is exactly the freedom + # this tool is checking has not been taken; measuring under it would make the + # parity result meaningless. + target_compile_options(bench_a2_planar PRIVATE + -Wall -Wextra -Wpedantic -Wstrict-aliasing -Wunreachable-code -Wno-unused-parameter + "$<$:-Og;-ggdb;-Werror>" + "$<$:-O3>" + ) +endif() + add_executable(run_tests run_tests.cpp test/allocation_tracking.cpp ${NAM_SOURCES}) # Compile run_tests without optimizations to ensure allocation tracking works correctly # Also ensure assertions are enabled (NDEBUG is not defined) so tests actually run diff --git a/tools/bench_a2_planar.cpp b/tools/bench_a2_planar.cpp new file mode 100644 index 00000000..2ceb636c --- /dev/null +++ b/tools/bench_a2_planar.cpp @@ -0,0 +1,339 @@ +// Head-to-head for the planar NEON A2 kernels against the reference A2 fast +// path: same model, same weights, same input, same process, in one binary. +// +// It checks before it times. Every run first renders the whole signal through +// both engines and compares the output bit for bit; if they differ it says so +// and reports no speed at all, because a speed number for a kernel that is not +// reproducing the reference is not worth having. +// +// Timing follows the shape that survived being wrong in earlier attempts: +// interference on a desktop machine is one-sided -- it can only make a pass +// slower -- so the estimate is the mean of the fastest 70% of passes rather +// than the mean or the median of all of them, and the fastest single pass is +// printed next to it so the spread is visible. +// +// Usage: +// bench_a2_planar [--buffer N] [--seconds S] [--warmup W] [--passes P] +// [--submodel widest|narrowest|] ... +// +// A .nam holding a SlimmableContainer is unwrapped and one submodel is +// measured; --submodel picks it by width, not by position, so a reordering in +// the trainer cannot silently change what is being measured. +// +// Only compiled when NAM_ENABLE_A2_FAST is defined. + +#if defined(NAM_ENABLE_A2_FAST) + + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + + #include "json.hpp" + + #include "NAM/dsp.h" + #include "NAM/wavenet/a2_fast.h" + #include "NAM/wavenet/a2_planar.h" + +using hr_clock = std::chrono::high_resolution_clock; + +namespace +{ + +struct Options +{ + int buffer_size = 64; + double seconds = 10.9; // one pass of audio + double warmup_seconds = 5.0; // discarded + int passes = 12; // timed + double accept_fraction = 0.7; + std::string submodel = "widest"; + std::vector model_paths; +}; + +Options parse_args(int argc, char** argv) +{ + Options o; + for (int i = 1; i < argc; i++) + { + std::string a = argv[i]; + if (a == "--buffer" && i + 1 < argc) + o.buffer_size = std::atoi(argv[++i]); + else if (a == "--seconds" && i + 1 < argc) + o.seconds = std::atof(argv[++i]); + else if (a == "--warmup" && i + 1 < argc) + o.warmup_seconds = std::atof(argv[++i]); + else if (a == "--passes" && i + 1 < argc) + o.passes = std::atoi(argv[++i]); + else if (a == "--submodel" && i + 1 < argc) + o.submodel = argv[++i]; + else if (a == "-h" || a == "--help") + { + std::cerr << "Usage: bench_a2_planar [--buffer N] [--seconds S] [--warmup W] [--passes P]\n" + << " [--submodel widest|narrowest|] ...\n"; + std::exit(0); + } + else + o.model_paths.push_back(std::move(a)); + } + return o; +} + +struct LoadedModel +{ + nlohmann::json config; + std::vector weights; + double sample_rate = 48000.0; + std::string path; + std::string note; // which submodel, when unwrapped +}; + +/// Pull one WaveNet out of a .nam, unwrapping a SlimmableContainer if that is +/// what it holds. Selection is by channel count, so it does not depend on the +/// order the submodels happen to be stored in. +LoadedModel load_nam(const std::string& path, const std::string& submodel) +{ + std::ifstream is(path); + if (!is) + throw std::runtime_error("Could not open " + path); + nlohmann::json j; + is >> j; + + const std::string arch = j.value("architecture", std::string()); + nlohmann::json wavenet = j; + std::string note; + + if (arch == "SlimmableContainer") + { + const auto& subs = j.at("config").at("submodels"); + if (!subs.is_array() || subs.empty()) + throw std::runtime_error(path + ": SlimmableContainer has no submodels"); + + int chosen = -1; + if (submodel == "widest" || submodel == "narrowest") + { + int best = -1; + for (size_t i = 0; i < subs.size(); i++) + { + const auto& m = subs[i].at("model"); + const int ch = m.at("config").at("layers")[0].value("channels", 0); + const bool better = (chosen < 0) || (submodel == "widest" ? ch > best : ch < best); + if (better) + { + best = ch; + chosen = static_cast(i); + } + } + } + else + { + chosen = std::atoi(submodel.c_str()); + if (chosen < 0 || chosen >= static_cast(subs.size())) + throw std::runtime_error(path + ": no submodel " + submodel); + } + + wavenet = subs[chosen].at("model"); + note = "submodel " + std::to_string(chosen) + " of " + std::to_string(subs.size()); + } + else if (arch != "WaveNet") + { + throw std::runtime_error(path + ": not a WaveNet or SlimmableContainer model"); + } + + LoadedModel m; + m.path = path; + m.note = note; + m.config = wavenet.at("config"); + m.weights = wavenet.at("weights").get>(); + if (wavenet.contains("sample_rate") && !wavenet["sample_rate"].is_null()) + m.sample_rate = wavenet["sample_rate"].get(); + return m; +} + +/// One full pass over the signal, in blocks. Returns wall time in milliseconds. +double run_pass(nam::DSP& dsp, const std::vector& input, std::vector& output, int buffer_size) +{ + const int total = static_cast(input.size()); + const auto t0 = hr_clock::now(); + int pos = 0; + while (pos < total) + { + const int n = std::min(buffer_size, total - pos); + const NAM_SAMPLE* in_ptr = input.data() + pos; + NAM_SAMPLE* out_ptr = output.data() + pos; + const NAM_SAMPLE* in_arr[] = {in_ptr}; + NAM_SAMPLE* out_arr[] = {out_ptr}; + dsp.process(const_cast(in_arr), out_arr, n); + pos += n; + } + const auto t1 = hr_clock::now(); + return std::chrono::duration(t1 - t0).count(); +} + +struct Timing +{ + double mean = 0.0; // of the fastest `accept_fraction` of passes + double fastest = 0.0; + double slowest = 0.0; + int kept = 0; +}; + +Timing summarise(std::vector times, double accept_fraction) +{ + Timing t; + if (times.empty()) + return t; + std::sort(times.begin(), times.end()); + t.fastest = times.front(); + t.slowest = times.back(); + t.kept = std::max(1, static_cast(times.size() * accept_fraction)); + double sum = 0.0; + for (int i = 0; i < t.kept; i++) + sum += times[i]; + t.mean = sum / t.kept; + return t; +} + +/// Bit-for-bit, not within a tolerance. Returns the index of the first +/// difference, or -1. +long long first_difference(const std::vector& a, const std::vector& b) +{ + if (std::memcmp(a.data(), b.data(), a.size() * sizeof(NAM_SAMPLE)) == 0) + return -1; + for (size_t i = 0; i < a.size(); i++) + if (a[i] != b[i]) + return static_cast(i); + return 0; // differing bits in equal values (a signed zero); still a difference +} + +bool bench_model(const LoadedModel& m, const Options& o) +{ + int channels = 0; + if (!nam::wavenet::a2_fast::is_a2_shape(m.config, &channels)) + { + std::cerr << "[skip] " << m.path << ": not an A2-shaped WaveNet\n"; + return true; + } + + auto reference = nam::wavenet::a2_fast::create_a2_fast_reference_model(channels, m.weights, m.sample_rate); + auto planar = nam::wavenet::a2_fast::create_a2_planar_model(channels, m.weights, m.sample_rate); + if (planar == nullptr) + { + std::cerr << "[skip] " << m.path << ": no planar kernel for " << channels << " channels on this target\n"; + return true; + } + + const int total = static_cast(o.seconds * m.sample_rate); + std::vector input(total); + for (int i = 0; i < total; i++) + { + const double t = static_cast(i) / m.sample_rate; + input[i] = + static_cast(0.25 * std::sin(2.0 * M_PI * 220.0 * t) + 0.10 * std::sin(2.0 * M_PI * 1230.0 * t) + + 0.05 * std::sin(2.0 * M_PI * 3170.0 * t)); + } + std::vector out_reference(total, static_cast(0)); + std::vector out_planar(total, static_cast(0)); + + reference->Reset(m.sample_rate, o.buffer_size); + planar->Reset(m.sample_rate, o.buffer_size); + + const std::string arch = (channels == 3) ? "A2 nano" : "A2 standard"; + std::cout << "\n== " << m.path << (m.note.empty() ? "" : (" [" + m.note + "]")) << "\n" + << " " << arch << ", " << channels << " channels, " << m.weights.size() << " weights; " << std::fixed + << std::setprecision(2) << (total / m.sample_rate) << " s of audio per pass at " + << static_cast(m.sample_rate) << " Hz, " << o.buffer_size << "-frame blocks\n"; + + // --- Parity, before anything is timed --------------------------------------- + run_pass(*reference, input, out_reference, o.buffer_size); + run_pass(*planar, input, out_planar, o.buffer_size); + const long long diff = first_difference(out_reference, out_planar); + if (diff >= 0) + { + std::cerr << " PARITY FAILED: first difference at sample " << diff << " of " << total + << " (reference=" << out_reference[static_cast(diff)] + << ", planar=" << out_planar[static_cast(diff)] << "). Not timing.\n"; + return false; + } + std::cout << " parity: bit-identical over all " << total << " frames\n"; + + // --- Warm up, then time ----------------------------------------------------- + const int warmup_passes = std::max(1, static_cast(o.warmup_seconds / o.seconds + 0.5)); + for (int i = 0; i < warmup_passes; i++) + { + run_pass(*reference, input, out_reference, o.buffer_size); + run_pass(*planar, input, out_planar, o.buffer_size); + } + + std::vector t_reference, t_planar; + t_reference.reserve(o.passes); + t_planar.reserve(o.passes); + for (int i = 0; i < o.passes; i++) + { + // Interleaved, so a slow patch on the machine lands on both. + t_reference.push_back(run_pass(*reference, input, out_reference, o.buffer_size)); + t_planar.push_back(run_pass(*planar, input, out_planar, o.buffer_size)); + } + + const Timing r = summarise(t_reference, o.accept_fraction); + const Timing p = summarise(t_planar, o.accept_fraction); + const double audio_ms = 1000.0 * total / m.sample_rate; + + std::cout << std::fixed << std::setprecision(2); + std::cout << " " << o.passes << " passes, mean of the fastest " << r.kept << "\n"; + std::cout << " mean/pass fastest slowest x real time\n"; + std::cout << " a2_fast " << std::setw(8) << r.mean << " ms " << std::setw(9) << r.fastest << " " + << std::setw(9) << r.slowest << " " << std::setw(8) << (audio_ms / r.mean) << "\n"; + std::cout << " planar NEON " << std::setw(8) << p.mean << " ms " << std::setw(9) << p.fastest << " " + << std::setw(9) << p.slowest << " " << std::setw(8) << (audio_ms / p.mean) << "\n"; + std::cout << std::setprecision(3) << " speedup: " << (r.mean / p.mean) << "x on the mean, " + << (r.fastest / p.fastest) << "x on the fastest pass\n"; + return true; +} + +} // namespace + +int main(int argc, char** argv) +{ + const Options o = parse_args(argc, argv); + if (o.model_paths.empty()) + { + std::cerr << "Usage: bench_a2_planar [options] ... (--help for options)\n"; + return 2; + } + + bool ok = true; + for (const auto& path : o.model_paths) + { + try + { + ok = bench_model(load_nam(path, o.submodel), o) && ok; + } + catch (const std::exception& e) + { + std::cerr << "[error] " << path << ": " << e.what() << "\n"; + ok = false; + } + } + return ok ? 0 : 1; +} + +#else // NAM_ENABLE_A2_FAST + + #include + +int main() +{ + std::cerr << "bench_a2_planar: built without NAM_ENABLE_A2_FAST\n"; + return 2; +} + +#endif // NAM_ENABLE_A2_FAST diff --git a/tools/run_tests.cpp b/tools/run_tests.cpp index 0c8cd932..366400a7 100644 --- a/tools/run_tests.cpp +++ b/tools/run_tests.cpp @@ -39,6 +39,7 @@ #include "test/test_render_slim.cpp" #include "test/test_slimmable_wavenet.cpp" #include "test/test_a2_fast.cpp" +#include "test/test_a2_planar.cpp" int main() { @@ -400,6 +401,12 @@ int main() test_a2_fast::test_cached_prewarm_full(); test_a2_fast::test_process_realtime_safe_lite(); test_a2_fast::test_process_realtime_safe_full(); + + // Planar NEON A2 kernels: bit-identity against the reference fast path. + // No-ops where the planar kernels are not built. + test_a2_planar::test_bit_identical_nano(); + test_a2_planar::test_bit_identical_standard(); + test_a2_planar::test_factory_selects_planar(); #endif std::cout << "Success!" << std::endl; diff --git a/tools/test/test_a2_planar.cpp b/tools/test/test_a2_planar.cpp new file mode 100644 index 00000000..163e77bb --- /dev/null +++ b/tools/test/test_a2_planar.cpp @@ -0,0 +1,230 @@ +// Bit-identity verification for the planar NEON A2 kernels. +// +// The claim these kernels make is stronger than "close enough": for the A2 nano +// and A2 standard shapes they produce the *same float32 bits* as the reference +// A2 fast path, sample for sample. This asserts exactly that -- memcmp, not a +// tolerance -- across a spread of block sizes, including ones that exercise the +// partial-tile and single-frame tails. +// +// Built only where the planar kernels exist (AArch64 with the A2 fast path on). +// Everywhere else the test bodies compile to nothing. + +#if defined(NAM_ENABLE_A2_FAST) + + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + + #include "json.hpp" + + #include "NAM/dsp.h" + #include "NAM/wavenet/a2_fast.h" + #include "NAM/wavenet/a2_planar.h" + +namespace test_a2_planar +{ + + #if defined(NAM_A2_PLANAR) + +namespace +{ + +nlohmann::json build_a2_config(int channels) +{ + using nlohmann::json; + + json activation = json::array(); + json gating_mode = json::array(); + json secondary = json::array(); + json kernel_sizes = json::array(); + json dilations = json::array(); + for (int i = 0; i < nam::wavenet::a2_fast::kNumLayers; i++) + { + activation.push_back({{"type", "LeakyReLU"}, {"negative_slope", nam::wavenet::a2_fast::kLeakySlope}}); + gating_mode.push_back("none"); + secondary.push_back(nullptr); + kernel_sizes.push_back(nam::wavenet::a2_fast::kKernelSizes[i]); + dilations.push_back(nam::wavenet::a2_fast::kDilations[i]); + } + + json film_inactive = {{"active", false}, {"shift", true}, {"groups", 1}}; + + json layer; + layer["input_size"] = 1; + layer["condition_size"] = 1; + layer["channels"] = channels; + layer["bottleneck"] = channels; + layer["kernel_sizes"] = kernel_sizes; + layer["dilations"] = dilations; + layer["activation"] = activation; + layer["gating_mode"] = gating_mode; + layer["secondary_activation"] = secondary; + layer["head"] = {{"out_channels", 1}, {"kernel_size", nam::wavenet::a2_fast::kHeadKernelSize}, {"bias", true}}; + layer["head1x1"] = {{"active", false}, {"out_channels", 1}, {"groups", 1}}; + layer["layer1x1"] = {{"active", true}, {"groups", 1}}; + layer["conv_pre_film"] = film_inactive; + layer["conv_post_film"] = film_inactive; + layer["input_mixin_pre_film"] = film_inactive; + layer["input_mixin_post_film"] = film_inactive; + layer["activation_pre_film"] = film_inactive; + layer["activation_post_film"] = film_inactive; + layer["layer1x1_post_film"] = film_inactive; + layer["head1x1_post_film"] = film_inactive; + layer["groups_input"] = 1; + layer["groups_input_mixin"] = 1; + + json config; + config["layers"] = json::array({layer}); + config["head_scale"] = 0.01f; + return config; +} + +int a2_weight_count(int channels) +{ + const int bn = channels; + int total = /*rechannel*/ channels; + for (int i = 0; i < nam::wavenet::a2_fast::kNumLayers; i++) + { + const int K = nam::wavenet::a2_fast::kKernelSizes[i]; + total += bn * channels * K + bn; // conv1d weights + bias + total += bn; // input mixin (no bias) + total += channels * bn + channels; // layer1x1 + bias + } + total += channels * nam::wavenet::a2_fast::kHeadKernelSize + 1; // head rechannel + bias + total += 1; // trailing head_scale + return total; +} + +std::vector make_deterministic_weights(int count, uint32_t seed) +{ + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-0.3f, 0.3f); + std::vector w(count); + for (auto& x : w) + x = dist(rng); + return w; +} + +std::vector make_test_input(int num_frames, double sample_rate) +{ + std::vector in(num_frames); + for (int i = 0; i < num_frames; i++) + { + const double t = static_cast(i) / sample_rate; + in[i] = static_cast(0.25 * std::sin(2.0 * M_PI * 220.0 * t) + 0.10 * std::sin(2.0 * M_PI * 1230.0 * t)); + } + return in; +} + +std::vector run_dsp(nam::DSP& dsp, const std::vector& input, int block_size) +{ + dsp.Reset(48000.0, block_size); // also prewarms + std::vector out(input.size(), static_cast(0)); + int pos = 0; + const int total = static_cast(input.size()); + while (pos < total) + { + const int n = std::min(block_size, total - pos); + const NAM_SAMPLE* in_ptr = input.data() + pos; + NAM_SAMPLE* out_ptr = out.data() + pos; + const NAM_SAMPLE* in_arr[] = {in_ptr}; + NAM_SAMPLE* out_arr[] = {out_ptr}; + dsp.process(const_cast(in_arr), out_arr, n); + pos += n; + } + return out; +} + +/// The whole point: identical bits, not a tolerance. +void assert_bit_identical(const std::vector& reference, const std::vector& planar, int channels, + int block_size) +{ + assert(reference.size() == planar.size()); + if (std::memcmp(reference.data(), planar.data(), reference.size() * sizeof(NAM_SAMPLE)) == 0) + return; + + size_t first = 0; + while (first < reference.size() && reference[first] == planar[first]) + first++; + std::cerr << "A2 planar kernel (channels=" << channels << ", block=" << block_size + << ") is not bit-identical to the reference: first difference at sample " << first + << " (reference=" << reference[first] << ", planar=" << planar[first] << ")" << std::endl; + assert(false); +} + +void check_channels(int channels) +{ + const auto config = build_a2_config(channels); + int detected = 0; + assert(nam::wavenet::a2_fast::is_a2_shape(config, &detected)); + assert(detected == channels); + + const auto weights = make_deterministic_weights(a2_weight_count(channels), 0xA2u + channels); + // Long enough that every layer's ring wraps and rewinds several times. + const auto input = make_test_input(20000, 48000.0); + + // Block sizes chosen to hit each path: below one vector, not a multiple of + // four, exactly one vector, between one vector and one tile, the tile widths + // themselves (32 for nano, 8 for standard), and well past them. + for (const int block_size : {1, 3, 4, 7, 8, 15, 16, 31, 32, 33, 64, 65, 128, 512}) + { + auto reference = nam::wavenet::a2_fast::create_a2_fast_reference_model(channels, weights, 48000.0); + auto planar = nam::wavenet::a2_fast::create_a2_planar_model(channels, weights, 48000.0); + assert(planar != nullptr); + assert(reference->GetPrewarmSamples() == planar->GetPrewarmSamples()); + + const auto out_reference = run_dsp(*reference, input, block_size); + const auto out_planar = run_dsp(*planar, input, block_size); + assert_bit_identical(out_reference, out_planar, channels, block_size); + } +} + +} // namespace + +void test_bit_identical_nano() +{ + check_channels(3); +} + +void test_bit_identical_standard() +{ + check_channels(8); +} + +/// The dispatcher must actually route to the planar kernel where it exists, +/// otherwise the tests above would be checking something nothing uses. +void test_factory_selects_planar() +{ + for (const int channels : {3, 8}) + { + const auto config = build_a2_config(channels); + const auto weights = make_deterministic_weights(a2_weight_count(channels), 0x5Eu + channels); + auto model_config = nam::wavenet::a2_fast::create_a2_fast_config(config, 48000.0); + auto from_factory = model_config->create(weights, 48000.0); + auto planar = nam::wavenet::a2_fast::create_a2_planar_model(channels, weights, 48000.0); + + const auto input = make_test_input(4096, 48000.0); + const auto out_factory = run_dsp(*from_factory, input, 64); + const auto out_planar = run_dsp(*planar, input, 64); + assert(std::memcmp(out_factory.data(), out_planar.data(), out_factory.size() * sizeof(NAM_SAMPLE)) == 0); + } +} + + #else // NAM_A2_PLANAR + +void test_bit_identical_nano() {} +void test_bit_identical_standard() {} +void test_factory_selects_planar() {} + + #endif // NAM_A2_PLANAR + +} // namespace test_a2_planar + +#endif // NAM_ENABLE_A2_FAST From cbacb7614bfe8cf87129937ec2f2c21118365386 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Sat, 8 Aug 2026 20:36:22 +0100 Subject: [PATCH 2/8] Gate the planar kernels to Apple Silicon, and fix the tool's build off it Two corrections, both about where this code is allowed to exist. bench_a2_planar.cpp was guarded on NAM_ENABLE_A2_FAST but called create_a2_planar_model unconditionally, so it failed to compile on any target without the planar kernels -- which is every non-AArch64 target, including the x86 Linux runners CI uses. Caught by cross-building for x86_64. It now builds everywhere and, where there is no planar kernel, prints that there is nothing to measure and exits 0. The target is still built on every platform on purpose: a tool that quietly disappears from some configurations is a tool nobody notices has stopped compiling. The activation gate was any AArch64 target. It is now Apple Silicon (__APPLE__ && __aarch64__). The kernels are very likely correct and faster on any AArch64 part, but they have only been built and measured on Apple Silicon, and two of the things they depend 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. Claiming a target nobody has run is not worth the reach. Verified on x86_64 (cross-built on this machine): every target builds, a2_planar.cpp.o contains no symbols at all, and the full test suite passes. Off Apple Silicon the only thing that changes anywhere is that two lines of A2FastConfig::create now live in a named function. --- NAM/wavenet/a2_fast.cpp | 2 +- NAM/wavenet/a2_planar.h | 29 +++++++++++++++++++++++------ tools/bench_a2_planar.cpp | 26 +++++++++++++++++--------- tools/test/test_a2_planar.cpp | 5 +++-- 4 files changed, 44 insertions(+), 18 deletions(-) diff --git a/NAM/wavenet/a2_fast.cpp b/NAM/wavenet/a2_fast.cpp index 337cf12f..7d76b8e7 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -774,7 +774,7 @@ struct A2FastConfig : public ModelConfig std::unique_ptr create(std::vector weights, double sampleRate) override { #if defined(NAM_A2_PLANAR) - // On AArch64, prefer the planar NEON kernels. They are bit-identical to the + // On Apple Silicon, prefer the planar NEON kernels. They are bit-identical to the // reference model below -- same float32 bits out, sample for sample -- so // this is a speed choice and nothing else. A channel count they do not cover // returns nullptr and falls through. diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h index 1b83b34d..d40ca520 100644 --- a/NAM/wavenet/a2_planar.h +++ b/NAM/wavenet/a2_planar.h @@ -1,6 +1,6 @@ #pragma once -// Planar NEON kernels for the A2 fast path (AArch64 only). +// Planar NEON kernels for the A2 fast path (Apple Silicon only). // // These are drop-in replacements for A2FastModel<3> and A2FastModel<8> that // produce **bit-identical** output: not "within a tolerance", not "below the @@ -12,14 +12,31 @@ // per-frame scalar reduction verbatim. Nothing is reassociated, which is what // makes the bit-identity claim hold rather than being a lucky accident. // -// Availability is decided here rather than at the call site: NAM_A2_PLANAR is -// defined only when the A2 fast path is built for AArch64. Everywhere else this -// header declares nothing and a2_fast keeps its existing behaviour. Define -// NAM_DISABLE_A2_PLANAR to opt out on AArch64 too (useful for A/B measurement). +// ----------------------------------------------------------------------------- +// Where this is active, and where it is not +// +// NAM_A2_PLANAR is defined only when the A2 fast path is being built for Apple +// Silicon. On every other target -- x86, and also every *other* AArch64 target +// -- this header declares nothing, a2_planar.cpp compiles to an object with no +// symbols, the call site in a2_fast.cpp is preprocessed away, and the A2 path +// is byte for byte the code that is there today. There is nothing to regress. +// +// The gate is __APPLE__ rather than plain __aarch64__ on purpose. The kernels +// are almost certainly correct and probably faster on any AArch64 part, but +// they have only been built and measured on Apple Silicon, and two things there +// are toolchain-dependent rather than architectural: the tile widths are M2 +// measurements, and bit-identity relies on the compiler contracting a*b+c into +// an FMA in a2_fast's own 3-channel branch, which clang and gcc do by default +// and MSVC at /fp:precise does not. Rather than claim a target nobody has run, +// the gate stops at the one that has been. +// +// NAM_DISABLE_A2_PLANAR opts out on Apple Silicon too, which is what makes an +// A/B measurement against the reference a one-flag change. +// ----------------------------------------------------------------------------- #if defined(NAM_ENABLE_A2_FAST) - #if (defined(__aarch64__) || defined(_M_ARM64)) && !defined(NAM_DISABLE_A2_PLANAR) + #if defined(__APPLE__) && defined(__aarch64__) && !defined(NAM_DISABLE_A2_PLANAR) #define NAM_A2_PLANAR 1 #endif diff --git a/tools/bench_a2_planar.cpp b/tools/bench_a2_planar.cpp index 2ceb636c..998bd7a4 100644 --- a/tools/bench_a2_planar.cpp +++ b/tools/bench_a2_planar.cpp @@ -20,9 +20,18 @@ // measured; --submodel picks it by width, not by position, so a reordering in // the trainer cannot silently change what is being measured. // -// Only compiled when NAM_ENABLE_A2_FAST is defined. +// There is only something to measure where the planar kernels exist, so on any +// other target this builds to a main() that says so and exits. The target is +// still built everywhere, deliberately: a tool that silently vanishes from some +// configurations is a tool nobody notices has stopped compiling. + +#include #if defined(NAM_ENABLE_A2_FAST) + #include "NAM/wavenet/a2_planar.h" // defines NAM_A2_PLANAR where it applies +#endif + +#if defined(NAM_A2_PLANAR) #include #include @@ -31,7 +40,6 @@ #include #include #include - #include #include #include #include @@ -41,7 +49,6 @@ #include "NAM/dsp.h" #include "NAM/wavenet/a2_fast.h" - #include "NAM/wavenet/a2_planar.h" using hr_clock = std::chrono::high_resolution_clock; @@ -326,14 +333,15 @@ int main(int argc, char** argv) return ok ? 0 : 1; } -#else // NAM_ENABLE_A2_FAST - - #include +#else // NAM_A2_PLANAR int main() { - std::cerr << "bench_a2_planar: built without NAM_ENABLE_A2_FAST\n"; - return 2; + // Not an error: there is simply no planar kernel in this build to compare + // against, either because NAM_ENABLE_A2_FAST is off, because the target is + // not Apple Silicon, or because NAM_DISABLE_A2_PLANAR was set. + std::cout << "bench_a2_planar: this build has no planar A2 kernel; nothing to measure.\n"; + return 0; } -#endif // NAM_ENABLE_A2_FAST +#endif // NAM_A2_PLANAR diff --git a/tools/test/test_a2_planar.cpp b/tools/test/test_a2_planar.cpp index 163e77bb..736b1092 100644 --- a/tools/test/test_a2_planar.cpp +++ b/tools/test/test_a2_planar.cpp @@ -6,8 +6,9 @@ // tolerance -- across a spread of block sizes, including ones that exercise the // partial-tile and single-frame tails. // -// Built only where the planar kernels exist (AArch64 with the A2 fast path on). -// Everywhere else the test bodies compile to nothing. +// Built only where the planar kernels exist (Apple Silicon with the A2 fast +// path on). Everywhere else the test bodies compile to nothing, so run_tests +// calls them unconditionally and they cost nothing on other targets. #if defined(NAM_ENABLE_A2_FAST) From 6c2a82143e92802aaa24fd412d3ab5e3ceb60501 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Mon, 10 Aug 2026 13:45:27 +0100 Subject: [PATCH 3/8] Widen the planar gate from Apple Silicon to AArch64 The gate was __APPLE__ && __aarch64__ because Apple Silicon was the only place these kernels had been built and measured. It is now __aarch64__. Bit-identity was the thing worth checking off Apple, since it leans on the compiler contracting a*b+c into an FMA inside a2_fast's own 3-channel branch -- a toolchain behaviour rather than an architectural one. It holds: both submodels bit-identical to a2_fast, max|diff| exactly zero over a full render, on a Cortex-A76 (Raspberry Pi 500, Ubuntu 24.04) under GCC 13, and on Neoverse N2 under GCC 14 and Clang 18. The speed holds too, with a different shape. M2: 2.47x on A2 standard, 2.00x on A2 nano. Cortex-A76: 2.13x and 2.94x. Still __aarch64__ rather than a spelling that also catches MSVC's _M_ARM64. MSVC at /fp:precise does not contract into an FMA, so the reference branch it would be compared against computes something else and bit-identity would not hold. clang-cl on ARM64 defines __aarch64__ and is unaffected. The tile widths remain M2 measurements. They affect speed only, never output, and the Cortex-A76's different profile suggests re-tuning per part would be worth someone's time. --- NAM/wavenet/a2_fast.cpp | 2 +- NAM/wavenet/a2_planar.h | 51 +++++++++++++++++++++++------------ tools/bench_a2_planar.cpp | 2 +- tools/test/test_a2_planar.cpp | 2 +- 4 files changed, 37 insertions(+), 20 deletions(-) diff --git a/NAM/wavenet/a2_fast.cpp b/NAM/wavenet/a2_fast.cpp index 7d76b8e7..337cf12f 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -774,7 +774,7 @@ struct A2FastConfig : public ModelConfig std::unique_ptr create(std::vector weights, double sampleRate) override { #if defined(NAM_A2_PLANAR) - // On Apple Silicon, prefer the planar NEON kernels. They are bit-identical to the + // On AArch64, prefer the planar NEON kernels. They are bit-identical to the // reference model below -- same float32 bits out, sample for sample -- so // this is a speed choice and nothing else. A channel count they do not cover // returns nullptr and falls through. diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h index d40ca520..4cf3ac5e 100644 --- a/NAM/wavenet/a2_planar.h +++ b/NAM/wavenet/a2_planar.h @@ -1,6 +1,6 @@ #pragma once -// Planar NEON kernels for the A2 fast path (Apple Silicon only). +// Planar NEON kernels for the A2 fast path (AArch64). // // These are drop-in replacements for A2FastModel<3> and A2FastModel<8> that // produce **bit-identical** output: not "within a tolerance", not "below the @@ -15,28 +15,45 @@ // ----------------------------------------------------------------------------- // Where this is active, and where it is not // -// NAM_A2_PLANAR is defined only when the A2 fast path is being built for Apple -// Silicon. On every other target -- x86, and also every *other* AArch64 target -// -- this header declares nothing, a2_planar.cpp compiles to an object with no -// symbols, the call site in a2_fast.cpp is preprocessed away, and the A2 path -// is byte for byte the code that is there today. There is nothing to regress. +// NAM_A2_PLANAR is defined only when the A2 fast path is being built for +// AArch64. On every other target -- x86 above all -- this header declares +// nothing, a2_planar.cpp compiles to an object with no symbols, the call site in +// a2_fast.cpp is preprocessed away, and the A2 path is byte for byte the code +// that is there today. There is nothing to regress. // -// The gate is __APPLE__ rather than plain __aarch64__ on purpose. The kernels -// are almost certainly correct and probably faster on any AArch64 part, but -// they have only been built and measured on Apple Silicon, and two things there -// are toolchain-dependent rather than architectural: the tile widths are M2 -// measurements, and bit-identity relies on the compiler contracting a*b+c into -// an FMA in a2_fast's own 3-channel branch, which clang and gcc do by default -// and MSVC at /fp:precise does not. Rather than claim a target nobody has run, -// the gate stops at the one that has been. +// The gate was __APPLE__ && __aarch64__ at first, because Apple Silicon was the +// only place these had been built and measured. It has since been widened to +// AArch64 generally, on evidence rather than optimism: // -// NAM_DISABLE_A2_PLANAR opts out on Apple Silicon too, which is what makes an -// A/B measurement against the reference a one-flag change. +// * Bit-identity holds off Apple. The property it leans on is the compiler +// contracting a*b+c into an FMA inside a2_fast's *own* 3-channel branch, +// which is a toolchain behaviour, not an architectural one. Checked on a +// Cortex-A76 (Raspberry Pi 500, Ubuntu 24.04) under GCC 13, and on Neoverse +// N2 under GCC 14 and Clang 18: both submodels bit-identical to a2_fast, +// max|diff| exactly zero, over a full render. +// +// * The speed holds too, though the shape of the win is not the same. On an +// M2: 2.47x on A2 standard and 2.00x on A2 nano. On a Cortex-A76: 2.13x and +// 2.94x. Faster on both parts, on both submodels. +// +// __aarch64__ specifically, rather than a spelling that would also catch MSVC's +// _M_ARM64. That is deliberate and is the one part of the old gate worth +// keeping: MSVC at /fp:precise does not contract a*b+c into an FMA, so the +// reference branch it would be compared against computes something else, and +// bit-identity -- the whole claim -- would not hold. clang-cl on ARM64 defines +// __aarch64__ and is fine. +// +// The tile widths remain M2 measurements. They affect speed only, never output, +// and the Cortex-A76's rather different profile suggests re-tuning them per part +// would be worth someone's time. +// +// NAM_DISABLE_A2_PLANAR opts out anywhere, which is what makes an A/B +// measurement against the reference a one-flag change. // ----------------------------------------------------------------------------- #if defined(NAM_ENABLE_A2_FAST) - #if defined(__APPLE__) && defined(__aarch64__) && !defined(NAM_DISABLE_A2_PLANAR) + #if defined(__aarch64__) && !defined(NAM_DISABLE_A2_PLANAR) #define NAM_A2_PLANAR 1 #endif diff --git a/tools/bench_a2_planar.cpp b/tools/bench_a2_planar.cpp index 998bd7a4..186e68f5 100644 --- a/tools/bench_a2_planar.cpp +++ b/tools/bench_a2_planar.cpp @@ -339,7 +339,7 @@ int main() { // Not an error: there is simply no planar kernel in this build to compare // against, either because NAM_ENABLE_A2_FAST is off, because the target is - // not Apple Silicon, or because NAM_DISABLE_A2_PLANAR was set. + // not AArch64, or because NAM_DISABLE_A2_PLANAR was set. std::cout << "bench_a2_planar: this build has no planar A2 kernel; nothing to measure.\n"; return 0; } diff --git a/tools/test/test_a2_planar.cpp b/tools/test/test_a2_planar.cpp index 736b1092..9a9eef21 100644 --- a/tools/test/test_a2_planar.cpp +++ b/tools/test/test_a2_planar.cpp @@ -6,7 +6,7 @@ // tolerance -- across a spread of block sizes, including ones that exercise the // partial-tile and single-frame tails. // -// Built only where the planar kernels exist (Apple Silicon with the A2 fast +// Built only where the planar kernels exist (AArch64 with the A2 fast // path on). Everywhere else the test bodies compile to nothing, so run_tests // calls them unconditionally and they cost nothing on other targets. From a2b71b73162fe77ae20ab89eeffe7bd493f5493d Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Tue, 1 Sep 2026 14:11:34 +0100 Subject: [PATCH 4/8] Widen the planar gate to 32-bit ARM, and re-tune the tiles for it The kernels were AArch64-only because that is where they had been built and measured. They now build for ARMv7-A with NEON and VFPv4, gated on __arm__ && __ARM_NEON && __ARM_FEATURE_FMA -- the FMA half being load-bearing, since without it Eigen computes a2_fast's own C=8 path with non-fused vmlaq_f32 and the reference stops being the reference. Measured on a Rockchip RK3288 (quad Cortex-A17), GCC 13.3, clock-pinned: both submodels bit-identical to a2_fast, max|diff| exactly zero, with A2 standard at 78.5% -> 57.8% of one core and A2 nano 12.32% -> 8.87%. Three things this needed beyond a recompile: * ARMv7 has no by-element FMA at all, so vfmaq_laneq_f32 and friends do not exist. LaneWeights hides the difference: lane-addressed on AArch64, broadcast-from-memory (vld1q_dup_f32) on ARMv7, which is the right shape anyway on a machine with 16 Q registers rather than 32. * The tile widths do not transfer. The C=3 ladder peaks at 32 frames on an M2 and at 8 on a Cortex-A17, where 32 runs slower than a2_fast. Tile 8 is the last rung whose accumulators fit in 16 registers, and the measured spill counts turn over exactly there. * The C=8 mixin is a product and then a separate add -- two roundings, because that is what Eigen does -- and on ARMv7 -ffp-contract=fast folds the pair back into one vfma. round_now blocks that locally, at the point a2_fast rounds twice, rather than by turning contraction off globally, which the C=3 branch depends on. AArch64 is unchanged in output and in speed: bit-identical as before, identical fmla and fmul counts in the object, and 2.126x -> 2.146x (A2 standard) and 2.927x -> 2.921x (A2 nano) measured back to back on a Cortex-A76, which is noise. The header records the three caveats the ARMv7 claim carries: bit-identity rests on FPSCR.FZ being set, and both the C=8 parity and the C=3 speedup are GCC claims. Co-Authored-By: Claude Opus 5 --- NAM/wavenet/a2_planar.cpp | 396 +++++++++++++++++++++++++++----------- NAM/wavenet/a2_planar.h | 90 +++++++-- 2 files changed, 365 insertions(+), 121 deletions(-) diff --git a/NAM/wavenet/a2_planar.cpp b/NAM/wavenet/a2_planar.cpp index 57be0f63..840362c0 100644 --- a/NAM/wavenet/a2_planar.cpp +++ b/NAM/wavenet/a2_planar.cpp @@ -73,6 +73,133 @@ namespace a2_fast namespace { +// ----------------------------------------------------------------------------- +// The two instruction sets, reconciled +// +// AArch64 multiply-accumulates against a *lane* of a loaded weight vector -- +// vfmaq_laneq_f32 and friends. ARMv7's VFPv4 has no VFMA-by-scalar encoding at +// all, so none of the by-element forms exist. The replacement is not the same +// thing spelled differently: on a machine with 16 Q registers rather than 32, +// the right move is not to hold the weight vector at all, but to broadcast each +// weight straight out of memory as it is used (vld1q_dup_f32, one instruction). +// +// Both spellings are exact, for three reasons: +// +// 1. A splat is a bit-copy. The multiplicand the FMA sees is the identical +// float32 the lane-addressed form would have seen; no rounding is added. +// 2. vfmaq_f32 is VFMA.F32 -- one rounding, exactly as vfmaq_laneq_f32 is. +// 3. Lane independence is untouched: each lane still runs one frame's chain in +// a2_fast's order, so nothing is reassociated. +// +// LaneWeights is what keeps the two shapes behind one call site, so the kernels +// below are written once. On AArch64 it holds the weight vectors and indexes +// them by lane, compiling to exactly the instructions this file emitted before +// ARMv7 was added -- that path is byte for byte what it was. On ARMv7 it holds +// only the pointer. +// ----------------------------------------------------------------------------- + + #if defined(NAM_A2_PLANAR_A32) +// The kernel's premise is that the accumulators stay in registers across the tap +// loop, which stops being true the moment a tile helper is left out of line: the +// accumulator arrays are passed by reference, so a call boundary spills them. +// GCC inlines these on its own. Clang, measured on this target, does not -- it +// leaves the helpers out of line and spends 164 loads and 39 stores to do 87 +// FMAs, which costs most of the win. Forcing the decision is free where the +// compiler was going to make it anyway. + #define NAM_A2_PLANAR_INLINE inline __attribute__((always_inline)) + #else + #define NAM_A2_PLANAR_INLINE inline + #endif + +/// A run of 4*NV weights, delivered to the multiplier the way the target +/// prefers. Every index is a template parameter, never a runtime value: on +/// AArch64 the lane is encoded in the instruction, and on ARMv7 the offset is +/// folded into the load. +template +struct LaneWeights +{ + static constexpr int kFloats = 4 * NV; + + const float* p; + #if !defined(NAM_A2_PLANAR_A32) + float32x4_t v[NV]; + #endif + + NAM_A2_PLANAR_INLINE explicit LaneWeights(const float* q) + : p(q) + { + #if !defined(NAM_A2_PLANAR_A32) + for (int u = 0; u < NV; u++) + v[u] = vld1q_f32(q + 4 * u); + #endif + } + + /// Weight I broadcast to every lane. + template + NAM_A2_PLANAR_INLINE float32x4_t dup() const + { + static_assert(I >= 0 && I < kFloats, "weight index out of range"); + #if defined(NAM_A2_PLANAR_A32) + return vld1q_dup_f32(p + I); + #else + return vdupq_laneq_f32(v[I / 4], I % 4); + #endif + } + + /// acc + s * w[I], one rounding. + template + NAM_A2_PLANAR_INLINE float32x4_t fma(float32x4_t acc, float32x4_t s) const + { + static_assert(I >= 0 && I < kFloats, "weight index out of range"); + #if defined(NAM_A2_PLANAR_A32) + return vfmaq_f32(acc, s, vld1q_dup_f32(p + I)); + #else + return vfmaq_laneq_f32(acc, s, v[I / 4], I % 4); + #endif + } + + /// s * w[I], one rounding. Seeds a chain where an FMA against zero would + /// otherwise be needed; both are one rounding of w*x, so this is exact either + /// way and just saves the init. + template + NAM_A2_PLANAR_INLINE float32x4_t mul(float32x4_t s) const + { + static_assert(I >= 0 && I < kFloats, "weight index out of range"); + #if defined(NAM_A2_PLANAR_A32) + return vmulq_f32(s, vld1q_dup_f32(p + I)); + #else + return vmulq_laneq_f32(s, v[I / 4], I % 4); + #endif + } +}; + +/// Force `v` to be a rounded value before it is used again. +/// +/// The C=8 layer body computes the mixin as a product and then a separate add -- +/// two roundings, because that is what Eigen does. Under -ffp-contract=fast +/// (GCC's default) nothing stops the compiler folding that pair into one vfma, +/// which is one rounding and different bits. +/// +/// This is measured rather than defensive, and only on ARMv7: an ordering probe +/// against Eigen under arm-linux-gnueabihf-g++ matches 224/224 at +/// -ffp-contract=off and 11/224 at =fast, with the conv stage unaffected either +/// way. AArch64 was checked and does not contract here, so the barrier is not +/// applied there -- those kernels are measured as they stand. +/// +/// Turning contraction off globally would be the wrong fix: a2_fast's C=3 branch +/// *depends* on the compiler contracting a*b+c, which is the whole bit-identity +/// premise at that width. So contraction stays on and is blocked here, locally, +/// at exactly the point a2_fast rounds twice. +/// +/// Compiles to nothing; "w" is the VFP/NEON register constraint on ARM. +NAM_A2_PLANAR_INLINE float32x4_t round_now(float32x4_t v) +{ + #if defined(NAM_A2_PLANAR_A32) + __asm__("" : "+w"(v)); + #endif + return v; +} + // ----------------------------------------------------------------------------- // Weights // @@ -180,7 +307,14 @@ PlanarWeights parse_weights(const std::vector& weights) /// two roundings whatever the compiler decides. inline float mul_rounded(float a, float b) { - return vget_lane_f32(vmul_f32(vdup_n_f32(a), vdup_n_f32(b)), 0); + float p = vget_lane_f32(vmul_f32(vdup_n_f32(a), vdup_n_f32(b)), 0); + #if defined(NAM_A2_PLANAR_A32) + // Routing through a vector register is enough on AArch64, where the pair is + // not contracted anyway. On ARMv7 the same barrier the vector path needs is + // applied here too, for the same reason and with the same cost: none. + __asm__("" : "+w"(p)); + #endif + return p; } /// Receptive field, counted exactly as A2FastModel counts it, so the planar @@ -206,6 +340,17 @@ int planar_prewarm_samples() // has no mirror to maintain, no masking on reads, and every read is contiguous; // it pays instead with an occasional large memmove, amortised over many blocks // by the 2*lookback sizing. +// +// ARMv7 re-ran that sweep and agrees at C=3, emphatically: linear+rewind is +// 1.378x against a2_fast where the eager mirror is 1.322x, the largest single +// gain available after the tile width. On a 32 KB L1D the mirror's per-block +// memcpy across all 23 layers is not merely overhead, it is eviction. +// +// At C=8 the same sweep put the lazy mirror marginally ahead -- 1.312x against +// linear's 1.309x -- and that difference is kept out of this file deliberately. +// It is 0.2%, against a run-to-run spread of 0.5% on the same board, so it is +// below what the measurement can resolve; carrying a second ring implementation +// into shipped code to chase it would be buying complexity with noise. // ----------------------------------------------------------------------------- template struct PlanarRing @@ -261,10 +406,28 @@ struct PlanarRing // a2_fast's 9 scalar FMAs per frame. // ============================================================================= -/// Frames per tile. Twelve accumulator registers at 32 (3 channels x 8 vectors), -/// which is where the sweep peaked: 8/16/32/64 measured 1.39x/1.56x/1.75x/1.46x +/// Frames per tile, swept separately on each architecture because the answer is +/// not close. +/// +/// AArch64: twelve accumulator registers at 32 (3 channels x 8 vectors), which +/// is where the sweep peaked -- 8/16/32/64 measured 1.39x/1.56x/1.75x/1.46x /// against a2_fast. 64 spills. +/// +/// ARMv7: the ladder inverts. 2/4/8/12/16/32 measured +/// 0.888x/1.116x/1.242x/1.086x/1.006x/0.954x, so AArch64's 32 is *slower than +/// a2_fast* here. Counting registers says why. The live set is +/// +/// z accumulators C*T/lanes + t accumulators C*T/lanes + input T/lanes + 1 +/// +/// which at C=3, tile 8, 4 lanes is 6 + 6 + 2 + 1 = 15 of the 16 Q registers +/// this machine has; tile 12 needs 22. The measured spill counts turn over at +/// exactly that step: 0.48 memory operations per FMA at tile 8, 1.13 at 12. +/// Tile 8 is the last rung that fits, and it is the peak. + #if defined(NAM_A2_PLANAR_A32) +constexpr int kNanoTile = 8; + #else constexpr int kNanoTile = 32; + #endif constexpr int kNanoVecs = kNanoTile / 4; /// One layer's weights, padded to four lanes so each group of three is a single @@ -460,40 +623,38 @@ class A2PlanarNano : public DSP /// registers across every tap instead of making a round trip to memory per /// tap per frame. template - inline void tile(const NanoLayer& P, const float* const* h, const int (&tapb)[K], int f0, float* const* d, - float* const* hs) + NAM_A2_PLANAR_INLINE void tile(const NanoLayer& P, const float* const* h, const int (&tapb)[K], int f0, + float* const* d, float* const* hs) { - const float32x4_t cb = vld1q_f32(P.conv_b.data()); + const LaneWeights<1> cb(P.conv_b.data()); float32x4_t a0[NVEC], a1[NVEC], a2[NVEC]; for (int v = 0; v < NVEC; v++) { - a0[v] = vdupq_laneq_f32(cb, 0); - a1[v] = vdupq_laneq_f32(cb, 1); - a2[v] = vdupq_laneq_f32(cb, 2); + a0[v] = cb.dup<0>(); + a1[v] = cb.dup<1>(); + a2[v] = cb.dup<2>(); } const float* cw = P.conv_w.data(); for (int k = 0; k < K; k++) { - const float* wk = cw + static_cast(k) * 12; - const float32x4_t A = vld1q_f32(wk); // w0 w1 w2 w3 - const float32x4_t B = vld1q_f32(wk + 4); // w4 w5 w6 w7 - const float32x4_t Cw = vld1q_f32(wk + 8); // w8 . . . + // Nine weights then three pad, per tap: [out i][in j] at j * 3 + i. + const LaneWeights<3> w(cw + static_cast(k) * 12); const int b = tapb[k] + f0; for (int v = 0; v < NVEC; v++) { const float32x4_t s0 = vld1q_f32(h[0] + b + 4 * v); const float32x4_t s1 = vld1q_f32(h[1] + b + 4 * v); const float32x4_t s2 = vld1q_f32(h[2] + b + 4 * v); - a0[v] = vfmaq_laneq_f32(a0[v], s0, A, 0); - a1[v] = vfmaq_laneq_f32(a1[v], s0, A, 1); - a2[v] = vfmaq_laneq_f32(a2[v], s0, A, 2); - a0[v] = vfmaq_laneq_f32(a0[v], s1, A, 3); - a1[v] = vfmaq_laneq_f32(a1[v], s1, B, 0); - a2[v] = vfmaq_laneq_f32(a2[v], s1, B, 1); - a0[v] = vfmaq_laneq_f32(a0[v], s2, B, 2); - a1[v] = vfmaq_laneq_f32(a1[v], s2, B, 3); - a2[v] = vfmaq_laneq_f32(a2[v], s2, Cw, 0); + a0[v] = w.fma<0>(a0[v], s0); + a1[v] = w.fma<1>(a1[v], s0); + a2[v] = w.fma<2>(a2[v], s0); + a0[v] = w.fma<3>(a0[v], s1); + a1[v] = w.fma<4>(a1[v], s1); + a2[v] = w.fma<5>(a2[v], s1); + a0[v] = w.fma<6>(a0[v], s2); + a1[v] = w.fma<7>(a1[v], s2); + a2[v] = w.fma<8>(a2[v], s2); } } @@ -505,10 +666,11 @@ class A2PlanarNano : public DSP /// reloading it here is cheaper than carrying it through the tap loop, which /// is what decides how wide the tile can usefully get. template - inline void post(const NanoLayer& P, const float* const* h, int last_tap, int f0, float32x4_t (&a0)[NVEC], - float32x4_t (&a1)[NVEC], float32x4_t (&a2)[NVEC], float* const* d, float* const* hs) + NAM_A2_PLANAR_INLINE void post(const NanoLayer& P, const float* const* h, int last_tap, int f0, + float32x4_t (&a0)[NVEC], float32x4_t (&a1)[NVEC], float32x4_t (&a2)[NVEC], + float* const* d, float* const* hs) { - const float32x4_t M = vld1q_f32(P.mixin_w.data()); + const LaneWeights<1> M(P.mixin_w.data()); const float32x4_t zero = vdupq_n_f32(0.0f); const float32x4_t slope = vdupq_n_f32(kLeakySlope); const float* cond = _cond.data(); @@ -516,9 +678,11 @@ class A2PlanarNano : public DSP for (int v = 0; v < NVEC; v++) { const float32x4_t cf = vld1q_f32(cond + f0 + 4 * v); - a0[v] = vfmaq_laneq_f32(a0[v], cf, M, 0); - a1[v] = vfmaq_laneq_f32(a1[v], cf, M, 1); - a2[v] = vfmaq_laneq_f32(a2[v], cf, M, 2); + // Contracted, unlike the C=8 mixin: a2_fast's 3-channel branch writes + // `a += mixin * cond` as one FMA, so one rounding is the correct order. + a0[v] = M.fma<0>(a0[v], cf); + a1[v] = M.fma<1>(a1[v], cf); + a2[v] = M.fma<2>(a2[v], cf); a0[v] = vbslq_f32(vcltq_f32(a0[v], zero), vmulq_f32(a0[v], slope), a0[v]); a1[v] = vbslq_f32(vcltq_f32(a1[v], zero), vmulq_f32(a1[v], slope), a1[v]); a2[v] = vbslq_f32(vcltq_f32(a2[v], zero), vmulq_f32(a2[v], slope), a2[v]); @@ -547,26 +711,25 @@ class A2PlanarNano : public DSP if constexpr (DoL1x1) { - const float32x4_t LA = vld1q_f32(P.l1x1_w.data()); // l0 l1 l2 l3 - const float32x4_t LB = vld1q_f32(P.l1x1_w.data() + 4); // l4 l5 l6 l7 - const float32x4_t LC = vld1q_f32(P.l1x1_w.data() + 8); // l8 . . . - const float32x4_t LBias = vld1q_f32(P.l1x1_b.data()); + // Nine weights then three pad: [out i][in j] at j * 3 + i, as the conv. + const LaneWeights<3> L(P.l1x1_w.data()); + const LaneWeights<1> LBias(P.l1x1_b.data()); for (int v = 0; v < NVEC; v++) { const int o = f0 + 4 * v; - float32x4_t o0 = vdupq_laneq_f32(LBias, 0); - float32x4_t o1 = vdupq_laneq_f32(LBias, 1); - float32x4_t o2 = vdupq_laneq_f32(LBias, 2); - o0 = vfmaq_laneq_f32(o0, a0[v], LA, 0); - o0 = vfmaq_laneq_f32(o0, a1[v], LA, 3); - o0 = vfmaq_laneq_f32(o0, a2[v], LB, 2); - o1 = vfmaq_laneq_f32(o1, a0[v], LA, 1); - o1 = vfmaq_laneq_f32(o1, a1[v], LB, 0); - o1 = vfmaq_laneq_f32(o1, a2[v], LB, 3); - o2 = vfmaq_laneq_f32(o2, a0[v], LA, 2); - o2 = vfmaq_laneq_f32(o2, a1[v], LB, 1); - o2 = vfmaq_laneq_f32(o2, a2[v], LC, 0); + float32x4_t o0 = LBias.dup<0>(); + float32x4_t o1 = LBias.dup<1>(); + float32x4_t o2 = LBias.dup<2>(); + o0 = L.fma<0>(o0, a0[v]); + o0 = L.fma<3>(o0, a1[v]); + o0 = L.fma<6>(o0, a2[v]); + o1 = L.fma<1>(o1, a0[v]); + o1 = L.fma<4>(o1, a1[v]); + o1 = L.fma<7>(o1, a2[v]); + o2 = L.fma<2>(o2, a0[v]); + o2 = L.fma<5>(o2, a1[v]); + o2 = L.fma<8>(o2, a2[v]); vst1q_f32(d[0] + o, vaddq_f32(vld1q_f32(h[0] + last_tap + o), o0)); vst1q_f32(d[1] + o, vaddq_f32(vld1q_f32(h[1] + last_tap + o), o1)); vst1q_f32(d[2] + o, vaddq_f32(vld1q_f32(h[2] + last_tap + o), o2)); @@ -587,21 +750,25 @@ class A2PlanarNano : public DSP const float s0 = h[0][tapb[k] + f]; const float s1 = h[1][tapb[k] + f]; const float s2 = h[2][tapb[k] + f]; - a[0] += wk[0] * s0; - a[1] += wk[1] * s0; - a[2] += wk[2] * s0; - a[0] += wk[3] * s1; - a[1] += wk[4] * s1; - a[2] += wk[5] * s1; - a[0] += wk[6] * s2; - a[1] += wk[7] * s2; - a[2] += wk[8] * s2; + // Spelled as fused rather than written `a += w * s` and left to + // -ffp-contract. Both compile to the same instruction where contraction is + // on, and this way the tail's exactness is a property of the source + // instead of a property of a flag. + a[0] = __builtin_fmaf(wk[0], s0, a[0]); + a[1] = __builtin_fmaf(wk[1], s0, a[1]); + a[2] = __builtin_fmaf(wk[2], s0, a[2]); + a[0] = __builtin_fmaf(wk[3], s1, a[0]); + a[1] = __builtin_fmaf(wk[4], s1, a[1]); + a[2] = __builtin_fmaf(wk[5], s1, a[2]); + a[0] = __builtin_fmaf(wk[6], s2, a[0]); + a[1] = __builtin_fmaf(wk[7], s2, a[1]); + a[2] = __builtin_fmaf(wk[8], s2, a[2]); } const float cf = _cond[f]; for (int c = 0; c < C; c++) { - a[c] += P.mixin_w[c] * cf; + a[c] = __builtin_fmaf(P.mixin_w[c], cf, a[c]); a[c] = (a[c] < 0.0f) ? a[c] * kLeakySlope : a[c]; if constexpr (StoreHead) hs[c][f] = 0.0f + a[c]; @@ -614,9 +781,9 @@ class A2PlanarNano : public DSP for (int c = 0; c < C; c++) { float o = P.l1x1_b[c]; - o += P.l1x1_w[0 + c] * a[0]; - o += P.l1x1_w[3 + c] * a[1]; - o += P.l1x1_w[6 + c] * a[2]; + o = __builtin_fmaf(P.l1x1_w[0 + c], a[0], o); + o = __builtin_fmaf(P.l1x1_w[3 + c], a[1], o); + o = __builtin_fmaf(P.l1x1_w[6 + c], a[2], o); d[c][f] = h[c][tapb[K - 1] + f] + o; } } @@ -647,10 +814,10 @@ class A2PlanarNano : public DSP float32x4_t y = vdupq_n_f32(_w.head_b); for (int k = 0; k < kHeadKernelSize; k++) { - const float32x4_t W = vld1q_f32(hw + static_cast(k) * 4); - y = vfmaq_laneq_f32(y, vld1q_f32(p[0] + hb[k] + f), W, 0); - y = vfmaq_laneq_f32(y, vld1q_f32(p[1] + hb[k] + f), W, 1); - y = vfmaq_laneq_f32(y, vld1q_f32(p[2] + hb[k] + f), W, 2); + const LaneWeights<1> W(hw + static_cast(k) * 4); + y = W.fma<0>(y, vld1q_f32(p[0] + hb[k] + f)); + y = W.fma<1>(y, vld1q_f32(p[1] + hb[k] + f)); + y = W.fma<2>(y, vld1q_f32(p[2] + hb[k] + f)); } vst1q_f32(out + f, vmulq_n_f32(y, scale)); } @@ -660,9 +827,9 @@ class A2PlanarNano : public DSP for (int k = 0; k < kHeadKernelSize; k++) { const float* w = hw + static_cast(k) * 4; - y += w[0] * p[0][hb[k] + f]; - y += w[1] * p[1][hb[k] + f]; - y += w[2] * p[2][hb[k] + f]; + y = __builtin_fmaf(w[0], p[0][hb[k] + f], y); + y = __builtin_fmaf(w[1], p[1][hb[k] + f], y); + y = __builtin_fmaf(w[2], p[2][hb[k] + f], y); } out[f] = y * scale; } @@ -696,6 +863,14 @@ class A2PlanarNano : public DSP /// Frames per conv tile. a2_fast's association needs the running total `z` and /// the current tap's partial `t` live at once, which is 2 x 8 x (tile/4) vector /// registers; the measured curve peaks at 8 and falls off at 16. +/// +/// Both architectures agree here, which is worth saying because at C=3 they do +/// not. The ARMv7 ladder over 2/4/8/12/16/32 is +/// 1.049x/1.126x/1.255x/1.135x/0.899x/0.822x -- the same peak, reached from a +/// very different place: tile 8 at C=8 needs 35 registers on a machine with 16 +/// and already spends 1.96 memory operations per FMA. It does not fit by a +/// factor of two and still wins, so what is being measured past that point is +/// how gracefully the spilling degrades, which no register count predicts. constexpr int kFullTile = 8; constexpr int kFullVecs = kFullTile / 4; @@ -703,7 +878,15 @@ constexpr int kFullVecs = kFullTile / 4; /// so it is latency-bound rather than throughput-bound; running eight chains at /// once costs nothing in registers and nothing in exactness, because each chain /// still covers its own frames in a2_fast's own order. +/// +/// "Costs nothing in registers" is an AArch64 statement. Eight chains is eight +/// live accumulators plus the two weight vectors, which on ARMv7 is ten of +/// sixteen before a single input is loaded, and the sweep there chose one chain. + #if defined(NAM_A2_PLANAR_A32) +constexpr int kFullHeadVecs = 1; + #else constexpr int kFullHeadVecs = 8; + #endif class A2PlanarFull : public DSP { @@ -877,8 +1060,8 @@ class A2PlanarFull : public DSP /// tap's partial -- both live at once, because that separation *is* a2_fast's /// association. z never reaches memory. template - inline void tile(const PlanarLayerWeights& L, const float* lt, const float* const* h, const int (&tapb)[K], int f0, - float* const* d, float* const* hs) + NAM_A2_PLANAR_INLINE void tile(const PlanarLayerWeights& L, const float* lt, const float* const* h, + const int (&tapb)[K], int f0, float* const* d, float* const* hs) { float32x4_t z[C][NVEC]; const float32x4_t zero = vdupq_n_f32(0.0f); @@ -898,10 +1081,7 @@ class A2PlanarFull : public DSP // rounding of w*x, so this is exact either way; it just saves the init. const auto do_j = [&](auto jc) { constexpr int j = decltype(jc)::value; - const float* wj = wk + j * C; // W(0..C-1, j), contiguous - float32x4_t wv[C / 4]; - for (int u = 0; u < C / 4; u++) - wv[u] = vld1q_f32(wj + 4 * u); + const LaneWeights wv(wk + j * C); // W(0..C-1, j), contiguous const float* hp = h[j] + base; for (int v = 0; v < NVEC; v++) { @@ -909,9 +1089,9 @@ class A2PlanarFull : public DSP const auto do_i = [&](auto ic) { constexpr int i = decltype(ic)::value; if constexpr (j == 0) - t[i][v] = vmulq_laneq_f32(s, wv[i / 4], i % 4); + t[i][v] = wv.template mul(s); else - t[i][v] = vfmaq_laneq_f32(t[i][v], s, wv[i / 4], i % 4); + t[i][v] = wv.template fma(t[i][v], s); }; [&](std::integer_sequence) { (do_i(std::integral_constant{}), ...); @@ -934,8 +1114,8 @@ class A2PlanarFull : public DSP /// residual -- in a2_fast's order, with the mixin's multiply and add rounded /// separately as Eigen rounds them. template - inline void post(const PlanarLayerWeights& L, const float* lt, const float* const* h, int last_tap, int f0, - float32x4_t (&z)[C][NVEC], float* const* d, float* const* hs) + NAM_A2_PLANAR_INLINE void post(const PlanarLayerWeights& L, const float* lt, const float* const* h, int last_tap, + int f0, float32x4_t (&z)[C][NVEC], float* const* d, float* const* hs) { const float32x4_t zero = vdupq_n_f32(0.0f); const float32x4_t slope = vdupq_n_f32(kLeakySlope); @@ -952,7 +1132,10 @@ class A2PlanarFull : public DSP for (int v = 0; v < NVEC; v++) { float32x4_t a = vaddq_f32(z[i][v], b); - a = vaddq_f32(a, vmulq_n_f32(cf[v], m)); // product then add, two roundings + // Product then add, two roundings, as Eigen rounds them. round_now is + // what stops the pair being contracted straight back into one vfma; see + // its definition for which target needs it and why. + a = vaddq_f32(a, round_now(vmulq_n_f32(cf[v], m))); z[i][v] = vbslq_f32(vcltq_f32(a, zero), vmulq_f32(a, slope), a); } } @@ -976,19 +1159,18 @@ class A2PlanarFull : public DSP for (int i = 0; i < C; i++) { const float32x4_t bi = vdupq_n_f32(L.l1x1_b[i]); - const float32x4_t la = vld1q_f32(lt + static_cast(i) * C); // L(i, 0..3) - const float32x4_t lb = vld1q_f32(lt + static_cast(i) * C + 4); // L(i, 4..7) + const LaneWeights lw(lt + static_cast(i) * C); // L(i, 0..C-1) for (int v = 0; v < NVEC; v++) { // u_i = sum over j in increasing order, from zero. - float32x4_t u = vmulq_laneq_f32(z[0][v], la, 0); - u = vfmaq_laneq_f32(u, z[1][v], la, 1); - u = vfmaq_laneq_f32(u, z[2][v], la, 2); - u = vfmaq_laneq_f32(u, z[3][v], la, 3); - u = vfmaq_laneq_f32(u, z[4][v], lb, 0); - u = vfmaq_laneq_f32(u, z[5][v], lb, 1); - u = vfmaq_laneq_f32(u, z[6][v], lb, 2); - u = vfmaq_laneq_f32(u, z[7][v], lb, 3); + float32x4_t u = lw.mul<0>(z[0][v]); + u = lw.fma<1>(u, z[1][v]); + u = lw.fma<2>(u, z[2][v]); + u = lw.fma<3>(u, z[3][v]); + u = lw.fma<4>(u, z[4][v]); + u = lw.fma<5>(u, z[5][v]); + u = lw.fma<6>(u, z[6][v]); + u = lw.fma<7>(u, z[7][v]); const float32x4_t prev = vld1q_f32(h[i] + last_tap + f0 + 4 * v); vst1q_f32(d[i] + f0 + 4 * v, vaddq_f32(vaddq_f32(prev, u), bi)); } @@ -1016,7 +1198,7 @@ class A2PlanarFull : public DSP { const float s = h[j][tapb[k] + f]; for (int i = 0; i < C; i++) - t[i] += wk[static_cast(j) * C + i] * s; + t[i] = __builtin_fmaf(wk[static_cast(j) * C + i], s, t[i]); } for (int i = 0; i < C; i++) z[i] += t[i]; @@ -1040,7 +1222,7 @@ class A2PlanarFull : public DSP { float u = 0.0f; for (int j = 0; j < C; j++) - u += L.l1x1_w[static_cast(j) * C + i] * z[j]; + u = __builtin_fmaf(L.l1x1_w[static_cast(j) * C + i], z[j], u); d[i][f] = (h[i][tapb[K - 1] + f] + u) + L.l1x1_b[i]; } } @@ -1073,20 +1255,19 @@ class A2PlanarFull : public DSP y[u] = bias; for (int k = 0; k < kHeadKernelSize; k++) { - const float32x4_t wa = vld1q_f32(hw + static_cast(k) * C); - const float32x4_t wb = vld1q_f32(hw + static_cast(k) * C + 4); + const LaneWeights w(hw + static_cast(k) * C); const int base = hb[k] + f; for (int u = 0; u < kFullHeadVecs; u++) { const int o = base + 4 * u; - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[0] + o), wa, 0); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[1] + o), wa, 1); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[2] + o), wa, 2); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[3] + o), wa, 3); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[4] + o), wb, 0); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[5] + o), wb, 1); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[6] + o), wb, 2); - y[u] = vfmaq_laneq_f32(y[u], vld1q_f32(p[7] + o), wb, 3); + y[u] = w.fma<0>(y[u], vld1q_f32(p[0] + o)); + y[u] = w.fma<1>(y[u], vld1q_f32(p[1] + o)); + y[u] = w.fma<2>(y[u], vld1q_f32(p[2] + o)); + y[u] = w.fma<3>(y[u], vld1q_f32(p[3] + o)); + y[u] = w.fma<4>(y[u], vld1q_f32(p[4] + o)); + y[u] = w.fma<5>(y[u], vld1q_f32(p[5] + o)); + y[u] = w.fma<6>(y[u], vld1q_f32(p[6] + o)); + y[u] = w.fma<7>(y[u], vld1q_f32(p[7] + o)); } } for (int u = 0; u < kFullHeadVecs; u++) @@ -1097,17 +1278,16 @@ class A2PlanarFull : public DSP float32x4_t y = bias; for (int k = 0; k < kHeadKernelSize; k++) { - const float32x4_t wa = vld1q_f32(hw + static_cast(k) * C); - const float32x4_t wb = vld1q_f32(hw + static_cast(k) * C + 4); + const LaneWeights w(hw + static_cast(k) * C); const int o = hb[k] + f; - y = vfmaq_laneq_f32(y, vld1q_f32(p[0] + o), wa, 0); - y = vfmaq_laneq_f32(y, vld1q_f32(p[1] + o), wa, 1); - y = vfmaq_laneq_f32(y, vld1q_f32(p[2] + o), wa, 2); - y = vfmaq_laneq_f32(y, vld1q_f32(p[3] + o), wa, 3); - y = vfmaq_laneq_f32(y, vld1q_f32(p[4] + o), wb, 0); - y = vfmaq_laneq_f32(y, vld1q_f32(p[5] + o), wb, 1); - y = vfmaq_laneq_f32(y, vld1q_f32(p[6] + o), wb, 2); - y = vfmaq_laneq_f32(y, vld1q_f32(p[7] + o), wb, 3); + y = w.fma<0>(y, vld1q_f32(p[0] + o)); + y = w.fma<1>(y, vld1q_f32(p[1] + o)); + y = w.fma<2>(y, vld1q_f32(p[2] + o)); + y = w.fma<3>(y, vld1q_f32(p[3] + o)); + y = w.fma<4>(y, vld1q_f32(p[4] + o)); + y = w.fma<5>(y, vld1q_f32(p[5] + o)); + y = w.fma<6>(y, vld1q_f32(p[6] + o)); + y = w.fma<7>(y, vld1q_f32(p[7] + o)); } vst1q_f32(out + f, vmulq_n_f32(y, scale)); } @@ -1118,7 +1298,7 @@ class A2PlanarFull : public DSP { const float* wk = hw + static_cast(k) * C; for (int b = 0; b < C; b++) - y += wk[b] * p[b][hb[k] + f]; + y = __builtin_fmaf(wk[b], p[b][hb[k] + f], y); } out[f] = y * scale; } diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h index 4cf3ac5e..31cd35bf 100644 --- a/NAM/wavenet/a2_planar.h +++ b/NAM/wavenet/a2_planar.h @@ -1,6 +1,6 @@ #pragma once -// Planar NEON kernels for the A2 fast path (AArch64). +// Planar NEON kernels for the A2 fast path (AArch64 and ARMv7). // // These are drop-in replacements for A2FastModel<3> and A2FastModel<8> that // produce **bit-identical** output: not "within a tolerance", not "below the @@ -15,15 +15,17 @@ // ----------------------------------------------------------------------------- // Where this is active, and where it is not // -// NAM_A2_PLANAR is defined only when the A2 fast path is being built for -// AArch64. On every other target -- x86 above all -- this header declares -// nothing, a2_planar.cpp compiles to an object with no symbols, the call site in -// a2_fast.cpp is preprocessed away, and the A2 path is byte for byte the code -// that is there today. There is nothing to regress. +// NAM_A2_PLANAR is defined only when the A2 fast path is being built for an ARM +// target that has fused multiply-add. On every other target -- x86 above all -- +// this header declares nothing, a2_planar.cpp compiles to an object with no +// symbols, the call site in a2_fast.cpp is preprocessed away, and the A2 path is +// byte for byte the code that is there today. There is nothing to regress. // // The gate was __APPLE__ && __aarch64__ at first, because Apple Silicon was the -// only place these had been built and measured. It has since been widened to -// AArch64 generally, on evidence rather than optimism: +// only place these had been built and measured. It has since been widened +// twice, each time on evidence rather than optimism. +// +// AArch64 generally: // // * Bit-identity holds off Apple. The property it leans on is the compiler // contracting a*b+c into an FMA inside a2_fast's *own* 3-channel branch, @@ -36,6 +38,17 @@ // M2: 2.47x on A2 standard and 2.00x on A2 nano. On a Cortex-A76: 2.13x and // 2.94x. Faster on both parts, on both submodels. // +// ARMv7-A with NEON and VFPv4 (see the caveats below, which are load-bearing): +// +// * Measured on a Rockchip RK3288 (quad Cortex-A17, ASUS Tinker Board), GCC +// 13.3, clock-pinned to 1416 MHz, over the same 523,808-frame render. Both +// submodels bit-identical to a2_fast, max|diff| exactly zero. +// +// * A2 standard 78.5% -> 57.8% of one core and A2 nano 12.32% -> 8.87% at +// 32-frame blocks (1.359x and 1.390x). Smaller than the AArch64 wins, and +// on this part that is the difference between one instance per core with +// nothing left over and one with room for the rest of a signal chain. +// // __aarch64__ specifically, rather than a spelling that would also catch MSVC's // _M_ARM64. That is deliberate and is the one part of the old gate worth // keeping: MSVC at /fp:precise does not contract a*b+c into an FMA, so the @@ -43,9 +56,53 @@ // bit-identity -- the whole claim -- would not hold. clang-cl on ARM64 defines // __aarch64__ and is fine. // -// The tile widths remain M2 measurements. They affect speed only, never output, -// and the Cortex-A76's rather different profile suggests re-tuning them per part -// would be worth someone's time. +// The ARMv7 arm of the gate additionally requires __ARM_NEON and +// __ARM_FEATURE_FMA, and both are load-bearing rather than defensive: +// +// * Without NEON there is no kernel at all. +// * Without FMA, Eigen stops defining EIGEN_VECTORIZE_FMA and selects the +// non-fused vmlaq_f32 for every pmadd, which changes a2_fast's *own* C=8 +// arithmetic. The kernels would then be bit-identical to a reference that +// is no longer there. -mfpu=neon alone reaches that state; -mfpu=neon-vfpv4 +// is what makes the comparison meaningful. +// +// ----------------------------------------------------------------------------- +// Three caveats on ARMv7, stated here rather than buried +// +// 1. **Bit-identity rests on FPSCR.FZ.** AArch32 Advanced SIMD is +// *unconditionally* flush-to-zero for single precision, while the VFP +// scalar code these kernels are compared against honours the FZ bit. The +// two agree bit-for-bit only when the host has FZ set -- which audio hosts +// usually do, and which a host application is under no obligation to do. If +// FZ is clear and a denormal reaches a layer, the NEON kernel and the +// scalar reference will differ. This is the one condition under which the +// claim above is false. AArch64 has no such caveat: there FPCR.FZ applies +// to scalar and vector alike, so both sides move together. +// +// 2. **The C=8 claim is a GCC claim.** Under clang 18.1.3 on this target +// Eigen's gebp_kernel emits non-fused vmla.f32, so a2_fast itself computes +// different bits (reference checksum -17.478711597881365 under GCC against +// -17.478718637490147 under clang) and these kernels land 127.5 dB from it +// rather than at zero. The general honest form of the claim is +// "bit-identical wherever the reference contracts its own FMAs". +// +// 3. **The C=3 speed is a GCC claim**, for an unrelated reason: clang declines +// to inline the tile helpers and spills the accumulators the kernel exists +// to keep in registers, which costs most of the win. Correctness is +// unaffected; only the speed is. +// +// ----------------------------------------------------------------------------- +// Tile widths are per-architecture, and the difference is not a small one +// +// The tile widths affect speed only, never output. They were swept separately on +// each architecture, and the AArch64 values are actively wrong on ARMv7: the C=3 +// ladder peaks at 32 frames on an M2 and at 8 on a Cortex-A17, with tile 32 +// running *slower than a2_fast* there. ARMv7 has 16 Q registers against +// AArch64's 32, and tile 8 is the last rung whose accumulators fit. +// +// Anyone porting these to a third architecture should re-sweep rather than +// inherit. See A32-PATH.md in the NAMBench repository for the full ladder and +// the spill counts underneath it. // // NAM_DISABLE_A2_PLANAR opts out anywhere, which is what makes an A/B // measurement against the reference a one-flag change. @@ -53,8 +110,15 @@ #if defined(NAM_ENABLE_A2_FAST) - #if defined(__aarch64__) && !defined(NAM_DISABLE_A2_PLANAR) - #define NAM_A2_PLANAR 1 + #if !defined(NAM_DISABLE_A2_PLANAR) + #if defined(__aarch64__) + #define NAM_A2_PLANAR 1 + #elif defined(__arm__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) + #define NAM_A2_PLANAR 1 + /// 32-bit ARM: 16 Q registers, no by-element FMA, narrower NEON datapath. + /// Selects the ARMv7 weight delivery and tile widths in a2_planar.cpp. + #define NAM_A2_PLANAR_A32 1 + #endif #endif #if defined(NAM_A2_PLANAR) From 99f8fff20e3064435d7509cd246c95857d4a7b2a Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Wed, 2 Sep 2026 21:48:50 +0100 Subject: [PATCH 5/8] Give ARMv7 its own C=8 conv loop shape The gate widened to 32-bit ARM in the previous commit, and the kernels were bit-identical there from the first build. They were also slow: 68.8% of one core on a Cortex-A17 against 57.4% for the lab kernel they were ported from, which is behind even a plain tile-8 kernel carrying none of the ring switches. The cause is the fold over generic lambdas that unrolls the C=8 conv's input and output channel loops. AArch64 needs that: each index has to be a compile-time value for vfmaq_laneq_f32. ARMv7 has no by-lane FMA at all, so it gains nothing from the fold and pays for it heavily -- with 16 Q registers the accumulators cannot stay resident at any useful tile width, and GCC schedules the spill traffic it cannot avoid far better for a plain loop nest than for a fold it must first decide to inline. Measured over the 523,808-frame render, fold form against loop form: 12.12 G instructions and 8.33 G memory accesses become 9.79 G and 6.34 G. Not stalls -- the fold form had the higher IPC of the two; it simply did a third more memory traffic. So ARMv7 gets the loop nest and AArch64 keeps the fold. The AArch64 object is byte-identical before and after this commit. On a Rockchip RK3288 at a pinned 1416 MHz, 32-frame blocks, both submodels bit-identical to a2_fast over the full render: A2 standard 78.79% -> 55.63% of one core (1.416x) A2 nano 12.29% -> 8.36% of one core (1.470x) Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01GJ6wuGZ9ndKjCDh7Zkjui3 --- NAM/wavenet/a2_planar.cpp | 37 ++++++++++++++++++++++++++++++++ NAM/wavenet/a2_planar.h | 45 ++++++++++++++++++++++++++++----------- 2 files changed, 69 insertions(+), 13 deletions(-) diff --git a/NAM/wavenet/a2_planar.cpp b/NAM/wavenet/a2_planar.cpp index 840362c0..af2a70c8 100644 --- a/NAM/wavenet/a2_planar.cpp +++ b/NAM/wavenet/a2_planar.cpp @@ -1076,6 +1076,42 @@ class A2PlanarFull : public DSP const int base = tapb[k] + f0; float32x4_t t[C][NVEC]; + #if defined(NAM_A2_PLANAR_A32) + // ARMv7: plain loops over j and i, and the weight broadcast from memory + // at the point of use. + // + // This is the same arithmetic in the same order as the AArch64 form + // below, written the way a 16-register machine wants it, and the + // difference it makes is not marginal. Over the reference render, with + // the fold form here: 12.12 G instructions, 8.33 G memory accesses, + // 10.58 G cycles, 68.7% of one core. With this one: 9.79 G, 6.34 G, + // 8.53 G, 55.4%. Nearly a third of the memory traffic was the fold. + // + // The accumulators are the whole reason. 16 Q registers cannot hold + // z[C][NVEC] and t[C][NVEC] at once at any useful tile width, so what + // decides this kernel is not whether it spills -- it must -- but how well + // the compiler schedules the traffic it cannot avoid. GCC does that far + // better for a loop nest it can see the shape of than for a fold of + // generic lambdas it has to decide to inline first. On AArch64, where the + // accumulators fit, the same fold costs nothing and buys the by-lane + // encoding, which is why both forms are here rather than one. + // + // Seeding j == 0 with a multiply instead of an FMA against zero is exact + // either way and saves the zeroing pass, but it is not done here: it + // needs j as a compile-time value, and having j as a loop counter is + // worth more than the pass it would save. + for (int i = 0; i < C; i++) + for (int v = 0; v < NVEC; v++) + t[i][v] = zero; + + for (int j = 0; j < C; j++) + for (int v = 0; v < NVEC; v++) + { + const float32x4_t s = vld1q_f32(h[j] + base + 4 * v); + for (int i = 0; i < C; i++) + t[i][v] = vfmaq_f32(t[i][v], s, vld1q_dup_f32(&wk[j * C + i])); + } + #else // Input channel j is unrolled at compile time so that j == 0 can seed the // partial with a multiply instead of an FMA against zero. Both are one // rounding of w*x, so this is exact either way; it just saves the init. @@ -1101,6 +1137,7 @@ class A2PlanarFull : public DSP [&](std::integer_sequence) { (do_j(std::integral_constant{}), ...); }(std::make_integer_sequence{}); + #endif for (int i = 0; i < C; i++) for (int v = 0; v < NVEC; v++) diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h index 31cd35bf..1d27989f 100644 --- a/NAM/wavenet/a2_planar.h +++ b/NAM/wavenet/a2_planar.h @@ -44,8 +44,8 @@ // 13.3, clock-pinned to 1416 MHz, over the same 523,808-frame render. Both // submodels bit-identical to a2_fast, max|diff| exactly zero. // -// * A2 standard 78.5% -> 57.8% of one core and A2 nano 12.32% -> 8.87% at -// 32-frame blocks (1.359x and 1.390x). Smaller than the AArch64 wins, and +// * A2 standard 78.79% -> 55.63% of one core and A2 nano 12.29% -> 8.36% at +// 32-frame blocks (1.416x and 1.470x). Smaller than the AArch64 wins, and // on this part that is the difference between one instance per core with // nothing left over and one with room for the rest of a signal chain. // @@ -92,17 +92,36 @@ // unaffected; only the speed is. // // ----------------------------------------------------------------------------- -// Tile widths are per-architecture, and the difference is not a small one -// -// The tile widths affect speed only, never output. They were swept separately on -// each architecture, and the AArch64 values are actively wrong on ARMv7: the C=3 -// ladder peaks at 32 frames on an M2 and at 8 on a Cortex-A17, with tile 32 -// running *slower than a2_fast* there. ARMv7 has 16 Q registers against -// AArch64's 32, and tile 8 is the last rung whose accumulators fit. -// -// Anyone porting these to a third architecture should re-sweep rather than -// inherit. See A32-PATH.md in the NAMBench repository for the full ladder and -// the spill counts underneath it. +// What is per-architecture here, and why none of it was inherited +// +// Two things vary by target, and neither changes the output -- only the speed. +// +// **Tile widths**, and the difference is not a small one. They were swept +// separately on each architecture, and the AArch64 values are actively wrong on +// ARMv7: the C=3 ladder peaks at 32 frames on an M2 and at 8 on a Cortex-A17, +// with tile 32 running *slower than a2_fast* there. ARMv7 has 16 Q registers +// against AArch64's 32, and tile 8 is the last rung whose accumulators fit. +// +// **The shape of the C=8 conv loop**, which is the less obvious one and cost +// more to find. The AArch64 form unrolls the input and output channels with a +// fold over generic lambdas, because it needs each index as a compile-time +// value for the by-lane FMA encoding. ARMv7 has no by-lane FMA, so it gains +// nothing from that -- and loses a great deal: with 16 registers the +// accumulators cannot stay resident, and GCC schedules the unavoidable spill +// traffic far better for a plain loop nest than for the fold. The two forms are +// 68.7% and 55.4% of one core on the same part. See the comment at the branch +// itself for the instruction and memory-access counts. +// +// The general lesson, stated because it was learned the expensive way: porting +// these kernels is not a matter of widening the gate and re-sweeping the tiles. +// The first ARMv7 build of this file was bit-identical, passed every +// conformance check, and was still slower than the kernel it was ported from by +// enough to lose most of the win -- and nothing but a measurement on the part +// said so. Re-measure against the reference on the target itself, and do not +// trust a figure carried over from a lab kernel of the same shape. +// +// See A32-PATH.md in the NAMBench repository for the full ladder and the spill +// counts underneath it. // // NAM_DISABLE_A2_PLANAR opts out anywhere, which is what makes an A/B // measurement against the reference a one-flag change. From d3814adb2f6e327b71a8df1e36c5515a772d2db4 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Wed, 2 Sep 2026 23:45:26 +0100 Subject: [PATCH 6/8] Summarise the four configurations at the top of the header The header already justified every per-architecture decision, but only in prose, spread over two sections and a caveat list. Anyone arriving at this file to find out what it actually does on their target had to read all of it to learn that there is one engine here rather than four, and to assemble the speed figures from three different paragraphs. Two small tables up front: what varies per target and channel count -- tile width, weight delivery, and the C=8 conv loop shape -- and what each combination is worth against a2_fast. The speed table also records something the prose did not. The M2 rows are carried from the Apple Silicon campaign these kernels came out of and have not been re-measured since; every other row is a direct measurement of the code as it stands. And the Cortex-A76 rows are at a 32-frame block, which is why they read 2.40x/2.87x against the 2.13x/2.94x quoted further down -- same code, different block size. Both figures were already in the file with nothing to say they were not in conflict. Comment-only: the AArch64 object is byte-identical, and both conformance suites still pass. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01GJ6wuGZ9ndKjCDh7Zkjui3 --- NAM/wavenet/a2_planar.h | 45 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/NAM/wavenet/a2_planar.h b/NAM/wavenet/a2_planar.h index 1d27989f..b662f549 100644 --- a/NAM/wavenet/a2_planar.h +++ b/NAM/wavenet/a2_planar.h @@ -25,6 +25,48 @@ // only place these had been built and measured. It has since been widened // twice, each time on evidence rather than optimism. // +// ----------------------------------------------------------------------------- +// At a glance +// +// The promoted kernel is a2_planar in every case: one engine, four +// configurations. What varies is the tile width, how a weight reaches the +// multiplier, and -- at C=8 -- the shape of the conv loop. +// +// target submodel C tile weights C=8 conv loop +// ------------------------------------------------------------------ +// AArch64 A2 nano 3 32 by-lane -- +// AArch64 A2 standard 8 8 by-lane fold over lambdas +// ARMv7 A2 nano 3 8 broadcast -- +// ARMv7 A2 standard 8 8 broadcast plain loop nest +// +// "by-lane" is vfmaq_laneq_f32, which encodes the lane in the instruction; +// "broadcast" is vld1q_dup_f32 at the point of use, because ARMv7 has no +// by-element FMA at all. Both sections below say why the remaining two columns +// differ, and neither difference was inherited -- each was measured on the part. +// +// The speed, all of it bit-identical to a2_fast over a full render: +// +// part submodel C vs a2_fast conditions +// ------------------------------------------------------------------ +// Apple M2 A2 nano 3 2.00x see note +// Apple M2 A2 standard 8 2.47x see note +// Cortex-A76 (Pi 500) A2 nano 3 2.87x block 32 +// Cortex-A76 (Pi 500) A2 standard 8 2.40x block 32 +// Cortex-A17 (RK3288) A2 nano 3 1.47x block 32, 1416 MHz +// Cortex-A17 (RK3288) A2 standard 8 1.42x block 32, 1416 MHz +// +// The M2 rows are carried from the Apple Silicon campaign these kernels came +// out of, and have not been re-measured since; every other row is a direct +// measurement of the code as it stands. The A76 rows are at a 32-frame block, +// which is why they do not match the 2.13x/2.94x quoted just below -- same +// code, different block size, not a discrepancy. +// +// Note the inversion. On the M2 the wide model wins bigger; on both ARM parts +// the narrow one does. That is the same architectural story as the tile widths +// further down, and the reason none of this travels between targets by +// assumption. +// ----------------------------------------------------------------------------- +// // AArch64 generally: // // * Bit-identity holds off Apple. The property it leans on is the compiler @@ -36,7 +78,8 @@ // // * The speed holds too, though the shape of the win is not the same. On an // M2: 2.47x on A2 standard and 2.00x on A2 nano. On a Cortex-A76: 2.13x and -// 2.94x. Faster on both parts, on both submodels. +// 2.94x. Faster on both parts, on both submodels. (At a 32-frame block the +// A76 reads 2.40x and 2.87x; see the table above.) // // ARMv7-A with NEON and VFPv4 (see the caveats below, which are load-bearing): // From 552c5ab54cb8c2027ba9567c3ae974884e5b68c5 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Fri, 11 Sep 2026 16:19:52 +0100 Subject: [PATCH 7/8] Give the planar models the cached-prewarm path A2FastModel gained upstream Between this branch's original base and current main, upstream added a cached-prewarm optimisation to A2FastModel: Reset() restores a cached steady-state instead of re-running the legacy silence-processing prewarm, verified by a new test asserting zero allocations on that path (A2FastModel::cached prewarm). The planar models derive from DSP directly, not from A2FastModel, so they simply inherited the old, allocating base-class prewarm() -- and that new test caught it after the rebase: 12 allocations where it expected zero. Same mechanism, mirrored for the planar ring layout: each ring's steady-state prewarm makes every column of every channel plane equal to the last one written, so caching that one column per plane is enough to rebuild the whole ring later without reprocessing silence. Added PrewarmFromCache / CacheStateAsPrewarmed to A2PlanarNano and A2PlanarFull, and an override of prewarm() that uses the cache once one exists, matching A2FastModel's own prewarm()/PrewarmFromCache()/CacheStateAsPrewarmed() shape. Full test suite passes after this, cross-built for AArch64, ARMv7 and x86_64. Co-Authored-By: Claude Sonnet 5 --- NAM/wavenet/a2_planar.cpp | 86 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/NAM/wavenet/a2_planar.cpp b/NAM/wavenet/a2_planar.cpp index af2a70c8..31aefc9d 100644 --- a/NAM/wavenet/a2_planar.cpp +++ b/NAM/wavenet/a2_planar.cpp @@ -4,6 +4,7 @@ #if defined(NAM_A2_PLANAR) + #include #include #include #include @@ -394,6 +395,31 @@ struct PlanarRing int tap(int lookback_frames, int n) const { return wpos - n - lookback_frames; } }; +// ----------------------------------------------------------------------------- +// Prewarm-state cache, planar equivalent of A2FastModel's. A ring holding +// steady-state silence has every column equal to the last one written, so +// caching that one column per channel plane is enough to rebuild the whole +// ring later without reprocessing silence. +// ----------------------------------------------------------------------------- +template +void CacheRingState(const PlanarRing& ring, Cache& cache) +{ + const int last_column = ring.wpos - 1; + for (int c = 0; c < C; c++) + cache[c] = ring.plane(c)[last_column]; +} + +template +void RestoreRingFromCache(PlanarRing& ring, const Cache& cache) +{ + for (int c = 0; c < C; c++) + { + float* p = ring.plane(c); + std::fill(p, p + ring.cap, cache[c]); + } + ring.wpos = ring.lookback; +} + // ============================================================================= // Channels == 3 (A2 nano) // @@ -490,6 +516,17 @@ class A2PlanarNano : public DSP int GetPrewarmSamples() override { return _prewarm_samples; } + void prewarm() override + { + if (_has_cached_prewarm_state) + { + PrewarmFromCache(); + return; + } + DSP::prewarm(); + CacheStateAsPrewarmed(); + } + void process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames) override { if (num_frames > GetMaxBufferSize()) @@ -835,6 +872,21 @@ class A2PlanarNano : public DSP } } + void PrewarmFromCache() + { + for (int li = 0; li < kNumLayers; li++) + RestoreRingFromCache(_rings[li], _cached_layer_state[li]); + RestoreRingFromCache(_head_ring, _cached_head_state); + } + + void CacheStateAsPrewarmed() + { + for (int li = 0; li < kNumLayers; li++) + CacheRingState(_rings[li], _cached_layer_state[li]); + CacheRingState(_head_ring, _cached_head_state); + _has_cached_prewarm_state = true; + } + PlanarWeights _w; int _prewarm_samples = 0; @@ -844,6 +896,10 @@ class A2PlanarNano : public DSP std::array _rings; Ring _head_ring; + std::array, kNumLayers> _cached_layer_state{}; + std::array _cached_head_state{}; + bool _has_cached_prewarm_state = false; + std::vector _layer_in; std::vector _head_sum; std::vector _cond; @@ -921,6 +977,17 @@ class A2PlanarFull : public DSP int GetPrewarmSamples() override { return _prewarm_samples; } + void prewarm() override + { + if (_has_cached_prewarm_state) + { + PrewarmFromCache(); + return; + } + DSP::prewarm(); + CacheStateAsPrewarmed(); + } + void process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames) override { if (num_frames > GetMaxBufferSize()) @@ -1341,6 +1408,21 @@ class A2PlanarFull : public DSP } } + void PrewarmFromCache() + { + for (int li = 0; li < kNumLayers; li++) + RestoreRingFromCache(_rings[li], _cached_layer_state[li]); + RestoreRingFromCache(_head_ring, _cached_head_state); + } + + void CacheStateAsPrewarmed() + { + for (int li = 0; li < kNumLayers; li++) + CacheRingState(_rings[li], _cached_layer_state[li]); + CacheRingState(_head_ring, _cached_head_state); + _has_cached_prewarm_state = true; + } + PlanarWeights _w; int _prewarm_samples = 0; @@ -1350,6 +1432,10 @@ class A2PlanarFull : public DSP std::array _rings; Ring _head_ring; + std::array, kNumLayers> _cached_layer_state{}; + std::array _cached_head_state{}; + bool _has_cached_prewarm_state = false; + std::vector _layer_in; std::vector _cond; std::vector _head_out; From 44412fad6ad135d51218785e15dc7418c29a2124 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Sat, 12 Sep 2026 11:34:20 +0100 Subject: [PATCH 8/8] Build the A2 kernels optimised for run_tests, so the parity test can hold test_a2_planar asserts memcmp equality against A2FastModel, but run_tests compiles every one of its sources at -O0 (upstream does this so the allocation tracking behaves). A2FastModel's 3-channel branch is a plain `a * b + c` chain, and its bit-identity premise is that the compiler contracts that into an FMA. GCC only contracts in its optimisers, so at -O0 the reference computes something the planar kernels -- intrinsics, so unconditionally fused -- cannot match: the test failed on the first sample of the first block size on every GCC target. Clang contracts during codegen, which is why the same test passed on Apple Silicon and hid this. Every shipping build is optimised, so -O0 was the one configuration in which the reference is not the code the claim is about. The kernels now build as an object library at -O3 (/O2 on MSVC) that only run_tests links, leaving every other test at the -O0 the allocation tracking wants. The allocation assertions that reach these kernels all require zero allocations, which optimisation cannot introduce. Verified: the full suite passes on an Apple M2 (clang), a Cortex-A76 and a Cortex-A17 (GCC 13). Before this, the latter two failed at channels=3, block=1. Co-Authored-By: Claude Opus 5 --- tools/CMakeLists.txt | 41 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index e430aac3..17875d2f 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -100,7 +100,46 @@ else() ) endif() -add_executable(run_tests run_tests.cpp test/allocation_tracking.cpp ${NAM_SOURCES}) +# The A2 kernels are built optimised even for run_tests, which is otherwise -O0. +# +# test_a2_planar asserts memcmp equality against A2FastModel, and A2FastModel's +# 3-channel branch is a plain `a * b + c` chain whose bit-identity premise is +# that the compiler contracts it into an FMA. GCC only contracts in its +# optimisers, so at -O0 the reference computes something the planar kernels -- +# which issue vfma unconditionally, being intrinsics -- cannot match, and the +# test fails on the very first sample. (Clang contracts during codegen, so it +# passes there either way; this is a GCC-only failure.) Every shipping build is +# optimised, so -O0 is the one configuration in which the reference is not the +# code the claim is about. +# +# Only these two translation units are raised, and only for this target, so +# every other test keeps the -O0 the allocation tracking was set up for. The +# allocation assertions that do reach these kernels all require *zero* +# allocations, and optimisation cannot introduce any. +set(A2_KERNEL_SOURCES "") +set(RUN_TESTS_NAM_SOURCES "") +foreach(_nam_src ${NAM_SOURCES}) + if(_nam_src MATCHES "/NAM/wavenet/a2_(fast|planar)\\.cpp$") + list(APPEND A2_KERNEL_SOURCES "${_nam_src}") + else() + list(APPEND RUN_TESTS_NAM_SOURCES "${_nam_src}") + endif() +endforeach() + +add_library(run_tests_a2_kernels OBJECT ${A2_KERNEL_SOURCES}) +target_compile_features(run_tests_a2_kernels PUBLIC cxx_std_20) +if(MSVC) + target_compile_options(run_tests_a2_kernels PRIVATE /O2) +else() + target_compile_options(run_tests_a2_kernels PRIVATE -O3) +endif() +# Keep assertions live here too, as they are in the rest of run_tests. +target_compile_options(run_tests_a2_kernels PRIVATE + $<$,$,$>:-UNDEBUG> +) + +add_executable(run_tests run_tests.cpp test/allocation_tracking.cpp ${RUN_TESTS_NAM_SOURCES} + $) # Compile run_tests without optimizations to ensure allocation tracking works correctly # Also ensure assertions are enabled (NDEBUG is not defined) so tests actually run set_target_properties(run_tests PROPERTIES COMPILE_OPTIONS "-O0")