nam_rs/models/linear_fft/process.rs
1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.
3
4//! Hot-path processing methods for the FFT convolution state.
5//!
6//! Separated from the state definition to keep the core struct and
7//! constructors in `linear_fft.rs` while isolating the RT-critical
8//! reset and tail block processing logic.
9
10use crate::math::common::Avx2Math;
11use crate::math::common::Avx512Math;
12use crate::math::common::Avx512VnniBf16Math;
13use crate::math::common::dispatch::InstructionSet;
14use crate::math::common::traits::SimdMath;
15
16impl super::LinearFftState {
17 /// Resets all runtime buffers to zero and re-initializes counters.
18 ///
19 /// This operation is allocation-free: it only zero-fills the existing
20 /// pre-allocated buffers. The pre-computed `h_fdl_*` spectra are
21 /// **not** modified (they depend only on the IR, which is static).
22 #[cold]
23 pub fn reset(&mut self) {
24 self.fdl_re.fill(0.0f32);
25 self.fdl_im.fill(0.0f32);
26 self.input_buf.fill(0.0f32);
27 self.fft_re.fill(0.0f32);
28 self.fft_im.fill(0.0f32);
29 self.acc_re.fill(0.0f32);
30 self.acc_im.fill(0.0f32);
31 self.output_buf.fill(0.0f32);
32 self.tail_output_buf.fill(0.0f32);
33 self.fdl_write_idx = self.num_partitions.saturating_sub(1);
34 self.sample_counter = 0;
35 }
36
37 /// Processes one block of tail convolution using overlap-save FFT.
38 ///
39 /// `input_window` must be a contiguous slice of the last `2P` input
40 /// samples (oldest to newest), typically obtained from the
41 /// `MirroredBuffer` via `history[write_pos - 2*P .. write_pos]`.
42 ///
43 /// After this call completes, `tail_output_buf` contains `P` valid
44 /// tail output samples ready for sequential per-sample consumption
45 /// via `sample_counter`.
46 ///
47 /// # Algorithm (overlap-save, zero-latency hybrid)
48 ///
49 /// 1. Compute forward RFFT of the `2P`-sample input window.
50 /// 2. Read past input spectra from the circular FDL (delays `P`,
51 /// `2P`, …, `K×P`), multiply by the corresponding pre-computed
52 /// tail IR spectra in the frequency domain using SIMD complex MAC.
53 /// 3. Store the new input spectrum in the FDL and advance the write
54 /// index (overwrites the oldest entry, now `K+1` blocks ago).
55 /// 4. Inverse RFFT the accumulated spectrum back to time domain.
56 /// 5. Extract the valid `P` output samples (overlap-save: indices
57 /// `P..2P-1`) into `tail_output_buf`.
58 ///
59 /// # RT-Safety
60 ///
61 /// Zero heap allocation, zero locks, zero panics in production
62 /// (debug assertions only). All buffers were pre-allocated at
63 /// construction time. The ISA for SIMD dispatch was captured once
64 /// at construction time — no runtime CPU feature checks on the hot
65 /// path.
66 pub fn process_tail_block(&mut self, input_window: &[f32]) {
67 let p = self.p;
68 let block_size = 2 * p;
69 let num_bins = self.num_bins;
70 let num_partitions = self.num_partitions;
71
72 debug_assert_eq!(input_window.len(), block_size);
73
74 if num_partitions == 0 {
75 return;
76 }
77
78 // ── Step 1: Copy input window and compute forward RFFT ──
79 self.input_buf[..block_size].copy_from_slice(input_window);
80 self.rfft
81 .process_forward(&self.input_buf, &mut self.fft_re, &mut self.fft_im);
82
83 // ── Step 2: Frequency-domain MAC ──
84 // The tail output for the NEXT block needs the current input spectrum
85 // (block B) for partition 0 (delay P) and FDL entries (blocks B-1,
86 // B-2, ...) for partitions 1..K-1 (delays 2P, 3P, ..., K×P).
87 // SAFETY: all slices have length num_bins, guaranteed by construction.
88 // ISA was captured at construction time.
89 self.acc_re[..num_bins].fill(0.0);
90 self.acc_im[..num_bins].fill(0.0);
91
92 // Partition 0 (delays P..2P−1): uses the current block's input spectrum.
93 unsafe {
94 match self.isa {
95 InstructionSet::Avx512VnniBf16 => Avx512VnniBf16Math::complex_mac_accumulate(
96 &self.h_fdl_re[..num_bins],
97 &self.h_fdl_im[..num_bins],
98 &self.fft_re[..num_bins],
99 &self.fft_im[..num_bins],
100 &mut self.acc_re[..num_bins],
101 &mut self.acc_im[..num_bins],
102 ),
103 InstructionSet::Avx512 => Avx512Math::complex_mac_accumulate(
104 &self.h_fdl_re[..num_bins],
105 &self.h_fdl_im[..num_bins],
106 &self.fft_re[..num_bins],
107 &self.fft_im[..num_bins],
108 &mut self.acc_re[..num_bins],
109 &mut self.acc_im[..num_bins],
110 ),
111 InstructionSet::Avx2 => Avx2Math::complex_mac_accumulate(
112 &self.h_fdl_re[..num_bins],
113 &self.h_fdl_im[..num_bins],
114 &self.fft_re[..num_bins],
115 &self.fft_im[..num_bins],
116 &mut self.acc_re[..num_bins],
117 &mut self.acc_im[..num_bins],
118 ),
119 }
120 }
121
122 // Partitions 1..K−1 (delays 2P..K×P): use past input spectra from FDL.
123 for k in 1..num_partitions {
124 let input_idx = (self.fdl_write_idx + num_partitions - k) % num_partitions;
125 let fdl_start = input_idx * num_bins;
126 let h_start = k * num_bins;
127
128 unsafe {
129 match self.isa {
130 InstructionSet::Avx512VnniBf16 => Avx512VnniBf16Math::complex_mac_accumulate(
131 &self.h_fdl_re[h_start..h_start + num_bins],
132 &self.h_fdl_im[h_start..h_start + num_bins],
133 &self.fdl_re[fdl_start..fdl_start + num_bins],
134 &self.fdl_im[fdl_start..fdl_start + num_bins],
135 &mut self.acc_re[..num_bins],
136 &mut self.acc_im[..num_bins],
137 ),
138 InstructionSet::Avx512 => Avx512Math::complex_mac_accumulate(
139 &self.h_fdl_re[h_start..h_start + num_bins],
140 &self.h_fdl_im[h_start..h_start + num_bins],
141 &self.fdl_re[fdl_start..fdl_start + num_bins],
142 &self.fdl_im[fdl_start..fdl_start + num_bins],
143 &mut self.acc_re[..num_bins],
144 &mut self.acc_im[..num_bins],
145 ),
146 InstructionSet::Avx2 => Avx2Math::complex_mac_accumulate(
147 &self.h_fdl_re[h_start..h_start + num_bins],
148 &self.h_fdl_im[h_start..h_start + num_bins],
149 &self.fdl_re[fdl_start..fdl_start + num_bins],
150 &self.fdl_im[fdl_start..fdl_start + num_bins],
151 &mut self.acc_re[..num_bins],
152 &mut self.acc_im[..num_bins],
153 ),
154 }
155 }
156 }
157
158 // ── Step 3: Store new input spectrum in FDL ──
159 {
160 let fdl_base = self.fdl_write_idx * num_bins;
161 self.fdl_re[fdl_base..fdl_base + num_bins].copy_from_slice(&self.fft_re);
162 self.fdl_im[fdl_base..fdl_base + num_bins].copy_from_slice(&self.fft_im);
163 }
164
165 // ── Step 4: Advance FDL write index ──
166 self.fdl_write_idx += 1;
167 if self.fdl_write_idx >= num_partitions {
168 self.fdl_write_idx = 0;
169 }
170
171 // ── Step 5: Inverse RFFT (complex → real, 2P samples) ──
172 self.rfft
173 .process_inverse(&mut self.acc_re, &mut self.acc_im, &mut self.output_buf);
174
175 // ── Step 6: Extract valid output (overlap-save: samples P..2P-1) ──
176 self.tail_output_buf[..p].copy_from_slice(&self.output_buf[p..block_size]);
177 }
178}