dspk_ops/spatial/
car.rs

1use cubecl::prelude::*;
2use cubecl::server::Handle;
3use cubecl::tune::{LocalTuner, Tunable, TunableSet, local_tuner};
4use dspk_core::compute::tune::{size_class, tune_id};
5use dspk_core::compute::LaunchGeometry;
6
7use crate::core::{buffer, DspFloat};
8pub use super::kernels::{
9    common_median_bitonic_kernel, common_median_kernel, common_median_vector_kernel, common_median_vector_unaligned_kernel, direct_car_kernel,
10    direct_car_vector_kernel, direct_car_vector_unaligned_kernel,
11};
12
13/// High-level host dispatcher for direct, single-pass Common Average Referencing (CAR).
14/// Computes and subtracts the common average directly in VRAM without auxiliary buffers, with the
15/// faster of [`CarKernel`]'s choices for this device (same output).
16pub fn execute_direct_car<F: DspFloat>(
17    client: &Client,
18    input: &Handle,
19    output: &Handle,
20    channels: usize,
21    samples: usize,
22) {
23    execute_direct_car_with::<F>(client, input, output, channels, samples, None);
24}
25
26/// Which common-average kernel runs (both give the same output).
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum CarKernel {
29    /// One unit per sample ([`direct_car_kernel`]).
30    Direct,
31    /// `width` consecutive samples per unit ([`direct_car_vector_kernel`]; SIMD on the CPU
32    /// runtime). Runs the widest width up to `width` that the runtime offers (else `Direct`); when
33    /// it does not divide the sample count, [`direct_car_vector_unaligned_kernel`] (lane by lane
34    /// loads and stores).
35    Vectorized { width: usize },
36}
37
38/// [`execute_direct_car`] with a fixed kernel (`None` = autotuned); for tests and benchmarks.
39pub fn execute_direct_car_with<F: DspFloat>(
40    client: &Client,
41    input: &Handle,
42    output: &Handle,
43    channels: usize,
44    samples: usize,
45    kernel: Option<CarKernel>,
46) {
47    if channels == 0 || samples == 0 {
48        return;
49    }
50    let inputs = CmrInputs { client: client.clone(), input: input.clone(), output: output.clone(), channels, samples };
51    match kernel {
52        Some(kernel) => inputs.run_car::<F>(kernel),
53        None => tuned_car::<F>(inputs),
54    }
55}
56
57impl CmrInputs {
58    fn run_car<F: DspFloat>(&self, kernel: CarKernel) {
59        let client = &self.client;
60        let total = self.channels * self.samples;
61        let width = match kernel {
62            CarKernel::Vectorized { width } => client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1 && w <= width).max(),
63            CarKernel::Direct => None,
64        };
65        // SAFETY: both buffers hold `channels · samples` values of `F`
66        let (input, output) = unsafe { (BufferArg::from_raw_parts(self.input.clone(), total), BufferArg::from_raw_parts(self.output.clone(), total)) };
67        match width {
68            None => {
69                let geom = LaunchGeometry::per_sample(client, self.samples);
70                direct_car_kernel::launch::<F>(client, geom.cube_count, geom.cube_dim, input, output, self.channels as u32, self.samples as u32)
71            }
72            Some(width) if self.samples % width == 0 => {
73                let vectors = self.samples / width;
74                let geom = LaunchGeometry::elementwise(client, vectors);
75                direct_car_vector_kernel::launch::<F>(client, geom.cube_count, geom.cube_dim, width, input, output, self.channels as u32, vectors as u32)
76            }
77            Some(width) => {
78                let units = self.samples.div_ceil(width);
79                let geom = LaunchGeometry::elementwise(client, units);
80                direct_car_vector_unaligned_kernel::launch::<F>(
81                    client, geom.cube_count, geom.cube_dim, width, input, output, self.channels as u32, self.samples as u32, units as u32,
82                )
83            }
84        }
85    }
86}
87
88/// [`CmrInputs::run_car`] with the kernel CubeCL's autotuner found fastest for this device,
89/// element type and size. Benchmarks write to a scratch output.
90fn tuned_car<F: DspFloat>(inputs: CmrInputs) {
91    static TUNER: LocalTuner<String, String> = local_tuner!("common-average-v2");
92    let client = inputs.client.clone();
93    let id = tune_id(&client);
94    let widths: Vec<usize> = client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1).collect();
95    let set = TUNER.init(&id, move || {
96        let key = |p: &CmrInputs| format!("{}-c{}-t{}-a{}", F::type_name(), size_class(p.channels), size_class(p.samples), p.samples.trailing_zeros().min(4));
97        let scratch = |_: &String, p: &CmrInputs| CmrInputs { output: buffer::empty::<F>(&p.client, p.channels * p.samples), ..p.clone() };
98        let set: TunableSet<String, CmrInputs, ()> = TunableSet::new(key, scratch);
99        let kernels = std::iter::once(CarKernel::Direct).chain(widths.iter().map(|&width| CarKernel::Vectorized { width }));
100        kernels.fold(set, |set, kernel| {
101            set.with(Tunable::new(&format!("{kernel:?}"), move |p: CmrInputs| {
102                p.run_car::<F>(kernel);
103                Ok::<_, String>(())
104            }))
105        })
106    });
107    TUNER.execute(&id, &client, set, inputs)
108}
109
110/// How [`execute_common_median`] finds each sample's median over channels. All are exact and give
111/// the same output; which is faster depends on the device, so the autotuner picks.
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub enum CmrKernel {
114    /// Rank counting ([`common_median_kernel`]): each value compared with every channel, no barriers.
115    Rank,
116    /// Bitonic sort in shared memory ([`common_median_bitonic_kernel`]): about `log₂² C / 2` steps
117    /// per sample, one barrier each.
118    Bitonic,
119    /// Rank counting on `width` consecutive samples per unit ([`common_median_vector_kernel`];
120    /// SIMD on the CPU runtime). Runs the widest width up to `width` that the runtime offers (else
121    /// `Rank`); when it does not divide the sample count,
122    /// [`common_median_vector_unaligned_kernel`] (lane by lane loads and stores).
123    Vectorized { width: usize },
124}
125
126/// One common-median call (cloned per autotune candidate).
127#[derive(Clone)]
128struct CmrInputs {
129    client: Client,
130    input: Handle,
131    output: Handle,
132    channels: usize,
133    samples: usize,
134}
135
136impl CmrInputs {
137    fn run<F: DspFloat>(&self, kernel: CmrKernel) {
138        let client = &self.client;
139        if let CmrKernel::Vectorized { width } = kernel {
140            let width = client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1 && w <= width).max();
141            if let Some(width) = width.filter(|&w| self.samples % w != 0) {
142                let total = self.channels * self.samples;
143                let units = self.samples.div_ceil(width);
144                let geom = LaunchGeometry::elementwise(client, units);
145                // SAFETY: both buffers hold `channels · samples` values
146                unsafe {
147                    common_median_vector_unaligned_kernel::launch::<F>(
148                        client,
149                        geom.cube_count,
150                        geom.cube_dim,
151                        width,
152                        BufferArg::from_raw_parts(self.input.clone(), total),
153                        BufferArg::from_raw_parts(self.output.clone(), total),
154                        self.channels as u32,
155                        self.samples as u32,
156                        units as u32,
157                    );
158                }
159                return;
160            }
161            if let Some(width) = width {
162                let total = self.channels * self.samples;
163                let vectors = self.samples / width;
164                let geom = LaunchGeometry::elementwise(client, vectors);
165                // SAFETY: both buffers hold `channels · samples` values (`total / width` vectors)
166                unsafe {
167                    common_median_vector_kernel::launch::<F>(
168                        client,
169                        geom.cube_count,
170                        geom.cube_dim,
171                        width,
172                        BufferArg::from_raw_parts(self.input.clone(), total),
173                        BufferArg::from_raw_parts(self.output.clone(), total),
174                        self.channels as u32,
175                        vectors as u32,
176                    );
177                }
178                return;
179            }
180            return self.run::<F>(CmrKernel::Rank);
181        }
182        let geom = LaunchGeometry::per_row(client, self.samples, self.channels);
183        let total = self.channels * self.samples;
184        let (c, s, units) = (self.channels as u32, self.samples as u32, geom.cube_dim.x);
185        // SAFETY: both buffers hold `channels · samples` values
186        let (input, output) = unsafe {
187            (BufferArg::from_raw_parts(self.input.clone(), total), BufferArg::from_raw_parts(self.output.clone(), total))
188        };
189        match kernel {
190            CmrKernel::Rank | CmrKernel::Vectorized { .. } => common_median_kernel::launch::<F>(client, geom.cube_count, geom.cube_dim, input, output, c, s, c, units),
191            CmrKernel::Bitonic => {
192                let padded = self.channels.next_power_of_two() as u32;
193                common_median_bitonic_kernel::launch::<F>(client, geom.cube_count, geom.cube_dim, input, output, c, s, padded, units)
194            }
195        }
196    }
197}
198
199/// [`CmrInputs::run`] with the kernel CubeCL's autotuner found fastest for this device, element type
200/// and problem size. Benchmarks write to a scratch output.
201fn tuned<F: DspFloat>(inputs: CmrInputs) {
202    static TUNER: LocalTuner<String, String> = local_tuner!("common-median-v3");
203    let client = inputs.client.clone();
204    let id = tune_id(&client);
205    let widths: Vec<usize> = client.io_optimized_vector_sizes(size_of::<F>()).filter(|&w| w > 1).collect();
206    // CubeCL's CPU runtime emulates barriers: Bitonic takes seconds there, even just to time it
207    let cpu = client.name() == "cpu";
208    let set = TUNER.init(&id, move || {
209        // `a`: the power of two dividing the sample count (up to 16): which widths run aligned
210        let key = |p: &CmrInputs| format!("{}-c{}-t{}-a{}", F::type_name(), size_class(p.channels), size_class(p.samples), p.samples.trailing_zeros().min(4));
211        let scratch = |_: &String, p: &CmrInputs| CmrInputs { output: buffer::empty::<F>(&p.client, p.channels * p.samples), ..p.clone() };
212        let set: TunableSet<String, CmrInputs, ()> = TunableSet::new(key, scratch);
213        let kernels = [CmrKernel::Rank, CmrKernel::Bitonic].into_iter().filter(|&k| !(k == CmrKernel::Bitonic && cpu)).chain(widths.iter().map(|&width| CmrKernel::Vectorized { width }));
214        kernels.fold(set, |set, kernel| {
215            set.with(Tunable::new(&format!("{kernel:?}"), move |p: CmrInputs| {
216                p.run::<F>(kernel);
217                Ok::<_, String>(())
218            }))
219        })
220    });
221    TUNER.execute(&id, &client, set, inputs)
222}
223
224/// Common median reference: subtracts the median over channels at every sample (SpikeInterface's
225/// `common_reference(operator="median")`), one cube per sample on the device, with the faster of
226/// [`CmrKernel`]'s exact methods for this device.
227pub fn execute_common_median<F: DspFloat>(
228    client: &Client,
229    input: &Handle,
230    output: &Handle,
231    channels: usize,
232    samples: usize,
233) {
234    execute_common_median_with::<F>(client, input, output, channels, samples, None);
235}
236
237/// [`execute_common_median`] with a fixed kernel (`None` = autotuned); for tests and benchmarks.
238pub fn execute_common_median_with<F: DspFloat>(
239    client: &Client,
240    input: &Handle,
241    output: &Handle,
242    channels: usize,
243    samples: usize,
244    kernel: Option<CmrKernel>,
245) {
246    if channels == 0 || samples == 0 {
247        return;
248    }
249    let inputs = CmrInputs { client: client.clone(), input: input.clone(), output: output.clone(), channels, samples };
250    match kernel {
251        Some(kernel) => inputs.run::<F>(kernel),
252        None => tuned::<F>(inputs),
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::core::buffer;
260
261    /// Exact against a sorted reference, in `F` (signed fractional values, so the key's sign
262    /// handling and both middle values of an even count are exercised).
263    fn medians_of<F: DspFloat + std::fmt::Debug>(client: &Client) {
264        // 304 samples: every vector width divides; 300: widths above 4 run unaligned; 301: all do
265        for (channels, samples) in [(1usize, 300usize), (4, 304), (5, 301), (37, 300), (384, 304)] {
266            let x: Vec<F> = (0..channels * samples).map(|i| crate::core::cast::<F>(((i * 7919) % 211) as f64 * 0.37 - 40.0)).collect();
267            let half = crate::core::cast::<F>(0.5);
268            let expected: Vec<F> = (0..channels * samples)
269                .map(|i| {
270                    let (c, s) = (i / samples, i % samples);
271                    let mut col: Vec<F> = (0..channels).map(|c| x[c * samples + s]).collect();
272                    col.sort_by(|a, b| a.partial_cmp(b).expect("finite"));
273                    let m = if channels % 2 == 1 { col[channels / 2] } else { (col[channels / 2 - 1] + col[channels / 2]) * half };
274                    x[c * samples + s] - m
275                })
276                .collect();
277            let vectorized = [2, 4, 8, 16].map(|width| Some(CmrKernel::Vectorized { width }));
278            for kernel in [Some(CmrKernel::Rank), Some(CmrKernel::Bitonic), None].into_iter().chain(vectorized) {
279                let out = buffer::empty::<F>(client, x.len());
280                execute_common_median_with::<F>(client, &buffer::upload(client, &x), &out, channels, samples, kernel);
281                let got = buffer::download::<F>(client, out);
282                assert_eq!(got, expected, "{channels} channels, {kernel:?}");
283            }
284        }
285    }
286
287    fn medians(client: &Client) {
288        medians_of::<f32>(client);
289        // float64 where the runtime has it (not WebGPU)
290        if client.properties().supports_type(f64::elem_type_native()) {
291            medians_of::<f64>(client);
292        }
293    }
294    runtime_test!(test_common_median_reference, medians);
295
296    /// The vectorized common average is bit-identical to `Direct` (sample counts every width
297    /// divides, some, none).
298    fn averages(client: &Client) {
299        for (channels, samples) in [(1usize, 300usize), (5, 304), (37, 301), (384, 304)] {
300            let x: Vec<f32> = (0..channels * samples).map(|i| ((i * 7919) % 211) as f32 * 0.37 - 40.0).collect();
301            let input = buffer::upload(client, &x);
302            let run = |kernel| {
303                let out = buffer::empty::<f32>(client, x.len());
304                execute_direct_car_with::<f32>(client, &input, &out, channels, samples, Some(kernel));
305                buffer::download::<f32>(client, out)
306            };
307            let direct = run(CarKernel::Direct);
308            for width in [2usize, 4, 16] {
309                assert_eq!(run(CarKernel::Vectorized { width }), direct, "{} {channels} ch {samples} samples width {width}", client.name());
310            }
311        }
312    }
313    runtime_test!(test_common_average_vectorized_matches_direct, averages);
314}
dspk_ops/spatial/car.rs

dspk_ops/spatial/
car.rs
