dspk_ops/math/
elementwise.rs

1//! One function of each sample: rectify (`|x|`), square, square root, base-10 logarithm. With
2//! filters they make the usual envelopes: rectify then low-pass (EMG envelope); square, moving
3//! average, square root (moving RMS); log of an energy (dB-like scales).
4
5use cubecl::prelude::*;
6use cubecl::server::Handle;
7use dspk_core::compute::LaunchGeometry;
8
9use super::kernels::elementwise_kernel;
10use crate::core::{cast, DspFloat};
11
12/// Smallest magnitude [`Elementwise::Log10`] takes the logarithm of (`log₁₀` of it: −12): zeros
13/// and silence give a finite floor instead of −∞.
14pub const LOG10_FLOOR: f32 = 1e-12;
15
16/// The function [`execute_elementwise`] applies to each sample.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum Elementwise {
19    /// `|x|` (full-wave rectification).
20    Abs,
21    /// `x²` (instantaneous power).
22    Square,
23    /// `√x`, negative values taken as 0.
24    Sqrt,
25    /// `log₁₀ |x|`, magnitudes below [`LOG10_FLOOR`] taken as it.
26    Log10,
27}
28
29impl Elementwise {
30    fn code(self) -> u32 {
31        match self {
32            Elementwise::Abs => 0,
33            Elementwise::Square => 1,
34            Elementwise::Sqrt => 2,
35            Elementwise::Log10 => 3,
36        }
37    }
38
39    /// The function on the host (references and tests).
40    pub fn apply(self, x: f64) -> f64 {
41        match self {
42            Elementwise::Abs => x.abs(),
43            Elementwise::Square => x * x,
44            Elementwise::Sqrt => x.max(0.0).sqrt(),
45            Elementwise::Log10 => x.abs().max(LOG10_FLOOR as f64).log10(),
46        }
47    }
48}
49
50/// `output = f(input)` over `len` values (any shape: element-wise).
51pub fn execute_elementwise<F: DspFloat>(client: &Client, input: &Handle, output: &Handle, len: usize, function: Elementwise) {
52    let geom = LaunchGeometry::elementwise(client, len);
53    unsafe {
54        elementwise_kernel::launch::<F>(
55            client,
56            geom.cube_count,
57            geom.cube_dim,
58            BufferArg::from_raw_parts(input.clone(), len),
59            BufferArg::from_raw_parts(output.clone(), len),
60            cast::<F>(LOG10_FLOOR as f64),
61            function.code(),
62        );
63    }
64}
dspk_ops/math/elementwise.rs

dspk_ops/math/
elementwise.rs
