dspk_ops/crossings/kernels/
mod.rs1use cubecl::prelude::*;
13use dspk_core::compute::{channel_position, sample_position};
14
15pub const CROSSING_UP: u32 = 0;
17pub const CROSSING_DOWN: u32 = 1;
18
19#[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#[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#[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#[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#[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}