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..4cf3ac5e --- /dev/null +++ b/NAM/wavenet/a2_planar.h @@ -0,0 +1,87 @@ +#pragma once + +// 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 +// 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. +// +// ----------------------------------------------------------------------------- +// 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. +// +// 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: +// +// * 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(__aarch64__) && !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..186e68f5 --- /dev/null +++ b/tools/bench_a2_planar.cpp @@ -0,0 +1,347 @@ +// 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. +// +// 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 + #include + #include + #include + #include + #include + #include + #include + #include + #include + + #include "json.hpp" + + #include "NAM/dsp.h" + #include "NAM/wavenet/a2_fast.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_A2_PLANAR + +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 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; +} + +#endif // NAM_A2_PLANAR 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..9a9eef21 --- /dev/null +++ b/tools/test/test_a2_planar.cpp @@ -0,0 +1,231 @@ +// 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, so run_tests +// calls them unconditionally and they cost nothing on other targets. + +#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