use cubecl::prelude::*;
use cubecl::server::Handle;
use dspk_core::compute::LaunchGeometry;
use super::kernels::elementwise_kernel;
use crate::core::{cast, DspFloat};
pub const LOG10_FLOOR: f32 = 1e-12;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Elementwise {
Abs,
Square,
Sqrt,
Log10,
}
impl Elementwise {
fn code(self) -> u32 {
match self {
Elementwise::Abs => 0,
Elementwise::Square => 1,
Elementwise::Sqrt => 2,
Elementwise::Log10 => 3,
}
}
pub fn apply(self, x: f64) -> f64 {
match self {
Elementwise::Abs => x.abs(),
Elementwise::Square => x * x,
Elementwise::Sqrt => x.max(0.0).sqrt(),
Elementwise::Log10 => x.abs().max(LOG10_FLOOR as f64).log10(),
}
}
}
pub fn execute_elementwise<F: DspFloat>(client: &Client, input: &Handle, output: &Handle, len: usize, function: Elementwise) {
let geom = LaunchGeometry::elementwise(client, len);
unsafe {
elementwise_kernel::launch::<F>(
client,
geom.cube_count,
geom.cube_dim,
BufferArg::from_raw_parts(input.clone(), len),
BufferArg::from_raw_parts(output.clone(), len),
cast::<F>(LOG10_FLOOR as f64),
function.code(),
);
}
}