diff --git a/crates/xtrain-cuda/build.rs b/crates/xtrain-cuda/build.rs index 1d27f02..9d29fa3 100644 --- a/crates/xtrain-cuda/build.rs +++ b/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"); } diff --git a/crates/xtrain-cuda/src/ffi.rs b/crates/xtrain-cuda/src/ffi.rs index 1f6bfc7..a13a3e3 100644 --- a/crates/xtrain-cuda/src/ffi.rs +++ b/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; diff --git a/crates/xtrain-tensor/src/tensor.rs b/crates/xtrain-tensor/src/tensor.rs index 98e7cdd..6cbb8d7 100644 --- a/crates/xtrain-tensor/src/tensor.rs +++ b/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 { diff --git a/csrc/ops/gemm.cu b/csrc/ops/gemm.cu new file mode 100644 index 0000000..ed7630c --- /dev/null +++ b/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<<>>(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<<>>(in, out, rows, cols); +} + +}