dspk_ops/spatial/kernels/
sparse.rs

1use cubecl::prelude::*;
2use dspk_core::compute::{channel_position, sample_position};
3
4/// Spatial projection with at most `width` nonzero weights per output channel (ELLPACK rows):
5/// `output[o, t] = Σ_k values[o, k] · input[indices[o, k], t]`. Padding entries carry weight 0.
6/// One unit per `(output channel, sample)`.
7#[cube(launch)]
8pub fn sparse_rows_multiply_kernel<F: Float>(
9    input: &[F],
10    values: &[F],
11    indices: &[u32],
12    output: &mut [F],
13    num_channels: u32,
14    num_samples: u32,
15    width: u32,
16) {
17    let t = sample_position();
18    let out_ch = channel_position();
19
20    if t < num_samples && out_ch < num_channels {
21        let row = out_ch * width;
22        let mut acc = F::new(0.0f32);
23        let mut k = 0u32;
24        while k < width {
25            let in_ch = indices[(row + k) as usize];
26            acc += values[(row + k) as usize] * input[(in_ch * num_samples + t) as usize];
27            k += 1u32;
28        }
29        output[(out_ch * num_samples + t) as usize] = acc;
30    }
31}
dspk_ops/spatial/kernels/sparse.rs

dspk_ops/spatial/kernels/
sparse.rs
