dspk_synapse/streaming/kernels/
template_reduce.rs

1use cubecl::prelude::*;
2use dspk_core::compute::element_position;
3use dspk_core::compute::LaunchGeometry;
4
5/// CubeCL kernel reducing extracted `[num_spikes, k_neighbors, snippet_len]` snippets in VRAM into
6/// per-channel `[channels, k_neighbors, snippet_len]` mean and second central moment ($M_2$).
7///
8/// Spikes arrive segmented by primary channel: `spike_order[segment_offsets[ch]..segment_offsets[ch + 1]]`
9/// lists the spikes of channel `ch`. One unit per `(channel, k_idx, sample_idx)` slot walks only its
10/// channel's segment (total work O(N·K·L)) with Welford updates, so no large sums cancel in f32.
11#[cube(launch)]
12pub fn reduce_channel_templates_kernel(
13    snippets: &[f32],
14    spike_order: &[u32],
15    segment_offsets: &[u32],
16    out_mean: &mut [f32],
17    out_m2: &mut [f32],
18    num_channels: u32,
19    elems_per_spike: u32,
20) {
21    let tid = element_position() as u32;
22    if tid < num_channels * elems_per_spike {
23        let ch = tid / elems_per_spike;
24        let rem = tid - ch * elems_per_spike;
25        let start = segment_offsets[ch as usize];
26        let end = segment_offsets[(ch + 1u32) as usize];
27
28        let mut mean = 0.0f32;
29        let mut m2 = 0.0f32;
30        let mut j = start;
31        while j < end {
32            let spike = spike_order[j as usize];
33            let val = snippets[(spike * elems_per_spike + rem) as usize];
34            let n = (j - start + 1u32) as f32;
35            let delta = val - mean;
36            mean += delta / n;
37            m2 += delta * (val - mean);
38            j += 1u32;
39        }
40
41        out_mean[tid as usize] = mean;
42        out_m2[tid as usize] = m2;
43    }
44}
45
46/// Per-channel batch moments downloaded from VRAM: `counts[channels]`, and `mean` / `m2` of
47/// shape `[channels, k_neighbors, snippet_len]`.
48pub struct BatchTemplateStats {
49    pub counts: Vec<u32>,
50    pub mean: Vec<f32>,
51    pub m2: Vec<f32>,
52}
53
54/// Host-side dispatcher executing [`fn@reduce_channel_templates_kernel`] on in-VRAM extracted snippets.
55/// `primary_channels[i]` is the primary channel of snippet `i`; snippets on channels
56/// `>= channels` are ignored.
57pub fn execute_reduce_templates_in_vram(
58    client: &Client,
59    snippets_handle: &cubecl::server::Handle,
60    primary_channels: &[u32],
61    channels: usize,
62    k_neighbors: usize,
63    snippet_len: usize,
64) -> BatchTemplateStats {
65    let num_spikes = primary_channels.len();
66    let elems_per_spike = k_neighbors * snippet_len;
67    let total_slots = channels * elems_per_spike;
68
69    // Counting sort of the snippets by primary channel.
70    let mut counts = vec![0u32; channels];
71    for &ch in primary_channels {
72        if let Some(c) = counts.get_mut(ch as usize) {
73            *c += 1;
74        }
75    }
76    let mut offsets = vec![0u32; channels + 1];
77    for ch in 0..channels {
78        offsets[ch + 1] = offsets[ch] + counts[ch];
79    }
80    let mut cursor = offsets.clone();
81    let mut order = vec![0u32; offsets[channels] as usize];
82    for (i, &ch) in primary_channels.iter().enumerate() {
83        let ch = ch as usize;
84        if ch < channels {
85            order[cursor[ch] as usize] = i as u32;
86            cursor[ch] += 1;
87        }
88    }
89    if order.is_empty() || total_slots == 0 {
90        return BatchTemplateStats { counts, mean: vec![0.0; total_slots], m2: vec![0.0; total_slots] };
91    }
92
93    let order_handle = client.create_from_slice(u32::as_bytes(&order));
94    let offsets_handle = client.create_from_slice(u32::as_bytes(&offsets));
95    let out_mean_handle = client.empty(total_slots * std::mem::size_of::<f32>());
96    let out_m2_handle = client.empty(total_slots * std::mem::size_of::<f32>());
97
98    let geom = LaunchGeometry::elementwise(client, total_slots);
99    unsafe {
100        reduce_channel_templates_kernel::launch(
101            client,
102            geom.cube_count,
103            geom.cube_dim,
104            BufferArg::from_raw_parts(snippets_handle.clone(), num_spikes * elems_per_spike),
105            BufferArg::from_raw_parts(order_handle, order.len()),
106            BufferArg::from_raw_parts(offsets_handle, channels + 1),
107            BufferArg::from_raw_parts(out_mean_handle.clone(), total_slots),
108            BufferArg::from_raw_parts(out_m2_handle.clone(), total_slots),
109            channels as u32,
110            elems_per_spike as u32,
111        );
112    }
113
114    let mean_bytes = client.read_one(out_mean_handle).expect("VRAM read mean");
115    let m2_bytes = client.read_one(out_m2_handle).expect("VRAM read m2");
116    BatchTemplateStats {
117        counts,
118        mean: f32::from_bytes(&mean_bytes).to_vec(),
119        m2: f32::from_bytes(&m2_bytes).to_vec(),
120    }
121}
dspk_synapse/streaming/kernels/template_reduce.rs

dspk_synapse/streaming/kernels/
template_reduce.rs
