diff --git a/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu b/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu index 5f50e42..2206941 100644 --- a/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu +++ b/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu @@ -67,6 +67,15 @@ __device__ __forceinline__ long long mul_i64(long long a, long long b) { return a * b; } +// Float/double tropical multiply in function form, so the shared kernel body +// (TROPICAL_GEMM_BODY) can take MUL_FN as a call for every scalar type instead +// of special-casing the `+`/`*` operator for floats. These inline to a bare +// add/mul -- identical codegen to the operator form. +__device__ __forceinline__ float add_f32(float a, float b) { return a + b; } +__device__ __forceinline__ float mul_f32(float a, float b) { return a * b; } +__device__ __forceinline__ double add_f64(double a, double b) { return a + b; } +__device__ __forceinline__ double mul_f64(double a, double b) { return a * b; } + // Drifted tropical-zero detection for argmax canonicalization. A no-contribution // output cell's value sits in "infinity territory" (past S/2) after the // guard-free add drifts it (`S + data`). Used ONLY at the O(M*N) write-out (not @@ -193,337 +202,113 @@ __device__ double atomicAddDouble(double* address, double val) { } // ============================================================================ -// F32 GEMM KERNEL MACRO +// SHARED FORWARD GEMM KERNEL BODY // ============================================================================ -// Block sizes for f32: 64x32x64, Thread sizes: 4x4 -// Generates: tropical_{semiring}_f32_nn - -#define TROPICAL_GEMM_F32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_OP) \ -extern "C" __global__ void KERNEL_NAME( \ - const float* __restrict__ A, \ - const float* __restrict__ B, \ - float* __restrict__ C, \ - int M, int N, int K \ -) { \ - const int BLOCK_SIZE_M = 64; \ - const int BLOCK_SIZE_K = 32; \ - const int BLOCK_SIZE_N = 64; \ - const int THREAD_SIZE_M = 4; \ - const int THREAD_SIZE_N = 4; \ - \ - const int bszm = BLOCK_SIZE_M / THREAD_SIZE_M; \ - const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ - const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ - \ - int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ - int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ - int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ - int BLOCK_IDY = blockIdx.x / DIM_GRID_X; \ - \ - const int tid = threadIdx.y * bszm + threadIdx.x; \ - \ - __shared__ float As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ - __shared__ float Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ - \ - float accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ - float regs_a[THREAD_SIZE_M]; \ - float regs_b[THREAD_SIZE_N]; \ - \ - _Pragma("unroll") \ - for (int i = 0; i < THREAD_SIZE_M * THREAD_SIZE_N; ++i) { \ - accum[i] = INIT_VAL; \ - } \ - \ - const int A_TILE_COL = tid / BLOCK_SIZE_M; \ - const int A_TILE_ROW = tid % BLOCK_SIZE_M; \ - const int B_TILE_COL = tid / BLOCK_SIZE_K; \ - const int B_TILE_ROW = tid % BLOCK_SIZE_K; \ - const int A_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_M; \ - const int B_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_K; \ - \ - for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) { \ - LOAD_A_TILE(A, INIT_VAL) \ - \ - LOAD_B_TILE(B, INIT_VAL) \ - \ - __syncthreads(); \ - \ - _Pragma("unroll") \ - for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ - k, BLOCK_SIZE_M)]; \ - } \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - regs_b[tn] = Bs[OFFSET_COL(k, threadIdx.y * THREAD_SIZE_N + tn,\ - BLOCK_SIZE_K)]; \ - } \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - float prod = regs_a[tm] MUL_OP regs_b[tn]; \ - int idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ - accum[idx] = COMPARE_FN(accum[idx], prod); \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_TILE(C) \ -} +// The tiled max-/min-plus / max-mul inner kernel is identical across every +// scalar type AND across the single-matrix and strided-batched entry points; +// it differs only in (1) the scalar TYPE, (2) the block tiling sizes, and +// (3) which operand base pointer each tile load/store reads. It lives here once +// so the algorithm has a single source of truth. The thin wrappers below supply +// the kernel signature and the operand base expressions: plain `A/B/C` for the +// single-matrix kernels, `A + blockIdx.z*strideA` etc. for the batched ones. +// +// COMPARE_FN and MUL_FN are both function-style (e.g. fmaxf / add_f32) so the +// body is type-agnostic; see the add_*/mul_* and max_*/min_* helpers above. +// A_SRC/B_SRC/C_DST are spliced into `SRC[...]` / `DST[...]`, so a batched base +// expression must be parenthesised by the caller. + +#define TROPICAL_GEMM_BODY(TYPE, BSM, BSK, BSN, INIT_VAL, COMPARE_FN, MUL_FN, A_SRC, B_SRC, C_DST) \ + const int BLOCK_SIZE_M = BSM; \ + const int BLOCK_SIZE_K = BSK; \ + const int BLOCK_SIZE_N = BSN; \ + const int THREAD_SIZE_M = 4; \ + const int THREAD_SIZE_N = 4; \ + const int bszm = BLOCK_SIZE_M / THREAD_SIZE_M; \ + const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ + const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ + int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ + int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ + int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ + int BLOCK_IDY = blockIdx.x / DIM_GRID_X; \ + const int tid = threadIdx.y * bszm + threadIdx.x; \ + __shared__ TYPE As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ + __shared__ TYPE Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ + TYPE accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ + TYPE regs_a[THREAD_SIZE_M]; \ + TYPE regs_b[THREAD_SIZE_N]; \ + _Pragma("unroll") \ + for (int i = 0; i < THREAD_SIZE_M * THREAD_SIZE_N; ++i) { \ + accum[i] = INIT_VAL; \ + } \ + const int A_TILE_COL = tid / BLOCK_SIZE_M; \ + const int A_TILE_ROW = tid % BLOCK_SIZE_M; \ + const int B_TILE_COL = tid / BLOCK_SIZE_K; \ + const int B_TILE_ROW = tid % BLOCK_SIZE_K; \ + const int A_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_M; \ + const int B_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_K; \ + for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) { \ + LOAD_A_TILE(A_SRC, INIT_VAL) \ + LOAD_B_TILE(B_SRC, INIT_VAL) \ + __syncthreads(); \ + _Pragma("unroll") \ + for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ + _Pragma("unroll") \ + for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ + regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, k, BLOCK_SIZE_M)]; \ + } \ + _Pragma("unroll") \ + for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ + regs_b[tn] = Bs[OFFSET_COL(k, threadIdx.y * THREAD_SIZE_N + tn, BLOCK_SIZE_K)]; \ + } \ + _Pragma("unroll") \ + for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ + _Pragma("unroll") \ + for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ + TYPE prod = MUL_FN(regs_a[tm], regs_b[tn]); \ + int idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ + accum[idx] = COMPARE_FN(accum[idx], prod); \ + } \ + } \ + } \ + __syncthreads(); \ + } \ + STORE_C_TILE(C_DST) // ============================================================================ -// F64 GEMM KERNEL MACRO +// SINGLE-MATRIX FORWARD GEMM KERNELS (blockIdx.x carries the M/N tiling) // ============================================================================ -// Block sizes for f64: 32x16x32, Thread sizes: 4x4 - -#define TROPICAL_GEMM_F64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_OP) \ -extern "C" __global__ void KERNEL_NAME( \ - const double* __restrict__ A, \ - const double* __restrict__ B, \ - double* __restrict__ C, \ - int M, int N, int K \ -) { \ - const int BLOCK_SIZE_M = 32; \ - const int BLOCK_SIZE_K = 16; \ - const int BLOCK_SIZE_N = 32; \ - const int THREAD_SIZE_M = 4; \ - const int THREAD_SIZE_N = 4; \ - \ - const int bszm = BLOCK_SIZE_M / THREAD_SIZE_M; \ - const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ - const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ - \ - int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ - int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ - int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ - int BLOCK_IDY = blockIdx.x / DIM_GRID_X; \ - \ - const int tid = threadIdx.y * bszm + threadIdx.x; \ - \ - __shared__ double As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ - __shared__ double Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ - \ - double accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ - double regs_a[THREAD_SIZE_M]; \ - double regs_b[THREAD_SIZE_N]; \ - \ - _Pragma("unroll") \ - for (int i = 0; i < THREAD_SIZE_M * THREAD_SIZE_N; ++i) { \ - accum[i] = INIT_VAL; \ - } \ - \ - const int A_TILE_COL = tid / BLOCK_SIZE_M; \ - const int A_TILE_ROW = tid % BLOCK_SIZE_M; \ - const int B_TILE_COL = tid / BLOCK_SIZE_K; \ - const int B_TILE_ROW = tid % BLOCK_SIZE_K; \ - const int A_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_M; \ - const int B_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_K; \ - \ - for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) { \ - LOAD_A_TILE(A, INIT_VAL) \ - \ - LOAD_B_TILE(B, INIT_VAL) \ - \ - __syncthreads(); \ - \ - _Pragma("unroll") \ - for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ - k, BLOCK_SIZE_M)]; \ - } \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - regs_b[tn] = Bs[OFFSET_COL(k, threadIdx.y * THREAD_SIZE_N + tn,\ - BLOCK_SIZE_K)]; \ - } \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - double prod = regs_a[tm] MUL_OP regs_b[tn]; \ - int idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ - accum[idx] = COMPARE_FN(accum[idx], prod); \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_TILE(C) \ +// Block sizes: f32/i32 = 64x32x64, f64/i64 = 32x16x32 (8-byte types halve the tile). + +#define TROPICAL_GEMM_F32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const float* __restrict__ A, const float* __restrict__ B, \ + float* __restrict__ C, int M, int N, int K \ +) { \ + TROPICAL_GEMM_BODY(float, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, A, B, C) \ } -// ============================================================================ -// I32 GEMM KERNEL MACRO -// ============================================================================ -// Block sizes for i32: 64x32x64, Thread sizes: 4x4 (same as f32) -// Multiply is a bare add (MUL_FN); the tropical zero is a large sentinel. - -#define TROPICAL_GEMM_I32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ -extern "C" __global__ void KERNEL_NAME( \ - const int* __restrict__ A, \ - const int* __restrict__ B, \ - int* __restrict__ C, \ - int M, int N, int K \ -) { \ - const int BLOCK_SIZE_M = 64; \ - const int BLOCK_SIZE_K = 32; \ - const int BLOCK_SIZE_N = 64; \ - const int THREAD_SIZE_M = 4; \ - const int THREAD_SIZE_N = 4; \ - \ - const int bszm = BLOCK_SIZE_M / THREAD_SIZE_M; \ - const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ - const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ - \ - int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ - int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ - int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ - int BLOCK_IDY = blockIdx.x / DIM_GRID_X; \ - \ - const int tid = threadIdx.y * bszm + threadIdx.x; \ - \ - __shared__ int As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ - __shared__ int Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ - \ - int accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ - int regs_a[THREAD_SIZE_M]; \ - int regs_b[THREAD_SIZE_N]; \ - \ - _Pragma("unroll") \ - for (int i = 0; i < THREAD_SIZE_M * THREAD_SIZE_N; ++i) { \ - accum[i] = INIT_VAL; \ - } \ - \ - const int A_TILE_COL = tid / BLOCK_SIZE_M; \ - const int A_TILE_ROW = tid % BLOCK_SIZE_M; \ - const int B_TILE_COL = tid / BLOCK_SIZE_K; \ - const int B_TILE_ROW = tid % BLOCK_SIZE_K; \ - const int A_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_M; \ - const int B_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_K; \ - \ - for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) { \ - LOAD_A_TILE(A, INIT_VAL) \ - \ - LOAD_B_TILE(B, INIT_VAL) \ - \ - __syncthreads(); \ - \ - _Pragma("unroll") \ - for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ - k, BLOCK_SIZE_M)]; \ - } \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - regs_b[tn] = Bs[OFFSET_COL(k, threadIdx.y * THREAD_SIZE_N + tn,\ - BLOCK_SIZE_K)]; \ - } \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - int prod = MUL_FN(regs_a[tm], regs_b[tn]); \ - int idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ - accum[idx] = COMPARE_FN(accum[idx], prod); \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_TILE(C) \ +#define TROPICAL_GEMM_F64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const double* __restrict__ A, const double* __restrict__ B, \ + double* __restrict__ C, int M, int N, int K \ +) { \ + TROPICAL_GEMM_BODY(double, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, A, B, C) \ } -// ============================================================================ -// I64 GEMM KERNEL MACRO -// ============================================================================ -// Block sizes for i64: 32x16x32, Thread sizes: 4x4 (same as f64) +#define TROPICAL_GEMM_I32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const int* __restrict__ A, const int* __restrict__ B, \ + int* __restrict__ C, int M, int N, int K \ +) { \ + TROPICAL_GEMM_BODY(int, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, A, B, C) \ +} -#define TROPICAL_GEMM_I64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ -extern "C" __global__ void KERNEL_NAME( \ - const long long* __restrict__ A, \ - const long long* __restrict__ B, \ - long long* __restrict__ C, \ - int M, int N, int K \ -) { \ - const int BLOCK_SIZE_M = 32; \ - const int BLOCK_SIZE_K = 16; \ - const int BLOCK_SIZE_N = 32; \ - const int THREAD_SIZE_M = 4; \ - const int THREAD_SIZE_N = 4; \ - \ - const int bszm = BLOCK_SIZE_M / THREAD_SIZE_M; \ - const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ - const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ - \ - int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ - int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ - int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ - int BLOCK_IDY = blockIdx.x / DIM_GRID_X; \ - \ - const int tid = threadIdx.y * bszm + threadIdx.x; \ - \ - __shared__ long long As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ - __shared__ long long Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ - \ - long long accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ - long long regs_a[THREAD_SIZE_M]; \ - long long regs_b[THREAD_SIZE_N]; \ - \ - _Pragma("unroll") \ - for (int i = 0; i < THREAD_SIZE_M * THREAD_SIZE_N; ++i) { \ - accum[i] = INIT_VAL; \ - } \ - \ - const int A_TILE_COL = tid / BLOCK_SIZE_M; \ - const int A_TILE_ROW = tid % BLOCK_SIZE_M; \ - const int B_TILE_COL = tid / BLOCK_SIZE_K; \ - const int B_TILE_ROW = tid % BLOCK_SIZE_K; \ - const int A_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_M; \ - const int B_TILE_COL_STRIDE = THREAD_NUM_PER_BLOCK / BLOCK_SIZE_K; \ - \ - for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) { \ - LOAD_A_TILE(A, INIT_VAL) \ - \ - LOAD_B_TILE(B, INIT_VAL) \ - \ - __syncthreads(); \ - \ - _Pragma("unroll") \ - for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ - k, BLOCK_SIZE_M)]; \ - } \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - regs_b[tn] = Bs[OFFSET_COL(k, threadIdx.y * THREAD_SIZE_N + tn,\ - BLOCK_SIZE_K)]; \ - } \ - _Pragma("unroll") \ - for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ - _Pragma("unroll") \ - for (int tn = 0; tn < THREAD_SIZE_N; ++tn) { \ - long long prod = MUL_FN(regs_a[tm], regs_b[tn]); \ - int idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ - accum[idx] = COMPARE_FN(accum[idx], prod); \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_TILE(C) \ +#define TROPICAL_GEMM_I64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const long long* __restrict__ A, const long long* __restrict__ B, \ + long long* __restrict__ C, int M, int N, int K \ +) { \ + TROPICAL_GEMM_BODY(long long, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, A, B, C) \ } // ============================================================================ @@ -927,14 +712,14 @@ extern "C" __global__ void KERNEL_NAME( \ // ============================================================================ // --- F32 Basic GEMM Kernels --- -TROPICAL_GEMM_F32(tropical_maxplus_f32_nn, NEG_INF_F32, fmaxf, +) -TROPICAL_GEMM_F32(tropical_minplus_f32_nn, INF_F32, fminf, +) -TROPICAL_GEMM_F32(tropical_maxmul_f32_nn, 0.0f, fmaxf, *) +TROPICAL_GEMM_F32(tropical_maxplus_f32_nn, NEG_INF_F32, fmaxf, add_f32) +TROPICAL_GEMM_F32(tropical_minplus_f32_nn, INF_F32, fminf, add_f32) +TROPICAL_GEMM_F32(tropical_maxmul_f32_nn, 0.0f, fmaxf, mul_f32) // --- F64 Basic GEMM Kernels --- -TROPICAL_GEMM_F64(tropical_maxplus_f64_nn, NEG_INF_F64, fmax, +) -TROPICAL_GEMM_F64(tropical_minplus_f64_nn, INF_F64, fmin, +) -TROPICAL_GEMM_F64(tropical_maxmul_f64_nn, 0.0, fmax, *) +TROPICAL_GEMM_F64(tropical_maxplus_f64_nn, NEG_INF_F64, fmax, add_f64) +TROPICAL_GEMM_F64(tropical_minplus_f64_nn, INF_F64, fmin, add_f64) +TROPICAL_GEMM_F64(tropical_maxmul_f64_nn, 0.0, fmax, mul_f64) // --- F32 GEMM with Argmax Kernels --- TROPICAL_GEMM_ARGMAX_F32(tropical_maxplus_f32_nn_with_argmax, NEG_INF_F32, >, +) @@ -1075,3 +860,78 @@ extern "C" __global__ void KERNEL_NAME( \ TROPICAL_GEMM_BATCHED_ARGMAX_F32(tropical_maxplus_f32_nn_batched_with_argmax, NEG_INF_F32, >, +) TROPICAL_GEMM_BATCHED_ARGMAX_F32(tropical_minplus_f32_nn_batched_with_argmax, INF_F32, <, +) TROPICAL_GEMM_BATCHED_ARGMAX_F32(tropical_maxmul_f32_nn_batched_with_argmax, 0.0f, >, *) + +// ============================================================================ +// STRIDED-BATCHED FORWARD GEMM KERNELS (blockIdx.z selects the batch element) +// ============================================================================ +// Same shared body as the single-matrix kernels; only the operand base pointers +// differ -- each advances by its per-batch stride. This replaces omeinsum's +// host-side per-slice clone_dtod loop with a single launch over already- +// contiguous device buffers. `(int)blockIdx.z` matches the original `int +// batch_idx = blockIdx.z` indexing; the host chunks the batch under the +// gridDim.z cap and pre-offsets each chunk's base, so blockIdx.z stays in range. + +#define TROPICAL_GEMM_BATCHED_F32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const float* __restrict__ A, const float* __restrict__ B, \ + float* __restrict__ C, int M, int N, int K, \ + int strideA, int strideB, int strideC \ +) { \ + TROPICAL_GEMM_BODY(float, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, \ + (A + (int)blockIdx.z * strideA), \ + (B + (int)blockIdx.z * strideB), \ + (C + (int)blockIdx.z * strideC)) \ +} + +#define TROPICAL_GEMM_BATCHED_F64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const double* __restrict__ A, const double* __restrict__ B, \ + double* __restrict__ C, int M, int N, int K, \ + int strideA, int strideB, int strideC \ +) { \ + TROPICAL_GEMM_BODY(double, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, \ + (A + (int)blockIdx.z * strideA), \ + (B + (int)blockIdx.z * strideB), \ + (C + (int)blockIdx.z * strideC)) \ +} + +#define TROPICAL_GEMM_BATCHED_I32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const int* __restrict__ A, const int* __restrict__ B, \ + int* __restrict__ C, int M, int N, int K, \ + int strideA, int strideB, int strideC \ +) { \ + TROPICAL_GEMM_BODY(int, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, \ + (A + (int)blockIdx.z * strideA), \ + (B + (int)blockIdx.z * strideB), \ + (C + (int)blockIdx.z * strideC)) \ +} + +#define TROPICAL_GEMM_BATCHED_I64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +extern "C" __global__ void KERNEL_NAME( \ + const long long* __restrict__ A, const long long* __restrict__ B, \ + long long* __restrict__ C, int M, int N, int K, \ + int strideA, int strideB, int strideC \ +) { \ + TROPICAL_GEMM_BODY(long long, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, \ + (A + (int)blockIdx.z * strideA), \ + (B + (int)blockIdx.z * strideB), \ + (C + (int)blockIdx.z * strideC)) \ +} + +// --- Forward batched GEMM kernel instantiations (mirror the non-batched set) --- +TROPICAL_GEMM_BATCHED_F32(tropical_maxplus_f32_nn_batched, NEG_INF_F32, fmaxf, add_f32) +TROPICAL_GEMM_BATCHED_F32(tropical_minplus_f32_nn_batched, INF_F32, fminf, add_f32) +TROPICAL_GEMM_BATCHED_F32(tropical_maxmul_f32_nn_batched, 0.0f, fmaxf, mul_f32) + +TROPICAL_GEMM_BATCHED_F64(tropical_maxplus_f64_nn_batched, NEG_INF_F64, fmax, add_f64) +TROPICAL_GEMM_BATCHED_F64(tropical_minplus_f64_nn_batched, INF_F64, fmin, add_f64) +TROPICAL_GEMM_BATCHED_F64(tropical_maxmul_f64_nn_batched, 0.0, fmax, mul_f64) + +TROPICAL_GEMM_BATCHED_I32(tropical_maxplus_i32_nn_batched, NEG_INF_I32, max_i32, add_i32) +TROPICAL_GEMM_BATCHED_I32(tropical_minplus_i32_nn_batched, INF_I32, min_i32, add_i32) +TROPICAL_GEMM_BATCHED_I32(tropical_maxmul_i32_nn_batched, 0, max_i32, mul_i32) + +TROPICAL_GEMM_BATCHED_I64(tropical_maxplus_i64_nn_batched, NEG_INF_I64, max_i64, add_i64) +TROPICAL_GEMM_BATCHED_I64(tropical_minplus_i64_nn_batched, INF_I64, min_i64, add_i64) +TROPICAL_GEMM_BATCHED_I64(tropical_maxmul_i64_nn_batched, 0LL, max_i64, mul_i64) diff --git a/crates/tropical-gemm-cuda/src/context.rs b/crates/tropical-gemm-cuda/src/context.rs index 0810c57..fd00951 100644 --- a/crates/tropical-gemm-cuda/src/context.rs +++ b/crates/tropical-gemm-cuda/src/context.rs @@ -63,6 +63,19 @@ const KERNEL_NAMES: &[&str] = &[ "tropical_maxplus_f32_nn_batched_with_argmax", "tropical_minplus_f32_nn_batched_with_argmax", "tropical_maxmul_f32_nn_batched_with_argmax", + // Forward batched GEMM kernels (no argmax): one launch, blockIdx.z = batch. + "tropical_maxplus_f32_nn_batched", + "tropical_minplus_f32_nn_batched", + "tropical_maxmul_f32_nn_batched", + "tropical_maxplus_f64_nn_batched", + "tropical_minplus_f64_nn_batched", + "tropical_maxmul_f64_nn_batched", + "tropical_maxplus_i32_nn_batched", + "tropical_minplus_i32_nn_batched", + "tropical_maxmul_i32_nn_batched", + "tropical_maxplus_i64_nn_batched", + "tropical_minplus_i64_nn_batched", + "tropical_maxmul_i64_nn_batched", ]; /// CUDA context for tropical GEMM operations. diff --git a/crates/tropical-gemm-cuda/src/kernels.rs b/crates/tropical-gemm-cuda/src/kernels.rs index ef839c3..23fb666 100644 --- a/crates/tropical-gemm-cuda/src/kernels.rs +++ b/crates/tropical-gemm-cuda/src/kernels.rs @@ -15,13 +15,36 @@ //! barrier (e.g. before timing) can call `ctx.stream().synchronize()`. use crate::context::CudaContext; +use crate::error::CudaError; use crate::error::Result; use crate::memory::{ ExternalGpuMatrix, ExternalGpuTensor3, GpuMatrix, GpuMatrixWithArgmax, GpuTensor3WithArgmax, }; -use cudarc::driver::{DeviceRepr, LaunchConfig, PushKernelArg, ValidAsZeroBits}; +use cudarc::driver::{CudaSlice, DeviceRepr, LaunchConfig, PushKernelArg, ValidAsZeroBits}; use tropical_gemm::types::{TropicalMaxMul, TropicalMaxPlus, TropicalMinPlus, TropicalSemiring}; +/// Maximum extent of `gridDim.z` on all CUDA compute capabilities (the batch +/// dimension of the strided-batched kernels maps to `blockIdx.z`). A single +/// launch with more than this many batch elements fails at launch time with +/// `CUDA_ERROR_INVALID_VALUE`, so batched launches are split into chunks of at +/// most this size (see `launch_kernel_batched_impl`). +const MAX_GRID_DIM_Z: usize = 65535; + +/// Convert a per-batch element stride to the `i32` the kernels take, failing +/// instead of truncating. The chunked launches derive operand base offsets from +/// the `usize` stride while the kernel reads the stride as `i32`; an `as i32` +/// cast that wrapped would desynchronise the two and silently corrupt +/// addressing. The kernel signature is `i32`, so a stride past `i32::MAX` is +/// unrepresentable on the device regardless — reject it at the boundary. +fn stride_to_i32(stride: usize, what: &str) -> Result { + i32::try_from(stride).map_err(|_| { + CudaError::DimensionMismatch(format!( + "batched GEMM {what} stride {stride} exceeds i32::MAX; \ + matrix too large for the strided-batched kernel" + )) + }) +} + /// Trait for types that can be computed on GPU. pub trait CudaKernel: TropicalSemiring where @@ -30,6 +53,9 @@ where /// Kernel function name. const KERNEL_NAME: &'static str; + /// Forward batched kernel function name (one launch, `blockIdx.z` = batch). + const BATCHED_KERNEL_NAME: &'static str; + /// Execute the tropical GEMM kernel. /// /// Computes C = A ⊗ B where ⊗ is tropical matrix multiplication. @@ -39,6 +65,26 @@ where b: &GpuMatrix, c: &mut GpuMatrix, ) -> Result<()>; + + /// Execute `batch` independent tropical GEMMs in a single launch. + /// + /// Computes `C[i] = A[i] ⊗ B[i]` for `i in 0..batch` over contiguous, + /// already-device-resident operands: `a` is `batch × m × k`, `b` is + /// `batch × k × n`, `c` is `batch × m × n`, each column-major per matrix + /// with contiguous per-batch stride (`m*k`, `k*n`, `m*n`). This is the + /// strided-batched replacement for a host-side per-slice loop: no per-slice + /// `clone_dtod`, no per-slice allocation, no reassembly copy. The kernel + /// fully writes every element of `c`, so `c` may be uninitialized on entry. + fn launch_gemm_batched( + ctx: &CudaContext, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + ) -> Result<()>; } /// Helper function to launch a CUDA kernel with given grid/block dimensions. @@ -85,12 +131,100 @@ fn launch_kernel_impl( Ok(()) } +/// Helper to launch a forward batched kernel over contiguous device buffers. +/// +/// `grid.z` selects the batch element; the per-batch strides are the contiguous +/// extents `m*k` / `k*n` / `m*n`. Operands are passed as borrowed `CudaSlice`s +/// — no copy, no allocation. Follows the same asynchronous launch contract as +/// [`launch_kernel_impl`]. +#[allow(clippy::too_many_arguments)] +fn launch_kernel_batched_impl( + ctx: &CudaContext, + kernel_name: &'static str, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + grid_xy: u32, + block: (u32, u32, u32), +) -> Result<()> { + // The raw-slice API loses GpuMatrix's dimension invariant, so validate the + // contiguous batched extents before handing pointers to the kernel. + let want = |dim: usize, len: usize, what: &str| -> Result<()> { + if len != dim { + return Err(CudaError::DimensionMismatch(format!( + "batched GEMM {what}: expected {dim} elements (batch={batch}, m={m}, k={k}, n={n}), got {len}" + ))); + } + Ok(()) + }; + want(batch * m * k, a.len(), "operand A")?; + want(batch * k * n, b.len(), "operand B")?; + want(batch * m * n, c.len(), "output C")?; + + let kernel = ctx.get_kernel(kernel_name)?; + + // Bind scalar kernel args to locals so they outlive the launch builder. + let m_i32 = m as i32; + let n_i32 = n as i32; + let k_i32 = k as i32; + // Per-batch element extents, used both to offset each chunk's operand base + // (as `usize`) and as the kernel's `i32` stride. Derive the `i32` form + // fallibly so a stride past `i32::MAX` can't truncate out of sync with the + // `usize` offsets below (a.len()/b.len()/c.len() already fit in `usize`). + let (sa, sb, sc) = (m * k, k * n, m * n); + let stride_a = stride_to_i32(sa, "operand A")?; + let stride_b = stride_to_i32(sb, "operand B")?; + let stride_c = stride_to_i32(sc, "output C")?; + + let stream = ctx.stream(); + // The batch maps to `blockIdx.z`, which CUDA caps at `MAX_GRID_DIM_Z`. Launch + // the batch in chunks of at most that size, offsetting every operand's base + // by the chunk start so `blockIdx.z ∈ [0, chunk)` indexes the correct slice. + // `grid_xy` from the caller is the x tiling; y is 1 and z is the chunk size. + let mut start = 0usize; + while start < batch { + let chunk = (batch - start).min(MAX_GRID_DIM_Z); + let a_view = a.slice(start * sa..(start + chunk) * sa); + let b_view = b.slice(start * sb..(start + chunk) * sb); + let mut c_view = c.slice_mut(start * sc..(start + chunk) * sc); + let cfg = LaunchConfig { + grid_dim: (grid_xy, 1, chunk as u32), + block_dim: block, + shared_mem_bytes: 0, + }; + let mut builder = stream.launch_builder(&kernel); + builder + .arg(&a_view) + .arg(&b_view) + .arg(&mut c_view) + .arg(&m_i32) + .arg(&n_i32) + .arg(&k_i32) + .arg(&stride_a) + .arg(&stride_b) + .arg(&stride_c); + unsafe { + builder.launch(cfg)?; + } + start += chunk; + } + + // Async: no per-launch device sync (stream-ordered; host reads sync in + // `to_host`). See the module-level "Asynchronous launch contract". + Ok(()) +} + /// Macro to implement CudaKernel for f32 types. macro_rules! impl_cuda_kernel_f32 { - ($($semiring:ty => $kernel_name:literal),* $(,)?) => { + ($($semiring:ty => $kernel_name:literal, $batched_name:literal);* $(;)?) => { $( impl CudaKernel for $semiring { const KERNEL_NAME: &'static str = $kernel_name; + const BATCHED_KERNEL_NAME: &'static str = $batched_name; fn launch_gemm( ctx: &CudaContext, @@ -102,6 +236,23 @@ macro_rules! impl_cuda_kernel_f32 { let block = CudaContext::block_dims_f32(); launch_kernel_impl(ctx, Self::KERNEL_NAME, a, b, c, grid, block) } + + fn launch_gemm_batched( + ctx: &CudaContext, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + ) -> Result<()> { + let grid_xy = CudaContext::grid_dims_f32(m, n).0; + let block = CudaContext::block_dims_f32(); + launch_kernel_batched_impl( + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, + ) + } } )* }; @@ -109,10 +260,11 @@ macro_rules! impl_cuda_kernel_f32 { /// Macro to implement CudaKernel for f64 types. macro_rules! impl_cuda_kernel_f64 { - ($($semiring:ty => $kernel_name:literal),* $(,)?) => { + ($($semiring:ty => $kernel_name:literal, $batched_name:literal);* $(;)?) => { $( impl CudaKernel for $semiring { const KERNEL_NAME: &'static str = $kernel_name; + const BATCHED_KERNEL_NAME: &'static str = $batched_name; fn launch_gemm( ctx: &CudaContext, @@ -124,30 +276,48 @@ macro_rules! impl_cuda_kernel_f64 { let block = CudaContext::block_dims_f64(); launch_kernel_impl(ctx, Self::KERNEL_NAME, a, b, c, grid, block) } + + fn launch_gemm_batched( + ctx: &CudaContext, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + ) -> Result<()> { + let grid_xy = CudaContext::grid_dims_f64(m, n).0; + let block = CudaContext::block_dims_f64(); + launch_kernel_batched_impl( + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, + ) + } } )* }; } impl_cuda_kernel_f32! { - TropicalMaxPlus => "tropical_maxplus_f32_nn", - TropicalMinPlus => "tropical_minplus_f32_nn", - TropicalMaxMul => "tropical_maxmul_f32_nn", + TropicalMaxPlus => "tropical_maxplus_f32_nn", "tropical_maxplus_f32_nn_batched"; + TropicalMinPlus => "tropical_minplus_f32_nn", "tropical_minplus_f32_nn_batched"; + TropicalMaxMul => "tropical_maxmul_f32_nn", "tropical_maxmul_f32_nn_batched"; } impl_cuda_kernel_f64! { - TropicalMaxPlus => "tropical_maxplus_f64_nn", - TropicalMinPlus => "tropical_minplus_f64_nn", - TropicalMaxMul => "tropical_maxmul_f64_nn", + TropicalMaxPlus => "tropical_maxplus_f64_nn", "tropical_maxplus_f64_nn_batched"; + TropicalMinPlus => "tropical_minplus_f64_nn", "tropical_minplus_f64_nn_batched"; + TropicalMaxMul => "tropical_maxmul_f64_nn", "tropical_maxmul_f64_nn_batched"; } /// Macro to implement CudaKernel for i32 types. /// Uses same block sizes as f32 (64x32x64) since int is 4 bytes. macro_rules! impl_cuda_kernel_i32 { - ($($semiring:ty => $kernel_name:literal),* $(,)?) => { + ($($semiring:ty => $kernel_name:literal, $batched_name:literal);* $(;)?) => { $( impl CudaKernel for $semiring { const KERNEL_NAME: &'static str = $kernel_name; + const BATCHED_KERNEL_NAME: &'static str = $batched_name; fn launch_gemm( ctx: &CudaContext, @@ -159,6 +329,23 @@ macro_rules! impl_cuda_kernel_i32 { let block = CudaContext::block_dims_f32(); launch_kernel_impl(ctx, Self::KERNEL_NAME, a, b, c, grid, block) } + + fn launch_gemm_batched( + ctx: &CudaContext, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + ) -> Result<()> { + let grid_xy = CudaContext::grid_dims_f32(m, n).0; + let block = CudaContext::block_dims_f32(); + launch_kernel_batched_impl( + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, + ) + } } )* }; @@ -167,10 +354,11 @@ macro_rules! impl_cuda_kernel_i32 { /// Macro to implement CudaKernel for i64 types. /// Uses same block sizes as f64 (32x16x32) since long long is 8 bytes. macro_rules! impl_cuda_kernel_i64 { - ($($semiring:ty => $kernel_name:literal),* $(,)?) => { + ($($semiring:ty => $kernel_name:literal, $batched_name:literal);* $(;)?) => { $( impl CudaKernel for $semiring { const KERNEL_NAME: &'static str = $kernel_name; + const BATCHED_KERNEL_NAME: &'static str = $batched_name; fn launch_gemm( ctx: &CudaContext, @@ -182,21 +370,38 @@ macro_rules! impl_cuda_kernel_i64 { let block = CudaContext::block_dims_f64(); launch_kernel_impl(ctx, Self::KERNEL_NAME, a, b, c, grid, block) } + + fn launch_gemm_batched( + ctx: &CudaContext, + a: &CudaSlice, + b: &CudaSlice, + c: &mut CudaSlice, + batch: usize, + m: usize, + k: usize, + n: usize, + ) -> Result<()> { + let grid_xy = CudaContext::grid_dims_f64(m, n).0; + let block = CudaContext::block_dims_f64(); + launch_kernel_batched_impl( + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, + ) + } } )* }; } impl_cuda_kernel_i32! { - TropicalMaxPlus => "tropical_maxplus_i32_nn", - TropicalMinPlus => "tropical_minplus_i32_nn", - TropicalMaxMul => "tropical_maxmul_i32_nn", + TropicalMaxPlus => "tropical_maxplus_i32_nn", "tropical_maxplus_i32_nn_batched"; + TropicalMinPlus => "tropical_minplus_i32_nn", "tropical_minplus_i32_nn_batched"; + TropicalMaxMul => "tropical_maxmul_i32_nn", "tropical_maxmul_i32_nn_batched"; } impl_cuda_kernel_i64! { - TropicalMaxPlus => "tropical_maxplus_i64_nn", - TropicalMinPlus => "tropical_minplus_i64_nn", - TropicalMaxMul => "tropical_maxmul_i64_nn", + TropicalMaxPlus => "tropical_maxplus_i64_nn", "tropical_maxplus_i64_nn_batched"; + TropicalMinPlus => "tropical_minplus_i64_nn", "tropical_minplus_i64_nn_batched"; + TropicalMaxMul => "tropical_maxmul_i64_nn", "tropical_maxmul_i64_nn_batched"; } // ============================================================================ @@ -548,50 +753,73 @@ pub unsafe fn launch_gemm_external_batched_with_argmax_f32( ) -> Result> { // Apply row-major → column-major trick: swap inputs and swap M↔N // Same as non-batched version, but for each batch - let mut c = GpuTensor3WithArgmax::::alloc(ctx, batch, m, n)?; + let c = GpuTensor3WithArgmax::::alloc(ctx, batch, m, n)?; - // Grid: (ceil(N/64) * ceil(M/64), 1, batch) with swapped M↔N - let grid_xy = ((n + 63) / 64) * ((m + 63) / 64); - let grid = (grid_xy as u32, 1, batch as u32); + // Grid: (ceil(N/64) * ceil(M/64), 1, batch) with swapped M↔N. Reuse the + // shared tile-count helper (keyed off the f32 block constants) rather than + // hardcoding 64, which would silently desync if the block size changed. + let grid_xy = CudaContext::grid_dims_f32(n, m).0; let block = CudaContext::block_dims_f32(); - let kernel = ctx.get_kernel(kernel_name)?; - let cfg = LaunchConfig { - grid_dim: grid, - block_dim: block, - shared_mem_bytes: 0, - }; - // Compute strides before borrowing mutably - let stride_a = a.stride() as i32; - let stride_b = b.stride() as i32; - let stride_c = c.tensor.stride() as i32; - - // Bind raw external pointers and scalar args to locals so they outlive the - // launch builder. - let b_ptr: u64 = b.device_ptr(); // B becomes "A" in kernel - let a_ptr: u64 = a.device_ptr(); // A becomes "B" in kernel + // Per-batch element strides. A/B carry the (possibly padded) stride declared + // by the external DLPack tensor; C/argmax are freshly allocated contiguous, + // so their stride is rows*cols. Derive the kernel's `i32` strides fallibly + // so a stride past `i32::MAX` can't truncate out of sync with the `usize` + // base-pointer offsets below. + let stride_a = a.stride(); + let stride_b = b.stride(); + let stride_c = c.tensor.stride(); + let stride_a_i32 = stride_to_i32(stride_a, "operand A")?; + let stride_b_i32 = stride_to_i32(stride_b, "operand B")?; + let stride_c_i32 = stride_to_i32(stride_c, "output C")?; let n_i32 = n as i32; // Swapped: N becomes "M" let m_i32 = m as i32; // Swapped: M becomes "N" let k_i32 = k as i32; - // Swap order: pass B first, then A, and swap M↔N + // Operand base device addresses (byte pointers). C and argmax are owned by + // `c`, which outlives every (async, stream-ordered) launch below. + let a_base: u64 = a.device_ptr(); + let b_base: u64 = b.device_ptr(); + let c_base: u64 = c.tensor.device_ptr(); + let am_base: u64 = c.argmax.device_ptr(); + let f32_bytes = std::mem::size_of::() as u64; + let idx_bytes = std::mem::size_of::() as u64; // ArgmaxIndex + + // The batch maps to `blockIdx.z` (CUDA cap `MAX_GRID_DIM_Z`). Launch in + // chunks of at most that size, advancing each operand base pointer by the + // chunk start so `blockIdx.z ∈ [0, chunk)` indexes the correct batch slice. + // Swap order: pass B first, then A, and swap M↔N (row-major→col-major). // Kernel signature: (A, B, C, argmax, M, N, K, strideA, strideB, strideC) // We pass: (B, A, C, argmax, N, M, K, strideB, strideA, strideC) let stream = ctx.stream(); - let mut builder = stream.launch_builder(&kernel); - builder - .arg(&b_ptr) - .arg(&a_ptr) - .arg(c.tensor.as_slice_mut()) - .arg(c.argmax.as_slice_mut()) - .arg(&n_i32) - .arg(&m_i32) - .arg(&k_i32) - .arg(&stride_b) // strideA (B's stride in our swap) - .arg(&stride_a) // strideB (A's stride in our swap) - .arg(&stride_c); // strideC - builder.launch(cfg)?; + let mut start = 0usize; + while start < batch { + let chunk = (batch - start).min(MAX_GRID_DIM_Z); + let a_ptr = a_base + (start * stride_a) as u64 * f32_bytes; + let b_ptr = b_base + (start * stride_b) as u64 * f32_bytes; + let c_ptr = c_base + (start * stride_c) as u64 * f32_bytes; + let am_ptr = am_base + (start * stride_c) as u64 * idx_bytes; + let cfg = LaunchConfig { + grid_dim: (grid_xy, 1, chunk as u32), + block_dim: block, + shared_mem_bytes: 0, + }; + let mut builder = stream.launch_builder(&kernel); + builder + .arg(&b_ptr) + .arg(&a_ptr) + .arg(&c_ptr) + .arg(&am_ptr) + .arg(&n_i32) + .arg(&m_i32) + .arg(&k_i32) + .arg(&stride_b_i32) // strideA (B's stride in our swap) + .arg(&stride_a_i32) // strideB (A's stride in our swap) + .arg(&stride_c_i32); // strideC + builder.launch(cfg)?; + start += chunk; + } // Async: no per-launch device sync (stream-ordered; host reads sync in // `to_host`). See the module-level "Asynchronous launch contract". diff --git a/crates/tropical-gemm-cuda/src/lib.rs b/crates/tropical-gemm-cuda/src/lib.rs index 418c473..ab3a7a8 100644 --- a/crates/tropical-gemm-cuda/src/lib.rs +++ b/crates/tropical-gemm-cuda/src/lib.rs @@ -1202,6 +1202,89 @@ mod tests { assert!((c[3] - 12.0).abs() < 1e-5, "C[1,1] = {}, expected 12", c[3]); } + /// CPU column-major max-plus reference: C[i,j] = max_k(A[i,k] + B[k,j]). + fn cpu_maxplus(a: &[f32], m: usize, k: usize, b: &[f32], n: usize) -> Vec { + let mut c = vec![f32::NEG_INFINITY; m * n]; + for j in 0..n { + for i in 0..m { + let mut acc = f32::NEG_INFINITY; + for p in 0..k { + acc = acc.max(a[p * m + i] + b[j * k + p]); + } + c[j * m + i] = acc; + } + } + c + } + + /// The forward batched kernel must match the trusted single-matrix path for + /// every batch element, including non-block-aligned sizes (edge tiles). This + /// is the correctness gate for `launch_gemm_batched`. + #[test] + fn test_tropical_gemm_batched_matches_single() { + use cudarc::driver::CudaSlice; + let Some(ctx) = cuda_context_or_skip() else { + return; + }; + let stream = ctx.stream(); + + // Cases chosen to exercise: a tiny hand-checkable shape, a single full + // tile, and a multi-block shape whose M/N are NOT multiples of the + // 64-wide block (so the guarded edge-tile load/store path runs). + for &(batch, m, k, n) in &[(2usize, 2usize, 3usize, 2usize), (3, 70, 33, 50)] { + // Build `batch` distinct operands, contiguous column-major per slice. + let mut a_all = vec![0f32; batch * m * k]; + let mut b_all = vec![0f32; batch * k * n]; + for bi in 0..batch { + for idx in 0..m * k { + a_all[bi * m * k + idx] = (((bi * 7 + idx * 3) % 13) as f32) - 6.0; + } + for idx in 0..k * n { + b_all[bi * k * n + idx] = (((bi * 5 + idx * 2) % 11) as f32) - 5.0; + } + } + + // Reference: trusted single-matrix GPU path + CPU, per slice. + let mut expected = vec![0f32; batch * m * n]; + for bi in 0..batch { + let a_i = &a_all[bi * m * k..(bi + 1) * m * k]; + let b_i = &b_all[bi * k * n..(bi + 1) * k * n]; + let ref_gpu = + tropical_matmul_gpu::>(a_i, m, k, b_i, n).unwrap(); + let ref_cpu = cpu_maxplus(a_i, m, k, b_i, n); + for idx in 0..m * n { + assert!( + (ref_gpu[idx] - ref_cpu[idx]).abs() < 1e-4, + "single-kernel disagrees with CPU ref at batch {bi} idx {idx}: {} vs {}", + ref_gpu[idx], + ref_cpu[idx] + ); + } + expected[bi * m * n..(bi + 1) * m * n].copy_from_slice(&ref_gpu); + } + + // Batched: one launch over the whole contiguous buffer, uninit output. + let a_dev: CudaSlice = stream.clone_htod(&a_all).unwrap(); + let b_dev: CudaSlice = stream.clone_htod(&b_all).unwrap(); + let mut c_dev = unsafe { stream.alloc::(batch * m * n) }.unwrap(); + as CudaKernel>::launch_gemm_batched( + ctx, &a_dev, &b_dev, &mut c_dev, batch, m, k, n, + ) + .unwrap(); + let got = stream.clone_dtoh(&c_dev).unwrap(); + stream.synchronize().unwrap(); + + for idx in 0..batch * m * n { + assert!( + (got[idx] - expected[idx]).abs() < 1e-4, + "batched != single at (batch,m,k,n)=({batch},{m},{k},{n}) idx {idx}: {} vs {}", + got[idx], + expected[idx] + ); + } + } + } + #[test] fn test_tropical_matmul_gpu_with_argmax_maxplus() { if cuda_context_or_skip().is_none() { @@ -1659,14 +1742,12 @@ mod tests { // 2 batches of 2x2 matrices, stored contiguously (column-major) let a = vec![ // Batch 0: [[1,2],[3,4]] col-major: [1,3,2,4] - 1.0f32, 3.0, 2.0, 4.0, - // Batch 1: [[5,6],[7,8]] col-major: [5,7,6,8] + 1.0f32, 3.0, 2.0, 4.0, // Batch 1: [[5,6],[7,8]] col-major: [5,7,6,8] 5.0, 7.0, 6.0, 8.0, ]; let b = vec![ // Batch 0: [[1,0],[0,1]] col-major: [1,0,0,1] - 1.0f32, 0.0, 0.0, 1.0, - // Batch 1: [[1,2],[3,4]] col-major: [1,3,2,4] + 1.0f32, 0.0, 0.0, 1.0, // Batch 1: [[1,2],[3,4]] col-major: [1,3,2,4] 1.0, 3.0, 2.0, 4.0, ];