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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user