//! CubeCL RMS normalization kernel with learned weight.
//!
//! Each work-unit processes one row: compute RMS, normalize, then scale by
//! the learned weight vector.
//!
//! 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::*;

/// RMSNorm kernel: one work-unit per row.
///
/// Layout (row-major, contiguous):
///   input:  [num_rows, hidden_size]
///   weight: [hidden_size]
///   output: [num_rows, hidden_size]
///
/// Launch grid: (num_rows) work-units.
///
/// `hidden_size` is `#[comptime]` — baked into the compiled kernel.
/// Epsilon is hardcoded to 1e-6 (standard for all transformer RMSNorm layers;
/// avoids f32: Hash requirement from CubeCL's kernel-ID caching system).
#[cube(launch)]
pub fn rmsnorm_kernel<F: Float>(
    input: &Tensor<F>,
    weight: &Tensor<F>,
    output: &mut Tensor<F>,
    #[comptime] hidden_size: u32,
) {
    let row = ABSOLUTE_POS;
    let offset = row * hidden_size;

    // Compute mean of squares
    let mut ss = F::new(0.0);
    for i in 0..hidden_size {
        let val = input[offset + i];
        ss += val * val;
    }
    // Multiply by 1/N — computed as a f32 constant at kernel-compile time
    // (comptime u32 → f32 division stays on CPU, avoids From<u32> trait bound)
    ss = ss * F::new(1.0_f32 / hidden_size as f32);

    // Compute 1/sqrt(ss + eps), eps=1e-6 (standard transformer epsilon)
    let rms = F::sqrt(ss + F::new(1e-6));

    // Normalize and apply weight
    for i in 0..hidden_size {
        output[offset + i] = (input[offset + i] / rms) * weight[i];
    }
}

/// Launch configuration for the RMSNorm kernel.
pub fn rmsnorm_launch_config(num_rows: u32) -> (CubeCount, CubeDim) {
    let block_size = 256u32;
    let cube_dim = CubeDim::new(block_size, 1, 1);
    let cube_count = CubeCount::Static((num_rows + 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_rmsnorm_launch_config_exact() {
        let (count, _dim) = rmsnorm_launch_config(256);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

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

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