dspk_ops/linalg/kernels/
covariance.rs

1use cubecl::prelude::*;
2use dspk_core::compute::element_position;
3
4/// Copies columns `0..cols` of `rows` strided rows into a split-major buffer `[splits, rows,
5/// split]`: `output[(s · rows + r) · split + j] = input[offset + r · row_stride + t] − mean[r]` with
6/// `t = s · split + j`, zero where `t ≥ cols` (the padded end of the last split; zero columns add
7/// nothing to a product with itself). `centred = false` skips the mean (`mean` is not read). The
8/// offset is applied here, so the input binds at the buffer's start. One unit per output element
9/// ([`fn@element_position`]).
10#[cube(launch)]
11#[allow(clippy::too_many_arguments)]
12pub fn split_rows_kernel<F: Float>(
13    input: &[F],
14    mean: &[F],
15    output: &mut [F],
16    offset: u32,
17    row_stride: u32,
18    cols: u32,
19    rows: u32,
20    split: u32,
21    total: u32,
22    #[comptime] centred: bool,
23) {
24    let e = element_position() as u32;
25    if e < total {
26        let j = e % split;
27        let sr = e / split;
28        let r = sr % rows;
29        let s = sr / rows;
30        let t = s * split + j;
31        let mut value = F::new(0.0f32);
32        if t < cols {
33            value = input[(offset + r * row_stride + t) as usize];
34            if comptime!(centred) {
35                value -= mean[r as usize];
36            }
37        }
38        output[e as usize] = value;
39    }
40}
41
42/// `out[e] = scale · Σ_s partial[s · len + e]` over the `slices` partial products; `accumulate`
43/// adds to `out` instead of overwriting it. One unit per element ([`fn@element_position`]).
44#[cube(launch)]
45pub fn sum_slices_kernel<F: Float + CubeElement + LaunchArg>(
46    partial: &[F],
47    out: &mut [F],
48    len: u32,
49    slices: u32,
50    scale: F,
51    #[comptime] accumulate: bool,
52) {
53    let e = element_position() as u32;
54    if e < len {
55        let mut acc = F::new(0.0f32);
56        let mut s = 0u32;
57        while s < slices {
58            acc += partial[(s * len + e) as usize];
59            s += 1u32;
60        }
61        let value = acc * scale;
62        if comptime!(accumulate) {
63            out[e as usize] += value;
64        } else {
65            out[e as usize] = value;
66        }
67    }
68}
dspk_ops/linalg/kernels/covariance.rs

dspk_ops/linalg/kernels/
covariance.rs
