Skip to main content

dspk_ops/crossings/kernels/
mod.rs

1//! Level-crossing kernels for [`super`]: count the crossings of each `(block, channel)`, write them
2//! compacted, the extremum of each stretch between crossings, and one column of a buffer.
3//!
4//! Values are compared in *signed* space, `s · x` with `s = +1` (rising: a level above the
5//! baseline) or `s = −1` (falling: a level below it), so one kernel serves both directions:
6//! - an **up** crossing at `t`: `s · x[t−1] < level ≤ s · x[t]`;
7//! - a **down** crossing at `t`: `s · x[t−1] ≥ release > s · x[t]`.
8//!
9//! At `t = 0` the previous value is `prev[ch]` (the last sample of the window before, or NaN for
10//! none: NaN compares false, so the first sample of a recording never crosses).
11
12use cubecl::prelude::*;
13use dspk_core::compute::{channel_position, sample_position};
14
15/// Kind ids written by [`write_crossings_kernel`].
16pub const CROSSING_UP: u32 = 0;
17pub const CROSSING_DOWN: u32 = 1;
18
19/// Crossing kind at local sample `t` of the row starting at `row`: [`CROSSING_UP`],
20/// [`CROSSING_DOWN`], or 2 for none.
21#[cube]
22fn crossing<F: Float>(trace: &[F], prev: &[F], ch: u32, row: u32, t: u32, sign: F, level: F, release: F) -> u32 {
23    let mut p = prev[ch as usize];
24    if t > 0u32 {
25        p = trace[(row + t - 1u32) as usize];
26    }
27    let sp = sign * p;
28    let sv = sign * trace[(row + t) as usize];
29    let mut kind = 2u32;
30    if sp < level && sv >= level {
31        kind = CROSSING_UP;
32    } else if sp >= release && sv < release {
33        kind = CROSSING_DOWN;
34    }
35    kind
36}
37
38/// Counts the crossings of each `(block, channel)` unit (contiguous blocks of `block` samples from
39/// `scan_start`) into `block_counts[ch · num_blocks + b]`.
40#[cube(launch)]
41pub fn count_crossings_kernel<F: Float + CubeElement + LaunchArg>(
42    trace: &[F],
43    prev: &[F],
44    levels: &[F],
45    releases: &[F],
46    block_counts: &mut [u32],
47    num_channels: u32,
48    num_samples: u32,
49    scan_start: u32,
50    scan_end: u32,
51    num_blocks: u32,
52    block: u32,
53    sign: F,
54) {
55    let b = sample_position();
56    let ch = channel_position();
57    if ch < num_channels && b < num_blocks {
58        let (level, release) = (levels[ch as usize], releases[ch as usize]);
59        let row = ch * num_samples;
60        let mut n = 0u32;
61        let mut t = scan_start + b * block;
62        let mut end = t + block;
63        if end > scan_end {
64            end = scan_end;
65        }
66        while t < end {
67            if crossing::<F>(trace, prev, ch, row, t, sign, level, release) < 2u32 {
68                n += 1u32;
69            }
70            t += 1u32;
71        }
72        block_counts[(ch * num_blocks + b) as usize] = n;
73    }
74}
75
76/// Writes each unit's crossings (local sample and kind) from `channel_bases[ch] +
77/// block_offsets[ch · num_blocks + b]` on: every channel's crossings come out contiguous and in
78/// time order (blocks are contiguous and in order).
79#[cube(launch)]
80pub fn write_crossings_kernel<F: Float + CubeElement + LaunchArg>(
81    trace: &[F],
82    prev: &[F],
83    levels: &[F],
84    releases: &[F],
85    block_offsets: &[u32],
86    channel_bases: &[u32],
87    out_indices: &mut [u32],
88    out_kinds: &mut [u32],
89    num_channels: u32,
90    num_samples: u32,
91    scan_start: u32,
92    scan_end: u32,
93    num_blocks: u32,
94    block: u32,
95    sign: F,
96) {
97    let b = sample_position();
98    let ch = channel_position();
99    if ch < num_channels && b < num_blocks {
100        let (level, release) = (levels[ch as usize], releases[ch as usize]);
101        let row = ch * num_samples;
102        let mut slot = channel_bases[ch as usize] + block_offsets[(ch * num_blocks + b) as usize];
103        let mut t = scan_start + b * block;
104        let mut end = t + block;
105        if end > scan_end {
106            end = scan_end;
107        }
108        while t < end {
109            let kind = crossing::<F>(trace, prev, ch, row, t, sign, level, release);
110            if kind < 2u32 {
111                out_indices[slot as usize] = t;
112                out_kinds[slot as usize] = kind;
113                slot += 1u32;
114            }
115            t += 1u32;
116        }
117    }
118}
119
120/// One unit per stretch `i`: the largest signed value `s · x` over local samples
121/// `starts[i]..ends[i]` of channel `channels[i]`, and where it is (the first such sample). An empty
122/// stretch gives `−3e38` at its start.
123#[cube(launch)]
124pub fn stretch_extrema_kernel<F: Float + CubeElement + LaunchArg>(
125    trace: &[F],
126    channels: &[u32],
127    starts: &[u32],
128    ends: &[u32],
129    out_values: &mut [F],
130    out_indices: &mut [u32],
131    num_stretches: u32,
132    num_samples: u32,
133    sign: F,
134) {
135    let i = ABSOLUTE_POS_X;
136    if i < num_stretches {
137        let row = channels[i as usize] * num_samples;
138        let start = starts[i as usize];
139        let mut best = F::new(-3.0e38f32);
140        let mut at = start;
141        let mut t = start;
142        while t < ends[i as usize] {
143            let v = sign * trace[(row + t) as usize];
144            if v > best {
145                best = v;
146                at = t;
147            }
148            t += 1u32;
149        }
150        out_values[i as usize] = best;
151        out_indices[i as usize] = at;
152    }
153}
154
155/// `out[ch] = trace[ch · num_samples + col]`: one column of a `[channels, samples]` buffer.
156#[cube(launch)]
157pub fn column_kernel<F: Float>(trace: &[F], out: &mut [F], num_channels: u32, num_samples: u32, col: u32) {
158    let ch = ABSOLUTE_POS_X;
159    if ch < num_channels {
160        out[ch as usize] = trace[(ch * num_samples + col) as usize];
161    }
162}