xserv/crates/xserv-kernels/src/gemm.rs

use std::ffi::c_void;
use xserv_cuda::error::{self, Result};
use xserv_tensor::{DType, Device, Tensor};

#[derive(Debug, Clone, Copy)]
pub enum GemmBackend {
    Naive,
    Tiled,
    CuBlas,
}

// --- FFI: custom CUDA kernels ---
unsafe extern "C" {
    fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
    fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
    fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
    fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
}

// --- FFI: cuBLAS ---
type CublasHandle = *mut c_void;

#[allow(non_upper_case_globals)]
const CUBLAS_OP_N: i32 = 0;

// cudaDataType
const CUDA_R_32F: i32 = 0;
const CUDA_R_16BF: i32 = 14;

// cublasComputeType
const CUBLAS_COMPUTE_32F: i32 = 68;

unsafe extern "C" {
    fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
    fn cublasDestroy_v2(handle: CublasHandle) -> i32;
    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
    fn cublasGemmEx(
        handle: CublasHandle,
        transa: i32, transb: i32,
        m: i32, n: i32, k: i32,
        alpha: *const c_void,
        a: *const c_void, a_type: i32, lda: i32,
        b: *const c_void, b_type: i32, ldb: i32,
        beta: *const c_void,
        c: *mut c_void, c_type: i32, ldc: i32,
        compute_type: i32,
        algo: i32,
    ) -> i32;
    fn cublasGemmStridedBatchedEx(
        handle: CublasHandle,
        transa: i32, transb: i32,
        m: i32, n: i32, k: i32,
        alpha: *const c_void,
        a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
        b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
        beta: *const c_void,
        c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
        batch_count: i32,
        compute_type: i32,
        algo: i32,
    ) -> i32;
}

pub struct CublasContext {
    handle: CublasHandle,
}

impl CublasContext {
    pub fn new() -> Result<Self> {
        let mut handle = std::ptr::null_mut();
        error::check(unsafe { cublasCreate_v2(&mut handle) })?;
        Ok(Self { handle })
    }
}

impl Drop for CublasContext {
    fn drop(&mut self) {
        if !self.handle.is_null() {
            unsafe { cublasDestroy_v2(self.handle) };
        }
    }
}

/// Matrix multiplication: C = A @ B
/// A: [M, K], B: [K, N], C: [M, N]
/// All tensors must be contiguous and on the same GPU.
pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
    assert_eq!(a.ndim(), 2);
    assert_eq!(b.ndim(), 2);
    assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
    assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
    assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
    assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");

    let m = a.shape()[0];
    let k = a.shape()[1];
    let n = b.shape()[1];
    let dtype = a.dtype();

    let c = Tensor::zeros(&[m, n], dtype, a.device());

    let a_ptr = a.data_ptr() as *const c_void;
    let b_ptr = b.data_ptr() as *const c_void;
    let c_ptr = c.data_ptr() as *mut c_void;
    let null_stream = std::ptr::null_mut();

    match backend {
        GemmBackend::Naive => {
            unsafe {
                match dtype {
                    DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
                    DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
                    _ => panic!("unsupported dtype for naive GEMM"),
                }
            }
            xserv_cuda::device::synchronize().unwrap();
        }
        GemmBackend::Tiled => {
            unsafe {
                match dtype {
                    DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
                    DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
                    _ => panic!("unsupported dtype for tiled GEMM"),
                }
            }
            xserv_cuda::device::synchronize().unwrap();
        }
        GemmBackend::CuBlas => {
            // cuBLAS uses column-major, but we have row-major tensors.
            // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
            // cuBLAS sees our row-major data as column-major transposed.
            let ctx = CublasContext::new().unwrap();
            let alpha = 1.0f32;
            let beta = 0.0f32;

            let (a_type, b_type, c_type) = match dtype {
                DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
                DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
                _ => panic!("unsupported dtype for cuBLAS GEMM"),
            };

            unsafe {
                cublasSetStream_v2(ctx.handle, null_stream);
                // Row-major trick: swap A/B and transpose flags
                // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
                error::check(cublasGemmEx(
                    ctx.handle,
                    CUBLAS_OP_N, CUBLAS_OP_N,
                    n as i32, m as i32, k as i32,
                    &alpha as *const f32 as *const c_void,
                    b_ptr, b_type, n as i32,  // B as col-major = B^T
                    a_ptr, a_type, k as i32,  // A as col-major = A^T
                    &beta as *const f32 as *const c_void,
                    c_ptr, c_type, n as i32,  // C as col-major = C^T
                    CUBLAS_COMPUTE_32F,
                    -1, // default algo
                )).expect("cuBLAS GEMM failed");
            }
            xserv_cuda::device::synchronize().unwrap();
        }
    }

    c
}

/// Batched matrix multiplication via cuBLAS: C[b] = A[b] @ B[b]
/// a: [..., M, K], b: [..., K, N] → [..., M, N]
/// Leading dimensions must match and tensors must be contiguous.
pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
    assert!(a.ndim() >= 2 && b.ndim() >= 2);
    assert_eq!(a.ndim(), b.ndim());
    assert!(a.is_contiguous() && b.is_contiguous());
    assert!(matches!(a.device(), Device::Cuda(_)));
    assert_eq!(a.dtype(), b.dtype());

    let ndim = a.ndim();
    let m = a.shape()[ndim - 2];
    let k = a.shape()[ndim - 1];
    let n = b.shape()[ndim - 1];
    assert_eq!(b.shape()[ndim - 2], k, "inner dimension mismatch");

    // Compute batch count from leading dimensions
    let batch: usize = a.shape()[..ndim - 2].iter().product();
    assert_eq!(
        b.shape()[..ndim - 2].iter().product::<usize>(),
        batch,
        "batch dimensions mismatch"
    );

    let mut out_shape: Vec<usize> = a.shape()[..ndim - 2].to_vec();
    out_shape.push(m);
    out_shape.push(n);
    let c = Tensor::zeros(&out_shape, a.dtype(), a.device());

    let dtype = a.dtype();
    let (a_type, b_type, c_type) = match dtype {
        DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
        DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
        _ => panic!("unsupported dtype for batched matmul"),
    };

    let alpha = 1.0f32;
    let beta = 0.0f32;
    // cuBLAS strides are in elements (not bytes)
    let stride_a = (m * k) as i64;
    let stride_b = (k * n) as i64;
    let stride_c = (m * n) as i64;

    let ctx = CublasContext::new().unwrap();
    unsafe {
        cublasSetStream_v2(ctx.handle, std::ptr::null_mut());
        // Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
        error::check(cublasGemmStridedBatchedEx(
            ctx.handle,
            CUBLAS_OP_N, CUBLAS_OP_N,
            n as i32, m as i32, k as i32,
            &alpha as *const f32 as *const c_void,
            b.data_ptr() as _, b_type, n as i32, stride_b,
            a.data_ptr() as _, a_type, k as i32, stride_a,
            &beta as *const f32 as *const c_void,
            c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
            batch as i32,
            CUBLAS_COMPUTE_32F,
            -1,
        )).expect("cuBLAS batched GEMM failed");
    }
    xserv_cuda::device::synchronize().unwrap();
    c
}