diff --git a/NAM/wavenet/a2_fast.cpp b/NAM/wavenet/a2_fast.cpp index 08d52e72..3d7cf0c6 100644 --- a/NAM/wavenet/a2_fast.cpp +++ b/NAM/wavenet/a2_fast.cpp @@ -154,6 +154,7 @@ class A2FastModel : public DSP void _head_ring_write(int num_frames); void _layer_forward(int layer_idx, const float* cond, int num_frames); void _head_forward(float* output, int num_frames); + void _process_in_chunks(NAM_SAMPLE** input, NAM_SAMPLE** output, int num_frames); // Compile-time-specialized per-layer kernel. KernelSize is lifted to a // template parameter so clang can fully unroll the tap loop and schedule @@ -733,7 +734,10 @@ template void A2FastModel::process(NAM_SAMPLE** input, NAM_SAMPLE** output, int num_frames) { if (num_frames > GetMaxBufferSize()) - SetMaxBufferSize(num_frames); + { + _process_in_chunks(input, output, num_frames); + return; + } const NAM_SAMPLE* in0 = input[0]; NAM_SAMPLE* out0 = output[0]; @@ -763,6 +767,24 @@ void A2FastModel::process(NAM_SAMPLE** input, NAM_SAMPLE** output, int out0[f] = static_cast(head_out[f]); } +// A call with more frames than the maximum buffer size, as consecutive calls of at most that size. Growing the +// buffers instead would allocate on the audio thread, and SetMaxBufferSize() resets every history. +template +void A2FastModel::_process_in_chunks(NAM_SAMPLE** input, NAM_SAMPLE** output, int num_frames) +{ + if (GetMaxBufferSize() == 0) + { + SetMaxBufferSize(num_frames); // Never Reset(): no state to keep, so size for this call + } + const int max_frames = GetMaxBufferSize(); + for (int offset = 0; offset < num_frames; offset += max_frames) + { + NAM_SAMPLE* input_chunk = input[0] + offset; + NAM_SAMPLE* output_chunk = output[0] + offset; + A2FastModel::process(&input_chunk, &output_chunk, std::min(max_frames, num_frames - offset)); + } +} + // ----------------------------------------------------------------------------- // A2FastConfig — wraps the constructed DSP behind the ModelConfig interface. // ----------------------------------------------------------------------------- diff --git a/tools/run_tests.cpp b/tools/run_tests.cpp index ee827818..b7ca6d7c 100644 --- a/tools/run_tests.cpp +++ b/tools/run_tests.cpp @@ -411,6 +411,8 @@ int main() test_a2_fast::test_cached_prewarm_full(); test_a2_fast::test_process_realtime_safe_lite(); test_a2_fast::test_process_realtime_safe_full(); + test_a2_fast::test_oversized_call_matches_consecutive_calls_lite(); + test_a2_fast::test_oversized_call_matches_consecutive_calls_full(); #endif std::cout << "Success!" << std::endl; diff --git a/tools/test/test_a2_fast.cpp b/tools/test/test_a2_fast.cpp index 914c57d9..d51076ab 100644 --- a/tools/test/test_a2_fast.cpp +++ b/tools/test/test_a2_fast.cpp @@ -405,7 +405,7 @@ void test_process_realtime_safe(int channels) // Exercise several block sizes all within a single pre-sized state so the // internal "num_frames > max_buffer_size" guard in process() never fires - // (which would legitimately reallocate). + // (test_oversized_call_matches_consecutive_calls covers that path). const int max_buffer = 256; fast_dsp->Reset(48000.0, max_buffer); @@ -453,6 +453,51 @@ void test_process_realtime_safe_full() test_process_realtime_safe(8); } +// A call with more frames than the maximum buffer size must give exactly what the host gets by making calls of at +// most that size, without allocating. The fast path used to grow its buffers instead, which allocated on the audio +// thread and reset every history mid-stream. +void test_oversized_call_matches_consecutive_calls(int channels) +{ + const int max_buffer = 32; + const int call_size = 100; // Not a multiple of max_buffer, so each call's last chunk is short + const int num_calls = 20; + const auto cfg = build_a2_config(channels); + const auto weights = make_deterministic_weights(a2_weight_count(channels), /*seed=*/0xA2FA500u + channels); + std::vector> dsps; + for (int i = 0; i < 2; i++) + { + std::vector w = weights; + dsps.push_back(nam::wavenet::a2_fast::create_a2_fast_config(cfg, 48000.0)->create(std::move(w), 48000.0)); + dsps.back()->Reset(48000.0, max_buffer); + } + auto& consecutive = *dsps[0]; + auto& oversized = *dsps[1]; + + const auto input = make_test_input(num_calls * call_size, 48000.0); + std::vector actual(call_size); + for (int call = 0; call < num_calls; call++) + { + std::vector x(input.begin() + call * call_size, input.begin() + (call + 1) * call_size); + const auto expected = process_dsp(consecutive, x, max_buffer); + NAM_SAMPLE* in_ptr = x.data(); + NAM_SAMPLE* out_ptr = actual.data(); + const std::string test_name = "A2FastModel<" + std::to_string(channels) + ">::process oversized"; + allocation_tracking::run_allocation_test_no_allocations( + nullptr, [&]() { oversized.process(&in_ptr, &out_ptr, call_size); }, nullptr, test_name.c_str()); + assert(actual == expected); + } +} + +void test_oversized_call_matches_consecutive_calls_lite() +{ + test_oversized_call_matches_consecutive_calls(3); +} + +void test_oversized_call_matches_consecutive_calls_full() +{ + test_oversized_call_matches_consecutive_calls(8); +} + } // namespace test_a2_fast #endif // NAM_ENABLE_A2_FAST