Skip to main content

dspk_ops/linalg/kernels/
subspace.rs

1use cubecl::prelude::*;
2use dspk_core::compute::element_position;
3
4/// `x[e] −= mean[e mod cols]`: rows of a contiguous `[rows, cols]` buffer centred on the column
5/// means, in place. One unit per element ([`fn@element_position`]).
6#[cube(launch)]
7pub fn centre_columns_kernel<F: Float>(x: &mut [F], mean: &[F], cols: u32, total: u32) {
8    let e = element_position() as u32;
9    if e < total {
10        x[e as usize] = x[e as usize] - mean[(e % cols) as usize];
11    }
12}
13
14/// `acc[e] += x[e]`. One unit per element ([`fn@element_position`]).
15#[cube(launch)]
16pub fn add_assign_kernel<F: Float>(acc: &mut [F], x: &[F], total: u32) {
17    let e = element_position() as u32;
18    if e < total {
19        acc[e as usize] = acc[e as usize] + x[e as usize];
20    }
21}
22