Cache A2 prewarm state

CMakeLists.txt

@@ -51,7 +51,7 @@ endif()
 set(NAM_DEPS_PATH "${CMAKE_CURRENT_SOURCE_DIR}/Dependencies")
 include_directories(SYSTEM "${NAM_DEPS_PATH}/eigen")
 
-# Build the specialized A2 fast-path WaveNet (A2 standard + A2 nano). When ON,
+# Build the specialized A2 fast-path WaveNet (A2-Full + A2-Lite). When ON,
 # models whose config matches the A2 shape signature are routed to a hand-optimized
 # WaveNet implementation instead of the generic one. When OFF, all models go
 # through the generic path.

NAM/conv1d.cpp

@@ -1,5 +1,6 @@
 #include "conv1d.h"
 #include "compiler.h"
+#include <cassert>
 #include <cstring>
 #include <stdexcept>
 
@@ -9,6 +10,7 @@ namespace nam
 
 void Conv1D::set_weights_(std::vector<float>::iterator& weights)
 {
+  _has_cached_prewarm_state = false;
   if (this->_is_depthwise)
   {
     // Depthwise convolution: one weight per channel per kernel tap
@@ -109,6 +111,10 @@ void Conv1D::set_size_(const int in_channels, const int out_channels, const int
   }
   else
     this->_bias.resize(0);
+
+  _cached_prewarm_state.resize(in_channels);
+  _cached_prewarm_state.setZero();
+  _has_cached_prewarm_state = false;
 }
 
 void Conv1D::set_size_and_weights_(const int in_channels, const int out_channels, const int kernel_size,
@@ -142,6 +148,17 @@ void Conv1D::SetMaxBufferSize(const int maxBufferSize)
   _output.setZero();
 }
 
+void Conv1D::PrewarmFromCache()
+{
+  assert(HasCachedPrewarmState());
+  _input_buffer.FillWithSample(_cached_prewarm_state);
+}
+
+void Conv1D::CacheStateAsPrewarmed()
+{
+  _input_buffer.CacheLastWrittenSample(_cached_prewarm_state);
+  _has_cached_prewarm_state = true;
+}
 
 void Conv1D::Process(const Eigen::MatrixXf& input, const int num_frames)
 {

NAM/conv1d.h

@@ -117,6 +117,15 @@ class Conv1D
   /// \return true if bias is present, false otherwise
   bool has_bias() const { return this->_bias.size() > 0; };
 
+  /// \brief Check whether a steady-state prewarm sample has been cached
+  bool HasCachedPrewarmState() const { return _has_cached_prewarm_state; }
+
+  /// \brief Restore the input history from the cached steady-state sample
+  void PrewarmFromCache();
+
+  /// \brief Cache the most recently written input sample as the steady prewarm state
+  void CacheStateAsPrewarmed();
+
 protected:
   // conv[kernel](cout, cin) - used for non-depthwise convolutions
   std::vector<Eigen::MatrixXf> _weight;
@@ -131,6 +140,8 @@ class Conv1D
 
 private:
   RingBuffer _input_buffer; // Ring buffer for input (channels x buffer_size)
+  Eigen::VectorXf _cached_prewarm_state;
+  bool _has_cached_prewarm_state = false;
   Eigen::MatrixXf _output; // Pre-allocated output buffer (out_channels x maxBufferSize)
   int _max_buffer_size = 0; // Stored maxBufferSize
 };

NAM/ring_buffer.cpp

@@ -61,6 +61,20 @@ void RingBuffer::Advance(const int num_frames)
   _write_pos += num_frames;
 }
 
+void RingBuffer::CacheLastWrittenSample(Eigen::VectorXf& destination) const
+{
+  assert(_write_pos > 0);
+  assert(destination.size() == _storage.rows());
+  destination = _storage.col(_write_pos - 1);
+}
+
+void RingBuffer::FillWithSample(const Eigen::VectorXf& sample)
+{
+  assert(sample.size() == _storage.rows());
+  _storage.colwise() = sample;
+  _write_pos = _max_lookback;
+}
+
 bool RingBuffer::NeedsRewind(const int num_frames) const
 {
   return _write_pos + num_frames > (long)_storage.cols();

NAM/ring_buffer.h

@@ -52,6 +52,12 @@ class RingBuffer
   /// \param max_lookback Maximum lookback distance
   void SetMaxLookback(const long max_lookback) { _max_lookback = max_lookback; }
 
+  /// \brief Copy the most recently written sample into pre-allocated storage
+  void CacheLastWrittenSample(Eigen::VectorXf& destination) const;
+
+  /// \brief Fill the buffer with a constant sample and restore its initial write position
+  void FillWithSample(const Eigen::VectorXf& sample);
+
 private:
   // Wrap buffer when approaching end (called automatically if needed)
   void Rewind();

NAM/wavenet/a2_fast.cpp

@@ -66,6 +66,7 @@ class A2FastModel : public DSP
   ~A2FastModel() override = default;
 
   void process(NAM_SAMPLE** input, NAM_SAMPLE** output, int num_frames) override;
+  void prewarm() override;
   int GetPrewarmSamples() override { return _prewarm_samples; }
 
 protected:
@@ -92,6 +93,7 @@ class A2FastModel : public DSP
 
     // Conv1D input history ring buffer, column-major (Channels rows).
     std::vector<float> history;
+    std::array<float, Channels> cached_prewarm_state{};
   #if NAM_A2_RING_MODE == 1
     // pow2 ring + tail mirror. Storage = (pow2_size + max_buffer_size) cols.
     // write_pos is kept in [0, pow2_size), reads use (pos & pow2_mask) and are
@@ -124,6 +126,7 @@ class A2FastModel : public DSP
 
   // Head ring buffer (Channels rows, col-major). Same ring layout as per-layer.
   std::vector<float> _head_history;
+  std::array<float, Channels> _cached_head_prewarm_state{};
   #if NAM_A2_RING_MODE == 1
   int _head_pow2_size = 0;
   int _head_pow2_mask = 0;
@@ -141,8 +144,12 @@ class A2FastModel : public DSP
   std::vector<float> _head_out; // float32 head output before writing to NAM_SAMPLE
 
   int _prewarm_samples = 0;
+  bool _has_cached_prewarm_state = false;
 
   void _load_weights(std::vector<float>& weights);
+  bool HasCachedPrewarmState() const { return _has_cached_prewarm_state; }
+  void PrewarmFromCache();
+  void CacheStateAsPrewarmed();
   void _ring_write(Layer& L, int num_frames);
   void _head_ring_write(int num_frames);
   void _layer_forward(int layer_idx, const float* cond, int num_frames);
@@ -329,6 +336,73 @@ void A2FastModel<Channels>::SetMaxBufferSize(int maxBufferSize)
   #endif
 }
 
+// -----------------------------------------------------------------------------
+// Prewarm-state cache
+//
+// Processing silence for a full receptive field leaves every convolution
+// history constant in time. Keep one Channels-wide column from each layer and
+// the head so later prewarms can rebuild the complete histories directly.
+// -----------------------------------------------------------------------------
+template <int Channels>
+void A2FastModel<Channels>::prewarm()
+{
+  if (HasCachedPrewarmState())
+  {
+    PrewarmFromCache();
+    return;
+  }
+
+  DSP::prewarm();
+  CacheStateAsPrewarmed();
+}
+
+template <int Channels>
+void A2FastModel<Channels>::PrewarmFromCache()
+{
+  for (auto& L : _layers)
+  {
+    const size_t columns = L.history.size() / Channels;
+    for (size_t column = 0; column < columns; column++)
+    {
+      std::copy(L.cached_prewarm_state.begin(), L.cached_prewarm_state.end(),
+                L.history.begin() + static_cast<std::ptrdiff_t>(column * Channels));
+    }
+    L.write_pos = L.max_lookback;
+  }
+
+  const size_t head_columns = _head_history.size() / Channels;
+  for (size_t column = 0; column < head_columns; column++)
+  {
+    std::copy(_cached_head_prewarm_state.begin(), _cached_head_prewarm_state.end(),
+              _head_history.begin() + static_cast<std::ptrdiff_t>(column * Channels));
+  }
+  _head_write_pos = kHeadKernelSize - 1;
+}
+
+template <int Channels>
+void A2FastModel<Channels>::CacheStateAsPrewarmed()
+{
+  for (auto& L : _layers)
+  {
+  #if NAM_A2_RING_MODE == 1
+    const int last_column = (L.write_pos - 1) & L.pow2_mask;
+  #else
+    const int last_column = L.write_pos - 1;
+  #endif
+    std::copy_n(L.history.begin() + static_cast<std::ptrdiff_t>(last_column * Channels), Channels,
+                L.cached_prewarm_state.begin());
+  }
+
+  #if NAM_A2_RING_MODE == 1
+  const int last_head_column = (_head_write_pos - 1) & _head_pow2_mask;
+  #else
+  const int last_head_column = _head_write_pos - 1;
+  #endif
+  std::copy_n(_head_history.begin() + static_cast<std::ptrdiff_t>(last_head_column * Channels), Channels,
+              _cached_head_prewarm_state.begin());
+  _has_cached_prewarm_state = true;
+}
+
 // -----------------------------------------------------------------------------
 // Ring-write helpers.
 //   Mode 1: pow2 + tail mirror. Constant-time per block (one short memcpy
@@ -430,14 +504,14 @@ void A2FastModel<Channels>::_layer_forward_k(Layer& L, const float* cond, int nu
 
   // Two conv strategies, dispatched at compile time on Channels:
   //
-  //   - Channels <= 4 (A2 nano): full-block tap-major. The z accumulator lives
+  //   - Channels <= 4 (A2-Lite): full-block tap-major. The z accumulator lives
   //     in the heap buffer across all taps, and for each tap the inner f-loop
   //     iterates over all num_frames. This gives clang frame-level
   //     parallelism — it vectorizes across 4 frames at a time, which matters
   //     more than weight-reload cost when the b-loop (3 wide) can't saturate
   //     NEON lanes on its own.
   //
-  //   - Channels >= 8 (A2 standard): frame-tiled tap-major with T=4. ztile
+  //   - Channels >= 8 (A2-Full): frame-tiled tap-major with T=4. ztile
   //     stays in NEON registers across all K taps, amortizing weight loads
   //     over 4 frames — equivalent to what a GEMM kernel does. Weight reuse
   //     matters here because the b-loop (8 wide) already saturates SIMD, so

NAM/wavenet/a2_fast.h

@@ -1,7 +1,7 @@
 #pragma once
 
-// Specialized WaveNet fast path for the A2 standard (Channels=8) and
-// A2 nano (Channels=3) models. Shares the same architecture shape; only
+// Specialized WaveNet fast path for the A2-Full (Channels=8) and
+// A2-Lite (Channels=3) models. Shares the same architecture shape; only
 // the channel count differs.
 //
 // When NAM_ENABLE_A2_FAST is defined at build time, wavenet::create_config
@@ -34,17 +34,17 @@ constexpr int kHeadKernelSize = 16;
 /// \brief LeakyReLU negative-slope used by every layer.
 constexpr float kLeakySlope = 0.01f;
 
-/// \brief Per-layer kernel sizes (fixed pattern shared by A2 standard + nano).
+/// \brief Per-layer kernel sizes (fixed pattern shared by A2-Full + A2-Lite).
 inline constexpr std::array<int, kNumLayers> kKernelSizes = {
   6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 15, 15, 6, 6, 6, 6, 6, 6, 6};
 
-/// \brief Per-layer dilations (fixed pattern shared by A2 standard + nano).
+/// \brief Per-layer dilations (fixed pattern shared by A2-Full + A2-Lite).
 inline constexpr std::array<int, kNumLayers> kDilations = {
   1, 3, 7, 17, 41, 101, 239, 1, 3, 7, 17, 41, 101, 239, 1, 13, 1, 3, 7, 17, 41, 101, 239};
 
 /// \brief Strict detector: returns true iff config matches the A2 shape.
 /// \param config   The "config" sub-object from a .nam WaveNet entry.
-/// \param channels Out-param set to 3 (A2 nano) or 8 (A2 standard) on match.
+/// \param channels Out-param set to 3 (A2-Lite) or 8 (A2-Full) on match.
 /// \return true if every architectural knob matches the A2 signature exactly.
 bool is_a2_shape(const nlohmann::json& config, int* channels);
 

NAM/wavenet/detail.h

@@ -224,6 +224,10 @@ class Layer
   /// \return Const reference to the internal Conv1D object
   const Conv1D& get_conv() const { return _conv; }
 
+  bool HasCachedPrewarmState() const { return _conv.HasCachedPrewarmState(); }
+  void PrewarmFromCache() { _conv.PrewarmFromCache(); }
+  void CacheStateAsPrewarmed() { _conv.CacheStateAsPrewarmed(); }
+
 private:
   // The dilated convolution at the front of the block
   Conv1D _conv;
@@ -335,6 +339,10 @@ class LayerArray
   /// \return Receptive field size
   long get_receptive_field() const;
 
+  bool HasCachedPrewarmState() const;
+  void PrewarmFromCache();
+  void CacheStateAsPrewarmed();
+
 private:
   // The rechannel before the layers
   Conv1x1 _rechannel;
@@ -376,6 +384,10 @@ class Head
 
   const Eigen::MatrixXf& get_last_output() const { return _convs.back().GetOutput(); }
 
+  bool HasCachedPrewarmState() const;
+  void PrewarmFromCache();
+  void CacheStateAsPrewarmed();
+
 private:
   std::vector<nam::Conv1D> _convs;
   std::vector<nam::activations::Activation::Ptr> _activations;