xtrain/crates/xtrain-cuda/src/cublas.rs

//! cuBLAS GEMM backend (Phase T7).
//!
//! The hand-written tiled kernel (csrc/ops/gemm.cu) is kept as the T3 learning
//! artifact + correctness oracle's counterpart, but the forward + both backward
//! matmuls now route through cuBLAS `Sgemm` — fp32, so the result is numerically
//! the same GEMM (only the rounding order changes), which is why the T3 tolerance
//! against cuBLAS holds unchanged.
//!
//! **Layout.** cuBLAS is column-major; our tensors are row-major. A row-major
//! `[r,c]` matrix handed to cuBLAS with leading dim `c` is read as its transpose
//! (col-major `[c,r]`). To get a row-major result `C[m,n] = opA(A)·opB(B)` we
//! compute the col-major transpose `Cᵀ[n,m] = opB(B)ᵀ·opA(A)ᵀ`; the bytes of
//! col-major `Cᵀ` are exactly row-major `C`. See [`sgemm`] for the index algebra.
//!
//! **Handle.** cuBLAS handle creation is expensive (T3's oracle made one per
//! call). We cache one handle per thread for the lifetime of the process.

#![cfg(not(no_cuda))]

use crate::ffi::{self, CublasHandle};
use std::cell::RefCell;
use std::ffi::c_void;

thread_local! {
    static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
}

/// Run `f` with the thread's cached cuBLAS handle, creating it on first use.
fn with_handle<R>(f: impl FnOnce(CublasHandle) -> R) -> R {
    HANDLE.with(|h| {
        let mut slot = h.borrow_mut();
        if slot.is_none() {
            let mut handle: CublasHandle = std::ptr::null_mut();
            let status = unsafe { ffi::cublasCreate_v2(&mut handle) };
            assert_eq!(status, 0, "cublasCreate failed: {status}");
            *slot = Some(handle);
        }
        f(slot.unwrap())
    })
}

/// Row-major single-precision GEMM: `C[m,n] = opA(A) · opB(B)` with
/// `C = alpha·(…) + beta·C`. `A`/`B`/`C` are device pointers to row-major fp32
/// matrices; `trans_a`/`trans_b` request the transpose of the *logical* operand.
///
/// `m,n,k` are the dims of the math (`opA(A)` is `[m,k]`, `opB(B)` is `[k,n]`).
/// The stored, untransposed shapes are: `A` is `[m,k]` (or `[k,m]` if `trans_a`),
/// `B` is `[k,n]` (or `[n,k]` if `trans_b`). Their row-major leading dims are the
/// stored column counts, derived below.
///
/// We ask cuBLAS for col-major `Cᵀ[n,m] = opB(B)ᵀ · opA(A)ᵀ`. Since a row-major
/// `[r,c]` buffer is col-major `[c,r]`, a row-major operand already *is* its own
/// transpose to cuBLAS — so `opB(B)ᵀ` over the row-major bytes of `B` is obtained
/// by passing `B` with the OPPOSITE op flag of what `opB` would suggest. Working
/// it through: first cuBLAS arg = `B` with op `trans_b ? N : T`, second = `A` with
/// op `trans_a ? N : T`, sizes (m=n, n=m, k=k).
#[allow(clippy::too_many_arguments)]
pub fn sgemm(
    trans_a: bool,
    trans_b: bool,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    a: *const f32,
    b: *const f32,
    beta: f32,
    c: *mut f32,
) {
    // Leading dims = stored (row-major) column count of each untransposed matrix.
    let lda = if trans_a { m } else { k }; // A stored [m,k] or [k,m]
    let ldb = if trans_b { k } else { n }; // B stored [k,n] or [n,k]
    let ldc = n; // Cᵀ is [n,m] col-major with ld n  (== row-major C[m,n])

    let op_b = if trans_b {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let op_a = if trans_a {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };

    with_handle(|handle| {
        let status = unsafe {
            ffi::cublasSgemm_v2(
                handle, op_b, op_a, n as i32, // rows of Cᵀ
                m as i32, // cols of Cᵀ
                k as i32, &alpha, b, ldb as i32, a, lda as i32, &beta, c, ldc as i32,
            )
        };
        assert_eq!(status, 0, "cublasSgemm failed: {status}");
    });
}

/// Strided-batched row-major SGEMM: for each `i` in `0..batch`,
/// `C_i[m,n] = alpha·opA(A_i)·opB(B_i) + beta·C_i`, where `A_i`/`B_i`/`C_i` are
/// consecutive matrices laid `stride_*` elements apart in one contiguous buffer.
/// Same row-major⟺col-major trick as [`sgemm`] (compute col-major `Cᵀ`), applied
/// per batch element. Used for the batched attention `QKᵀ` / `PV` GEMMs (and their
/// backwards), so the whole attention runs as 2 batched-GEMM launches, not a
/// per-(batch,head) Python loop. `A`/`B`/`C` are device pointers to the first
/// matrix; strides are in ELEMENTS.
#[allow(clippy::too_many_arguments)]
pub fn sgemm_strided_batched(
    trans_a: bool,
    trans_b: bool,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    a: *const f32,
    stride_a: usize,
    b: *const f32,
    stride_b: usize,
    beta: f32,
    c: *mut f32,
    stride_c: usize,
    batch: usize,
) {
    let lda = if trans_a { m } else { k };
    let ldb = if trans_b { k } else { n };
    let ldc = n;
    let op_a = if trans_a {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let op_b = if trans_b {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };

    with_handle(|handle| {
        let status = unsafe {
            ffi::cublasSgemmStridedBatched(
                handle,
                op_b,
                op_a,
                n as i32,
                m as i32,
                k as i32,
                &alpha,
                b,
                ldb as i32,
                stride_b as i64,
                a,
                lda as i32,
                stride_a as i64,
                &beta,
                c,
                ldc as i32,
                stride_c as i64,
                batch as i32,
            )
        };
        assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
    });
}

/// bf16 row-major GEMM `C[m,n] = opA(A)·opB(B)` via `cublasGemmEx`: bf16 in/out,
/// **fp32 accumulation** (`CUBLAS_COMPUTE_32F`) — the standard AMP matmul (Phase
/// T12). `a`/`b`/`c` are device pointers to row-major **bf16** matrices; the
/// row-major⟺col-major transpose algebra is identical to [`sgemm`] (we compute
/// the col-major `Cᵀ`). `alpha`/`beta` are fp32 host scalars (compute is fp32).
#[allow(clippy::too_many_arguments)]
pub fn gemm_ex(
    trans_a: bool,
    trans_b: bool,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    a: *const c_void,
    b: *const c_void,
    beta: f32,
    c: *mut c_void,
) {
    let lda = if trans_a { m } else { k };
    let ldb = if trans_b { k } else { n };
    let ldc = n;
    let op_a = if trans_a {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let op_b = if trans_b {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let bf16 = ffi::CUDA_R_16BF;

    with_handle(|handle| {
        let status = unsafe {
            ffi::cublasGemmEx(
                handle,
                op_b,
                op_a,
                n as i32,
                m as i32,
                k as i32,
                &alpha as *const f32 as *const c_void,
                b,
                bf16,
                ldb as i32,
                a,
                bf16,
                lda as i32,
                &beta as *const f32 as *const c_void,
                c,
                bf16,
                ldc as i32,
                ffi::CUBLAS_COMPUTE_32F,
                ffi::CUBLAS_GEMM_DEFAULT,
            )
        };
        assert_eq!(status, 0, "cublasGemmEx failed: {status}");
    });
}

/// Strided-batched bf16 GEMM (Phase T12) — the [`gemm_ex`] analogue of
/// [`sgemm_strided_batched`] for the batched attention GEMMs. bf16 in/out, fp32
/// accumulation; strides are in ELEMENTS.
#[allow(clippy::too_many_arguments)]
pub fn gemm_ex_strided_batched(
    trans_a: bool,
    trans_b: bool,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    a: *const c_void,
    stride_a: usize,
    b: *const c_void,
    stride_b: usize,
    beta: f32,
    c: *mut c_void,
    stride_c: usize,
    batch: usize,
) {
    let lda = if trans_a { m } else { k };
    let ldb = if trans_b { k } else { n };
    let ldc = n;
    let op_a = if trans_a {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let op_b = if trans_b {
        ffi::CUBLAS_OP_T
    } else {
        ffi::CUBLAS_OP_N
    };
    let bf16 = ffi::CUDA_R_16BF;

    with_handle(|handle| {
        let status = unsafe {
            ffi::cublasGemmStridedBatchedEx(
                handle,
                op_b,
                op_a,
                n as i32,
                m as i32,
                k as i32,
                &alpha as *const f32 as *const c_void,
                b,
                bf16,
                ldb as i32,
                stride_b as i64,
                a,
                bf16,
                lda as i32,
                stride_a as i64,
                &beta as *const f32 as *const c_void,
                c,
                bf16,
                ldc as i32,
                stride_c as i64,
                batch as i32,
                ffi::CUBLAS_COMPUTE_32F,
                ffi::CUBLAS_GEMM_DEFAULT,
            )
        };
        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
    });
}