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

use std::cell::RefCell;
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);
    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
}

// --- FFI: cuBLAS ---
pub 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) };
        }
    }
}

thread_local! {
    static CUBLAS_CTX: RefCell<CublasContext> = RefCell::new(
        CublasContext::new().expect("failed to create thread-local cuBLAS handle")
    );
}

/// Borrow the thread-local cuBLAS handle for the duration of a closure.
fn with_cublas<F, R>(f: F) -> R
where
    F: FnOnce(CublasHandle) -> R,
{
    CUBLAS_CTX.with(|cell| {
        let ctx = cell.borrow();
        f(ctx.handle)
    })
}

/// Get the thread-local cuBLAS handle for use with dispatch module.
pub fn cublas_handle() -> CublasHandle {
    CUBLAS_CTX.with(|cell| {
        cell.borrow().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();

    // All backends (naive, tiled, cuBLAS with beta=0, custom GEMV) fully
    // overwrite every element of C, so we skip the cudaMemset.
    let c = Tensor::empty(&[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"),
                }
            }
        }
        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"),
                }
            }
        }
        GemmBackend::CuBlas => {
            // Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode)
            if m == 1 && dtype == DType::BF16 {
                let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
                unsafe {
                    launch_gemv_bf16(
                        a_ptr, b_ptr, c_ptr,
                        fp32_buf.as_mut_ptr() as *mut c_void,
                        k as i32, n as i32,
                        null_stream,
                    );
                }
                // fp32_buf returned to caching allocator pool on drop
            } else {
                // 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 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"),
                };

                with_cublas(|handle| unsafe {
                    cublasSetStream_v2(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(
                        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");
                });
            }
        }
    }

    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);
    // cuBLAS with beta=0 fully overwrites every element of C.
    let c = Tensor::empty(&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;

    with_cublas(|handle| unsafe {
        cublasSetStream_v2(handle, std::ptr::null_mut());
        // Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
        error::check(cublasGemmStridedBatchedEx(
            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");
    });
    c
}