dspk_ops/math/kernels/
elementwise.rs

1use cubecl::prelude::*;
2use dspk_core::compute::element_position;
3
4/// One function of each sample (`op`, fixed at compile time): `0` `|x|`, `1` `x²`, `2` `√max(x, 0)`,
5/// `3` `log₁₀ max(|x|, floor)`.
6#[cube(launch)]
7pub fn elementwise_kernel<F: Float + CubeElement + LaunchArg>(input: &[F], output: &mut [F], floor: F, #[comptime] op: u32) {
8    let i = element_position();
9    if i < input.len() {
10        let x = input[i];
11        if comptime!(op == 0) {
12            output[i] = F::abs(x);
13        } else if comptime!(op == 1) {
14            output[i] = x * x;
15        } else if comptime!(op == 2) {
16            output[i] = F::sqrt(F::max(x, F::new(0.0f32)));
17        } else {
18            output[i] = F::ln(F::max(F::abs(x), floor)) * F::new(core::f32::consts::LOG10_E);
19        }
20    }
21}

dspk_ops/math/kernels/elementwise.rs

dspk_ops/math/kernels/
elementwise.rs