From 688ef872c788fa00b59894d467bc5854e19025d5 Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Sat, 8 Aug 2026 20:01:25 +0100 Subject: [PATCH 1/3] 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 ee72ab60..17f27dc2 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 @@ -698,11 +699,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); } }; @@ -909,6 +914,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 7bc1b0b7..a4c2ecb8 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 492fb676..81ff653d 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -62,6 +62,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 5699bf78..632cb429 100644 --- a/tools/run_tests.cpp +++ b/tools/run_tests.cpp @@ -36,6 +36,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() { @@ -368,6 +369,12 @@ int main() test_a2_fast::test_prewarm_matches_generic_standard(); test_a2_fast::test_process_realtime_safe_nano(); test_a2_fast::test_process_realtime_safe_standard(); + + // 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 a4de938b4dfa5d3f7e2375067e64f95fe3b92e7f Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Sat, 8 Aug 2026 20:36:22 +0100 Subject: [PATCH 2/3] 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 17f27dc2..5a95e244 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -700,7 +700,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 ca349f6cfe07200aea9e6767cd29025d7152b33a Mon Sep 17 00:00:00 2001 From: Rik Hemsley Date: Mon, 10 Aug 2026 13:45:27 +0100 Subject: [PATCH 3/3] 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 5a95e244..17f27dc2 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -700,7 +700,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.