diff --git a/.gitignore b/.gitignore index 34ee36ee..3cf9a49d 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,8 @@ docs/_build/ *.DS_Store example_audio/output.wav + +# Local build and benchmark artifacts +build/ +build_nofused/ +benchmark_reports/tmp_models/ diff --git a/CMakeLists.txt b/CMakeLists.txt index d21f36a2..fef2b615 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -60,6 +60,16 @@ if(NAM_ENABLE_A2_FAST) add_compile_definitions(NAM_ENABLE_A2_FAST) endif() +# Build the fused NEON WaveNet engine. When ON, standard-shape WaveNet models +# (e.g. the A1 standard/lite/feather family and A2 standard) are routed to a +# fused, register-tiled NEON implementation on AArch64 (Apple Silicon etc.). +# On other architectures the flag is harmless: the shape detector returns +# false and all models use the generic path. +option(NAM_ENABLE_FUSED "Build the fused NEON WaveNet engine" ON) +if(NAM_ENABLE_FUSED) + add_compile_definitions(NAM_ENABLE_FUSED) +endif() + add_subdirectory(tools) #file(MAKE_DIRECTORY build/tools) diff --git a/NAM/wavenet/fused.cpp b/NAM/wavenet/fused.cpp new file mode 100644 index 00000000..75494801 --- /dev/null +++ b/NAM/wavenet/fused.cpp @@ -0,0 +1,1078 @@ +#if defined(NAM_ENABLE_FUSED) + + #include "fused.h" + + #include + #include + #include + #include + #include + #include + #include + #include + #include + #include + + #include "../activations.h" + #include "../dsp.h" + + #if defined(__aarch64__) || defined(_M_ARM64) + #define NAM_FUSED_KERNELS 1 + #include + #endif + +namespace nam +{ +namespace wavenet +{ +namespace fused +{ + +namespace +{ + +// ============================================================================= +// Shape spec: everything the engine needs, parsed once from the .nam JSON. +// ============================================================================= + +struct LayerSpec +{ + int kernel_size = 0; + int dilation = 0; + activations::ActivationConfig activation{}; +}; + +struct ArraySpec +{ + int input_size = 0; + int channels = 0; + int head_size = 0; + int head_kernel_size = 1; + bool head_bias = false; + bool layer1x1_active = true; + std::vector layers; +}; + +struct ShapeSpec +{ + std::vector arrays; +}; + +constexpr int kMaxChannels = 32; +// Cap on conv/head kernel sizes (bounds the per-block tap pointer table). +constexpr int kMaxKernelSize = 64; + +bool film_inactive(const nlohmann::json& layer, const char* key) +{ + auto it = layer.find(key); + if (it == layer.end() || it->is_null()) + return true; + if (it->is_boolean()) + return !it->get(); + if (it->is_object()) + return !it->value("active", false); + return false; +} + +bool gating_all_none(const nlohmann::json& la, size_t num_layers) +{ + auto gm_it = la.find("gating_mode"); + if (gm_it != la.end() && !gm_it->is_null()) + { + if (gm_it->is_string()) + return gm_it->get() == "none"; + if (!gm_it->is_array() || gm_it->size() != num_layers) + return false; + for (const auto& e : *gm_it) + { + if (!e.is_string() || e.get() != "none") + return false; + } + return true; + } + auto g_it = la.find("gated"); + if (g_it != la.end() && !g_it->is_null()) + return g_it->is_boolean() && !g_it->get(); + return true; // absent defaults to NONE +} + +/// Strict parser. Returns false (never throws) if the config doesn't match a +/// supported shape. On success fills *spec when non-null. +bool parse_spec(const nlohmann::json& config, ShapeSpec* spec) noexcept +{ + try + { + // Mono in, no condition DSP, no post-stack head. + if (config.value("in_channels", 1) != 1) + return false; + if (auto it = config.find("condition_dsp"); it != config.end() && !it->is_null()) + return false; + if (auto it = config.find("head"); it != config.end() && !it->is_null()) + return false; + if (auto it = config.find("head_scale"); it == config.end() || !it->is_number()) + return false; + + auto layers_it = config.find("layers"); + if (layers_it == config.end() || !layers_it->is_array() || layers_it->empty()) + return false; + + ShapeSpec local; + int prev_channels = -1; + int prev_head_size = -1; + for (const auto& la : *layers_it) + { + ArraySpec as; + + as.input_size = la.value("input_size", 0); + if (local.arrays.empty()) + { + if (as.input_size != 1) + return false; + } + else if (as.input_size != prev_channels) + return false; + + if (la.value("condition_size", 0) != 1) + return false; + + as.channels = la.value("channels", 0); + const int bottleneck = la.value("bottleneck", as.channels); + if (as.channels <= 0 || bottleneck != as.channels) + return false; + if (as.channels % 4 != 0 || as.channels > kMaxChannels) + return false; + if (!local.arrays.empty() && as.channels != prev_head_size) + return false; + + if (la.value("groups_input", 1) != 1 || la.value("groups_input_mixin", 1) != 1) + return false; + + // layer1x1: default active with groups=1 + if (auto it = la.find("layer1x1"); it != la.end() && !it->is_null()) + { + if (!it->is_object()) + return false; + as.layer1x1_active = it->value("active", false); + if (as.layer1x1_active && it->value("groups", 1) != 1) + return false; + // Generic path requires bottleneck == channels when layer1x1 inactive; + // that already holds here. + } + + // head1x1 must be inactive + if (auto it = la.find("head1x1"); it != la.end() && it->is_object() && it->value("active", false)) + return false; + + // No FiLM anywhere + for (const char* key : {"conv_pre_film", "conv_post_film", "input_mixin_pre_film", "input_mixin_post_film", + "activation_pre_film", "activation_post_film", "layer1x1_post_film", "head1x1_post_film"}) + { + if (!film_inactive(la, key)) + return false; + } + + // Not slimmable + if (auto it = la.find("slimmable"); it != la.end() && !it->is_null()) + return false; + + // Dilations & kernel sizes + auto dil_it = la.find("dilations"); + if (dil_it == la.end() || !dil_it->is_array() || dil_it->empty()) + return false; + const size_t num_layers = dil_it->size(); + + std::vector kernel_sizes; + const bool has_ks = la.find("kernel_size") != la.end(); + const bool has_kss = la.find("kernel_sizes") != la.end(); + if (has_ks == has_kss) // exactly one must be present + return false; + if (has_ks) + { + if (!la["kernel_size"].is_number_integer()) + return false; + kernel_sizes.assign(num_layers, la["kernel_size"].get()); + } + else + { + if (!la["kernel_sizes"].is_array() || la["kernel_sizes"].size() != num_layers) + return false; + for (const auto& k : la["kernel_sizes"]) + { + if (!k.is_number_integer()) + return false; + kernel_sizes.push_back(k.get()); + } + } + + // Activations (single or per-layer). Any config that + // ActivationConfig::from_json accepts is fine: the engine falls back to + // the generic Activation object for types without NEON kernels. + std::vector act_configs; + const auto& act_json = la.at("activation"); + if (act_json.is_array()) + { + if (act_json.size() != num_layers) + return false; + for (const auto& aj : act_json) + act_configs.push_back(activations::ActivationConfig::from_json(aj)); + } + else + { + act_configs.assign(num_layers, activations::ActivationConfig::from_json(act_json)); + } + + if (!gating_all_none(la, num_layers)) + return false; + + // Head rechannel: nested "head" object or legacy head_size/head_bias. + if (auto it = la.find("head"); it != la.end() && !it->is_null()) + { + if (!it->is_object()) + return false; + as.head_size = it->value("out_channels", 0); + as.head_kernel_size = it->value("kernel_size", 0); + as.head_bias = it->value("bias", false); + } + else if (la.find("head_size") != la.end()) + { + as.head_size = la["head_size"].get(); + as.head_kernel_size = 1; + as.head_bias = la.at("head_bias").get(); + } + else + return false; + if (as.head_size <= 0 || as.head_kernel_size < 1 || as.head_kernel_size > kMaxKernelSize) + return false; + + for (size_t i = 0; i < num_layers; i++) + { + LayerSpec ls; + ls.kernel_size = kernel_sizes[i]; + ls.dilation = (*dil_it)[i].get(); + ls.activation = act_configs[i]; + if (ls.kernel_size < 1 || ls.kernel_size > kMaxKernelSize || ls.dilation < 1) + return false; + as.layers.push_back(std::move(ls)); + } + + prev_channels = as.channels; + prev_head_size = as.head_size; + local.arrays.push_back(std::move(as)); + } + + if (spec) + *spec = std::move(local); + return true; + } + catch (const std::exception&) + { + return false; + } +} + +} // namespace + + #if defined(NAM_FUSED_KERNELS) + +namespace +{ + +// ============================================================================= +// NEON helpers +// ============================================================================= + +/// Vectorized version of activations::fast_tanh (same rational approximation). +/// Uses vdivq_f32: on Apple M-series the FP divider is a separate unit, so the +/// division overlaps with the FMA work and measures faster than a +/// reciprocal-refinement sequence (which competes with the FMAs) — and it's +/// exact, matching the scalar implementation up to FMA contraction. +inline float32x4_t vfast_tanh_f32(float32x4_t x) +{ + const float32x4_t a = vdupq_n_f32(2.45550750702956f); + const float32x4_t b = vdupq_n_f32(0.893229853513558f); + const float32x4_t c = vdupq_n_f32(0.821226666969744f); + const float32x4_t d = vdupq_n_f32(2.44506634652299f); + const float32x4_t e = vdupq_n_f32(0.814642734961073f); + const float32x4_t ax = vabsq_f32(x); + const float32x4_t x2 = vmulq_f32(x, x); + const float32x4_t num = vmulq_f32(x, vaddq_f32(vfmaq_f32(a, a, ax), vmulq_f32(vfmaq_f32(b, c, ax), x2))); + const float32x4_t den = vfmaq_f32(d, vaddq_f32(d, x2), vabsq_f32(vfmaq_f32(x, e, vmulq_f32(x, ax)))); + return vdivq_f32(num, den); +} + +// ============================================================================= +// Activation dispatch +// ============================================================================= + +enum class ActKind +{ + FastTanh, + ReLU, + LeakyReLU, + Hardtanh, + Softsign, + Other // exact Tanh, Sigmoid, PReLU, LUT, ... -> Activation::apply() +}; + +struct ResolvedActivation +{ + ActKind kind = ActKind::Other; + float leaky_slope = 0.01f; + activations::Activation::Ptr ptr; // always valid; used for ActKind::Other +}; + +/// Resolve the activation exactly the way the generic Layer would (through the +/// registry, so enable_fast_tanh()/LUT substitutions are honored), then pick a +/// NEON kernel only when the resolved object is a known implementation. +ResolvedActivation resolve_activation(const activations::ActivationConfig& config) +{ + ResolvedActivation out; + out.ptr = activations::Activation::get_activation(config); + if (out.ptr == nullptr) + throw std::runtime_error("FusedWaveNet: unsupported activation"); + + using activations::ActivationType; + switch (config.type) + { + case ActivationType::Tanh: + case ActivationType::Fasttanh: + // "Tanh" resolves to the fast-tanh singleton when enable_fast_tanh() is + // active, or to a LUT if one was installed. Only claim the NEON kernel + // when the resolved object really is the fast-tanh implementation. + if (dynamic_cast(out.ptr.get()) != nullptr) + out.kind = ActKind::FastTanh; + break; + case ActivationType::ReLU: out.kind = ActKind::ReLU; break; + case ActivationType::LeakyReLU: + out.kind = ActKind::LeakyReLU; + out.leaky_slope = config.negative_slope.value_or(0.01f); + break; + case ActivationType::Hardtanh: out.kind = ActKind::Hardtanh; break; + case ActivationType::Softsign: out.kind = ActKind::Softsign; break; + default: break; + } + return out; +} + +/// Apply activation to a contiguous buffer. n is always a multiple of 4 +/// (channels are multiples of 4). +void apply_activation(const ResolvedActivation& act, float* p, int n) +{ + assert(n % 4 == 0); + switch (act.kind) + { + case ActKind::FastTanh: + for (int i = 0; i < n; i += 4) + vst1q_f32(p + i, vfast_tanh_f32(vld1q_f32(p + i))); + break; + case ActKind::ReLU: + { + const float32x4_t zero = vdupq_n_f32(0.0f); + for (int i = 0; i < n; i += 4) + vst1q_f32(p + i, vmaxq_f32(vld1q_f32(p + i), zero)); + break; + } + case ActKind::LeakyReLU: + { + const float32x4_t zero = vdupq_n_f32(0.0f); + const float32x4_t slope = vdupq_n_f32(act.leaky_slope); + for (int i = 0; i < n; i += 4) + { + const float32x4_t x = vld1q_f32(p + i); + vst1q_f32(p + i, vbslq_f32(vcltq_f32(x, zero), vmulq_f32(x, slope), x)); + } + break; + } + case ActKind::Hardtanh: + { + const float32x4_t lo = vdupq_n_f32(-1.0f); + const float32x4_t hi = vdupq_n_f32(1.0f); + for (int i = 0; i < n; i += 4) + vst1q_f32(p + i, vminq_f32(vmaxq_f32(vld1q_f32(p + i), lo), hi)); + break; + } + case ActKind::Softsign: + { + const float32x4_t one = vdupq_n_f32(1.0f); + for (int i = 0; i < n; i += 4) + { + const float32x4_t x = vld1q_f32(p + i); + vst1q_f32(p + i, vdivq_f32(x, vaddq_f32(one, vabsq_f32(x)))); + } + break; + } + case ActKind::Other: act.ptr->apply(p, n); break; + } +} + +// ============================================================================= +// Ring buffer: power-of-2 capacity with a mirrored tail so any read of up to +// max_buffer_size frames starting inside [0, pow2) is one contiguous span. +// Same scheme as a2_fast ring mode 1 (constant per-block work, no rewind +// spikes on the audio thread). +// ============================================================================= + +int next_pow2(int v) +{ + int p = 1; + while (p < v) + p <<= 1; + return p; +} + +struct Ring +{ + std::vector data; + int channels = 0; + int pow2 = 0; + int mask = 0; + int wpos = 0; + int mbs = 0; + + void reset(int channels_, int max_lookback, int max_buffer) + { + channels = channels_; + mbs = max_buffer; + pow2 = next_pow2(max_lookback + max_buffer); + mask = pow2 - 1; + data.assign(static_cast(channels) * (pow2 + max_buffer), 0.0f); + wpos = max_lookback & mask; + } + + void write(const float* src, int n) + { + float* h = data.data(); + const int first = std::min(n, pow2 - wpos); + std::memcpy(h + static_cast(wpos) * channels, src, static_cast(first) * channels * sizeof(float)); + if (first < n) + { + std::memcpy(h, src + static_cast(first) * channels, + static_cast(n - first) * channels * sizeof(float)); + } + wpos = (wpos + n) & mask; + } + + /// Refresh the mirrored tail so that reads of up to `cols` frames past the + /// end of the ring see cols [0, cols). Called by tap() only for the blocks + /// where a read actually wraps, which keeps the common case copy-free. + void mirror(int cols) + { + float* h = data.data(); + std::memcpy(h + static_cast(pow2) * channels, h, static_cast(cols) * channels * sizeof(float)); + } + + /// Pointer to the first frame of an n-frame read looking back + /// `lookback_frames` from the start of the block just written. Ensures the + /// mirrored tail covers the read when it wraps past the end of the ring. + const float* tap(int lookback_frames, int n) + { + const int base = (wpos - n - lookback_frames) & mask; + const int overflow = base + n - pow2; + if (overflow > 0) + mirror(overflow); + return data.data() + static_cast(base) * channels; + } +}; + +// ============================================================================= +// Register-tiled NEON kernels, templated on Q = channels / 4. +// +// Layouts (all column-major, channels fastest): +// conv weights: per tap k, (C out x C in): w[k*C*C + j*C + i] +// 1x1 weights: (C out x C in): w[j*C + i] +// buffers: (C x frames): buf[f*C + c] +// ============================================================================= + +/// One tile of the dilated conv: acc = bias; for each tap, acc += W_k * x_k; +/// then acc += mixin * cond. T frames per tile, accumulators live in registers +/// across all taps. +template +inline void conv_lane(float32x4_t (&acc)[Q][T], const float32x4_t (&iv)[T], const float* wc) +{ + constexpr int C = 4 * Q; + float32x4_t w[Q]; + for (int q = 0; q < Q; q++) + w[q] = vld1q_f32(wc + LANE * C + 4 * q); + for (int t = 0; t < T; t++) + for (int q = 0; q < Q; q++) + acc[q][t] = vfmaq_laneq_f32(acc[q][t], w[q], iv[t], LANE); +} + +template +inline void conv_tile(const float* const* tap_src, int num_taps, const float* W, const float* bias, + const float* mixin, const float* cond, float* z, int f0) +{ + constexpr int C = 4 * Q; + float32x4_t acc[Q][T]; + for (int q = 0; q < Q; q++) + { + const float32x4_t b = vld1q_f32(bias + 4 * q); + for (int t = 0; t < T; t++) + acc[q][t] = b; + } + for (int k = 0; k < num_taps; k++) + { + const float* src = tap_src[k] + static_cast(f0) * C; + const float* wk = W + static_cast(k) * C * C; + for (int j4 = 0; j4 < Q; j4++) + { + float32x4_t iv[T]; + for (int t = 0; t < T; t++) + iv[t] = vld1q_f32(src + t * C + j4 * 4); + const float* wc = wk + static_cast(j4 * 4) * C; + conv_lane(acc, iv, wc); + conv_lane(acc, iv, wc); + conv_lane(acc, iv, wc); + conv_lane(acc, iv, wc); + } + } + for (int t = 0; t < T; t++) + { + const float32x4_t cf = vdupq_n_f32(cond[f0 + t]); + float* zc = z + static_cast(f0 + t) * C; + for (int q = 0; q < Q; q++) + vst1q_f32(zc + 4 * q, vfmaq_f32(acc[q][t], vld1q_f32(mixin + 4 * q), cf)); + } +} + +template +void conv_block(const float* const* tap_src, int num_taps, const float* W, const float* bias, const float* mixin, + const float* cond, float* z, int num_frames) +{ + constexpr int T = (Q <= 4) ? 4 : 2; + int f = 0; + for (; f + T <= num_frames; f += T) + conv_tile(tap_src, num_taps, W, bias, mixin, cond, z, f); + for (; f < num_frames; f++) + conv_tile(tap_src, num_taps, W, bias, mixin, cond, z, f); +} + +/// Post-activation tail: head_sum += a; if layer1x1 active, +/// lin += L * a + l1x1_bias (in-place residual for the next layer). +template +inline void tail_lane(float32x4_t (&out)[Q][T], const float32x4_t (&a)[Q][T], const float* lc, int j4) +{ + constexpr int C = 4 * Q; + float32x4_t w[Q]; + for (int q = 0; q < Q; q++) + w[q] = vld1q_f32(lc + LANE * C + 4 * q); + for (int t = 0; t < T; t++) + for (int q = 0; q < Q; q++) + out[q][t] = vfmaq_laneq_f32(out[q][t], w[q], a[j4][t], LANE); +} + +template +inline void tail_tile(const float* z, const float* L, const float* lb, bool has_l1x1, float* head_sum, float* lin, + int f0) +{ + constexpr int C = 4 * Q; + float32x4_t a[Q][T]; + for (int t = 0; t < T; t++) + { + const float* zc = z + static_cast(f0 + t) * C; + for (int q = 0; q < Q; q++) + a[q][t] = vld1q_f32(zc + 4 * q); + } + for (int t = 0; t < T; t++) + { + float* hc = head_sum + static_cast(f0 + t) * C; + for (int q = 0; q < Q; q++) + vst1q_f32(hc + 4 * q, vaddq_f32(vld1q_f32(hc + 4 * q), a[q][t])); + } + if (!has_l1x1) + return; + float32x4_t out[Q][T]; + for (int q = 0; q < Q; q++) + { + const float32x4_t b = vld1q_f32(lb + 4 * q); + for (int t = 0; t < T; t++) + out[q][t] = b; + } + for (int j4 = 0; j4 < Q; j4++) + { + const float* lc = L + static_cast(j4 * 4) * C; + tail_lane(out, a, lc, j4); + tail_lane(out, a, lc, j4); + tail_lane(out, a, lc, j4); + tail_lane(out, a, lc, j4); + } + for (int t = 0; t < T; t++) + { + float* lc = lin + static_cast(f0 + t) * C; + for (int q = 0; q < Q; q++) + vst1q_f32(lc + 4 * q, vaddq_f32(vld1q_f32(lc + 4 * q), out[q][t])); + } +} + +template +void tail_block(const float* z, const float* L, const float* lb, bool has_l1x1, float* head_sum, float* lin, + int num_frames) +{ + constexpr int T = (Q <= 4) ? 2 : 1; + int f = 0; + for (; f + T <= num_frames; f += T) + tail_tile(z, L, lb, has_l1x1, head_sum, lin, f); + for (; f < num_frames; f++) + tail_tile(z, L, lb, has_l1x1, head_sum, lin, f); +} + +/// Rechannel (1x1, no bias): out(C x N) = W(C x Cin) * x(Cin x N). +/// Cin is a runtime value (1 for the first array, previous channels after). +template +void rechannel_block(const float* W, const float* x, int c_in, float* out, int num_frames) +{ + constexpr int C = 4 * Q; + for (int f = 0; f < num_frames; f++) + { + const float* xc = x + static_cast(f) * c_in; + float32x4_t acc[Q]; + for (int q = 0; q < Q; q++) + acc[q] = vdupq_n_f32(0.0f); + for (int j = 0; j < c_in; j++) + { + const float32x4_t s = vdupq_n_f32(xc[j]); + const float* wc = W + static_cast(j) * C; + for (int q = 0; q < Q; q++) + acc[q] = vfmaq_f32(acc[q], vld1q_f32(wc + 4 * q), s); + } + float* oc = out + static_cast(f) * C; + for (int q = 0; q < Q; q++) + vst1q_f32(oc + 4 * q, acc[q]); + } +} + +/// Head rechannel conv: out(H x N) = sum_k W_k(H x C) * x_k + bias. +/// Weights row-major per tap: w[k*H*C + h*C + j]. H is small (often 8 or 1). +void head_conv_block(const float* const* tap_src, int num_taps, const float* W, const float* bias, int H, int C, + float* out, int num_frames) +{ + const int Cq = C / 4; + for (int f = 0; f < num_frames; f++) + { + float* oc = out + static_cast(f) * H; + for (int h = 0; h < H; h++) + { + float32x4_t acc = vdupq_n_f32(0.0f); + for (int k = 0; k < num_taps; k++) + { + const float* src = tap_src[k] + static_cast(f) * C; + const float* w = W + (static_cast(k) * H + h) * C; + for (int q = 0; q < Cq; q++) + acc = vfmaq_f32(acc, vld1q_f32(w + 4 * q), vld1q_f32(src + 4 * q)); + } + float y = vaddvq_f32(acc); + if (bias != nullptr) + y += bias[h]; + oc[h] = y; + } + } +} + +// Per-channel-count dispatch table. +struct KernelFns +{ + void (*conv)(const float* const*, int, const float*, const float*, const float*, const float*, float*, int) = + nullptr; + void (*tail)(const float*, const float*, const float*, bool, float*, float*, int) = nullptr; + void (*rechannel)(const float*, const float*, int, float*, int) = nullptr; +}; + +template +KernelFns make_kernel_fns() +{ + KernelFns fns; + fns.conv = &conv_block; + fns.tail = &tail_block; + fns.rechannel = &rechannel_block; + return fns; +} + +KernelFns kernel_fns_for_channels(int channels) +{ + switch (channels / 4) + { + case 1: return make_kernel_fns<1>(); + case 2: return make_kernel_fns<2>(); + case 3: return make_kernel_fns<3>(); + case 4: return make_kernel_fns<4>(); + case 5: return make_kernel_fns<5>(); + case 6: return make_kernel_fns<6>(); + case 7: return make_kernel_fns<7>(); + case 8: return make_kernel_fns<8>(); + default: throw std::runtime_error("FusedWaveNet: unsupported channel count " + std::to_string(channels)); + } +} + +// ============================================================================= +// FusedWaveNet +// ============================================================================= + +class FusedWaveNet : public DSP +{ +public: + FusedWaveNet(const ShapeSpec& spec, std::vector weights, double expected_sample_rate) + : DSP(/*in_channels=*/1, /*out_channels=*/spec.arrays.back().head_size, expected_sample_rate) + { + for (const auto& as : spec.arrays) + { + Array arr; + arr.c_in = as.input_size; + arr.channels = as.channels; + arr.head_size = as.head_size; + arr.head_kernel_size = as.head_kernel_size; + arr.head_bias_active = as.head_bias; + arr.fns = kernel_fns_for_channels(as.channels); + for (const auto& ls : as.layers) + { + Layer layer; + layer.kernel_size = ls.kernel_size; + layer.dilation = ls.dilation; + layer.max_lookback = (ls.kernel_size - 1) * ls.dilation; + layer.has_l1x1 = as.layer1x1_active; + layer.act = resolve_activation(ls.activation); + arr.layers.push_back(std::move(layer)); + } + _arrays.push_back(std::move(arr)); + } + + // The residual output of the very last layer feeds nothing (the generic + // path computes and then never reads it), so skip its layer1x1 GEMM. Its + // weights are still consumed by _load_weights. + if (!_arrays.empty() && !_arrays.back().layers.empty()) + _arrays.back().layers.back().skip_l1x1_output = true; + + _load_weights(weights); + + int prewarm = 1; + for (const auto& arr : _arrays) + { + for (const auto& layer : arr.layers) + prewarm += layer.max_lookback; + prewarm += arr.head_kernel_size - 1; + } + _prewarm_samples = prewarm; + } + + ~FusedWaveNet() override = default; + + void process(NAM_SAMPLE** input, NAM_SAMPLE** output, int num_frames) override + { + if (num_frames > GetMaxBufferSize()) + SetMaxBufferSize(num_frames); + const int N = num_frames; + + // Condition buffer (float copy of the input). + float* cond = _cond.data(); + const NAM_SAMPLE* in0 = input[0]; + for (int f = 0; f < N; f++) + cond[f] = static_cast(in0[f]); + + float* lin = _lin_a.data(); + float* alt = _lin_b.data(); + const float* prev_head_out = nullptr; + int prev_channels = 0; + + const float* tap_src[kMaxTaps]; + + for (auto& arr : _arrays) + { + const int C = arr.channels; + + // Rechannel into the layer input buffer. + if (prev_head_out == nullptr) + { + arr.fns.rechannel(arr.rechannel_w.data(), cond, 1, lin, N); + } + else + { + arr.fns.rechannel(arr.rechannel_w.data(), lin, prev_channels, alt, N); + std::swap(lin, alt); + } + + // Head accumulator: zero for the first array, copy of the previous + // array's head output after (its head_size == this array's channels). + if (prev_head_out == nullptr) + std::memset(_head_sum.data(), 0, static_cast(C) * N * sizeof(float)); + else + std::memcpy(_head_sum.data(), prev_head_out, static_cast(C) * N * sizeof(float)); + + for (auto& layer : arr.layers) + { + const int K = layer.kernel_size; + if (layer.max_lookback > 0) + { + layer.ring.write(lin, N); + for (int k = 0; k < K; k++) + tap_src[k] = layer.ring.tap((K - 1 - k) * layer.dilation, N); + } + else + { + tap_src[0] = lin; // K == 1: reads only the current block + } + arr.fns.conv(tap_src, K, layer.conv_w.data(), layer.conv_b.data(), layer.mixin_w.data(), cond, _z.data(), N); + apply_activation(layer.act, _z.data(), C * N); + arr.fns.tail(_z.data(), layer.l1x1_w.data(), layer.l1x1_b.data(), + layer.has_l1x1 && !layer.skip_l1x1_output, _head_sum.data(), lin, N); + } + + // Head rechannel conv over the accumulated head signal. + const int Kh = arr.head_kernel_size; + if (Kh > 1) + { + arr.head_ring.write(_head_sum.data(), N); + for (int k = 0; k < Kh; k++) + tap_src[k] = arr.head_ring.tap(Kh - 1 - k, N); + } + else + { + tap_src[0] = _head_sum.data(); + } + head_conv_block(tap_src, Kh, arr.head_w.data(), arr.head_bias_active ? arr.head_b.data() : nullptr, + arr.head_size, C, arr.head_out.data(), N); + + prev_head_out = arr.head_out.data(); + prev_channels = C; + } + + // Output: head_scale * last head output. + const int out_channels = NumOutputChannels(); + const float scale = _head_scale; + const float* ho = _arrays.back().head_out.data(); + if (out_channels == 1) + { + NAM_SAMPLE* out0 = output[0]; + for (int f = 0; f < N; f++) + out0[f] = static_cast(scale * ho[f]); + } + else + { + for (int ch = 0; ch < out_channels; ch++) + { + NAM_SAMPLE* oc = output[ch]; + for (int f = 0; f < N; f++) + oc[f] = static_cast(scale * ho[static_cast(f) * out_channels + ch]); + } + } + } + +protected: + void SetMaxBufferSize(int maxBufferSize) override + { + DSP::SetMaxBufferSize(maxBufferSize); + int max_channels = 0; + for (auto& arr : _arrays) + { + max_channels = std::max(max_channels, arr.channels); + for (auto& layer : arr.layers) + { + if (layer.max_lookback > 0) + layer.ring.reset(arr.channels, layer.max_lookback, maxBufferSize); + } + if (arr.head_kernel_size > 1) + arr.head_ring.reset(arr.channels, arr.head_kernel_size - 1, maxBufferSize); + arr.head_out.assign(static_cast(arr.head_size) * maxBufferSize, 0.0f); + } + const size_t buf = static_cast(max_channels) * maxBufferSize; + _cond.assign(static_cast(maxBufferSize), 0.0f); + _lin_a.assign(buf, 0.0f); + _lin_b.assign(buf, 0.0f); + _z.assign(buf, 0.0f); + _head_sum.assign(buf, 0.0f); + } + + int GetPrewarmSamples() override { return _prewarm_samples; } + +private: + static constexpr int kMaxTaps = kMaxKernelSize; + + struct Layer + { + int kernel_size = 0; + int dilation = 0; + int max_lookback = 0; + bool has_l1x1 = true; + // True for the final layer of the final array: its residual output is + // never consumed, so the layer1x1 GEMM is skipped (weights still loaded). + bool skip_l1x1_output = false; + ResolvedActivation act; + // Dilated conv (C -> C), column-major per tap + bias. + std::vector conv_w; + std::vector conv_b; + // Input mixin (1 -> C), no bias. + std::vector mixin_w; + // layer1x1 (C -> C) column-major + bias (empty when !has_l1x1). + std::vector l1x1_w; + std::vector l1x1_b; + Ring ring; + }; + + struct Array + { + int c_in = 0; + int channels = 0; + int head_size = 0; + int head_kernel_size = 1; + bool head_bias_active = false; + KernelFns fns; + // Rechannel (c_in -> C), column-major, no bias. + std::vector rechannel_w; + std::vector layers; + // Head rechannel (C -> head_size), row-major per tap + optional bias. + std::vector head_w; + std::vector head_b; + Ring head_ring; + std::vector head_out; + }; + + /// Consumes weights in exactly the order the generic WaveNet does: + /// per array: rechannel; per layer: conv (+bias), mixin, layer1x1 (+bias); + /// head rechannel (+bias). Trailing value: head_scale. + void _load_weights(std::vector& weights) + { + auto it = weights.begin(); + const auto end = weights.end(); + auto take = [&]() -> float { + if (it == end) + throw std::runtime_error("FusedWaveNet: weight stream exhausted"); + return *it++; + }; + + for (auto& arr : _arrays) + { + const int C = arr.channels; + // Rechannel Conv1x1 (c_in -> C, no bias): for i in C: for j in c_in. + arr.rechannel_w.assign(static_cast(C) * arr.c_in, 0.0f); + for (int i = 0; i < C; i++) + for (int j = 0; j < arr.c_in; j++) + arr.rechannel_w[static_cast(j) * C + i] = take(); + + for (auto& layer : arr.layers) + { + const int K = layer.kernel_size; + // Conv1D (C -> C, K taps, bias): for i: for j: for k; then bias. + layer.conv_w.assign(static_cast(K) * C * C, 0.0f); + layer.conv_b.assign(static_cast(C), 0.0f); + for (int i = 0; i < C; i++) + for (int j = 0; j < C; j++) + for (int k = 0; k < K; k++) + layer.conv_w[(static_cast(k) * C + j) * C + i] = take(); + for (int i = 0; i < C; i++) + layer.conv_b[i] = take(); + // Input mixin Conv1x1 (1 -> C, no bias). + layer.mixin_w.assign(static_cast(C), 0.0f); + for (int i = 0; i < C; i++) + layer.mixin_w[i] = take(); + // layer1x1 (C -> C, bias) when active. + if (layer.has_l1x1) + { + layer.l1x1_w.assign(static_cast(C) * C, 0.0f); + layer.l1x1_b.assign(static_cast(C), 0.0f); + for (int i = 0; i < C; i++) + for (int j = 0; j < C; j++) + layer.l1x1_w[static_cast(j) * C + i] = take(); + for (int i = 0; i < C; i++) + layer.l1x1_b[i] = take(); + } + else + { + // Kernels never read these when has_l1x1 is false, but keep them + // allocated so .data() stays valid. + layer.l1x1_w.assign(static_cast(C) * C, 0.0f); + layer.l1x1_b.assign(static_cast(C), 0.0f); + } + } + + // Head rechannel Conv1D (C -> head_size, Kh taps): for i: for j: for k. + const int H = arr.head_size; + const int Kh = arr.head_kernel_size; + arr.head_w.assign(static_cast(Kh) * H * C, 0.0f); + for (int i = 0; i < H; i++) + for (int j = 0; j < C; j++) + for (int k = 0; k < Kh; k++) + arr.head_w[(static_cast(k) * H + i) * C + j] = take(); + arr.head_b.assign(static_cast(H), 0.0f); + if (arr.head_bias_active) + for (int i = 0; i < H; i++) + arr.head_b[i] = take(); + } + + _head_scale = take(); + + if (it != end) + { + throw std::runtime_error("FusedWaveNet: weight stream has " + std::to_string(std::distance(it, end)) + + " unconsumed values"); + } + } + + std::vector _arrays; + float _head_scale = 1.0f; + int _prewarm_samples = 0; + + // Working buffers (sized in SetMaxBufferSize; no allocation in process()). + std::vector _cond; + std::vector _lin_a; + std::vector _lin_b; + std::vector _z; + std::vector _head_sum; +}; + +} // namespace + + #endif // NAM_FUSED_KERNELS + +// ============================================================================= +// Public API +// ============================================================================= + +namespace +{ + +struct FusedConfig : public ModelConfig +{ + ShapeSpec spec; + + std::unique_ptr create(std::vector weights, double sampleRate) override + { + #if defined(NAM_FUSED_KERNELS) + return std::make_unique(spec, std::move(weights), sampleRate); + #else + (void)weights; + (void)sampleRate; + throw std::runtime_error("FusedConfig::create called on a build without NEON kernels"); + #endif + } +}; + +} // namespace + +bool available() +{ + #if defined(NAM_FUSED_KERNELS) + return true; + #else + return false; + #endif +} + +bool is_fused_shape(const nlohmann::json& config) +{ + if (!available()) + return false; + return parse_spec(config, nullptr); +} + +std::unique_ptr create_fused_config(const nlohmann::json& config, double sampleRate) +{ + (void)sampleRate; + auto out = std::make_unique(); + if (!available() || !parse_spec(config, &out->spec)) + throw std::runtime_error("create_fused_config: config does not match a fused-engine shape"); + return out; +} + +} // namespace fused +} // namespace wavenet +} // namespace nam + +#endif // NAM_ENABLE_FUSED diff --git a/NAM/wavenet/fused.h b/NAM/wavenet/fused.h new file mode 100644 index 00000000..45b84932 --- /dev/null +++ b/NAM/wavenet/fused.h @@ -0,0 +1,67 @@ +#pragma once + +// Fused WaveNet fast path for AArch64 (Apple Silicon and other ARMv8+ CPUs). +// +// Covers the common "standard shape" WaveNet family — including the classic +// A1 standard/lite/feather models (channels 16/8/4, kernel 3) and the A2 +// standard model (channels 8) — with hand-written NEON kernels: +// +// - The dilated conv, input mixin, bias, activation, head accumulation, +// layer1x1 and residual add are fused into three passes over an +// L1-resident block buffer (vs ~8 buffer passes + 4 Eigen GEMM calls per +// layer in the generic path). +// - GEMMs are register-tiled with vfmaq_laneq_f32 and compile-time channel +// counts; measured at ~90% of fp32 FMA peak on Apple M-series for the +// 16-channel A1 layer (Eigen's dynamic GEMM reaches ~40% and allocates). +// - Conv history lives in power-of-2 ring buffers with a mirrored tail so +// every tap read is one contiguous span (no rewind memmove spikes). +// - Activations use vectorized NEON implementations where a known +// activation is in effect (fast tanh, ReLU, LeakyReLU, Hardtanh, +// Softsign); anything else (exact Tanh, Sigmoid, PReLU, LUTs) falls back +// to the exact same Activation object the generic path would call, so +// semantics never diverge. +// +// When NAM_ENABLE_FUSED is defined at build time, wavenet::create_config +// consults is_fused_shape() on every incoming WaveNet config and, on match, +// instantiates the fused engine instead of the generic WaveNet. The detector +// only matches on AArch64 builds; other platforms always use the generic +// path. +// +// Shape requirements (checked strictly by is_fused_shape): +// - in_channels == 1, no condition DSP, no post-stack head +// - each layer array: condition_size == 1, bottleneck == channels, +// channels a multiple of 4 (max 32), all groups == 1, gating "none", +// no FiLM, head1x1 inactive, layer1x1 groups == 1 +// - arbitrary kernel sizes, dilations, activations, head kernel size/bias + +#if defined(NAM_ENABLE_FUSED) + + #include + + #include "../model_config.h" + #include "json.hpp" + +namespace nam +{ +namespace wavenet +{ +namespace fused +{ + +/// \brief Whether this build has the fused NEON kernels (AArch64 only). +bool available(); + +/// \brief Strict detector: returns true iff the config matches a shape the +/// fused engine supports *and* the kernels are available in this build. +/// \param config The "config" sub-object from a .nam WaveNet entry. +bool is_fused_shape(const nlohmann::json& config); + +/// \brief Build a ModelConfig that instantiates the fused engine. +/// \pre is_fused_shape(config) returned true. +std::unique_ptr create_fused_config(const nlohmann::json& config, double sampleRate); + +} // namespace fused +} // namespace wavenet +} // namespace nam + +#endif // NAM_ENABLE_FUSED diff --git a/NAM/wavenet/model.cpp b/NAM/wavenet/model.cpp index 7fd9ed84..43dda259 100644 --- a/NAM/wavenet/model.cpp +++ b/NAM/wavenet/model.cpp @@ -15,6 +15,9 @@ #if defined(NAM_ENABLE_A2_FAST) #include "a2_fast.h" #endif +#if defined(NAM_ENABLE_FUSED) + #include "fused.h" +#endif // detail::Head (WaveNet post-stack head) ===================================== @@ -1236,6 +1239,14 @@ std::unique_ptr nam::wavenet::create_config(const nlohmann::js if (config_is_slimmable_wavenet(config)) return nam::slimmable_wavenet::create_config(config, sampleRate); +#if defined(NAM_ENABLE_FUSED) + // Fused NEON engine (AArch64). Checked before the A2 fast path: it covers + // shapes with channels % 4 == 0 (including A2 standard), while A2 nano + // (channels == 3) falls through to the A2 fast path below. + if (nam::wavenet::fused::is_fused_shape(config)) + return nam::wavenet::fused::create_fused_config(config, sampleRate); +#endif + #if defined(NAM_ENABLE_A2_FAST) if (int a2_channels = 0; nam::wavenet::a2_fast::is_a2_shape(config, &a2_channels)) return nam::wavenet::a2_fast::create_a2_fast_config(config, sampleRate); diff --git a/benchmark_reports/benchmark_report_20260712_192034.txt b/benchmark_reports/benchmark_report_20260712_192034.txt new file mode 100644 index 00000000..c460c5cf --- /dev/null +++ b/benchmark_reports/benchmark_report_20260712_192034.txt @@ -0,0 +1,99 @@ +================================================================ +NeuralAmpModelerCore Benchmark Report +================================================================ + +Date: 2026-07-12 18:20:34 UTC +Machine: Apple M2 +Arch: arm64 +OS: macOS 27.0 +CPU cores: 8 +RAM: 16.0 GB + +Build: Release, A2 fast ON +Git: 9c7b185 (v0.5.3) + +================================================================ +1. benchmodel — example models (10 runs) +================================================================ +Model Median Min Max Mean RTF +-------------------------------------------------------------------------------- +lstm.nam 8.193ms 7.992ms 8.713ms 8.258ms 244.1x +my_model.nam 38.725ms 38.534ms 40.635ms 39.138ms 51.6x +slimmable_container.nam 28.457ms 28.089ms 28.876ms 28.435ms 70.2x +slimmable_wavenet.nam 19.040ms 18.760ms 19.417ms 19.055ms 105.0x +wavenet.nam 6.057ms 5.908ms 6.396ms 6.128ms 330.1x +wavenet_a1_standard.nam 39.060ms 38.506ms 39.318ms 38.996ms 51.2x +wavenet_a2_max.nam 79.765ms 77.769ms 87.159ms 80.769ms 25.0x +wavenet_condition_dsp.nam 13.724ms 13.433ms 14.049ms 13.721ms 145.7x + +================================================================ +2. bench_a2_fast — A2 fast vs generic +================================================================ + +== /Users/rik/code.local/nam/NeuralAmpModelerPlugin/.claude/worktrees/neural-amp-apple-optimization-c3adc0/NeuralAmpModelerCore/benchmark_reports/tmp_models/a2_nano.nam (A2 nano, Channels=3) == + audio=2000ms @ 48000Hz, block=64, iters=10 (15000 blocks timed each) + per-block audio deadline: 1333.3 us + min p50 p99 p99.9 max mean + fast (us): 5.88 6.33 9.79 35.04 94.08 6.48 + generic (us): 42.21 43.83 69.92 166.83 285.38 45.64 + RTF (p50): fast=210.54x generic=30.42x speedup=6.92x + +== /Users/rik/code.local/nam/NeuralAmpModelerPlugin/.claude/worktrees/neural-amp-apple-optimization-c3adc0/NeuralAmpModelerCore/benchmark_reports/tmp_models/a2_standard.nam (A2 standard, Channels=8) == + audio=2000ms @ 48000Hz, block=64, iters=10 (15000 blocks timed each) + per-block audio deadline: 1333.3 us + min p50 p99 p99.9 max mean + fast (us): 44.96 45.33 71.75 164.04 376.46 47.24 + generic (us): 51.50 52.54 84.46 199.34 327.08 55.00 + RTF (p50): fast=29.41x generic=25.38x speedup=1.16x + +================================================================ +3. benchmodel_bufsize — buffer size sweep +================================================================ + +--- wavenet_a1_standard.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 40578.0 49.2x +32 40104.5 49.8x +64 39258.8 50.9x +128 38384.8 52.1x +256 37782.0 52.9x +512 54677.8 36.5x +1024 39633.0 50.4x + +--- wavenet.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 9860.0 202.8x +32 6589.2 303.5x +64 5164.3 387.2x +128 15677.1 127.5x +256 4095.3 488.3x +512 4020.0 497.5x +1024 3833.4 521.7x + +--- lstm.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 8514.8 234.8x +32 8621.8 231.9x +64 8319.2 240.4x +128 8570.4 233.3x +256 8271.4 241.7x +512 8622.4 231.9x +1024 8394.7 238.2x + +--- a2_standard.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 43189.8 46.3x +32 40768.2 49.0x +64 40225.9 49.7x +128 39783.5 50.2x +256 39106.1 51.1x +512 38534.6 51.9x +1024 38618.5 51.7x + +================================================================ +Report complete. +================================================================ diff --git a/benchmark_reports/benchmark_report_20260712_apple_m2_baseline.txt b/benchmark_reports/benchmark_report_20260712_apple_m2_baseline.txt new file mode 100644 index 00000000..240f68e9 --- /dev/null +++ b/benchmark_reports/benchmark_report_20260712_apple_m2_baseline.txt @@ -0,0 +1,140 @@ +================================================================ +NeuralAmpModelerCore Benchmark Report +================================================================ + +Date: 2026-07-12 17:24 UTC +Machine: Apple M2 +Arch: arm64 +OS: macOS 27.0 (build 26A5378j) +CPU cores: 8 +RAM: 16.0 GB + +Build: + Directory: build/ + Type: Release (-Ofast) + Compiler: Apple clang version 21.0.0 (clang-2100.3.20.102) + CMake: cmake version 4.3.4 + A2 fast: ON (NAM_ENABLE_A2_FAST) + +Git: + NeuralAmpModelerCore: 9c7b185 (v0.5.3) + +Benchmark methodology: + benchmodel: 2s audio @ 48kHz, buffer=64, fast tanh enabled, 10 runs + benchmodel_bufsize: 2s audio, 5 iterations per buffer size, fast tanh enabled + bench_a2_fast: 2s audio, buffer=64, 10 iterations, per-block timing (us) + +Build commands: + cd NeuralAmpModelerCore/build + cmake .. -DCMAKE_BUILD_TYPE=Release + cmake --build . --target benchmodel benchmodel_bufsize bench_a2_fast -j8 + +================================================================ +1. benchmodel — example models (median wall time for 2s audio) +================================================================ + +Model Median Min Max Mean RTF +-------------------------------------------------------------------------------- +lstm.nam 8.017ms 7.658ms 8.111ms 7.964ms 249.4x +my_model.nam 83.582ms 82.760ms 86.642ms 84.118ms 23.9x +slimmable_container.nam 29.819ms 27.618ms 33.222ms 29.555ms 67.0x +slimmable_wavenet.nam 20.491ms 20.024ms 21.453ms 20.584ms 97.6x +wavenet.nam 6.387ms 6.226ms 6.709ms 6.434ms 313.1x +wavenet_a1_standard.nam 86.197ms 84.139ms 91.861ms 86.765ms 23.2x +wavenet_a2_max.nam 77.875ms 77.392ms 115.734ms 81.670ms 25.6x +wavenet_condition_dsp.nam 12.734ms 12.570ms 12.935ms 12.742ms 157.0x + + RTF = 2000ms audio / processing time. Higher is faster. + All models process 2 seconds of audio (96,000 samples @ 48kHz) in buffer=64 chunks. + +================================================================ +2. bench_a2_fast — A2 fast path vs generic WaveNet +================================================================ + +Models generated with A2 shape signature (channels=3 nano, channels=8 standard). +Note: wavenet_a2_max.nam is a feature superset test case and does NOT match the +A2 fast-path shape detector; use purpose-built A2-shaped models for this benchmark. + +== a2_nano.nam (A2 nano, Channels=3) == + audio=2000ms @ 48000Hz, block=64, iters=10 (15000 blocks timed each) + per-block audio deadline: 1333.3 us + min p50 p99 p99.9 max mean + fast (us): 5.88 6.12 10.08 34.79 122.62 6.51 + generic (us): 42.42 43.71 76.46 167.00 267.62 46.54 + RTF (p50): fast=217.69x generic=30.51x speedup=7.14x + +== a2_standard.nam (A2 standard, Channels=8) == + audio=2000ms @ 48000Hz, block=64, iters=10 (15000 blocks timed each) + per-block audio deadline: 1333.3 us + min p50 p99 p99.9 max mean + fast (us): 44.62 45.17 76.50 152.92 208.46 47.92 + generic (us): 51.62 53.04 86.88 166.96 209.08 56.22 + RTF (p50): fast=29.52x generic=25.14x speedup=1.17x + +================================================================ +3. benchmodel_bufsize — buffer size sweep +================================================================ + +--- wavenet_a1_standard.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 108719.0 18.3x +32 88881.5 22.5x +64 81707.5 24.4x +128 76863.0 26.0x +256 72917.6 27.4x +512 69474.1 28.7x +1024 66738.2 29.9x + +--- wavenet.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 9814.2 203.7x +32 6769.9 295.4x +64 5321.7 375.8x +128 4446.5 449.7x +256 4188.6 477.4x +512 4030.2 496.2x +1024 3921.3 510.0x + +--- lstm.nam --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 8675.0 230.5x +32 8978.6 222.7x +64 9325.2 214.4x +128 8758.4 228.3x +256 8582.6 233.0x +512 8532.6 234.3x +1024 8610.8 232.2x + +--- a2_standard.nam (via dispatcher, not bench_a2_fast) --- +Buffer Avg Time (us) RTF +---------------------------------------------- +16 109809.0 18.2x +32 89100.8 22.4x +64 78363.0 25.5x +128 73456.6 27.2x +256 70778.1 28.2x +512 68142.7 29.3x +1024 68512.8 29.1x + +================================================================ +Summary +================================================================ + +Fastest model (benchmodel): wavenet.nam at ~6.4ms for 2s audio (313x RTF) +Slowest model (benchmodel): wavenet_a1_standard.nam at ~86ms (23x RTF) + +A2 fast path speedup: + - A2 nano (3 ch): 7.1x faster than generic at p50 + - A2 standard (8 ch): 1.2x faster than generic at p50 + +Buffer size effect: + - Small models (wavenet.nam) benefit strongly from larger buffers + - Large models (wavenet_a1, a2_standard) are relatively insensitive to buffer size + - LSTM performance is flat across buffer sizes (stateful, fixed per-sample cost) + +================================================================ +Report complete. +================================================================ diff --git a/benchmark_reports/generate_a2_models.py b/benchmark_reports/generate_a2_models.py new file mode 100755 index 00000000..4a6e741c --- /dev/null +++ b/benchmark_reports/generate_a2_models.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Generate A2-shaped .nam files for bench_a2_fast benchmarking.""" + +import json +import random +from pathlib import Path + +K_NUM_LAYERS = 23 +K_HEAD_KERNEL_SIZE = 16 +K_LEAKY_SLOPE = 0.01 +K_KERNEL_SIZES = [6] * 14 + [15, 15] + [6] * 7 +K_DILATIONS = [1, 3, 7, 17, 41, 101, 239, 1, 3, 7, 17, 41, 101, 239, 1, 13, 1, 3, 7, 17, 41, 101, 239] + + +def build_a2_config(channels: int) -> dict: + film_inactive = {"active": False, "shift": True, "groups": 1} + layer = { + "input_size": 1, + "condition_size": 1, + "channels": channels, + "bottleneck": channels, + "kernel_sizes": K_KERNEL_SIZES, + "dilations": K_DILATIONS, + "activation": [{"type": "LeakyReLU", "negative_slope": K_LEAKY_SLOPE} for _ in range(K_NUM_LAYERS)], + "gating_mode": ["none"] * K_NUM_LAYERS, + "secondary_activation": [None] * K_NUM_LAYERS, + "head": {"out_channels": 1, "kernel_size": K_HEAD_KERNEL_SIZE, "bias": True}, + "head1x1": {"active": False, "out_channels": 1, "groups": 1}, + "layer1x1": {"active": True, "groups": 1}, + "conv_pre_film": film_inactive, + "conv_post_film": film_inactive, + "input_mixin_pre_film": film_inactive, + "input_mixin_post_film": film_inactive, + "activation_pre_film": film_inactive, + "activation_post_film": film_inactive, + "layer1x1_post_film": film_inactive, + "head1x1_post_film": film_inactive, + "groups_input": 1, + "groups_input_mixin": 1, + } + return {"layers": [layer], "head_scale": 0.01} + + +def a2_weight_count(channels: int) -> int: + bn = channels + total = channels + for i in range(K_NUM_LAYERS): + k = K_KERNEL_SIZES[i] + total += bn * channels * k + bn # conv1d + total += bn # input mixin + total += channels * bn + channels # layer1x1 + total += channels * K_HEAD_KERNEL_SIZE + 1 # head rechannel + total += 1 # trailing head_scale + return total + + +def main() -> None: + out_dir = Path(__file__).resolve().parent / "tmp_models" + out_dir.mkdir(parents=True, exist_ok=True) + rng = random.Random(42) + + for channels, name in [(3, "a2_nano"), (8, "a2_standard")]: + n = a2_weight_count(channels) + weights = [rng.uniform(-0.3, 0.3) for _ in range(n)] + nam = { + "version": "0.6.0", + "architecture": "WaveNet", + "config": build_a2_config(channels), + "weights": weights, + "sample_rate": 48000, + "metadata": {"name": f"A2 benchmark {name}"}, + } + path = out_dir / f"{name}.nam" + path.write_text(json.dumps(nam)) + print(f"Wrote {path} ({n} weights)") + + +if __name__ == "__main__": + main() diff --git a/benchmark_reports/run_benchmarks.sh b/benchmark_reports/run_benchmarks.sh new file mode 100755 index 00000000..72720f38 --- /dev/null +++ b/benchmark_reports/run_benchmarks.sh @@ -0,0 +1,124 @@ +#!/bin/bash +# Run NeuralAmpModelerCore benchmarks and save a timestamped report. +# Usage: ./benchmark_reports/run_benchmarks.sh +# Prerequisite: build bench tools first: +# cd build && cmake .. -DCMAKE_BUILD_TYPE=Release && cmake --build . --target benchmodel benchmodel_bufsize bench_a2_fast -j$(sysctl -n hw.ncpu) + +set -euo pipefail +export PATH="/usr/bin:/bin:/usr/sbin:/sbin:/opt/homebrew/bin:${PATH:-}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" +cd "$ROOT" + +BENCH=./build/tools/benchmodel +BENCH_BUF=./build/tools/benchmodel_bufsize +BENCH_A2=./build/tools/bench_a2_fast +NUM_RUNS=10 +REPORT="$SCRIPT_DIR/benchmark_report_$(date +%Y%m%d_%H%M%S).txt" + +for exe in "$BENCH" "$BENCH_BUF" "$BENCH_A2"; do + if [[ ! -x "$exe" ]]; then + echo "Missing $exe — build first (see script header)." >&2 + exit 1 + fi +done + +# Generate A2-shaped test models if missing +if [[ ! -f "$SCRIPT_DIR/tmp_models/a2_nano.nam" ]]; then + python3 "$SCRIPT_DIR/generate_a2_models.py" +fi + +calc_stats() { + local f="$1" n min max mean median + n=$(wc -l < "$f" | tr -d ' ') + min=$(sort -n "$f" | head -1) + max=$(sort -n "$f" | tail -1) + mean=$(awk '{s+=$1} END{printf "%.3f", s/NR}' "$f") + if [ $((n % 2)) -eq 0 ]; then + local mid1=$((n / 2)) mid2=$((n / 2 + 1)) + local v1 v2 + v1=$(sort -n "$f" | sed -n "${mid1}p") + v2=$(sort -n "$f" | sed -n "${mid2}p") + median=$(echo "scale=3; ($v1 + $v2) / 2" | bc) + else + local mid=$((n / 2 + 1)) + median=$(sort -n "$f" | sed -n "${mid}p") + fi + echo "$median $min $max $mean" +} + +extract_ms() { + echo "$1" | grep -E "^[0-9]+\.[0-9]+ms$" | head -1 | sed 's/ms$//' +} + +{ +echo "================================================================" +echo "NeuralAmpModelerCore Benchmark Report" +echo "================================================================" +echo "" +echo "Date: $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "Machine: $(sysctl -n machdep.cpu.brand_string 2>/dev/null || echo unknown)" +echo "Arch: $(uname -m)" +echo "OS: $(sw_vers -productName 2>/dev/null) $(sw_vers -productVersion 2>/dev/null)" +echo "CPU cores: $(sysctl -n hw.ncpu 2>/dev/null)" +echo "RAM: $(sysctl -n hw.memsize 2>/dev/null | awk '{printf "%.1f GB", $1/1073741824}')" +echo "" +echo "Build: Release, A2 fast ON" +echo "Git: $(git rev-parse --short HEAD) ($(git describe --tags --always 2>/dev/null))" +echo "" + +echo "================================================================" +echo "1. benchmodel — example models ($NUM_RUNS runs)" +echo "================================================================" +printf "%-35s %10s %10s %10s %10s %8s\n" "Model" "Median" "Min" "Max" "Mean" "RTF" +echo "--------------------------------------------------------------------------------" +for model in example_models/*.nam; do + name=$(basename "$model") + tmp=$(mktemp) + for _ in $(seq 1 "$NUM_RUNS"); do + out=$("$BENCH" "$model" 2>&1) + extract_ms "$out" >> "$tmp" + done + read -r med min max mean <<< "$(calc_stats "$tmp")" + rtf=$(echo "scale=1; 2000 / $med" | bc) + printf "%-35s %9.3fms %9.3fms %9.3fms %9.3fms %7sx\n" "$name" "$med" "$min" "$max" "$mean" "$rtf" + rm -f "$tmp" +done + +echo "" +echo "================================================================" +echo "2. bench_a2_fast — A2 fast vs generic" +echo "================================================================" +"$BENCH_A2" "$SCRIPT_DIR/tmp_models/a2_nano.nam" "$SCRIPT_DIR/tmp_models/a2_standard.nam" 2>&1 + +echo "" +echo "================================================================" +echo "3. benchmodel_bufsize — buffer size sweep" +echo "================================================================" +for model_label in \ + "wavenet_a1_standard.nam:example_models/wavenet_a1_standard.nam" \ + "wavenet.nam:example_models/wavenet.nam" \ + "lstm.nam:example_models/lstm.nam" \ + "a2_standard.nam:$SCRIPT_DIR/tmp_models/a2_standard.nam"; do + label="${model_label%%:*}" + path="${model_label#*:}" + echo "" + echo "--- $label ---" + printf "%-12s %15s %12s\n" "Buffer" "Avg Time (us)" "RTF" + echo "----------------------------------------------" + for bufsize in 16 32 64 128 256 512 1024; do + result=$("$BENCH_BUF" "$path" "$bufsize" 5 2>&1) + us=$(echo "$result" | grep -E "^[0-9]+,[0-9.]+$" | cut -d, -f2) + rtf=$(echo "scale=1; 2000000 / $us" | bc) + printf "%-12s %15.1f %11sx\n" "$bufsize" "$us" "$rtf" + done +done + +echo "" +echo "================================================================" +echo "Report complete." +echo "================================================================" +} > "$REPORT" + +echo "Report saved to: $REPORT" diff --git a/docs/fused-engine.md b/docs/fused-engine.md new file mode 100644 index 00000000..73404d6b --- /dev/null +++ b/docs/fused-engine.md @@ -0,0 +1,121 @@ +# Fused NEON WaveNet engine + +`NAM/wavenet/fused.{h,cpp}` — an AArch64 (Apple Silicon, and any ARMv8+ CPU) +fast path for the common "standard shape" WaveNet family. Enabled with the +`NAM_ENABLE_FUSED` compile definition (CMake option `NAM_ENABLE_FUSED`, ON by +default; also defined by the plugin's mac/iOS/win build configs). + +## Why + +Profiling the generic WaveNet on an Apple M2 (`wavenet_a1_standard.nam`, +48 kHz, 64-frame blocks) showed: + +- **~65%** of wall time inside Eigen's dynamic-size GEMM machinery — and over + a quarter of *that* was packing/blocking overhead rather than math. Eigen's + GEMM also heap-allocates its packing buffers on the audio thread. +- **~8%** scalar fast-tanh, ~20% spread across ring-buffer writes, + `setZero`, and the many per-layer buffer-to-buffer passes + (conv → mixin-add → activation → head accumulate → 1x1 → residual add). + +The matrices are tiny (16×16 at most per tap), so a general-purpose GEMM is +the wrong tool: a register-tiled NEON microkernel with compile-time channel +counts reaches ~90% of the core's fp32 FMA peak on these shapes, where Eigen +reaches ~40%. Two further strategies were measured and rejected: + +- **Accelerate / AMX** (`cblas_sgemm`): ties the NEON kernel at these sizes + (932 vs 934 ns for a 16-ch, k=3, 64-frame layer conv including the im2col + gather) — the AMX only pulls ahead on much larger matrices. Not worth the + framework dependency or the (unspecified) real-time behavior of BLAS. +- **BNNS / Metal / ANE**: dispatch latency and/or conversion overhead make + them unsuitable for 64-sample (1.3 ms) real-time blocks. + +## What it does + +Per layer, the engine runs three passes over one L1-resident block buffer +instead of the generic path's ~8 passes + 4 Eigen GEMM calls: + +1. **conv pass** — all K dilated-conv taps, conv bias, and the input mixin + fused into one register-tiled kernel (`conv_block`, Q = channels/4). + Accumulators for a 4-frame tile stay in NEON registers across all taps; + weights stream through `vfmaq_laneq_f32`. +2. **activation pass** — vectorized NEON implementations for fast-tanh + (same rational approximation as `activations::fast_tanh`; `vdivq_f32` + measured *faster* than reciprocal-refinement on M-series because the FP + divider is a separate unit), ReLU, LeakyReLU, Hardtanh, Softsign. Any + other activation (exact Tanh, Sigmoid, PReLU, LUTs) calls the exact same + `Activation` object the generic path would use, so semantics never + diverge — including `enable_fast_tanh()` and LUT substitutions, which are + resolved through the registry at model-construction time in both paths. +3. **tail pass** — head accumulation, layer1x1 GEMM + bias, and the residual + add into the (in-place) layer input buffer, again register-tiled. + +Additional structure: + +- Conv history lives in **power-of-2 ring buffers with a mirrored tail**, so + every tap read is a single contiguous span. The mirror is only refreshed + for the (rare) blocks where a read actually wraps — constant, small + per-block cost, no rewind `memmove` latency spikes. +- The layer1x1 of the final layer of the final array is skipped (its output + is provably never consumed; the generic path computes and discards it). +- Everything is preallocated in `SetMaxBufferSize`; `process()` performs no + heap allocation (verified by an allocation-tracking test). + +## Supported shapes + +The strict detector (`is_fused_shape`) routes a model here only when: + +- `in_channels == 1`, no condition DSP, no post-stack head +- every layer array has `condition_size == 1`, `bottleneck == channels`, + channels a multiple of 4 (≤ 32), all groups == 1, gating `none`, no FiLM, + `head1x1` inactive, `layer1x1` groups == 1 (active or inactive) +- arbitrary kernel sizes (≤ 64), dilations, activations, head kernel + size/bias, any number of layer arrays + +This covers the classic A1 standard/lite family (16/8 and 8/4 channels, +kernel 3) and the A2 standard model (8 channels). A2 nano (3 channels) is +not a multiple of 4 and falls through to the existing `a2_fast` scalar path, +which is already excellent for it. Anything else (feather/nano 2-channel +arrays, gated/FiLM/grouped models) uses the generic path unchanged — the +fused engine never regresses unsupported models, it just declines them. + +Dispatch order in `wavenet::create_config`: slimmable → **fused** → a2_fast +→ generic. + +## Correctness + +- `tools/test/test_fused.cpp` compares fused vs generic outputs (tolerance + 5e-5, the same as the a2_fast tests) across: A1 standard 16/8 with exact + tanh and fast tanh, 8/4, channels 12/20/32, ReLU / LeakyReLU / Hardtanh / + Softsign / Sigmoid, A2-like mixed kernel sizes with head kernel 16, + layer1x1-inactive arrays, and block sizes 64/256/23 (the odd size + exercises the remainder tiles and ring wrap paths). Plus detector + negative tests and a zero-allocation real-time-safety test. +- End-to-end `render` comparison on the real trained models + (`wavenet_a1_standard.nam`, `my_model.nam`): max sample difference ~1e-6, + RMS error −127 dB relative to the signal. The only sources of deviation + are FMA contraction/ordering inside the conv sums. + +## Measured results (Apple M2, macOS, buffer = 64, 48 kHz) + +| Model | generic | fused | speedup | +|---|---|---|---| +| `wavenet_a1_standard.nam` (16/8) | ~86 ms / 2 s (23x RT) | ~39 ms / 2 s (~51x RT) | **~2.2x** | +| `a2_standard.nam` (8 ch) | ~83 ms / 2 s (24x RT) | ~40 ms / 2 s (~50x RT) | **~2.1x** (1.8x vs `a2_fast`) | + +The conv kernel itself runs at ~105 GFLOP/s fp32 — ~93% of a single M2 +P-core's theoretical FMA peak — so the remaining per-block time is dominated +by the (irreducible at fp32) tanh evaluation and the memory traffic floor. + +## Future directions considered + +- **fp16 / bf16**: `FMLAL`-style fp16 multiplies with fp32 accumulation give + the same 4 MACs/instruction as fp32 FMA on Apple cores — no compute win, + only bandwidth (we are compute-bound). Full fp16 accumulation doubles + throughput but the 11-bit mantissa is a real audio-quality risk across a + 48-term accumulation; not pursued. +- **SME/SME2 (M4+)**: the streaming-SVE matrix unit could substantially beat + NEON on these GEMMs; needs M4 hardware to develop/validate. +- **Multithreading**: layer arrays are sequential, so parallelism would have + to split channels within a layer; with 64-sample deadlines the sync jitter + is a poor trade for a plugin. Apple's audio-workgroup API would be the + right vehicle if ever needed. diff --git a/tools/benchmodel.cpp b/tools/benchmodel.cpp index b2630cf6..99ef72a7 100644 --- a/tools/benchmodel.cpp +++ b/tools/benchmodel.cpp @@ -25,13 +25,29 @@ int main(int argc, char* argv[]) double slimValue = -1.0; bool hasSlim = false; bool useFastTanh = true; + double seconds = 2.0; std::vector positionalArgs; positionalArgs.push_back(argv[0]); for (int i = 1; i < argc; i++) { std::string arg(argv[i]); - if (arg == "--slim") + if (arg == "--seconds") + { + if (i + 1 >= argc) + { + std::cerr << "Error: --seconds requires a positive number\n"; + return 1; + } + seconds = std::strtod(argv[i + 1], nullptr); + if (seconds <= 0.0) + { + std::cerr << "Error: --seconds must be positive\n"; + return 1; + } + i++; // skip the value + } + else if (arg == "--slim") { if (i + 1 >= argc) { @@ -102,7 +118,7 @@ int main(int argc, char* argv[]) size_t bufferSize = AUDIO_BUFFER_SIZE; model->Reset(model->GetExpectedSampleRate(), bufferSize); - size_t numBuffers = (48000 / bufferSize) * 2; + size_t numBuffers = static_cast((48000 / bufferSize) * seconds); // Allocate multi-channel buffers const int in_channels = model->NumInputChannels(); diff --git a/tools/run_tests.cpp b/tools/run_tests.cpp index 5699bf78..7292d879 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_fused.cpp" int main() { @@ -370,6 +371,21 @@ int main() test_a2_fast::test_process_realtime_safe_standard(); #endif +#if defined(NAM_ENABLE_FUSED) + // Fused NEON WaveNet engine: detector coverage + numerical match against + // generic across shapes, activations, and block sizes. + test_fused::test_detector_declines_non_fused_shapes(); + test_fused::test_detector_accepts_a1_legacy(); + test_fused::test_a1_standard_matches_generic(); + test_fused::test_a1_standard_fast_tanh_matches_generic(); + test_fused::test_a1_lite_feather_match_generic(); + test_fused::test_activations_match_generic(); + test_fused::test_mixed_kernel_sizes_and_head_kernel_match_generic(); + test_fused::test_layer1x1_inactive_matches_generic(); + test_fused::test_channels_12_20_match_generic(); + test_fused::test_process_is_realtime_safe(); +#endif + std::cout << "Success!" << std::endl; #ifdef ADDASSERT std::cerr << "===============================================================" << std::endl; diff --git a/tools/test/test_fused.cpp b/tools/test/test_fused.cpp new file mode 100644 index 00000000..0f5011df --- /dev/null +++ b/tools/test/test_fused.cpp @@ -0,0 +1,400 @@ +// Numerical verification for the fused NEON WaveNet engine: +// for configs that match the fused shape, the fused engine must produce the +// same output (within FMA-reordering tolerance) as the generic WaveNet on the +// same input and weights. +// +// Built only when NAM_ENABLE_FUSED is defined at compile time. On builds +// without NEON kernels (non-AArch64), the numerical tests are skipped and the +// detector is verified to decline everything. + +#if defined(NAM_ENABLE_FUSED) + + #include + #include + #include + #include + #include + #include + #include + #include + + #include "json.hpp" + + #include "NAM/dsp.h" + #include "NAM/wavenet/fused.h" + #include "NAM/wavenet/model.h" + + #include "allocation_tracking.h" + +namespace test_fused +{ +namespace +{ + +struct ArrayShape +{ + int input_size; + int channels; + int head_size; + int head_kernel_size; + bool head_bias; + bool layer1x1_active; + std::vector kernel_sizes; + std::vector dilations; + nlohmann::json activation; // single config applied to all layers +}; + +nlohmann::json build_config(const std::vector& arrays) +{ + using nlohmann::json; + json layers = json::array(); + for (const auto& a : arrays) + { + json la; + la["input_size"] = a.input_size; + la["condition_size"] = 1; + la["channels"] = a.channels; + la["bottleneck"] = a.channels; + la["kernel_sizes"] = a.kernel_sizes; + la["dilations"] = a.dilations; + la["activation"] = a.activation; + la["gating_mode"] = "none"; + la["head"] = {{"out_channels", a.head_size}, {"kernel_size", a.head_kernel_size}, {"bias", a.head_bias}}; + la["layer1x1"] = {{"active", a.layer1x1_active}, {"groups", 1}}; + la["head1x1"] = {{"active", false}, {"out_channels", 1}, {"groups", 1}}; + layers.push_back(la); + } + json config; + config["layers"] = layers; + config["head_scale"] = 0.02f; + return config; +} + +/// The classic "A1 standard" shape (what most captures in the wild use), +/// expressed with the legacy schema fields (kernel_size + head_size/head_bias) +/// to make sure the detector handles old .nam files. +nlohmann::json build_a1_legacy_config(int ch0, int ch1) +{ + using nlohmann::json; + json config; + json la0, la1; + const json dil = json::array({1, 2, 4, 8, 16, 32, 64, 128, 256, 512}); + la0["input_size"] = 1; + la0["condition_size"] = 1; + la0["channels"] = ch0; + la0["kernel_size"] = 3; + la0["dilations"] = dil; + la0["activation"] = "Tanh"; + la0["gated"] = false; + la0["head_size"] = ch1; + la0["head_bias"] = false; + la1["input_size"] = ch0; + la1["condition_size"] = 1; + la1["channels"] = ch1; + la1["kernel_size"] = 3; + la1["dilations"] = dil; + la1["activation"] = "Tanh"; + la1["gated"] = false; + la1["head_size"] = 1; + la1["head_bias"] = true; + config["layers"] = json::array({la0, la1}); + config["head_scale"] = 0.02f; + return config; +} + +int count_weights(const nlohmann::json& config) +{ + int total = 0; + int prev_channels = 1; + for (const auto& la : config["layers"]) + { + const int C = la["channels"].get(); + const int in_size = la["input_size"].get(); + (void)prev_channels; + total += C * in_size; // rechannel + const auto& dil = la["dilations"]; + for (size_t i = 0; i < dil.size(); i++) + { + int K; + if (la.find("kernel_sizes") != la.end()) + K = la["kernel_sizes"][i].get(); + else + K = la["kernel_size"].get(); + total += C * C * K + C; // conv + bias + total += C; // mixin + bool l1x1 = true; + if (la.find("layer1x1") != la.end()) + l1x1 = la["layer1x1"]["active"].get(); + if (l1x1) + total += C * C + C; + } + int H, Kh; + bool hb; + if (la.find("head") != la.end()) + { + H = la["head"]["out_channels"].get(); + Kh = la["head"]["kernel_size"].get(); + hb = la["head"]["bias"].get(); + } + else + { + H = la["head_size"].get(); + Kh = 1; + hb = la["head_bias"].get(); + } + total += H * C * Kh + (hb ? H : 0); + prev_channels = C; + } + total += 1; // trailing head_scale + return total; +} + +std::vector make_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_input(int num_frames) +{ + std::vector in(num_frames); + for (int i = 0; i < num_frames; i++) + { + const double t = static_cast(i) / 48000.0; + in[i] = static_cast(0.4 * std::sin(2.0 * M_PI * 110.0 * t) + 0.15 * std::sin(2.0 * M_PI * 997.0 * t)); + } + return in; +} + +std::vector run_dsp(nam::DSP& dsp, const std::vector& input, int block_size) +{ + dsp.Reset(48000.0, block_size); + 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; +} + +void compare(const std::vector& generic, const std::vector& fast, const std::string& what, + int block_size, double tol) +{ + assert(generic.size() == fast.size()); + double max_diff = 0.0; + size_t max_i = 0; + for (size_t i = 0; i < generic.size(); i++) + { + const double d = std::fabs(static_cast(generic[i]) - static_cast(fast[i])); + if (d > max_diff) + { + max_diff = d; + max_i = i; + } + } + if (!(max_diff < tol)) + { + std::cerr << "FusedWaveNet diverges from generic WaveNet for " << what << " (block=" << block_size + << "): max |diff| = " << max_diff << " at i=" << max_i << " (generic=" << generic[max_i] + << ", fused=" << fast[max_i] << ")" << std::endl; + assert(false); + } +} + +void check_matches_generic(const nlohmann::json& config, const std::string& what, uint32_t seed, + double tol = 5e-5) +{ + assert(nam::wavenet::fused::is_fused_shape(config)); + + const int weight_count = count_weights(config); + const auto weights = make_weights(weight_count, seed); + + auto fused_cfg = nam::wavenet::fused::create_fused_config(config, 48000.0); + std::vector w_fast = weights; + auto fused_dsp = fused_cfg->create(std::move(w_fast), 48000.0); + + // Reference: parse_config_json directly, bypassing the dispatcher shortcut. + auto generic_cfg = nam::wavenet::parse_config_json(config, 48000.0); + std::vector w_gen = weights; + auto generic_dsp = generic_cfg.create(std::move(w_gen), 48000.0); + + const int total = 4096; + const auto input = make_input(total); + // Include a non-multiple-of-4 block size to exercise the remainder tiles. + for (int block : {64, 256, 23}) + { + const auto out_generic = run_dsp(*generic_dsp, input, block); + const auto out_fused = run_dsp(*fused_dsp, input, block); + compare(out_generic, out_fused, what, block, tol); + } +} + +} // namespace + +void test_detector_declines_non_fused_shapes() +{ + // Gated layers must be declined. + { + auto cfg = build_config({{1, 8, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + cfg["layers"][0]["gating_mode"] = "gated"; + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } + // Channels not a multiple of 4. + { + auto cfg = build_config({{1, 6, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } + // FiLM active. + { + auto cfg = build_config({{1, 8, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + cfg["layers"][0]["conv_post_film"] = {{"active", true}, {"shift", true}, {"groups", 1}}; + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } + // head1x1 active. + { + auto cfg = build_config({{1, 8, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + cfg["layers"][0]["head1x1"] = {{"active", true}, {"out_channels", 4}, {"groups", 1}}; + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } + // Bottleneck != channels. + { + auto cfg = build_config({{1, 8, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + cfg["layers"][0]["bottleneck"] = 4; + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } + // Grouped layer1x1. + { + auto cfg = build_config({{1, 8, 1, 1, true, true, std::vector(4, 3), {1, 2, 4, 8}, "Tanh"}}); + cfg["layers"][0]["layer1x1"] = {{"active", true}, {"groups", 2}}; + assert(!nam::wavenet::fused::is_fused_shape(cfg)); + } +} + +void test_detector_accepts_a1_legacy() +{ + if (!nam::wavenet::fused::available()) + return; + assert(nam::wavenet::fused::is_fused_shape(build_a1_legacy_config(16, 8))); +} + +void test_a1_standard_matches_generic() +{ + if (!nam::wavenet::fused::available()) + return; + nam::activations::Activation::disable_fast_tanh(); + check_matches_generic(build_a1_legacy_config(16, 8), "A1 standard (16/8, exact tanh)", 0xF00D0001u); +} + +void test_a1_standard_fast_tanh_matches_generic() +{ + if (!nam::wavenet::fused::available()) + return; + nam::activations::Activation::enable_fast_tanh(); + check_matches_generic(build_a1_legacy_config(16, 8), "A1 standard (16/8, fast tanh)", 0xF00D0002u); + nam::activations::Activation::disable_fast_tanh(); +} + +void test_a1_lite_feather_match_generic() +{ + if (!nam::wavenet::fused::available()) + return; + nam::activations::Activation::enable_fast_tanh(); + check_matches_generic(build_a1_legacy_config(8, 4), "A1 lite-like (8/4)", 0xF00D0003u); + nam::activations::Activation::disable_fast_tanh(); +} + +void test_activations_match_generic() +{ + if (!nam::wavenet::fused::available()) + return; + const std::vector dil = {1, 2, 4, 8, 16}; + const std::vector ks(5, 3); + uint32_t seed = 0xF00D0100u; + for (const nlohmann::json act : {nlohmann::json("ReLU"), nlohmann::json("Hardtanh"), nlohmann::json("Softsign"), + nlohmann::json("Sigmoid"), nlohmann::json({{"type", "LeakyReLU"}, + {"negative_slope", 0.02}})}) + { + auto cfg = build_config({{1, 8, 1, 1, true, true, ks, dil, act}}); + check_matches_generic(cfg, "activation " + act.dump(), seed++); + } +} + +void test_mixed_kernel_sizes_and_head_kernel_match_generic() +{ + if (!nam::wavenet::fused::available()) + return; + // A2-standard-like: kernel sizes 6 and 15, head kernel 16, LeakyReLU. + std::vector ks = {6, 6, 6, 6, 15, 6, 6}; + std::vector dil = {1, 3, 7, 17, 1, 13, 41}; + nlohmann::json act = {{"type", "LeakyReLU"}, {"negative_slope", 0.01}}; + auto cfg = build_config({{1, 8, 1, 16, true, true, ks, dil, act}}); + check_matches_generic(cfg, "A2-like (k=6/15, head k=16)", 0xF00D0200u); +} + +void test_layer1x1_inactive_matches_generic() +{ + if (!nam::wavenet::fused::available()) + return; + auto cfg = build_config({{1, 8, 4, 1, false, false, {3, 3, 3}, {1, 2, 4}, "Tanh"}, + {8, 4, 1, 1, true, true, {3, 3, 3}, {1, 2, 4}, "Tanh"}}); + check_matches_generic(cfg, "layer1x1 inactive in first array", 0xF00D0300u); +} + +void test_channels_12_20_match_generic() +{ + if (!nam::wavenet::fused::available()) + return; + nam::activations::Activation::enable_fast_tanh(); + auto cfg12 = build_config({{1, 12, 1, 1, true, true, {3, 3, 3, 3}, {1, 2, 4, 8}, "Tanh"}}); + check_matches_generic(cfg12, "channels=12", 0xF00D0400u); + auto cfg20 = build_config({{1, 20, 1, 1, true, true, {3, 3, 3, 3}, {1, 2, 4, 8}, "Tanh"}}); + check_matches_generic(cfg20, "channels=20", 0xF00D0401u); + auto cfg32 = build_config({{1, 32, 1, 1, true, true, {3, 3}, {1, 2}, "Tanh"}}); + check_matches_generic(cfg32, "channels=32", 0xF00D0402u); + nam::activations::Activation::disable_fast_tanh(); +} + +void test_process_is_realtime_safe() +{ + if (!nam::wavenet::fused::available()) + return; + auto cfg = build_a1_legacy_config(16, 8); + const auto weights = make_weights(count_weights(cfg), 0xF00D0500u); + auto fused_cfg = nam::wavenet::fused::create_fused_config(cfg, 48000.0); + std::vector w = weights; + auto dsp = fused_cfg->create(std::move(w), 48000.0); + const int block = 64; + dsp->Reset(48000.0, block); + + const auto input = make_input(block); + std::vector out(block, 0.0); + const NAM_SAMPLE* in_arr[] = {input.data()}; + NAM_SAMPLE* out_arr[] = {out.data()}; + + // One call to settle anything lazy, then track allocations across many + // process() calls: expected zero allocations and zero deallocations. + dsp->process(const_cast(in_arr), out_arr, block); + allocation_tracking::run_allocation_test( + nullptr, + [&]() { + for (int i = 0; i < 32; i++) + dsp->process(const_cast(in_arr), out_arr, block); + }, + nullptr, /*expected_allocations=*/0, /*expected_deallocations=*/0, "FusedWaveNet::process realtime safety"); +} + +} // namespace test_fused + +#endif // NAM_ENABLE_FUSED