//! CubeCL softmax kernel.
//!
//! Numerically stable softmax over the last dimension. Each work-unit processes
//! one row: find max, compute exp(x - max), normalize by sum.
//!
//! 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::*;

/// Softmax kernel: one work-unit per row.
///
/// Layout (row-major, contiguous):
///   input: [num_rows, dim_size] — `input[row * dim_size + i]`
///   output: [num_rows, dim_size] — `output[row * dim_size + i]`
///
/// Launch grid: (num_rows) work-units.
///
/// The `dim_size` parameter is marked `#[comptime]` so it is baked
/// into the compiled kernel as a constant (loop bounds).
#[cube(launch)]
pub fn softmax_kernel<F: Float>(
    input: &Tensor<F>,
    output: &mut Tensor<F>,
    #[comptime] dim_size: u32,
) {
    let row = ABSOLUTE_POS;
    let offset = row * dim_size;

    // Find max for numerical stability
    let mut max_val = F::new(-3.4028235e38); // f32::MIN (NEG_INFINITY not available in CubeCL)
    for i in 0..dim_size {
        let val = input[offset + i];
        if val > max_val {
            max_val = val;
        }
    }

    // Compute exp(x - max) and running sum
    let mut sum = F::new(0.0);
    for i in 0..dim_size {
        let val = F::exp(input[offset + i] - max_val);
        output[offset + i] = val;
        sum += val;
    }

    // Normalize by sum
    for i in 0..dim_size {
        output[offset + i] = output[offset + i] / sum;
    }
}

/// Launch configuration for the softmax kernel.
///
/// Returns `(cube_count, cube_dim)` suitable for passing to the generated
/// `softmax_kernel::launch` function.
///
/// One thread per row, launched in 1D blocks of 256 threads.
pub fn softmax_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_softmax_launch_config_exact() {
        let (count, _dim) = softmax_launch_config(256);
        assert_eq!(cube_count_xyz(count), (1, 1, 1));
    }

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

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

    #[test]
    fn test_softmax_launch_config_large() {
        let (count, _dim) = softmax_launch_config(1024);
        assert_eq!(cube_count_xyz(count), (4, 1, 1));
    }
}
