//! CubeCL SiLU-gated element-wise multiply kernel.
//!
//! Computes `gate = silu(gate) * up` element-wise.
//! SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x)).
//!
//! One work-unit per element. The gate tensor is modified in-place (output
//! overwrites gate).
//!
//! Generic over `cubecl::prelude::Float`, so the same kernel compiles for
//! f32, f16, bf16 on any CubeCL-supported backend (CUDA, WGPU, etc.).

use cubecl::prelude::*;

/// SiLU-gated multiply kernel: gate[i] = silu(gate[i]) * up[i].
///
/// Layout: both `gate` and `up` are flat 1D tensors of length `len`.
///
/// Launch grid: (len) work-units. Each computes one element.
///
/// `len` is `#[comptime]` — baked into the compiled kernel for bounds checking.
#[cube(launch)]
pub fn silu_mul_kernel<F: Float>(
    gate: &Tensor<F>,
    up: &Tensor<F>,
    output: &mut Tensor<F>,
    #[comptime] len: u32,
) {
    let idx = ABSOLUTE_POS;
    if idx < len {
        let g = gate[idx];
        // SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
        let sigmoid = F::new(1.0) / (F::new(1.0) + F::exp(F::new(0.0) - g));
        let silu = g * sigmoid;
        output[idx] = silu * up[idx];
    }
}

/// Launch configuration for the SiLU-gated multiply kernel.
///
/// Returns `(cube_count, cube_dim)` suitable for passing to the generated
/// `silu_mul_kernel::launch` function.
///
/// One thread per element, launched in 1D blocks of 256 threads.
pub fn silu_mul_launch_config(len: u32) -> (CubeCount, CubeDim) {
    let block_size = 256u32;
    let cube_dim = CubeDim::new(block_size, 1, 1);
    let cube_count = CubeCount::Static((len + block_size - 1) / block_size, 1, 1);
    (cube_count, cube_dim)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn cube_count_xyz(c: CubeCount) -> (u32, u32, u32) {
        match c {
            CubeCount::Static(x, y, z) => (x, y, z),
            _ => panic!("expected CubeCount::Static"),
        }
    }

    #[test]
    fn test_silu_mul_launch_config_exact() {
        let (count, _dim) = silu_mul_launch_config(256);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

    #[test]
    fn test_silu_mul_launch_config_non_aligned() {
        let (count, _dim) = silu_mul_launch_config(257);
        assert_eq!(cube_count_xyz(count), (2, 1, 1));
    }

    #[test]
    fn test_silu_mul_launch_config_small() {
        let (count, _dim) = silu_mul_launch_config(1);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

    #[test]
    fn test_silu_mul_launch_config_large() {
        let (count, _dim) = silu_mul_launch_config(4096);
        assert_eq!(cube_count_xyz(count), (16, 1, 1));
    }
}
