Planar NEON kernels, >= 2x faster on Apple Silicon & Raspberry Pi, bit-identical to a2_fast

NAM/wavenet/a2_fast.cpp

@@ -9,6 +9,7 @@
   #endif
 
   #include "a2_fast.h"
+  #include "a2_planar.h"
 
   #include <algorithm>
   #include <array>
@@ -772,11 +773,15 @@ struct A2FastConfig : public ModelConfig
 
   std::unique_ptr<DSP> create(std::vector<float> weights, double sampleRate) override
   {
-    if (channels == 3)
-      return std::make_unique<A2FastModel<3>>(std::move(weights), sampleRate);
-    if (channels == 8)
-      return std::make_unique<A2FastModel<8>>(std::move(weights), sampleRate);
-    throw std::runtime_error("A2FastConfig: unsupported channel count " + std::to_string(channels));
+  #if defined(NAM_A2_PLANAR)
+    // On AArch64, prefer the planar NEON kernels. They are bit-identical to the
+    // reference model below -- same float32 bits out, sample for sample -- so
+    // this is a speed choice and nothing else. A channel count they do not cover
+    // returns nullptr and falls through.
+    if (auto planar = create_a2_planar_model(channels, weights, sampleRate))
+      return planar;
+  #endif
+    return create_a2_fast_reference_model(channels, std::move(weights), sampleRate);
   }
 };
 
@@ -983,6 +988,15 @@ bool is_a2_shape(const nlohmann::json& config, int* channels)
   return true;
 }
 
+std::unique_ptr<DSP> create_a2_fast_reference_model(int channels, std::vector<float> weights, double sampleRate)
+{
+  if (channels == 3)
+    return std::make_unique<A2FastModel<3>>(std::move(weights), sampleRate);
+  if (channels == 8)
+    return std::make_unique<A2FastModel<8>>(std::move(weights), sampleRate);
+  throw std::runtime_error("create_a2_fast_reference_model: unsupported channel count " + std::to_string(channels));
+}
+
 std::unique_ptr<ModelConfig> create_a2_fast_config(const nlohmann::json& config, double sampleRate)
 {
   (void)sampleRate;

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<ModelConfig> 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<DSP> create_a2_fast_reference_model(int channels, std::vector<float> weights, double sampleRate);
+
 } // namespace a2_fast
 } // namespace wavenet
 } // namespace nam