//! CubeCL matrix multiplication kernel.
//!
//! Naive matmul: one work-unit per output element. Each unit computes the dot
//! product of one row of A and one column of B.
//!
//! This is a correctness-first implementation. Future tasks will add tiling,
//! shared memory, and vectorized loads for production performance.
//!
//! 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::*;

/// Naive matrix multiplication kernel: C = A * B.
///
/// Layouts (row-major, contiguous):
///   A: [M, K]  — `a[row * K + i]`
///   B: [K, N]  — `b[i * N + col]`
///   C: [M, N]  — `out[row * N + col]`
///
/// Launch grid: (M, N) work-units. Each computes one element of C.
///
/// The `m`, `n`, `k` parameters are marked `#[comptime]` so they are baked
/// into the compiled kernel as constants (loop bounds, bounds-checks).
#[cube(launch)]
pub fn matmul_kernel<F: Float>(
    a: &Tensor<F>,
    b: &Tensor<F>,
    out: &mut Tensor<F>,
    #[comptime] m: u32,
    #[comptime] n: u32,
    #[comptime] k: u32,
) {
    let row = ABSOLUTE_POS_X;
    let col = ABSOLUTE_POS_Y;

    if row < m && col < n {
        let mut sum = F::new(0.0);
        for i in 0..k {
            sum += a[row * k + i] * b[i * n + col];
        }
        out[row * n + col] = sum;
    }
}

/// Launch configuration for the matmul kernel.
///
/// Returns `(cube_count, cube_dim)` suitable for passing to the generated
/// `matmul_kernel::launch` function.
///
/// The default tile is 16x16 threads per cube (256 threads = one warp on
/// NVIDIA, two warps on AMD). The cube count is ceil(M/16) x ceil(N/16).
pub fn matmul_launch_config(m: u32, n: u32) -> (CubeCount, CubeDim) {
    let tile = 16u32;
    let cube_dim = CubeDim::new(tile, tile, 1);
    let cube_count = CubeCount::Static(
        (m + tile - 1) / tile,
        (n + tile - 1) / tile,
        1,
    );
    (cube_count, cube_dim)
}

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

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

    /// Verify the launch config math (ceiling division).
    #[test]
    fn test_launch_config_exact() {
        let (count, dim) = matmul_launch_config(32, 48);
        // CubeDim doesn't impl PartialEq either, just check the count
        let _ = dim; // 16x16x1
        // 32/16 = 2, 48/16 = 3
        assert_eq!(cube_count_xyz(count), (2, 3, 1));
    }

    #[test]
    fn test_launch_config_non_aligned() {
        let (count, _dim) = matmul_launch_config(33, 17);
        // ceil(33/16)=3, ceil(17/16)=2
        assert_eq!(cube_count_xyz(count), (3, 2, 1));
    }

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

    #[test]
    fn test_launch_config_single_row() {
        let (count, _dim) = matmul_launch_config(1, 1024);
        // ceil(1/16)=1, ceil(1024/16)=64
        assert_eq!(cube_count_xyz(count), (1, 64, 1));
    }

    #[test]
    fn test_launch_config_single_col() {
        let (count, _dim) = matmul_launch_config(512, 1);
        // ceil(512/16)=32, ceil(1/16)=1
        assert_eq!(cube_count_xyz(count), (32, 1, 1));
    }
}
