fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32::MAX (#63) - #64
Merged
Merged
Conversation
…::MAX (#63) The strided-batched kernel rejected any operand stride >= i32::MAX, panicking on large matrices that fit in VRAM (e.g. miso's max-plus contraction at sc_target=32, where the output C stride m*n = 2^31). The ceiling was the kernels' 32-bit addressing, not memory: stride params were `int`, the batch base used `(int)blockIdx.z * stride`, and the within-matrix offset `OFFSET_COL(row,col,ld) = col*ld+row` overflowed `int` once a single matrix held >= 2^31 elements. Widen the GLOBAL-addressing path to 64-bit while leaving the hot inner loop (shared-memory As/Bs + accumulators, all tile-sized) at i32: - Add OFFSET_COL64 (`(long long)(col)*(ld)+(row)` -- cast wraps the multiply, not the finished sum) and use it for every global load/store in the shared LOAD/STORE tile macros. This fixes batched, single-matrix, and argmax forward kernels at once since they share those macros. - Batched stride params int -> long long; batch base `(long long)blockIdx.z * stride`. - Backward kernels: `total`/`idx` -> 64-bit, scatter offsets `i + (long long)k*M` and `k + (long long)j*K`. - Rust: stride_to_i32 -> stride_to_i64 with i64 kernel args (ABI stays in sync: i64 <-> long long, 8 bytes); backward grid_size computed in u64 via div_ceil (`total as u32` truncated when m*n >= 2^32). Public signatures unchanged. Add three #[ignore] regression tests crossing 2^31 (batched forward, single matrix, backward scatter); they sample device-side / zero bulk inputs on-device to keep host memory small. Validated on a 6x A800 80GB box (CUDA 12.1): all three pass and the full 57-test normal-size suite stays green (no regression from the shared-macro widening). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #64 +/- ##
=======================================
Coverage 96.29% 96.29%
=======================================
Files 19 19
Lines 918 918
=======================================
Hits 884 884
Misses 34 34 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Extract the duplicated column-major operand build and per-cell sample/verify shared by the two forward regression tests into `build_big_operands` / `assert_big_maxplus_cells` (+ `big_a_val`/`big_b_val`), and replace the three runtime `m*n > i32::MAX` asserts with one compile-time `const _` assert on BIG_DIM. Behavior-preserving test-only cleanup (the kernel fix is unchanged); net -8 lines, ~50 lines of copy-paste removed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #63.
Problem
The strided-batched tropical GEMM panicked on large operands that fit in VRAM:
2147483648 = 2³¹is one pasti32::MAX. This is a kernel addressing ceiling, not OOM — the same family as thegridDim.zcap fixed in #62. It blocked GPU tropical contraction exactly in the large-instance regime where the GPU backend matters (e.g. miso's max-plus contraction atsc_target=32, whose output C stridem·n = 2³¹).Three 32-bit sites caused it, all on the global-addressing path:
int,(int)blockIdx.z * stride,OFFSET_COL(row,col,ld) = col*ld+rowoverflowedintonce a single matrix held ≥ 2³¹ elements (so widening only the batch stride is insufficient).Fix
Widen the global-addressing path to 64-bit; leave the hot inner loop (shared-memory
As/Bs+ accumulators, all tile-sized) ati32so there's no cost where it counts.OFFSET_COL64=(long long)(col)*(ld)+(row)— the cast wraps the multiply (a singleIMAD.WIDEon the device), not the already-overflowed sum. Used for every global load/store in the shared tile macros, so batched, single-matrix, and argmax forward kernels are all fixed at once.int → long long; batch base(long long)blockIdx.z * stride.total/idx→ 64-bit; scatter offsetsi + (long long)k*M,k + (long long)j*K.stride_to_i32 → stride_to_i64withi64kernel args (ABI stays in sync:i64↔long long, 8 bytes; cudarc passes each arg by pointer so only size must match). Backwardgrid_sizenow computed in u64 viadiv_ceil—total as u32truncated whenm·n ≥ 2³². Public signatures unchanged, so downstream crates are unaffected.Tests
Adds three
#[ignore]regression tests crossing 2³¹ (46341² ≈ 2.147e9): batched-forward (the exact repro), single-matrix offset, and backward scatter. They sample device-side / zero bulk inputs on-device to keep host memory small.Validation (real hardware)
Run on a 6× A800 80GB box, CUDA 12.1 (
CUDARC_CUDA_VERSION=12010):m·n−1 = 2,147,488,280(well pasti32::MAX) computes the correct value, exactly what i32 indexing would corrupt.CI skips CUDA, so the large-allocation tests stay
#[ignore]d (they need ≥ 8.6 GB / 16 GB VRAM).🤖 Generated with Claude Code