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>
This commit is contained in:
2026-06-15 15:26:51 +08:00
parent 9ca98efd98
commit 08c88bf360
4 changed files with 221 additions and 0 deletions

View File

@@ -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;