dspk-ops 0.0.3

Signal-processing operations for dspk, written once in CubeCL: filters, resampling, linear algebra, statistics, peak finding, spatial operators
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

use cubecl::prelude::*;
use dspk_core::compute::{row_position, row_share};

/// Merges unit `b`'s `(count, mean, m2)` into unit `a`'s (Chan et al. parallel variance).
#[cube]
fn merge_moments<F: Float>(count: &mut Shared<[F]>, mean: &mut Shared<[F]>, m2: &mut Shared<[F]>, a: usize, b: usize) {
    let (na, nb) = (count[a], count[b]);
    let n = na + nb;
    if nb > F::new(0.0f32) {
        let delta = mean[b] - mean[a];
        let m2_b = m2[b];
        mean[a] += delta * nb / n;
        m2[a] += m2_b + delta * delta * na * nb / n;
        count[a] = n;
    }
}

/// Mean and population standard deviation (`/ cols`) of every row. `units` is the cube's x size
/// (a power of two).
#[cube(launch)]
pub fn row_mean_std_kernel<F: Float>(
    input: &[F],
    out_mean: &mut [F],
    out_std: &mut [F],
    rows: u32,
    cols: u32,
    #[comptime] units: u32,
) {
    let row = row_position();
    // `row` is uniform across the cube, so every unit of a cube takes the same branch
    if row < rows {
        let unit = UNIT_POS_X;
        let base = (row * cols) as usize;

        // Welford over this unit's share of the row
        let mut n = F::new(0.0f32);
        let mut mean = F::new(0.0f32);
        let mut m2 = F::new(0.0f32);
        let share = row_share(cols, units);
        let mut col = share.start;
        while col < share.end {
            let x = input[base + col as usize];
            n += F::new(1.0f32);
            let delta = x - mean;
            mean += delta / n;
            m2 += delta * (x - mean);
            col += share.step;
        }

        let mut count_s = Shared::<[F]>::new_slice(comptime!(units as usize));
        let mut mean_s = Shared::<[F]>::new_slice(comptime!(units as usize));
        let mut m2_s = Shared::<[F]>::new_slice(comptime!(units as usize));
        count_s[unit as usize] = n;
        mean_s[unit as usize] = mean;
        m2_s[unit as usize] = m2;
        sync_cube();

        let stride = RuntimeCell::<u32>::new(units / 2u32);
        while stride.read() > 0u32 {
            let s = stride.read();
            if unit < s {
                merge_moments::<F>(&mut count_s, &mut mean_s, &mut m2_s, unit as usize, (unit + s) as usize);
            }
            sync_cube();
            stride.store(s / 2u32);
        }

        if unit == 0u32 {
            out_mean[row as usize] = mean_s[0];
            let total = F::max(count_s[0], F::new(1.0f32));
            out_std[row as usize] = F::sqrt(m2_s[0] / total);
        }
    }
}