gemm: tiled F32 forward + transpose + backward (dA/dB)

Hand-written tiled GEMM (csrc/ops/gemm.cu, TILE_SIZE=32, FP32 accumulate,
boundary-masked) plus an out-of-place transpose kernel. Wire both through
xtrain-cuda FFI (no_cuda-gated) and expose at the tensor level:
Tensor::matmul, transpose_2d, and matmul_backward computing
dA = dC·Bᵀ and dB = Aᵀ·dC by materializing transposes and reusing the
forward. Also declare cuBLAS sgemm FFI + link cublas, used only as a
correctness reference in tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-cuda/build.rs

@@ -22,6 +22,8 @@ fn main() {
     println!("cargo:rustc-link-search=native={cuda_path}/lib64");
     println!("cargo:rustc-link-lib=dylib=cudart");
     println!("cargo:rustc-link-lib=dylib=cuda");
+    // cuBLAS is used only as a correctness reference for the hand-written GEMM.
+    println!("cargo:rustc-link-lib=dylib=cublas");
 
     cc::Build::new()
         .cuda(true)
@@ -29,6 +31,7 @@ fn main() {
         .flag("-gencode=arch=compute_120,code=sm_120")
         .file("../../csrc/test/vecadd.cu")
         .file("../../csrc/ops/elementwise.cu")
+        .file("../../csrc/ops/gemm.cu")
         .compile("xtrain_cuda_kernels");
 }
 

crates/xtrain-cuda/src/ffi.rs

@@ -39,4 +39,56 @@ unsafe extern "C" {
         n: i32,
         stream: CudaStream,
     );
+
+    // Tiled GEMM: C = A @ B, row-major F32. A:[M,K] B:[K,N] C:[M,N]
+    // (csrc/ops/gemm.cu).
+    pub fn launch_gemm_tiled_f32(
+        a: *const f32,
+        b: *const f32,
+        c: *mut f32,
+        m: i32,
+        n: i32,
+        k: i32,
+        stream: CudaStream,
+    );
+
+    // Out-of-place 2D transpose: out[j,i] = in[i,j]. in:[rows,cols] row-major,
+    // out:[cols,rows] row-major (csrc/ops/gemm.cu).
+    pub fn launch_transpose_f32(
+        input: *const f32,
+        out: *mut f32,
+        rows: i32,
+        cols: i32,
+        stream: CudaStream,
+    );
 }
+
+// cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in
+// tests. Declared (and linked, see build.rs) only when CUDA is compiled in.
+#[cfg(not(no_cuda))]
+pub type CublasHandle = *mut c_void;
+
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
+    pub fn cublasDestroy_v2(handle: CublasHandle) -> i32;
+    pub fn cublasSgemm_v2(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const f32,
+        a: *const f32,
+        lda: i32,
+        b: *const f32,
+        ldb: i32,
+        beta: *const f32,
+        c: *mut f32,
+        ldc: i32,
+    ) -> i32;
+}
+
+#[cfg(not(no_cuda))]
+pub const CUBLAS_OP_N: i32 = 0;

crates/xtrain-tensor/src/tensor.rs

@@ -156,6 +156,97 @@ impl Tensor {
         xtrain_cuda::device::synchronize().expect("scale kernel sync failed");
         out
     }
+
+    // --- GEMM (the T3 kernels) ---
+
+    /// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
+    ///
+    /// Runs the tiled `gemm_tiled_f32` CUDA kernel. Requires contiguous F32
+    /// tensors on the same GPU. Available only when CUDA is compiled in.
+    #[cfg(not(no_cuda))]
+    pub fn matmul(&self, other: &Tensor) -> Self {
+        assert_eq!(self.dtype, DType::F32, "matmul only supports F32");
+        assert_eq!(other.dtype, DType::F32, "matmul only supports F32");
+        assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
+        assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
+        assert_eq!(
+            self.shape[1], other.shape[0],
+            "inner dimension mismatch: {:?} @ {:?}",
+            self.shape, other.shape
+        );
+        assert!(
+            self.is_contiguous() && other.is_contiguous(),
+            "matmul requires contiguous tensors"
+        );
+        assert_eq!(self.device(), other.device(), "matmul device mismatch");
+        assert!(
+            matches!(self.device(), Device::Cuda(_)),
+            "matmul requires CUDA tensors"
+        );
+
+        let m = self.shape[0];
+        let k = self.shape[1];
+        let n = other.shape[1];
+        let out = Tensor::zeros(&[m, n], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_gemm_tiled_f32(
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                m as i32,
+                n as i32,
+                k as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("matmul kernel sync failed");
+        out
+    }
+
+    /// Out-of-place 2D transpose: returns a new contiguous tensor `out[j,i] =
+    /// self[i,j]`. Requires a contiguous F32 CUDA tensor.
+    #[cfg(not(no_cuda))]
+    pub fn transpose_2d(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "transpose only supports F32");
+        assert_eq!(self.ndim(), 2, "transpose_2d requires 2D tensor");
+        assert!(self.is_contiguous(), "transpose requires contiguous tensor");
+        assert!(
+            matches!(self.device(), Device::Cuda(_)),
+            "transpose requires a CUDA tensor"
+        );
+
+        let rows = self.shape[0];
+        let cols = self.shape[1];
+        let out = Tensor::zeros(&[cols, rows], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_transpose_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("transpose kernel sync failed");
+        out
+    }
+
+    /// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
+    /// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
+    /// ([K,N]). All tensors contiguous F32 on the same GPU.
+    #[cfg(not(no_cuda))]
+    pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
+        assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
+        assert_eq!(b.ndim(), 2, "matmul_backward requires 2D B");
+        assert_eq!(dc.ndim(), 2, "matmul_backward requires 2D dC");
+        assert_eq!(a.shape[1], b.shape[0], "A/B inner dim mismatch");
+        assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
+        assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");
+
+        let da = dc.matmul(&b.transpose_2d()); // [M,N] @ [N,K] = [M,K]
+        let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N]
+        (da, db)
+    }
 }
 
 impl std::fmt::Debug for Tensor {

csrc/ops/gemm.cu

@@ -0,0 +1,75 @@
+extern "C" {
+
+// Tiled GEMM (shared memory). C = A @ B, all row-major F32.
+// A: [M, K], B: [K, N], C: [M, N].
+// Each block computes a TILE_SIZE x TILE_SIZE tile of C, cooperatively loading
+// tiles of A and B into shared memory. FP32 accumulation.
+#define TILE_SIZE 32
+
+__global__ void gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+void launch_gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(A, B, C, M, N, K);
+}
+
+// Out-of-place transpose: out[j, i] = in[i, j].
+// in: [rows, cols] row-major, out: [cols, rows] row-major.
+__global__ void transpose_f32(const float* in, float* out, int rows, int cols) {
+    int col = blockIdx.x * blockDim.x + threadIdx.x; // over cols of `in`
+    int row = blockIdx.y * blockDim.y + threadIdx.y; // over rows of `in`
+    if (row < rows && col < cols) {
+        out[col * rows + row] = in[row * cols + col];
+    }
+}
+
+void launch_transpose_f32(const float* in, float* out, int rows, int cols, void* stream) {
+    dim3 block(16, 16);
+    dim3 grid((cols + block.x - 1) / block.x, (rows + block.y - 1) / block.y);
+    transpose_f32<<<grid, block, 0, (cudaStream_t)stream>>>(in, out, rows, cols);
+}
+
+}