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>

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);
+}
+
+}