[REFACTOR] Enhance SlimmableWavenet model management

.github/workflows/build.yml

@@ -5,10 +5,10 @@ on:
 
 jobs:
   build-ubuntu:
-    name: Build Ubuntu
+    name: Build Ubuntu (libstdc++)
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3.3.0
+    - uses: actions/checkout@v4
       with:
         submodules: recursive
 
@@ -55,3 +55,49 @@ jobs:
         ./build_inline/tools/benchmodel ./example_models/lstm.nam
         ./build_inline/tools/render ./example_models/wavenet.nam ./example_audio/input.wav ./example_audio/output.wav
 
+  build-ubuntu-libcxx:
+    name: Build Ubuntu (libc++)
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        submodules: recursive
+
+    # If libc++-dev is removed or renamed on a future ubuntu-latest image, pin versioned
+    # packages (e.g. libc++-18-dev libc++abi-18-dev) or add the LLVM apt repository.
+    - name: Install Clang and libc++
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y clang libc++-dev libc++abi-dev
+
+    - name: Create build directories
+      run: mkdir -p build_libcxx build_inline_libcxx
+
+    - name: Build Tools
+      working-directory: ${{github.workspace}}/build_libcxx
+      env:
+        CXX: clang++
+      run: |
+        cmake .. -DCMAKE_BUILD_TYPE=Debug -DNAM_USE_LIBCXX_LINUX=ON
+        cmake --build . -j4
+
+    - name: Build Tools (Inline GEMM)
+      working-directory: ${{github.workspace}}/build_inline_libcxx
+      env:
+        CXX: clang++
+      run: |
+        cmake .. -DCMAKE_BUILD_TYPE=Debug -DNAM_USE_LIBCXX_LINUX=ON -DCMAKE_CXX_FLAGS="-DNAM_USE_INLINE_GEMM"
+        cmake --build . -j4
+
+    - name: Run tests
+      working-directory: ${{github.workspace}}
+      run: |
+        ./build_libcxx/tools/run_tests
+        ./build_libcxx/tools/benchmodel ./example_models/wavenet.nam
+        ./build_libcxx/tools/benchmodel ./example_models/lstm.nam
+        ./build_libcxx/tools/render ./example_models/wavenet.nam ./example_audio/input.wav ./example_audio/output.wav
+        ./build_inline_libcxx/tools/run_tests
+        ./build_inline_libcxx/tools/benchmodel ./example_models/wavenet.nam
+        ./build_inline_libcxx/tools/benchmodel ./example_models/lstm.nam
+        ./build_inline_libcxx/tools/render ./example_models/wavenet.nam ./example_audio/input.wav ./example_audio/output.wav
+

CMakeLists.txt

@@ -14,10 +14,17 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Darwin")
 	set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -stdlib=libc++")
 endif()
 
+option(NAM_USE_LIBCXX_LINUX "Use libc++ on Linux (Clang + libc++-dev packages; tests SlimmableWavenet _LIBCPP_VERSION path)" OFF)
+
 if (CMAKE_SYSTEM_NAME STREQUAL "Darwin")
 	include_directories(SYSTEM /usr/local/include)
 elseif (CMAKE_SYSTEM_NAME STREQUAL "Linux")
-	link_libraries(stdc++fs)
+	if (NAM_USE_LIBCXX_LINUX)
+		set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -stdlib=libc++")
+		# Clang links libc++/libc++abi; do not pull in libstdc++fs (libstdc++).
+	else()
+		link_libraries(stdc++fs)
+	endif()
 elseif (CMAKE_SYSTEM_NAME STREQUAL "Windows")
 	add_compile_definitions(NOMINMAX WIN32_LEAN_AND_MEAN)
 else()

NAM/wavenet/slimmable.cpp

@@ -291,6 +291,49 @@ bool is_full_size(const std::vector<wavenet::LayerArrayParams>& params, const st
 
 } // anonymous namespace
 
+#ifdef _LIBCPP_VERSION
+void SlimmableWavenet::_pending_clear_release()
+{
+  std::atomic_store_explicit(&_pending_staged, std::shared_ptr<StagedSlimModel>{}, std::memory_order_release);
+}
+
+std::shared_ptr<SlimmableWavenet::StagedSlimModel> SlimmableWavenet::_pending_load_acquire() const
+{
+  return std::atomic_load_explicit(&_pending_staged, std::memory_order_acquire);
+}
+
+void SlimmableWavenet::_pending_store_release(std::shared_ptr<StagedSlimModel> p)
+{
+  std::atomic_store_explicit(&_pending_staged, std::move(p), std::memory_order_release);
+}
+
+std::shared_ptr<SlimmableWavenet::StagedSlimModel> SlimmableWavenet::_pending_exchange_take_acq_rel()
+{
+  return std::atomic_exchange_explicit(
+    &_pending_staged, std::shared_ptr<StagedSlimModel>{}, std::memory_order_acq_rel);
+}
+#else
+void SlimmableWavenet::_pending_clear_release()
+{
+  _pending_staged.store({}, std::memory_order_release);
+}
+
+std::shared_ptr<SlimmableWavenet::StagedSlimModel> SlimmableWavenet::_pending_load_acquire() const
+{
+  return _pending_staged.load(std::memory_order_acquire);
+}
+
+void SlimmableWavenet::_pending_store_release(std::shared_ptr<StagedSlimModel> p)
+{
+  _pending_staged.store(std::move(p), std::memory_order_release);
+}
+
+std::shared_ptr<SlimmableWavenet::StagedSlimModel> SlimmableWavenet::_pending_exchange_take_acq_rel()
+{
+  return _pending_staged.exchange({}, std::memory_order_acq_rel);
+}
+#endif
+
 // ============================================================================
 // SlimmableWavenet
 // ============================================================================
@@ -344,11 +387,8 @@ SlimmableWavenet::SlimmableWavenet(std::vector<wavenet::LayerArrayParams> origin
   _rebuild_model(full_channels);
 }
 
-void SlimmableWavenet::_rebuild_model(const std::vector<int>& target_channels)
+std::unique_ptr<DSP> SlimmableWavenet::_create_wavenet_for_channels(const std::vector<int>& target_channels)
 {
-  if (target_channels == _current_channels && _active_model)
-    return;
-
   std::vector<float> weights;
   std::vector<wavenet::LayerArrayParams> modified_params;
   const std::vector<wavenet::LayerArrayParams>* params_ptr;
@@ -371,16 +411,55 @@ void SlimmableWavenet::_rebuild_model(const std::vector<int>& target_channels)
     condition_dsp = get_dsp(_condition_dsp_json);
 
   double sampleRate = _current_sample_rate > 0 ? _current_sample_rate : GetExpectedSampleRate();
-  _active_model = std::make_unique<wavenet::WaveNet>(_in_channels, *params_ptr, _head_scale, _with_head, std::nullopt,
-                                                     std::move(weights), std::move(condition_dsp), sampleRate);
+  return std::make_unique<wavenet::WaveNet>(_in_channels, *params_ptr, _head_scale, _with_head, std::nullopt,
+                                            std::move(weights), std::move(condition_dsp), sampleRate);
+}
+
+void SlimmableWavenet::_rebuild_model(const std::vector<int>& target_channels)
+{
+  if (target_channels == _current_channels && _active_model)
+    return;
+
+  _pending_clear_release();
+
+  _active_model = std::shared_ptr<DSP>(_create_wavenet_for_channels(target_channels));
   _current_channels = target_channels;
 
   if (_current_buffer_size > 0)
     _active_model->Reset(_current_sample_rate, _current_buffer_size);
 }
 
+void SlimmableWavenet::_stage_rebuild_model(const std::vector<int>& target_channels)
+{
+  if (target_channels == _current_channels && _active_model)
+  {
+    _pending_clear_release();
+    return;
+  }
+
+  if (auto pending = _pending_load_acquire())
+  {
+    if (pending->channels == target_channels)
+      return;
+  }
+
+  auto pack = std::make_shared<StagedSlimModel>();
+  pack->model = std::shared_ptr<DSP>(_create_wavenet_for_channels(target_channels));
+  pack->channels = target_channels;
+
+  if (_current_buffer_size > 0)
+    pack->model->Reset(_current_sample_rate, _current_buffer_size);
+
+  _pending_store_release(std::move(pack));
+}
+
 void SlimmableWavenet::process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames)
 {
+  if (auto pack = _pending_exchange_take_acq_rel())
+  {
+    _active_model = std::move(pack->model);
+    _current_channels = std::move(pack->channels);
+  }
   if (_active_model)
     _active_model->process(input, output, num_frames);
 }
@@ -389,6 +468,8 @@ void SlimmableWavenet::prewarm()
 {
   if (_active_model)
     _active_model->prewarm();
+  if (auto pending = _pending_load_acquire())
+    pending->model->prewarm();
 }
 
 void SlimmableWavenet::Reset(const double sampleRate, const int maxBufferSize)
@@ -397,6 +478,8 @@ void SlimmableWavenet::Reset(const double sampleRate, const int maxBufferSize)
   _current_buffer_size = maxBufferSize;
   if (_active_model)
     _active_model->Reset(sampleRate, maxBufferSize);
+  if (auto pending = _pending_load_acquire())
+    pending->model->Reset(sampleRate, maxBufferSize);
 }
 
 void SlimmableWavenet::SetSlimmableSize(const double val)
@@ -413,7 +496,7 @@ void SlimmableWavenet::SetSlimmableSize(const double val)
       target[i] = ratio_to_channels(val, allowed);
   }
 
-  _rebuild_model(target);
+  _stage_rebuild_model(target);
 }
 
 // ============================================================================

NAM/wavenet/slimmable.h

@@ -3,6 +3,12 @@
 #include <memory>
 #include <vector>
 
+#ifdef _LIBCPP_VERSION
+// libc++: std::atomic<std::shared_ptr<T>> is not viable; staging uses deprecated atomic_* free functions.
+#else
+#include <atomic>
+#endif
+
 #include "../dsp.h"
 #include "json.hpp"
 #include "../model_config.h"
@@ -19,7 +25,10 @@ namespace slimmable_wavenet
 /// Stores the full WaveNet LayerArrayParams and weights. Each layer array has its
 /// own allowed_channels list (from the "slimmable" config field). On SetSlimmableSize(),
 /// maps the ratio to a channel count per array, extracts a weight subset, builds
-/// modified LayerArrayParams, and reconstructs the WaveNet.
+/// modified LayerArrayParams, and constructs a replacement WaveNet published to a staging
+/// slot (release). process() takes the slot (acquire/release) and installs the new WaveNet
+/// before DSP. libc++ uses std::atomic_* free functions on shared_ptr; other STLs use
+/// std::atomic<std::shared_ptr<…>>.
 class SlimmableWavenet : public DSP, public SlimmableModel
 {
 public:
@@ -52,12 +61,34 @@ class SlimmableWavenet : public DSP, public SlimmableModel
   bool _with_head;
   nlohmann::json _condition_dsp_json;
   std::vector<float> _full_weights;
-  std::unique_ptr<DSP> _active_model;
+  /// Shared ownership so a staged model can be moved onto the active slot without
+  /// transferring from a concurrent writer.
+  std::shared_ptr<DSP> _active_model;
+
+  struct StagedSlimModel
+  {
+    std::shared_ptr<DSP> model;
+    std::vector<int> channels;
+  };
+#ifdef _LIBCPP_VERSION
+  /// Staged model; synchronized via deprecated std::atomic_* overloads for shared_ptr only.
+  std::shared_ptr<StagedSlimModel> _pending_staged;
+#else
+  std::atomic<std::shared_ptr<StagedSlimModel>> _pending_staged;
+#endif
+
   std::vector<int> _current_channels;
   int _current_buffer_size = 0;
   double _current_sample_rate = 0.0;
 
+  std::unique_ptr<DSP> _create_wavenet_for_channels(const std::vector<int>& target_channels);
   void _rebuild_model(const std::vector<int>& target_channels);
+  void _stage_rebuild_model(const std::vector<int>& target_channels);
+
+  void _pending_clear_release();
+  std::shared_ptr<StagedSlimModel> _pending_load_acquire() const;
+  void _pending_store_release(std::shared_ptr<StagedSlimModel> p);
+  std::shared_ptr<StagedSlimModel> _pending_exchange_take_acq_rel();
 };
 
 // Config / registration