dspk_ops/resampler/
poly.rs

1//! Rational resampling (`scipy.signal.resample_poly`).
2
3use cubecl::prelude::*;
4use cubecl::server::Handle;
5use cubecl::tune::{LocalTuner, Tunable, TunableSet, local_tuner};
6use dspk_core::compute::LaunchGeometry;
7use dspk_core::compute::tune::{size_class, tune_id};
8
9use super::design::{firwin, FirWindow};
10use super::kernels::{upfirdn_kernel, upfirdn_vector_kernel, upfirdn_vector_unaligned_kernel};
11use crate::core::{buffer, cast_all, DspFloat, EdgeMode};
12
13/// `resample_poly`'s padding (`padtype="constant"`, zeros).
14pub const RESAMPLE_POLY_DEFAULT_EDGE: EdgeMode = EdgeMode::Zeros;
15
16/// Window `resample_poly` designs its filter with by default.
17pub const RESAMPLE_POLY_WINDOW: FirWindow = FirWindow::Kaiser { beta: 5.0 };
18
19/// Half-length of the designed filter per unit of the larger rate (`half_len = 10 · max(up, down)`).
20pub const RESAMPLE_POLY_HALF_LEN_PER_RATE: usize = 10;
21
22/// The anti-aliasing filter of [`execute_resample_poly`].
23#[derive(Debug, Clone, PartialEq)]
24pub enum ResampleFilter {
25    /// Designed with `firwin(2 · half_len + 1, 1 / max(up, down), window)`.
26    Window(FirWindow),
27    /// Explicit taps at the up-sampled rate (odd length; centred).
28    Taps(Vec<f64>),
29}
30
31/// Which polyphase kernel runs (every choice gives the same output).
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum ResampleKernel {
34    /// One unit per output.
35    Direct,
36    /// `width` consecutive outputs per unit (`Vector<F, width>`; SIMD on the CPU runtime). Runs the
37    /// widest width up to `width` that the runtime offers (else `Direct`); output rows it does not
38    /// divide store their outputs lane by lane.
39    Vectorized { width: usize },
40}
41
42impl Default for ResampleFilter {
43    fn default() -> Self {
44        ResampleFilter::Window(RESAMPLE_POLY_WINDOW)
45    }
46}
47
48fn gcd(a: usize, b: usize) -> usize {
49    if b == 0 { a } else { gcd(b, a % b) }
50}
51
52/// Samples out of `samples` resampled by `up / down` (`ceil(samples · up / down)`).
53pub fn resample_poly_len(samples: usize, up: usize, down: usize) -> usize {
54    let g = gcd(up, down).max(1);
55    (samples * (up / g)).div_ceil(down / g)
56}
57
58/// Resamples every channel of a `[channels, samples]` buffer of `F` by `up / down` into `output`
59/// (`[channels, resample_poly_len(samples, up, down)]`), as `scipy.signal.resample_poly`: the filter is
60/// scaled by `up` and aligned so output sample 0 sits on input sample 0. Returns the output length.
61///
62/// # Panics
63/// If `up` or `down` is zero.
64#[allow(clippy::too_many_arguments)]
65pub fn execute_resample_poly<F: DspFloat>(
66    client: &Client,
67    input: &Handle,
68    output: &Handle,
69    channels: usize,
70    samples: usize,
71    up: usize,
72    down: usize,
73    filter: &ResampleFilter,
74    edge: EdgeMode,
75) -> usize {
76    execute_resample_poly_with::<F>(client, input, output, channels, samples, up, down, filter, edge, None)
77}
78
79/// [`execute_resample_poly`] with a fixed kernel (`None` = autotuned); for tests and benchmarks.
80#[allow(clippy::too_many_arguments)]
81pub fn execute_resample_poly_with<F: DspFloat>(
82    client: &Client,
83    input: &Handle,
84    output: &Handle,
85    channels: usize,
86    samples: usize,
87    up: usize,
88    down: usize,
89    filter: &ResampleFilter,
90    edge: EdgeMode,
91    kernel: Option<ResampleKernel>,
92) -> usize {
93    assert!(up > 0 && down > 0, "resample_poly factors must be positive");
94    let g = gcd(up, down);
95    let (up, down) = (up / g, down / g);
96    let out_len = resample_poly_len(samples, up, down);
97    if channels == 0 || out_len == 0 {
98        return out_len;
99    }
100
101    let (mut taps, half_len) = match filter {
102        ResampleFilter::Window(window) => {
103            let half_len = RESAMPLE_POLY_HALF_LEN_PER_RATE * up.max(down);
104            (firwin(2 * half_len + 1, 1.0 / up.max(down) as f64, *window), half_len)
105        }
106        ResampleFilter::Taps(taps) => (taps.clone(), (taps.len().max(1) - 1) / 2),
107    };
108    for t in &mut taps {
109        *t *= up as f64;
110    }
111    // scipy's alignment: pad the filter by `pre_pad` zeros and drop the first `pre_remove` outputs
112    let pre_pad = down - half_len % down;
113    let pre_remove = (half_len + pre_pad) / down;
114    let start = (pre_remove * down) as i64 - pre_pad as i64;
115    // Farthest the taps reach outside the input, for non-zero edge modes
116    let s_min = (start - (taps.len() as i64 - 1)).div_euclid(up as i64);
117    let s_max = ((out_len as i64 - 1) * down as i64 + start).div_euclid(up as i64);
118    let ext_pad = (-s_min).max(s_max - (samples as i64 - 1)).max(0) as usize;
119
120    let inputs = PolyInputs {
121        client: client.clone(),
122        input: input.clone(),
123        output: output.clone(),
124        taps,
125        channels,
126        samples,
127        out_len,
128        up,
129        down,
130        start,
131        ext_pad,
132        edge,
133    };
134    match kernel {
135        Some(kernel) => inputs.run::<F>(kernel),
136        None => tuned::<F>(inputs),
137    }
138    out_len
139}
140
141/// One `resample_poly` call after design and alignment (cloned per autotune candidate).
142#[derive(Clone)]
143struct PolyInputs {
144    client: Client,
145    input: Handle,
146    output: Handle,
147    /// Scaled by `up`.
148    taps: Vec<f64>,
149    channels: usize,
150    samples: usize,
151    out_len: usize,
152    up: usize,
153    down: usize,
154    start: i64,
155    ext_pad: usize,
156    edge: EdgeMode,
157}
158
159impl PolyInputs {
160    fn run<F: DspFloat>(&self, kernel: ResampleKernel) {
161        let client = &self.client;
162        let width = match kernel {
163            ResampleKernel::Vectorized { width } => client
164                .io_optimized_vector_sizes(size_of::<F>())
165                .filter(|&w| w > 1 && w <= width)
166                .max(),
167            ResampleKernel::Direct => None,
168        };
169        let (in_total, out_total) = (self.channels * self.samples, self.channels * self.out_len);
170        let args = (self.samples as u32, self.out_len as u32, self.taps.len() as u32, self.up as u32, self.down as u32, self.start as i32, self.ext_pad as u32, self.edge.id());
171        match width {
172            None => {
173                let taps_h = buffer::upload(client, &cast_all::<F>(&self.taps));
174                let geom = LaunchGeometry::channels_samples(client, self.channels, self.out_len);
175                // SAFETY: input holds `in_total`, output `out_total` values of `F`
176                unsafe {
177                    upfirdn_kernel::launch::<F>(
178                        client,
179                        geom.cube_count,
180                        geom.cube_dim,
181                        BufferArg::from_raw_parts(self.input.clone(), in_total),
182                        BufferArg::from_raw_parts(taps_h, self.taps.len()),
183                        BufferArg::from_raw_parts(self.output.clone(), out_total),
184                        self.channels as u32,
185                        args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7,
186                    );
187                }
188            }
189            Some(width) => {
190                let table = buffer::upload(client, &cast_all::<F>(&poly_table(&self.taps, self.down, width)));
191                let table_len = (self.taps.len() + (width - 1) * self.down) * width;
192                // SAFETY: as above; `table` holds `table_len` values (`table_len / width` vectors)
193                let (input, table, output) = unsafe {
194                    (
195                        BufferArg::from_raw_parts(self.input.clone(), in_total),
196                        BufferArg::from_raw_parts(table, table_len),
197                        BufferArg::from_raw_parts(self.output.clone(), out_total),
198                    )
199                };
200                if self.out_len % width == 0 {
201                    let geom = LaunchGeometry::elementwise(client, out_total / width);
202                    upfirdn_vector_kernel::launch::<F>(
203                        client, geom.cube_count, geom.cube_dim, width, input, table, output,
204                        args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7,
205                    );
206                } else {
207                    // Rows the width does not divide: outputs stored lane by lane
208                    let per_row = self.out_len.div_ceil(width);
209                    let units = self.channels * per_row;
210                    let geom = LaunchGeometry::elementwise(client, units);
211                    upfirdn_vector_unaligned_kernel::launch::<F>(
212                        client, geom.cube_count, geom.cube_dim, width, input, table, output,
213                        args.0, args.1, args.2, args.3, args.4, args.5, args.6, per_row as u32, units as u32, args.7,
214                    );
215                }
216            }
217        }
218    }
219}
220
221/// The tap table of [`upfirdn_vector_kernel`]: `K + (width − 1) · down` vectors, lane `l` of vector
222/// `d + reach` holding `taps[d + l · down]` (zero outside), `reach = (width − 1) · down`.
223fn poly_table(taps: &[f64], down: usize, width: usize) -> Vec<f64> {
224    let reach = (width - 1) * down;
225    let mut table = Vec::with_capacity((taps.len() + reach) * width);
226    for e in 0..taps.len() + reach {
227        for l in 0..width {
228            let j = e as isize - reach as isize + (l * down) as isize;
229            table.push(if (0..taps.len() as isize).contains(&j) { taps[j as usize] } else { 0.0 });
230        }
231    }
232    table
233}
234
235/// [`PolyInputs::run`] with the kernel CubeCL's autotuner found fastest for this device, element
236/// type, rates, filter length and size. Benchmarks write to a scratch output.
237fn tuned<F: DspFloat>(inputs: PolyInputs) {
238    static TUNER: LocalTuner<String, String> = local_tuner!("resample-poly-v2");
239    let client = inputs.client.clone();
240    let id = tune_id(&client);
241    let widths: Vec<usize> = client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1).collect();
242    let set = TUNER.init(&id, move || {
243        let key = |p: &PolyInputs| {
244            // `a`: the power of two dividing the output row (up to 16): which widths run aligned
245            format!(
246                "{}-u{}-d{}-k{}-c{}-t{}-a{}",
247                F::type_name(),
248                p.up,
249                p.down,
250                p.taps.len(),
251                size_class(p.channels),
252                size_class(p.out_len),
253                p.out_len.trailing_zeros().min(4)
254            )
255        };
256        let scratch = |_: &String, p: &PolyInputs| PolyInputs { output: buffer::empty::<F>(&p.client, p.channels * p.out_len), ..p.clone() };
257        let set: TunableSet<String, PolyInputs, ()> = TunableSet::new(key, scratch);
258        let kernels = std::iter::once(ResampleKernel::Direct).chain(widths.iter().map(|&width| ResampleKernel::Vectorized { width }));
259        kernels.fold(set, |set, kernel| {
260            set.with(Tunable::new(&format!("{kernel:?}"), move |p: PolyInputs| {
261                p.run::<F>(kernel);
262                Ok::<_, String>(())
263            }))
264        })
265    });
266    TUNER.execute(&id, &client, set, inputs)
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    /// Host `resample_poly` straight from its definition: up-sample with zeros, convolve with the
274    /// padded filter, down-sample, keep `[pre_remove, pre_remove + out_len)`.
275    fn host_resample_poly(x: &[f64], up: usize, down: usize) -> Vec<f64> {
276        let half_len = RESAMPLE_POLY_HALF_LEN_PER_RATE * up.max(down);
277        let mut h = firwin(2 * half_len + 1, 1.0 / up.max(down) as f64, RESAMPLE_POLY_WINDOW);
278        h.iter_mut().for_each(|t| *t *= up as f64);
279        let pre_pad = down - half_len % down;
280        let pre_remove = (half_len + pre_pad) / down;
281        let h: Vec<f64> = std::iter::repeat_n(0.0, pre_pad).chain(h).collect();
282        let xu: Vec<f64> = (0..x.len() * up).map(|i| if i % up == 0 { x[i / up] } else { 0.0 }).collect();
283        let out_len = resample_poly_len(x.len(), up, down);
284        (0..out_len)
285            .map(|m| {
286                let n = (m + pre_remove) * down;
287                (0..h.len()).filter(|&k| k <= n && n - k < xu.len()).map(|k| h[k] * xu[n - k]).sum()
288            })
289            .collect()
290    }
291
292    fn matches_host(client: &Client) {
293        let samples = 997;
294        let x: Vec<f64> = (0..samples).map(|i| (i as f64 * 0.05).sin() * 10.0 + ((i * 7919) % 31) as f64 * 0.1).collect();
295        let input = buffer::upload(client, &x.iter().map(|v| *v as f32).collect::<Vec<_>>());
296        for (up, down) in [(1usize, 3usize), (3, 1), (2, 3), (160, 147)] {
297            let out_len = resample_poly_len(samples, up, down);
298            let output = buffer::empty::<f32>(client, out_len);
299            let n = execute_resample_poly::<f32>(client, &input, &output, 1, samples, up, down, &ResampleFilter::default(), RESAMPLE_POLY_DEFAULT_EDGE);
300            assert_eq!(n, out_len);
301            let got = buffer::download::<f32>(client, output);
302            let want = host_resample_poly(&x, up, down);
303            for (m, (g, w)) in got.iter().zip(&want).enumerate() {
304                assert!((*g as f64 - w).abs() < 1e-3, "{} {up}/{down} sample {m}: {g} vs {w}", client.name());
305            }
306        }
307    }
308    runtime_test!(test_resample_poly_matches_definition, matches_host);
309
310    /// The vectorized kernel is bit-identical to `Direct` (f32 / f64, rates up and down, every
311    /// edge, output rows divisible by the width or not).
312    fn vectorized_matches_direct_of<F: DspFloat + std::fmt::Debug>(client: &Client) {
313        let channels = 3;
314        for samples in [96usize, 997, 1_024] {
315            let x: Vec<F> = (0..channels * samples).map(|i| crate::core::cast::<F>(((i * 7919) % 211) as f64 * 0.37 - 40.0)).collect();
316            let input = buffer::upload(client, &x);
317            for (up, down) in [(1usize, 3usize), (3, 1), (2, 3), (1, 2), (160, 147)] {
318                for edge in [EdgeMode::Zeros, EdgeMode::Odd, EdgeMode::Reflect, EdgeMode::Nearest] {
319                    let out_len = resample_poly_len(samples, up, down);
320                    let run = |kernel| {
321                        let out = buffer::empty::<F>(client, channels * out_len);
322                        execute_resample_poly_with::<F>(client, &input, &out, channels, samples, up, down, &ResampleFilter::default(), edge, Some(kernel));
323                        buffer::download::<F>(client, out)
324                    };
325                    let direct = run(ResampleKernel::Direct);
326                    for width in [2usize, 4, 8, 16] {
327                        assert_eq!(run(ResampleKernel::Vectorized { width }), direct, "{} {} {samples} samples {up}/{down} {edge:?} width {width}", client.name(), F::type_name());
328                    }
329                }
330            }
331        }
332    }
333
334    fn vectorized_matches_direct(client: &Client) {
335        vectorized_matches_direct_of::<f32>(client);
336        if client.properties().supports_type(f64::elem_type_native()) {
337            vectorized_matches_direct_of::<f64>(client);
338        }
339    }
340    runtime_test!(test_resample_poly_vectorized_matches_direct, vectorized_matches_direct);
341}
dspk_ops/resampler/poly.rs

dspk_ops/resampler/
poly.rs
