diff --git a/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu b/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu index e6bbd36..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,16 +202,125 @@ __device__ double atomicAddDouble(double* address, double val) { } // ============================================================================ -// F32 GEMM KERNEL MACRO +// SHARED FORWARD GEMM KERNEL BODY +// ============================================================================ +// 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) + +// ============================================================================ +// SINGLE-MATRIX FORWARD GEMM KERNELS (blockIdx.x carries the M/N tiling) +// ============================================================================ +// 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) \ +} + +#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) \ +} + +#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 \ +) { \ + TROPICAL_GEMM_BODY(long long, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, A, B, C) \ +} + +// ============================================================================ +// F32 GEMM WITH ARGMAX KERNEL MACRO // ============================================================================ -// 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) \ +#define TROPICAL_GEMM_ARGMAX_F32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ extern "C" __global__ void KERNEL_NAME( \ const float* __restrict__ A, \ const float* __restrict__ B, \ float* __restrict__ C, \ + unsigned int* __restrict__ argmax, \ int M, int N, int K \ ) { \ const int BLOCK_SIZE_M = 64; \ @@ -226,12 +344,14 @@ extern "C" __global__ void KERNEL_NAME( \ __shared__ float Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ \ float accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ + int accum_idx[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; \ + accum_idx[i] = 0; \ } \ \ const int A_TILE_COL = tid / BLOCK_SIZE_M; \ @@ -250,6 +370,7 @@ extern "C" __global__ void KERNEL_NAME( \ \ _Pragma("unroll") \ for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ + int global_k = tile_idx + k; \ _Pragma("unroll") \ for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ @@ -266,26 +387,29 @@ extern "C" __global__ void KERNEL_NAME( \ 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); \ + if (prod COMPARE_OP accum[idx]) { \ + accum[idx] = prod; \ + accum_idx[idx] = global_k; \ + } \ } \ } \ } \ __syncthreads(); \ } \ \ - STORE_C_TILE(C) \ + STORE_C_ARGMAX_TILE(C, argmax) \ } // ============================================================================ -// F64 GEMM KERNEL MACRO +// F64 GEMM WITH ARGMAX KERNEL MACRO // ============================================================================ -// Block sizes for f64: 32x16x32, Thread sizes: 4x4 -#define TROPICAL_GEMM_F64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_OP) \ +#define TROPICAL_GEMM_ARGMAX_F64(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ extern "C" __global__ void KERNEL_NAME( \ const double* __restrict__ A, \ const double* __restrict__ B, \ double* __restrict__ C, \ + unsigned int* __restrict__ argmax, \ int M, int N, int K \ ) { \ const int BLOCK_SIZE_M = 32; \ @@ -309,12 +433,14 @@ extern "C" __global__ void KERNEL_NAME( \ __shared__ double Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ \ double accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ + int accum_idx[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; \ + accum_idx[i] = 0; \ } \ \ const int A_TILE_COL = tid / BLOCK_SIZE_M; \ @@ -333,6 +459,7 @@ extern "C" __global__ void KERNEL_NAME( \ \ _Pragma("unroll") \ for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ + int global_k = tile_idx + k; \ _Pragma("unroll") \ for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ @@ -349,27 +476,29 @@ extern "C" __global__ void KERNEL_NAME( \ 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); \ + if (prod COMPARE_OP accum[idx]) { \ + accum[idx] = prod; \ + accum_idx[idx] = global_k; \ + } \ } \ } \ } \ __syncthreads(); \ } \ \ - STORE_C_TILE(C) \ + STORE_C_ARGMAX_TILE(C, argmax) \ } // ============================================================================ -// I32 GEMM KERNEL MACRO +// I32 GEMM WITH ARGMAX 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) \ +#define TROPICAL_GEMM_ARGMAX_I32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_FN, ZERO_FN) \ extern "C" __global__ void KERNEL_NAME( \ const int* __restrict__ A, \ const int* __restrict__ B, \ int* __restrict__ C, \ + unsigned int* __restrict__ argmax, \ int M, int N, int K \ ) { \ const int BLOCK_SIZE_M = 64; \ @@ -393,12 +522,14 @@ extern "C" __global__ void KERNEL_NAME( \ __shared__ int Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ \ int accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ + int accum_idx[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; \ + accum_idx[i] = 0; \ } \ \ const int A_TILE_COL = tid / BLOCK_SIZE_M; \ @@ -417,6 +548,7 @@ extern "C" __global__ void KERNEL_NAME( \ \ _Pragma("unroll") \ for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ + int global_k = tile_idx + k; \ _Pragma("unroll") \ for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ @@ -433,26 +565,29 @@ extern "C" __global__ void KERNEL_NAME( \ 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); \ + if (prod COMPARE_OP accum[idx]) { \ + accum[idx] = prod; \ + accum_idx[idx] = global_k; \ + } \ } \ } \ } \ __syncthreads(); \ } \ \ - STORE_C_TILE(C) \ + STORE_C_ARGMAX_ZERO_TILE(C, argmax, ZERO_FN) \ } // ============================================================================ -// I64 GEMM KERNEL MACRO +// I64 GEMM WITH ARGMAX KERNEL MACRO // ============================================================================ -// Block sizes for i64: 32x16x32, Thread sizes: 4x4 (same as f64) -#define TROPICAL_GEMM_I64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ +#define TROPICAL_GEMM_ARGMAX_I64(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_FN, ZERO_FN) \ extern "C" __global__ void KERNEL_NAME( \ const long long* __restrict__ A, \ const long long* __restrict__ B, \ long long* __restrict__ C, \ + unsigned int* __restrict__ argmax, \ int M, int N, int K \ ) { \ const int BLOCK_SIZE_M = 32; \ @@ -476,12 +611,14 @@ extern "C" __global__ void KERNEL_NAME( \ __shared__ long long Bs[BLOCK_SIZE_K * BLOCK_SIZE_N]; \ \ long long accum[THREAD_SIZE_M * THREAD_SIZE_N]; \ + int accum_idx[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; \ + accum_idx[i] = 0; \ } \ \ const int A_TILE_COL = tid / BLOCK_SIZE_M; \ @@ -500,6 +637,7 @@ extern "C" __global__ void KERNEL_NAME( \ \ _Pragma("unroll") \ for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ + int global_k = tile_idx + k; \ _Pragma("unroll") \ for (int tm = 0; tm < THREAD_SIZE_M; ++tm) { \ regs_a[tm] = As[OFFSET_COL(threadIdx.x * THREAD_SIZE_M + tm, \ @@ -516,27 +654,124 @@ extern "C" __global__ void KERNEL_NAME( \ 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); \ + if (prod COMPARE_OP accum[idx]) { \ + accum[idx] = prod; \ + accum_idx[idx] = global_k; \ + } \ } \ } \ } \ __syncthreads(); \ } \ \ - STORE_C_TILE(C) \ + STORE_C_ARGMAX_ZERO_TILE(C, argmax, ZERO_FN) \ } // ============================================================================ -// F32 GEMM WITH ARGMAX KERNEL MACRO +// BACKWARD PASS KERNEL MACROS // ============================================================================ -#define TROPICAL_GEMM_ARGMAX_F32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ +#define TROPICAL_BACKWARD_A(KERNEL_NAME, TYPE, ATOMIC_ADD) \ +extern "C" __global__ void KERNEL_NAME( \ + const TYPE* __restrict__ grad_c, \ + const unsigned int* __restrict__ argmax, \ + TYPE* __restrict__ grad_a, \ + int M, int N, int K \ +) { \ + int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + int total = M * N; \ + if (idx < total) { \ + int i = idx % M; \ + int k = (int)argmax[idx]; \ + if (k >= 0 && k < K) { \ + ATOMIC_ADD(&grad_a[i + k * M], grad_c[idx]); \ + } \ + } \ +} + +#define TROPICAL_BACKWARD_B(KERNEL_NAME, TYPE, ATOMIC_ADD) \ +extern "C" __global__ void KERNEL_NAME( \ + const TYPE* __restrict__ grad_c, \ + const unsigned int* __restrict__ argmax, \ + TYPE* __restrict__ grad_b, \ + int M, int N, int K \ +) { \ + int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + int total = M * N; \ + if (idx < total) { \ + int j = idx / M; \ + int k = (int)argmax[idx]; \ + if (k >= 0 && k < K) { \ + ATOMIC_ADD(&grad_b[k + j * K], grad_c[idx]); \ + } \ + } \ +} + +// ============================================================================ +// KERNEL INSTANTIATIONS +// ============================================================================ + +// --- F32 Basic GEMM Kernels --- +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, 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, >, +) +TROPICAL_GEMM_ARGMAX_F32(tropical_minplus_f32_nn_with_argmax, INF_F32, <, +) +TROPICAL_GEMM_ARGMAX_F32(tropical_maxmul_f32_nn_with_argmax, 0.0f, >, *) + +// --- F64 GEMM with Argmax Kernels --- +TROPICAL_GEMM_ARGMAX_F64(tropical_maxplus_f64_nn_with_argmax, NEG_INF_F64, >, +) +TROPICAL_GEMM_ARGMAX_F64(tropical_minplus_f64_nn_with_argmax, INF_F64, <, +) +TROPICAL_GEMM_ARGMAX_F64(tropical_maxmul_f64_nn_with_argmax, 0.0, >, *) + +// --- I32 Basic GEMM Kernels --- +TROPICAL_GEMM_I32(tropical_maxplus_i32_nn, NEG_INF_I32, max_i32, add_i32) +TROPICAL_GEMM_I32(tropical_minplus_i32_nn, INF_I32, min_i32, add_i32) +TROPICAL_GEMM_I32(tropical_maxmul_i32_nn, 0, max_i32, mul_i32) + +// --- I64 Basic GEMM Kernels --- +TROPICAL_GEMM_I64(tropical_maxplus_i64_nn, NEG_INF_I64, max_i64, add_i64) +TROPICAL_GEMM_I64(tropical_minplus_i64_nn, INF_I64, min_i64, add_i64) +TROPICAL_GEMM_I64(tropical_maxmul_i64_nn, 0LL, max_i64, mul_i64) + +// --- I32 GEMM with Argmax Kernels --- +// ZERO_FN canonicalizes a drifted tropical-zero cell's argmax to 0 at write-out. +TROPICAL_GEMM_ARGMAX_I32(tropical_maxplus_i32_nn_with_argmax, NEG_INF_I32, >, add_i32, zero_maxplus_i32) +TROPICAL_GEMM_ARGMAX_I32(tropical_minplus_i32_nn_with_argmax, INF_I32, <, add_i32, zero_minplus_i32) +TROPICAL_GEMM_ARGMAX_I32(tropical_maxmul_i32_nn_with_argmax, 0, >, mul_i32, zero_never_i32) + +// --- I64 GEMM with Argmax Kernels --- +TROPICAL_GEMM_ARGMAX_I64(tropical_maxplus_i64_nn_with_argmax, NEG_INF_I64, >, add_i64, zero_maxplus_i64) +TROPICAL_GEMM_ARGMAX_I64(tropical_minplus_i64_nn_with_argmax, INF_I64, <, add_i64, zero_minplus_i64) +TROPICAL_GEMM_ARGMAX_I64(tropical_maxmul_i64_nn_with_argmax, 0LL, >, mul_i64, zero_never_i64) + +// --- Backward Pass Kernels (float/double only, no integer gradients) --- +TROPICAL_BACKWARD_A(tropical_backward_a_f32, float, atomicAdd) +TROPICAL_BACKWARD_B(tropical_backward_b_f32, float, atomicAdd) +TROPICAL_BACKWARD_A(tropical_backward_a_f64, double, atomicAddDouble) +TROPICAL_BACKWARD_B(tropical_backward_b_f64, double, atomicAddDouble) + +// ============================================================================ +// BATCHED F32 GEMM WITH ARGMAX KERNEL MACRO +// ============================================================================ +// Strided batched GEMM: processes batch_size independent GEMMs +// Uses blockIdx.z for batch index, strides for memory offsets + +#define TROPICAL_GEMM_BATCHED_ARGMAX_F32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ extern "C" __global__ void KERNEL_NAME( \ const float* __restrict__ A, \ const float* __restrict__ B, \ float* __restrict__ C, \ unsigned int* __restrict__ argmax, \ - int M, int N, int K \ + int M, int N, int K, \ + int strideA, int strideB, int strideC \ ) { \ const int BLOCK_SIZE_M = 64; \ const int BLOCK_SIZE_K = 32; \ @@ -548,11 +783,17 @@ extern "C" __global__ void KERNEL_NAME( \ const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ \ + int batch_idx = blockIdx.z; \ 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 float* A_batch = A + batch_idx * strideA; \ + const float* B_batch = B + batch_idx * strideB; \ + float* C_batch = C + batch_idx * strideC; \ + unsigned int* argmax_batch = argmax + batch_idx * strideC; \ + \ const int tid = threadIdx.y * bszm + threadIdx.x; \ \ __shared__ float As[BLOCK_SIZE_M * BLOCK_SIZE_K]; \ @@ -577,9 +818,9 @@ extern "C" __global__ void KERNEL_NAME( \ 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_A_TILE(A_batch, INIT_VAL) \ \ - LOAD_B_TILE(B, INIT_VAL) \ + LOAD_B_TILE(B_batch, INIT_VAL) \ \ __syncthreads(); \ \ @@ -612,824 +853,80 @@ extern "C" __global__ void KERNEL_NAME( \ __syncthreads(); \ } \ \ - STORE_C_ARGMAX_TILE(C, argmax) \ + STORE_C_ARGMAX_TILE(C_batch, argmax_batch) \ } +// --- Batched F32 GEMM with Argmax Kernels --- +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, >, *) + // ============================================================================ -// F64 GEMM WITH ARGMAX KERNEL MACRO +// 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_ARGMAX_F64(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ -extern "C" __global__ void KERNEL_NAME( \ - const double* __restrict__ A, \ - const double* __restrict__ B, \ - double* __restrict__ C, \ - unsigned int* __restrict__ argmax, \ - 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]; \ - int accum_idx[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; \ - accum_idx[i] = 0; \ - } \ - \ - 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) { \ - int global_k = tile_idx + 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); \ - if (prod COMPARE_OP accum[idx]) { \ - accum[idx] = prod; \ - accum_idx[idx] = global_k; \ - } \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_ARGMAX_TILE(C, argmax) \ +#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)) \ } -// ============================================================================ -// I32 GEMM WITH ARGMAX KERNEL MACRO -// ============================================================================ - -#define TROPICAL_GEMM_ARGMAX_I32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_FN, ZERO_FN) \ -extern "C" __global__ void KERNEL_NAME( \ - const int* __restrict__ A, \ - const int* __restrict__ B, \ - int* __restrict__ C, \ - unsigned int* __restrict__ argmax, \ - 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 accum_idx[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; \ - accum_idx[i] = 0; \ - } \ - \ - 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) { \ - int global_k = tile_idx + 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); \ - if (prod COMPARE_OP accum[idx]) { \ - accum[idx] = prod; \ - accum_idx[idx] = global_k; \ - } \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_ARGMAX_ZERO_TILE(C, argmax, ZERO_FN) \ -} - -// ============================================================================ -// I64 GEMM WITH ARGMAX KERNEL MACRO -// ============================================================================ - -#define TROPICAL_GEMM_ARGMAX_I64(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_FN, ZERO_FN) \ -extern "C" __global__ void KERNEL_NAME( \ - const long long* __restrict__ A, \ - const long long* __restrict__ B, \ - long long* __restrict__ C, \ - unsigned int* __restrict__ argmax, \ - 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]; \ - int accum_idx[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; \ - accum_idx[i] = 0; \ - } \ - \ - 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) { \ - int global_k = tile_idx + 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); \ - if (prod COMPARE_OP accum[idx]) { \ - accum[idx] = prod; \ - accum_idx[idx] = global_k; \ - } \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_ARGMAX_ZERO_TILE(C, argmax, ZERO_FN) \ -} - -// ============================================================================ -// BACKWARD PASS KERNEL MACROS -// ============================================================================ - -#define TROPICAL_BACKWARD_A(KERNEL_NAME, TYPE, ATOMIC_ADD) \ -extern "C" __global__ void KERNEL_NAME( \ - const TYPE* __restrict__ grad_c, \ - const unsigned int* __restrict__ argmax, \ - TYPE* __restrict__ grad_a, \ - int M, int N, int K \ -) { \ - int idx = blockIdx.x * blockDim.x + threadIdx.x; \ - int total = M * N; \ - if (idx < total) { \ - int i = idx % M; \ - int k = (int)argmax[idx]; \ - if (k >= 0 && k < K) { \ - ATOMIC_ADD(&grad_a[i + k * M], grad_c[idx]); \ - } \ - } \ -} - -#define TROPICAL_BACKWARD_B(KERNEL_NAME, TYPE, ATOMIC_ADD) \ -extern "C" __global__ void KERNEL_NAME( \ - const TYPE* __restrict__ grad_c, \ - const unsigned int* __restrict__ argmax, \ - TYPE* __restrict__ grad_b, \ - int M, int N, int K \ -) { \ - int idx = blockIdx.x * blockDim.x + threadIdx.x; \ - int total = M * N; \ - if (idx < total) { \ - int j = idx / M; \ - int k = (int)argmax[idx]; \ - if (k >= 0 && k < K) { \ - ATOMIC_ADD(&grad_b[k + j * K], grad_c[idx]); \ - } \ - } \ -} - -// ============================================================================ -// KERNEL INSTANTIATIONS -// ============================================================================ - -// --- 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, *) - -// --- 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, *) - -// --- F32 GEMM with Argmax Kernels --- -TROPICAL_GEMM_ARGMAX_F32(tropical_maxplus_f32_nn_with_argmax, NEG_INF_F32, >, +) -TROPICAL_GEMM_ARGMAX_F32(tropical_minplus_f32_nn_with_argmax, INF_F32, <, +) -TROPICAL_GEMM_ARGMAX_F32(tropical_maxmul_f32_nn_with_argmax, 0.0f, >, *) - -// --- F64 GEMM with Argmax Kernels --- -TROPICAL_GEMM_ARGMAX_F64(tropical_maxplus_f64_nn_with_argmax, NEG_INF_F64, >, +) -TROPICAL_GEMM_ARGMAX_F64(tropical_minplus_f64_nn_with_argmax, INF_F64, <, +) -TROPICAL_GEMM_ARGMAX_F64(tropical_maxmul_f64_nn_with_argmax, 0.0, >, *) - -// --- I32 Basic GEMM Kernels --- -TROPICAL_GEMM_I32(tropical_maxplus_i32_nn, NEG_INF_I32, max_i32, add_i32) -TROPICAL_GEMM_I32(tropical_minplus_i32_nn, INF_I32, min_i32, add_i32) -TROPICAL_GEMM_I32(tropical_maxmul_i32_nn, 0, max_i32, mul_i32) - -// --- I64 Basic GEMM Kernels --- -TROPICAL_GEMM_I64(tropical_maxplus_i64_nn, NEG_INF_I64, max_i64, add_i64) -TROPICAL_GEMM_I64(tropical_minplus_i64_nn, INF_I64, min_i64, add_i64) -TROPICAL_GEMM_I64(tropical_maxmul_i64_nn, 0LL, max_i64, mul_i64) - -// --- I32 GEMM with Argmax Kernels --- -// ZERO_FN canonicalizes a drifted tropical-zero cell's argmax to 0 at write-out. -TROPICAL_GEMM_ARGMAX_I32(tropical_maxplus_i32_nn_with_argmax, NEG_INF_I32, >, add_i32, zero_maxplus_i32) -TROPICAL_GEMM_ARGMAX_I32(tropical_minplus_i32_nn_with_argmax, INF_I32, <, add_i32, zero_minplus_i32) -TROPICAL_GEMM_ARGMAX_I32(tropical_maxmul_i32_nn_with_argmax, 0, >, mul_i32, zero_never_i32) - -// --- I64 GEMM with Argmax Kernels --- -TROPICAL_GEMM_ARGMAX_I64(tropical_maxplus_i64_nn_with_argmax, NEG_INF_I64, >, add_i64, zero_maxplus_i64) -TROPICAL_GEMM_ARGMAX_I64(tropical_minplus_i64_nn_with_argmax, INF_I64, <, add_i64, zero_minplus_i64) -TROPICAL_GEMM_ARGMAX_I64(tropical_maxmul_i64_nn_with_argmax, 0LL, >, mul_i64, zero_never_i64) - -// --- Backward Pass Kernels (float/double only, no integer gradients) --- -TROPICAL_BACKWARD_A(tropical_backward_a_f32, float, atomicAdd) -TROPICAL_BACKWARD_B(tropical_backward_b_f32, float, atomicAdd) -TROPICAL_BACKWARD_A(tropical_backward_a_f64, double, atomicAddDouble) -TROPICAL_BACKWARD_B(tropical_backward_b_f64, double, atomicAddDouble) - -// ============================================================================ -// BATCHED F32 GEMM WITH ARGMAX KERNEL MACRO -// ============================================================================ -// Strided batched GEMM: processes batch_size independent GEMMs -// Uses blockIdx.z for batch index, strides for memory offsets - -#define TROPICAL_GEMM_BATCHED_ARGMAX_F32(KERNEL_NAME, INIT_VAL, COMPARE_OP, MUL_OP) \ -extern "C" __global__ void KERNEL_NAME( \ - const float* __restrict__ A, \ - const float* __restrict__ B, \ - float* __restrict__ C, \ - unsigned int* __restrict__ argmax, \ - int M, int N, int K, \ - int strideA, int strideB, int strideC \ -) { \ - 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 batch_idx = blockIdx.z; \ - 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 float* A_batch = A + batch_idx * strideA; \ - const float* B_batch = B + batch_idx * strideB; \ - float* C_batch = C + batch_idx * strideC; \ - unsigned int* argmax_batch = argmax + batch_idx * strideC; \ - \ - 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]; \ - int accum_idx[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; \ - accum_idx[i] = 0; \ - } \ - \ - 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_batch, INIT_VAL) \ - \ - LOAD_B_TILE(B_batch, INIT_VAL) \ - \ - __syncthreads(); \ - \ - _Pragma("unroll") \ - for (int k = 0; k < BLOCK_SIZE_K; ++k) { \ - int global_k = tile_idx + 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); \ - if (prod COMPARE_OP accum[idx]) { \ - accum[idx] = prod; \ - accum_idx[idx] = global_k; \ - } \ - } \ - } \ - } \ - __syncthreads(); \ - } \ - \ - STORE_C_ARGMAX_TILE(C_batch, argmax_batch) \ -} - -// --- Batched F32 GEMM with Argmax Kernels --- -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, >, *) - -// ============================================================================ -// FORWARD BATCHED GEMM KERNEL MACROS (no argmax) -// ============================================================================ -// Strided batched forward GEMM: processes `batch` independent GEMMs in one -// launch, blockIdx.z selecting the batch element and strideA/strideB/strideC -// the per-batch offsets. Body is identical to the single-matrix forward macro -// (TROPICAL_GEMM_{F32,F64,I32,I64}) except operands/output are taken from the -// per-batch sub-buffer. This replaces omeinsum's host-side per-slice -// clone_dtod loop with a single launch over already-contiguous device buffers. - -#define TROPICAL_GEMM_BATCHED_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, \ - int strideA, int strideB, int strideC \ -) { \ - 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; \ - \ - int batch_idx = blockIdx.z; \ - const float* A_batch = A + batch_idx * strideA; \ - const float* B_batch = B + batch_idx * strideB; \ - float* C_batch = C + batch_idx * strideC; \ - \ - 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_batch, INIT_VAL) \ - \ - LOAD_B_TILE(B_batch, 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_batch) \ -} - -#define TROPICAL_GEMM_BATCHED_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, \ - int strideA, int strideB, int strideC \ -) { \ - 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; \ - \ - int batch_idx = blockIdx.z; \ - const double* A_batch = A + batch_idx * strideA; \ - const double* B_batch = B + batch_idx * strideB; \ - double* C_batch = C + batch_idx * strideC; \ - \ - 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_batch, INIT_VAL) \ - \ - LOAD_B_TILE(B_batch, 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_batch) \ -} - -#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 \ -) { \ - 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; \ - \ - int batch_idx = blockIdx.z; \ - const int* A_batch = A + batch_idx * strideA; \ - const int* B_batch = B + batch_idx * strideB; \ - int* C_batch = C + batch_idx * strideC; \ - \ - 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_batch, INIT_VAL) \ - \ - LOAD_B_TILE(B_batch, 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_batch) \ +#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 \ -) { \ - 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; \ - \ - int batch_idx = blockIdx.z; \ - const long long* A_batch = A + batch_idx * strideA; \ - const long long* B_batch = B + batch_idx * strideB; \ - long long* C_batch = C + batch_idx * strideC; \ - \ - 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_batch, INIT_VAL) \ - \ - LOAD_B_TILE(B_batch, 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_batch) \ +#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, +) -TROPICAL_GEMM_BATCHED_F32(tropical_minplus_f32_nn_batched, INF_F32, fminf, +) -TROPICAL_GEMM_BATCHED_F32(tropical_maxmul_f32_nn_batched, 0.0f, fmaxf, *) +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, +) -TROPICAL_GEMM_BATCHED_F64(tropical_minplus_f64_nn_batched, INF_F64, fmin, +) -TROPICAL_GEMM_BATCHED_F64(tropical_maxmul_f64_nn_batched, 0.0, fmax, *) +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) diff --git a/crates/tropical-gemm-cuda/src/kernels.rs b/crates/tropical-gemm-cuda/src/kernels.rs index 2eca6ec..23fb666 100644 --- a/crates/tropical-gemm-cuda/src/kernels.rs +++ b/crates/tropical-gemm-cuda/src/kernels.rs @@ -23,6 +23,28 @@ use crate::memory::{ 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 @@ -126,7 +148,7 @@ fn launch_kernel_batched_impl m: usize, k: usize, n: usize, - grid: (u32, u32, u32), + grid_xy: u32, block: (u32, u32, u32), ) -> Result<()> { // The raw-slice API loses GpuMatrix's dimension invariant, so validate the @@ -144,34 +166,51 @@ fn launch_kernel_batched_impl want(batch * m * n, c.len(), "output C")?; let kernel = ctx.get_kernel(kernel_name)?; - let cfg = LaunchConfig { - grid_dim: grid, - block_dim: block, - shared_mem_bytes: 0, - }; // 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; - let stride_a = (m * k) as i32; - let stride_b = (k * n) as i32; - let stride_c = (m * n) 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(); - let mut builder = stream.launch_builder(&kernel); - builder - .arg(a) - .arg(b) - .arg(c) - .arg(&m_i32) - .arg(&n_i32) - .arg(&k_i32) - .arg(&stride_a) - .arg(&stride_b) - .arg(&stride_c); - unsafe { - builder.launch(cfg)?; + // 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 @@ -208,11 +247,10 @@ macro_rules! impl_cuda_kernel_f32 { k: usize, n: usize, ) -> Result<()> { - let (grid_xy, _, _) = CudaContext::grid_dims_f32(m, n); - let grid = (grid_xy, 1, batch as u32); + 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, block, + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, ) } } @@ -249,11 +287,10 @@ macro_rules! impl_cuda_kernel_f64 { k: usize, n: usize, ) -> Result<()> { - let (grid_xy, _, _) = CudaContext::grid_dims_f64(m, n); - let grid = (grid_xy, 1, batch as u32); + 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, block, + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, ) } } @@ -303,11 +340,10 @@ macro_rules! impl_cuda_kernel_i32 { k: usize, n: usize, ) -> Result<()> { - let (grid_xy, _, _) = CudaContext::grid_dims_f32(m, n); - let grid = (grid_xy, 1, batch as u32); + 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, block, + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, ) } } @@ -345,11 +381,10 @@ macro_rules! impl_cuda_kernel_i64 { k: usize, n: usize, ) -> Result<()> { - let (grid_xy, _, _) = CudaContext::grid_dims_f64(m, n); - let grid = (grid_xy, 1, batch as u32); + 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, block, + ctx, Self::BATCHED_KERNEL_NAME, a, b, c, batch, m, k, n, grid_xy, block, ) } } @@ -718,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".