diff --git a/src/asr/diar/diar_pipeline.cpp b/src/asr/diar/diar_pipeline.cpp index 4303dce..0d3894e 100644 --- a/src/asr/diar/diar_pipeline.cpp +++ b/src/asr/diar/diar_pipeline.cpp @@ -264,19 +264,18 @@ DiarStream::run_one_chunk(bool force, bool final_flush) { // Trim consumed buffers. The next window starts at mel_consumed_ - // lc_mel_max; its first FE sample sits n_fft/2 before that frame's hop - // position. + // position. compact_front erases lazily, so a whole-file push costs + // amortized O(N) moves instead of one full-buffer memmove per chunk. The + // buffers keep holding [base, produced), so every index below stays valid. const int64_t keep_mel = std::max(mel_consumed_ - lc_mel_max, 0); - if (keep_mel > mel_base_) { - mel_buf_.erase( - mel_buf_.begin(), mel_buf_.begin() + (keep_mel - mel_base_) * m_.fe().n_mels()); + if (keep_mel > mel_base_ && + compact_front(mel_buf_, static_cast(keep_mel - mel_base_) * m_.fe().n_mels())) mel_base_ = keep_mel; - } const int64_t keep_sample = std::max(keep_mel * m_.fe().hop_length() - m_.fe().n_fft() / 2, 0); - if (keep_sample > static_cast(audio_base_)) { - audio_buf_.erase(audio_buf_.begin(), audio_buf_.begin() + (keep_sample - audio_base_)); + if (keep_sample > static_cast(audio_base_) && + compact_front(audio_buf_, static_cast(keep_sample) - audio_base_)) audio_base_ = static_cast(keep_sample); - } return true; } diff --git a/src/asr/diar/diar_pipeline.h b/src/asr/diar/diar_pipeline.h index a753b7e..99222e5 100644 --- a/src/asr/diar/diar_pipeline.h +++ b/src/asr/diar/diar_pipeline.h @@ -12,6 +12,7 @@ // serializes on the BackendManager mutex. #pragma once +#include #include #include #include @@ -117,6 +118,18 @@ std::vector diar_segments_from_probs( const float* probs, int64_t n_frames, int n_spk, double sec_per_frame, const DiarSegmentationCfg& cfg); +// Drops the first `dead` elements of `buf`, but only once they are at least as +// many as the elements after them, so trimming after every chunk costs +// amortized O(1) moves per element. Returns whether it erased. +template +bool +compact_front(std::vector& buf, size_t dead) { + if (dead < buf.size() - dead) + return false; + buf.erase(buf.begin(), buf.begin() + static_cast(dead)); + return true; +} + // Per-stream streaming state + timeline. class DiarStream { public: diff --git a/src/asr/features/fe.cpp b/src/asr/features/fe.cpp index 08a85c2..382a535 100644 --- a/src/asr/features/fe.cpp +++ b/src/asr/features/fe.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include "nvtx_utils.h" @@ -36,40 +37,74 @@ mel_to_hz(float m) { return 700.0f * (std::exp(m / 1127.0f) - 1.0f); } -// Iterative radix-2 FFT over real and imaginary arrays of equal length. +// Bit-reversal swap list and, for each stage of length len, its len/2 twiddles +// stored at offset len/2 - 1. Twiddles are generated by the same per-block +// recurrence (cur *= w, restarted at 1 for every block) the butterflies used to +// run inline, so every butterfly sees bit-identical twiddle values. static void -fft_radix2_inplace(std::vector& re, std::vector& im) { - const int n = static_cast(re.size()); +make_fft_tables( + int n, std::vector>& swaps, std::vector& tw_re, + std::vector& tw_im) { + swaps.clear(); + tw_re.clear(); + tw_im.clear(); if (n <= 1) return; for (int i = 1, j = 0; i < n; i++) { int bit = n >> 1; for (; j & bit; bit >>= 1) j ^= bit; j ^= bit; - if (i < j) { - std::swap(re[i], re[j]); - std::swap(im[i], im[j]); - } + if (i < j) + swaps.emplace_back(i, j); } + tw_re.assign(static_cast(n - 1), 0.0f); + tw_im.assign(static_cast(n - 1), 0.0f); for (int len = 2; len <= n; len <<= 1) { const float ang = -2.0f * kPi / len; const float wre = std::cos(ang); const float wim = std::sin(ang); + float cur_re = 1.0f, cur_im = 0.0f; + const int off = len / 2 - 1; + for (int k = 0; k < len / 2; k++) { + tw_re[off + k] = cur_re; + tw_im[off + k] = cur_im; + const float nre = cur_re * wre - cur_im * wim; + const float nim = cur_re * wim + cur_im * wre; + cur_re = nre; + cur_im = nim; + } + } +} + +// Iterative radix-2 FFT over real and imaginary arrays of equal length, using +// tables from make_fft_tables for that length. +static void +fft_radix2_inplace( + const std::vector>& swaps, const std::vector& tw_re_all, + const std::vector& tw_im_all, std::vector& re, std::vector& im) { + const int n = static_cast(re.size()); + if (n <= 1) + return; + for (const auto& s : swaps) { + std::swap(re[s.first], re[s.second]); + std::swap(im[s.first], im[s.second]); + } + for (int len = 2; len <= n; len <<= 1) { + const int half = len / 2; + const float* tw_re = tw_re_all.data() + (half - 1); + const float* tw_im = tw_im_all.data() + (half - 1); for (int i = 0; i < n; i += len) { - float cur_re = 1.0f, cur_im = 0.0f; - for (int k = 0; k < len / 2; k++) { + for (int k = 0; k < half; k++) { + const float cur_re = tw_re[k]; + const float cur_im = tw_im[k]; const float ur = re[i + k]; const float ui = im[i + k]; - const float vr = re[i + k + len / 2] * cur_re - im[i + k + len / 2] * cur_im; - const float vi = re[i + k + len / 2] * cur_im + im[i + k + len / 2] * cur_re; + const float vr = re[i + k + half] * cur_re - im[i + k + half] * cur_im; + const float vi = re[i + k + half] * cur_im + im[i + k + half] * cur_re; re[i + k] = ur + vr; im[i + k] = ui + vi; - re[i + k + len / 2] = ur - vr; - im[i + k + len / 2] = ui - vi; - const float nre = cur_re * wre - cur_im * wim; - const float nim = cur_re * wim + cur_im * wre; - cur_re = nre; - cur_im = nim; + re[i + k + half] = ur - vr; + im[i + k + half] = ui - vi; } } } @@ -266,6 +301,7 @@ MelSpectrogramExtractor::MelSpectrogramExtractor( } init_window(); init_mel_basis(); + init_cpu_tables(); if (bm_ != nullptr && bm_->gpu_backend_handle() != nullptr) { module_ = std::make_unique(cfg_, window_, mel_basis_); @@ -295,6 +331,27 @@ MelSpectrogramExtractor::batch_metrics() const { return batcher_ ? batcher_->metrics() : nemo_speech::asr::BatchMetrics{}; } +void +MelSpectrogramExtractor::init_cpu_tables() { + make_fft_tables(cfg_.n_fft, fft_swaps_, fft_tw_re_, fft_tw_im_); + // Triangular filterbank rows are nonzero only over a short contiguous span. + // Restricting each row's dot product to [kbeg, kend) between its first and + // last exactly-nonzero entry keeps the accumulation order, and the skipped + // terms add exact zeros, so the output is bit-identical to the dense product. + const int n_bins = cfg_.n_fft / 2 + 1; + mel_kbeg_.assign(cfg_.n_mels, 0); + mel_kend_.assign(cfg_.n_mels, 0); + for (int m = 0; m < cfg_.n_mels; m++) { + const float* row = mel_basis_.data() + static_cast(m) * n_bins; + int b = 0; + while (b < n_bins && row[b] == 0.0f) b++; + int e = n_bins; + while (e > b && row[e - 1] == 0.0f) e--; + mel_kbeg_[m] = b; + mel_kend_[m] = e; + } +} + void MelSpectrogramExtractor::init_window() { const int win = win_length(); @@ -352,6 +409,7 @@ MelSpectrogramExtractor::set_mel_basis(const float* fb, int n_mels, int n_bins) std::to_string(cfg_.n_mels) + ", " + std::to_string(expected_bins) + ")"); } mel_basis_.assign(fb, fb + (size_t)n_mels * (size_t)n_bins); + init_cpu_tables(); if (module_ && session_) { module_->update_mel_basis(session_.get(), fb, n_mels, n_bins); } @@ -466,15 +524,16 @@ MelSpectrogramExtractor::compute_padded( re[woff + i] = sample * window_[i]; } - fft_radix2_inplace(re, im); + fft_radix2_inplace(fft_swaps_, fft_tw_re_, fft_tw_im_, re, im); - // Power spectrum -> mel + // Power spectrum (once per frame, in place in re) -> mel. re is + // refilled at the start of every frame, so it can hold the power here. + float* power = re.data(); + for (int k = 0; k < n_bins; k++) power[k] = re[k] * re[k] + im[k] * im[k]; for (int m = 0; m < cfg_.n_mels; m++) { + const float* row = mel_basis_.data() + static_cast(m) * n_bins; float acc = 0.0f; - for (int k = 0; k < n_bins; k++) { - const float power = re[k] * re[k] + im[k] * im[k]; - acc += mel_basis_[m * n_bins + k] * power; - } + for (int k = mel_kbeg_[m]; k < mel_kend_[m]; k++) acc += row[k] * power[k]; features[static_cast(m) + static_cast(f) * cfg_.n_mels] = std::log(acc + cfg_.log_zero_guard); } diff --git a/src/asr/features/fe.h b/src/asr/features/fe.h index dcc6c19..b38d70c 100644 --- a/src/asr/features/fe.h +++ b/src/asr/features/fe.h @@ -9,6 +9,7 @@ #include #include #include +#include #include #include "batching.h" @@ -122,8 +123,14 @@ class MelSpectrogramExtractor { std::vector window_; // Hann window of length win_length std::vector mel_basis_; // (n_mels, n_fft/2 + 1) row-major + // CPU fast-path tables, built once and rebuilt whenever mel_basis_ changes. + std::vector> fft_swaps_; // bit-reversal swap list + std::vector fft_tw_re_, fft_tw_im_; // per-stage FFT twiddles + std::vector mel_kbeg_, mel_kend_; // nonzero span of each mel row + void init_window(); void init_mel_basis(); + void init_cpu_tables(); // GPU log-mel for one window - runs through session_->run(). Output // layout matches the CPU path: `out` is resized to (n_mels * n_frames), diff --git a/tests/cpp/asr/CMakeLists.txt b/tests/cpp/asr/CMakeLists.txt index 2a7b054..a922ed3 100644 --- a/tests/cpp/asr/CMakeLists.txt +++ b/tests/cpp/asr/CMakeLists.txt @@ -55,6 +55,10 @@ add_executable(test_diar_state test_diar_state.cpp) target_link_libraries(test_diar_state PRIVATE nemo_speech_asr) add_test(NAME diar_state COMMAND test_diar_state) +add_executable(test_mel_features test_mel_features.cpp) +target_link_libraries(test_mel_features PRIVATE nemo_speech_asr) +add_test(NAME mel_features COMMAND test_mel_features) + add_executable(test_diar_recognizer test_diar_recognizer.cpp) target_link_libraries(test_diar_recognizer PRIVATE nemo_speech_asr) diff --git a/tests/cpp/asr/test_diar_state.cpp b/tests/cpp/asr/test_diar_state.cpp index 0ffd95a..2c3e410 100644 --- a/tests/cpp/asr/test_diar_state.cpp +++ b/tests/cpp/asr/test_diar_state.cpp @@ -1,5 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#include #include #include #include @@ -213,6 +214,56 @@ test_learned_silence_embedding() { return state.n_sil_frames() == 0 && state.mean_sil_emb() == learned; } +// Replays DiarStream's buffer trim: `pushed` samples arrive `push` at a time, +// each chunk consumes `chunk` and keeps `lookback` before the next one. Every +// sample stores its global index, so a wrong base shows up as a wrong value. +// Returns the samples moved by erases; `peak` gets the largest buffer seen. +size_t +replay_trim(size_t pushed, size_t push, size_t chunk, size_t lookback, size_t& peak) { + std::vector buf; + size_t base = 0, produced = 0, consumed = 0, moved = 0; + peak = 0; + while (consumed + chunk <= pushed) { + while (produced < consumed + chunk) { + const size_t n = std::min(push, pushed - produced); + for (size_t i = 0; i < n; i++) buf.push_back(static_cast(produced + i)); + produced += n; + } + peak = std::max(peak, buf.size()); + require( + buf.size() == produced - base && buf[consumed - base] == static_cast(consumed), + "compaction: the buffer holds [base, produced)"); + consumed += chunk; + const size_t keep = consumed > lookback ? consumed - lookback : 0; + if (keep > base && compact_front(buf, keep - base)) { + moved += buf.size(); + base = keep; + } + } + return moved; +} + +void +test_buffer_compaction() { + std::vector buf{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + require(!compact_front(buf, 4) && buf.size() == 10, "compaction: a short dead prefix stays"); + require( + compact_front(buf, 5) && buf == std::vector({5, 6, 7, 8, 9}), + "compaction: a dead prefix as long as the live tail is erased"); + + // A whole 10 minute file pushed at once: erasing after every 160 ms chunk + // would move about 1.8e10 samples; the lazy rule must stay linear. + const size_t file = 600 * 16000, chunk = 2560, lookback = 16000; + size_t peak = 0; + require( + replay_trim(file, file, chunk, lookback, peak) <= file, + "compaction: a whole-file push moves each sample at most once on average"); + + // Live audio pushed 160 ms at a time: the buffer must stay bounded. + replay_trim(file, chunk, chunk, lookback, peak); + require(peak <= 2 * (lookback + 2 * chunk), "compaction: a live stream's buffer stays bounded"); +} + } // namespace int @@ -225,6 +276,7 @@ main() { test_geometry_and_word_cadence(); test_finite_parity(); test_invalid_rotary_geometry(); + test_buffer_compaction(); if (!test_learned_silence_embedding()) { std::fprintf(stderr, "[FAIL] learned silence embedding was not retained\n"); return 1; diff --git a/tests/cpp/asr/test_mel_features.cpp b/tests/cpp/asr/test_mel_features.cpp new file mode 100644 index 0000000..4642d20 --- /dev/null +++ b/tests/cpp/asr/test_mel_features.cpp @@ -0,0 +1,193 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// The CPU log-mel path must match a plain reference: an FFT with its twiddles +// computed inline and a dense filterbank product. Covers the generated +// filterbank and one set with set_mel_basis, including rows the sparse span has +// to handle at its edges. Builds that fuse multiply-adds (GCC by default, and +// clang on arm64) may fuse them differently here and in the library, so the +// comparison allows 1e-4, ten times below what the bugs this guards against +// produce. + +#include +#include +#include +#include +#include +#include + +#include "fe.h" +#include "numeric_parity.h" + +namespace { + +constexpr float kPi = 3.14159265358979323846f; + +void +reference_fft(std::vector& re, std::vector& im) { + const int n = static_cast(re.size()); + for (int i = 1, j = 0; i < n; i++) { + int bit = n >> 1; + for (; j & bit; bit >>= 1) j ^= bit; + j ^= bit; + if (i < j) { + std::swap(re[i], re[j]); + std::swap(im[i], im[j]); + } + } + for (int len = 2; len <= n; len <<= 1) { + const float ang = -2.0f * kPi / len; + const float wre = std::cos(ang); + const float wim = std::sin(ang); + for (int i = 0; i < n; i += len) { + float cur_re = 1.0f, cur_im = 0.0f; + for (int k = 0; k < len / 2; k++) { + const float ur = re[i + k]; + const float ui = im[i + k]; + const float vr = re[i + k + len / 2] * cur_re - im[i + k + len / 2] * cur_im; + const float vi = re[i + k + len / 2] * cur_im + im[i + k + len / 2] * cur_re; + re[i + k] = ur + vr; + im[i + k] = ui + vi; + re[i + k + len / 2] = ur - vr; + im[i + k + len / 2] = ui - vi; + const float nre = cur_re * wre - cur_im * wim; + const float nim = cur_re * wim + cur_im * wre; + cur_re = nre; + cur_im = nim; + } + } + } +} + +// The filterbank MelSpectrogramExtractor generates when a model has none. +std::vector +generated_basis(const MelSpecConfig& cfg) { + const int n_bins = cfg.n_fft / 2 + 1; + const float fmax = cfg.fmax > 0.0f ? cfg.fmax : 0.5f * cfg.sample_rate; + auto hz_to_mel = [](float f) { return 1127.0f * std::log(1.0f + f / 700.0f); }; + auto mel_to_hz = [](float m) { return 700.0f * (std::exp(m / 1127.0f) - 1.0f); }; + const float mel_min = hz_to_mel(cfg.fmin); + const float mel_max = hz_to_mel(fmax); + std::vector points(cfg.n_mels + 2); + for (int i = 0; i < cfg.n_mels + 2; i++) + points[i] = mel_to_hz(mel_min + (mel_max - mel_min) * i / (cfg.n_mels + 1)); + std::vector basis(static_cast(cfg.n_mels) * n_bins, 0.0f); + for (int m = 0; m < cfg.n_mels; m++) { + for (int k = 0; k < n_bins; k++) { + const float f = static_cast(k) * cfg.sample_rate / cfg.n_fft; + float w = 0.0f; + if (f >= points[m] && f <= points[m + 1]) + w = (f - points[m]) / (points[m + 1] - points[m] + 1e-12f); + else if (f >= points[m + 1] && f <= points[m + 2]) + w = (points[m + 2] - f) / (points[m + 2] - points[m + 1] + 1e-12f); + basis[static_cast(m) * n_bins + k] = std::max(0.0f, w); + } + } + return basis; +} + +// Unnormalized log-mel with reflect_left=true, as compute() produces it. +std::vector +reference_log_mel( + const MelSpecConfig& cfg, const std::vector& basis, const std::vector& audio) { + const int win = static_cast(cfg.window_size * cfg.sample_rate + 0.5f); + const int hop = static_cast(cfg.window_stride * cfg.sample_rate + 0.5f); + const int n_fft = cfg.n_fft; + const int n_bins = n_fft / 2 + 1; + std::vector window(win); + const float denom = cfg.hann_periodic ? static_cast(win) : static_cast(win - 1); + for (int i = 0; i < win; i++) window[i] = 0.5f * (1.0f - std::cos(2.0f * kPi * i / denom)); + + std::vector pre(audio); + for (size_t i = pre.size() - 1; i > 0; --i) pre[i] = pre[i] - cfg.preemph * pre[i - 1]; + std::vector padded(n_fft / 2, 0.0f); + padded.insert(padded.end(), pre.begin(), pre.end()); + padded.insert(padded.end(), n_fft / 2, 0.0f); + + const int n_frames = static_cast((padded.size() - n_fft) / hop + 1); + std::vector features(static_cast(cfg.n_mels) * n_frames); + std::vector re(n_fft), im(n_fft); + const int woff = cfg.stft_center_window ? (n_fft - win) / 2 : 0; + for (int f = 0; f < n_frames; f++) { + std::fill(re.begin(), re.end(), 0.0f); + std::fill(im.begin(), im.end(), 0.0f); + for (int i = 0; i < win; i++) re[woff + i] = padded[f * hop + woff + i] * window[i]; + reference_fft(re, im); + for (int m = 0; m < cfg.n_mels; m++) { + float acc = 0.0f; + for (int k = 0; k < n_bins; k++) { + const float power = re[k] * re[k] + im[k] * im[k]; + acc += basis[static_cast(m) * n_bins + k] * power; + } + features[static_cast(m) + static_cast(f) * cfg.n_mels] = + std::log(acc + cfg.log_zero_guard); + } + } + return features; +} + +constexpr float kTolerance = 1e-4f; + +std::vector +noise(size_t n, unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-0.5f, 0.5f); + std::vector audio(n); + for (float& s : audio) s = dist(rng); + return audio; +} + +void +check_config(const MelSpecConfig& cfg, const char* name) { + MelSpectrogramExtractor fe(cfg); + const std::vector audio = noise(20800, 7); + std::vector features; + int n_frames = 0; + + fe.compute(audio.data(), audio.size(), features, n_frames, true, false); + float diff = finite_max_abs_diff(features, reference_log_mel(cfg, generated_basis(cfg), audio)); + if (!(diff <= kTolerance)) { + std::fprintf(stderr, "[FAIL] %s: generated filterbank off by %g\n", name, diff); + throw std::runtime_error("mel features: generated filterbank"); + } + + // Replace the basis on the same extractor: the nonzero spans must follow. + const int n_bins = cfg.n_fft / 2 + 1; + std::vector basis = generated_basis(cfg); + std::reverse(basis.begin(), basis.end()); + std::fill_n(basis.begin(), n_bins, 0.0f); // an all-zero row + std::fill_n(basis.begin() + n_bins, n_bins, 0.25f); // a dense row + basis[2 * static_cast(n_bins) + n_bins - 1] = 1.0f; // nonzero in the last bin + basis[3 * static_cast(n_bins)] = 1.0f; // nonzero in the first bin + basis[3 * static_cast(n_bins) + n_bins / 2] = 0.5f; // with a gap in between + fe.set_mel_basis(basis.data(), cfg.n_mels, n_bins); + fe.compute(audio.data(), audio.size(), features, n_frames, true, false); + diff = finite_max_abs_diff(features, reference_log_mel(cfg, basis, audio)); + if (!(diff <= kTolerance)) { + std::fprintf(stderr, "[FAIL] %s: set_mel_basis output off by %g\n", name, diff); + throw std::runtime_error("mel features: set_mel_basis"); + } +} + +} // namespace + +int +main() { + MelSpecConfig asr; // legacy ASR: periodic Hann, right-aligned window + check_config(asr, "asr"); + + MelSpecConfig nemo; // NeMo parity: symmetric Hann, centered window + nemo.n_mels = 128; + nemo.hann_periodic = false; + nemo.stft_center_window = true; + check_config(nemo, "nemo"); + + MelSpecConfig wide; + wide.n_fft = 1024; + wide.n_mels = 64; + wide.fmin = 20.0f; + wide.fmax = 7600.0f; + check_config(wide, "n_fft 1024"); + + std::printf("[PASS] CPU log-mel matches the reference\n"); + return 0; +}