Skip to content

fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32::MAX (#63) #64

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
isPANN merged 2 commits into from
Jun 16, 2026
Merged

fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32::MAX (#63) #64

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2b0a91d
fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32…
isPANN Jun 16, 2026
0467db3
test(cuda): dedup the #63 regression tests
isPANN Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 47 additions & 37 deletions crates/tropical-gemm-cuda/kernels/tropical_gemm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,16 @@
#define INF_I64 (1LL << 60)
#define NEG_INF_I64 (-(1LL << 60))

// Memory layout helpers
// Memory layout helpers.
// OFFSET_COL is `int` arithmetic: use it ONLY for shared-memory tiles (As/Bs)
// and per-thread accumulator indices, whose leading dimension is a tile size, so
// the product can never reach 2^31. OFFSET_COL64 forces the multiply into 64-bit
// (32x32->64, a single IMAD.WIDE on the device) for GLOBAL-memory accesses, where
// a single matrix with >= 2^31 elements (col*ld) would otherwise overflow `int`
// and corrupt addressing (issue #63). The cast wraps the multiply itself, not the
// finished sum -- `(long long)(col*ld+row)` would already have overflowed.
#define OFFSET_COL(row, col, ld) ((col) * (ld) + (row))
#define OFFSET_COL64(row, col, ld) ((long long)(col) * (ld) + (row))

// Integer max/min functions
__device__ __forceinline__ int max_i32(int a, int b) { return a > b ? a : b; }
Expand Down Expand Up @@ -131,9 +139,9 @@ __device__ double atomicAddDouble(double* address, double val) {
int col = A_TILE_COL + i + tile_idx; \
int dst = OFFSET_COL(A_TILE_ROW, i + A_TILE_COL, BLOCK_SIZE_M); \
if (BLOCK_IDX == DIM_GRID_X - 1 || tile_idx >= K - BLOCK_SIZE_K) \
As[dst] = (row < M && col < K) ? SRC[OFFSET_COL(row, col, M)] : (PAD); \
As[dst] = (row < M && col < K) ? SRC[OFFSET_COL64(row, col, M)] : (PAD); \
else \
As[dst] = SRC[OFFSET_COL(row, col, M)]; \
As[dst] = SRC[OFFSET_COL64(row, col, M)]; \
}

// B tile -> Bs. B is column-major with leading dimension K.
Expand All @@ -144,9 +152,9 @@ __device__ double atomicAddDouble(double* address, double val) {
int col = BLOCK_SIZE_N * BLOCK_IDY + i + B_TILE_COL; \
int dst = OFFSET_COL(B_TILE_ROW, i + B_TILE_COL, BLOCK_SIZE_K); \
if (tile_idx >= K - BLOCK_SIZE_K || BLOCK_IDY == DIM_GRID_Y - 1) \
Bs[dst] = (row < K && col < N) ? SRC[OFFSET_COL(row, col, K)] : (PAD); \
Bs[dst] = (row < K && col < N) ? SRC[OFFSET_COL64(row, col, K)] : (PAD); \
else \
Bs[dst] = SRC[OFFSET_COL(row, col, K)]; \
Bs[dst] = SRC[OFFSET_COL64(row, col, K)]; \
}

// Edge block iff last block-row or last block-col; interior blocks are fully in
Expand All @@ -162,7 +170,7 @@ __device__ double atomicAddDouble(double* address, double val) {
int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \
int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \
if (!TILE_IS_EDGE || (row < M && col < N)) \
DST[OFFSET_COL(row, col, M)] = accum[OFFSET_COL(tm, tn, THREAD_SIZE_M)]; \
DST[OFFSET_COL64(row, col, M)] = accum[OFFSET_COL(tm, tn, THREAD_SIZE_M)]; \
} \
}

Expand All @@ -175,7 +183,7 @@ __device__ double atomicAddDouble(double* address, double val) {
int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \
int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \
if (!TILE_IS_EDGE || (row < M && col < N)) { \
int out_idx = OFFSET_COL(row, col, M); \
long long out_idx = OFFSET_COL64(row, col, M); \
int local_idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \
DST[out_idx] = accum[local_idx]; \
ARGMAX_DST[out_idx] = accum_idx[local_idx]; \
Expand All @@ -193,7 +201,7 @@ __device__ double atomicAddDouble(double* address, double val) {
int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \
int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \
if (!TILE_IS_EDGE || (row < M && col < N)) { \
int out_idx = OFFSET_COL(row, col, M); \
long long out_idx = OFFSET_COL64(row, col, M); \
int local_idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \
DST[out_idx] = accum[local_idx]; \
ARGMAX_DST[out_idx] = (ZERO_FN)(accum[local_idx]) ? 0 : accum_idx[local_idx]; \
Expand Down Expand Up @@ -678,13 +686,13 @@ extern "C" __global__ void KERNEL_NAME( \
TYPE* __restrict__ grad_a, \
int M, int N, int K \
) { \
int idx = blockIdx.x * blockDim.x + threadIdx.x; \
int total = M * N; \
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; \
long long total = (long long)M * N; \
if (idx < total) { \
int i = idx % M; \
int i = (int)(idx % M); \
int k = (int)argmax[idx]; \
if (k >= 0 && k < K) { \
ATOMIC_ADD(&grad_a[i + k * M], grad_c[idx]); \
ATOMIC_ADD(&grad_a[i + (long long)k * M], grad_c[idx]); \
} \
} \
}
Expand All @@ -696,13 +704,13 @@ extern "C" __global__ void KERNEL_NAME( \
TYPE* __restrict__ grad_b, \
int M, int N, int K \
) { \
int idx = blockIdx.x * blockDim.x + threadIdx.x; \
int total = M * N; \
long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; \
long long total = (long long)M * N; \
if (idx < total) { \
int j = idx / M; \
int j = (int)(idx / M); \
int k = (int)argmax[idx]; \
if (k >= 0 && k < K) { \
ATOMIC_ADD(&grad_b[k + j * K], grad_c[idx]); \
ATOMIC_ADD(&grad_b[k + (long long)j * K], grad_c[idx]); \
} \
} \
}
Expand Down Expand Up @@ -771,7 +779,7 @@ extern "C" __global__ void KERNEL_NAME( \
float* __restrict__ C, \
unsigned int* __restrict__ argmax, \
int M, int N, int K, \
int strideA, int strideB, int strideC \
long long strideA, long long strideB, long long strideC \
) { \
const int BLOCK_SIZE_M = 64; \
const int BLOCK_SIZE_K = 32; \
Expand All @@ -783,7 +791,7 @@ 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; \
long long 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; \
Expand Down Expand Up @@ -867,56 +875,58 @@ TROPICAL_GEMM_BATCHED_ARGMAX_F32(tropical_maxmul_f32_nn_batched_with_argmax, 0.
// 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.
// contiguous device buffers. The per-batch stride is `long long` and the base
// offset `(long long)blockIdx.z * stride` is computed in 64-bit, so a single
// matrix of >= 2^31 elements addresses correctly (issue #63); 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 \
long long strideA, long long strideB, long long 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)) \
(A + (long long)blockIdx.z * strideA), \
(B + (long long)blockIdx.z * strideB), \
(C + (long long)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 \
long long strideA, long long strideB, long long 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)) \
(A + (long long)blockIdx.z * strideA), \
(B + (long long)blockIdx.z * strideB), \
(C + (long long)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 \
long long strideA, long long strideB, long long 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)) \
(A + (long long)blockIdx.z * strideA), \
(B + (long long)blockIdx.z * strideB), \
(C + (long long)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 \
long long strideA, long long strideB, long long 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)) \
(A + (long long)blockIdx.z * strideA), \
(B + (long long)blockIdx.z * strideB), \
(C + (long long)blockIdx.z * strideC)) \
}

// --- Forward batched GEMM kernel instantiations (mirror the non-batched set) ---
Expand Down
47 changes: 23 additions & 24 deletions crates/tropical-gemm-cuda/src/kernels.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,17 @@ use tropical_gemm::types::{TropicalMaxMul, TropicalMaxPlus, TropicalMinPlus, Tro
/// 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
/// Convert a per-batch element stride to the `i64` 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> {
i32::try_from(stride).map_err(|_| {
/// the `usize` stride while the kernel reads the stride as `i64` (`long long`);
/// keeping the conversion fallible guarantees the two can never desynchronise.
/// The kernels do 64-bit global addressing (issue #63), so there is no longer an
/// `i32::MAX` ceiling — only a `usize` value past `i64::MAX` (≥ 2⁶³ elements,
/// physically impossible to allocate) would fail here.
fn stride_to_i64(stride: usize, what: &str) -> Result<i64> {
i64::try_from(stride).map_err(|_| {
CudaError::DimensionMismatch(format!(
"batched GEMM {what} stride {stride} exceeds i32::MAX; \
matrix too large for the strided-batched kernel"
"batched GEMM {what} stride {stride} exceeds i64::MAX"
))
})
}
Expand Down Expand Up @@ -172,13 +172,13 @@ fn launch_kernel_batched_impl<T: DeviceRepr + ValidAsZeroBits + Default + Clone>
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`).
// (as `usize`) and as the kernel's `i64` stride. Derive the `i64` form
// fallibly so it stays in 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 stride_a = stride_to_i64(sa, "operand A")?;
let stride_b = stride_to_i64(sb, "operand B")?;
let stride_c = stride_to_i64(sc, "output C")?;

let stream = ctx.stream();
// The batch maps to `blockIdx.z`, which CUDA caps at `MAX_GRID_DIM_Z`. Launch
Expand Down Expand Up @@ -764,15 +764,14 @@ pub unsafe fn launch_gemm_external_batched_with_argmax_f32(

// 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.
// so their stride is rows*cols. Derive the kernel's `i64` strides fallibly
// so they stay in 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 stride_a_i64 = stride_to_i64(stride_a, "operand A")?;
let stride_b_i64 = stride_to_i64(stride_b, "operand B")?;
let stride_c_i64 = stride_to_i64(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;
Expand Down Expand Up @@ -814,9 +813,9 @@ pub unsafe fn launch_gemm_external_batched_with_argmax_f32(
.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
.arg(&stride_b_i64) // strideA (B's stride in our swap)
.arg(&stride_a_i64) // strideB (A's stride in our swap)
.arg(&stride_c_i64); // strideC
builder.launch(cfg)?;
start += chunk;
}
Expand Down
Loading
Loading