feat(cuda): forward strided-batched tropical GEMM (one launch over the batch) - #62
Merged
Conversation
…dx.z=batch)
Adds TROPICAL_GEMM_BATCHED_{F32,F64,I32,I64} forward kernels (no argmax),
mirroring the single-matrix forward macros plus blockIdx.z batch indexing and
per-batch strides, and wires them through CudaKernel::launch_gemm_batched +
launch_kernel_batched_impl (non-swapped column-major convention, raw CudaSlice
operands so callers feed already-contiguous device buffers with no copy).
Lets the omeinsum device path run a whole node's batch in ONE launch instead of
a host-side per-slice clone_dtod loop, and the output may be allocated
uninitialized since the kernel fully writes every element. Correctness test
test_tropical_gemm_batched_matches_single checks the batched kernel against the
trusted single-matrix path + a CPU max-plus reference, incl. a non-block-aligned
shape (edge tiles).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
) * fix(cuda): chunk batched tropical GEMM over gridDim.z 65535 cap The strided-batched kernels map the batch dimension to blockIdx.z, whose grid extent CUDA caps at 65535 on all compute capabilities. A contraction whose batch exceeds that launched with grid.z > 65535 and failed at launch with CUDA_ERROR_INVALID_VALUE. Large tensor-network contractions (space complexity >= ~30) routinely exceed it, so the failure surfaced as an intermittent panic (the batch/m/n split depends on the contraction order). Split batched launches into chunks of at most 65535 batch elements, advancing each operand's base by the chunk start so blockIdx.z stays in [0, chunk). Fixes both the value path (launch_kernel_batched_impl, used by all four scalar-type macro impls) and the argmax path (launch_gemm_external_batched_with_argmax_f32). No kernel (.cu) change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * cuda: reject batched strides past i32::MAX instead of truncating Address Copilot review on PR #61. The chunked batched launches derive operand base offsets from the usize per-batch stride, while the kernels read the stride as i32. An `as i32` cast that wrapped would desynchronise the two and silently corrupt addressing. Add a shared `stride_to_i32` helper that converts fallibly and returns DimensionMismatch on overflow, and apply it in both the value path (launch_kernel_batched_impl) and the external argmax path (launch_gemm_external_batched_with_argmax_f32). The kernel signature is i32, so a stride past i32::MAX is unrepresentable on-device regardless; fail loudly at the boundary. Also correct the argmax-path stride comment: external A/B carry the (possibly padded) DLPack stride, not rows*cols. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor(cuda): dedupe GEMM kernel macros and batched launch wiring Quality cleanup of the strided-batched GEMM work (no behavior change): kernels/tropical_gemm.cu - Factor the tiled max-/min-plus / max-mul inner kernel into one shared TROPICAL_GEMM_BODY macro. It was previously copy-pasted 8x (4 scalar types x {single-matrix, batched}); the 8 kernel families are now thin wrappers that supply only the signature and operand base expressions (`A/B/C` vs `A + blockIdx.z*strideA`). Net -505 lines. - Add add_f32/mul_f32/add_f64/mul_f64 helpers so the body takes MUL_FN in function form for every type (symmetric with the int helpers) instead of special-casing the float `+`/`*` operator. src/kernels.rs - launch_gemm_external_batched_with_argmax_f32: reuse CudaContext:: grid_dims_f32 instead of re-deriving the tile count with a hardcoded 64 (which would silently desync if the block size changed). - launch_kernel_batched_impl: collapse the `grid: (u32,u32,u32)` parameter to a single `grid_xy: u32`. The caller-passed z (`batch as u32`) was dead (recomputed per chunk as the chunk size) and y was always 1. Verified on A800 (sm_80, CUDA 12.1): 57/57 lib tests pass. PTX/SASS diff of all 43 kernels vs the pre-refactor source — 31 non-batched kernels are byte-identical; the 12 batched kernels have identical instruction histograms and identical register/smem usage (only benign prologue scheduling differs). No correctness or performance regression. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #62 +/- ##
=======================================
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:
|
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
1 participant
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.
Rebased the
perf/forward-batched-gemmstack onto currentmain(which alreadyhas the async-launch change from #60) so it merges cleanly. The duplicate
924abb2async commit is dropped here; this branch carries only theforward-batched feature, its gridDim.z fix, and the macro/launch cleanup.
What's included
blockIdx.zselecting the batch element and per-batch strides, replacing omeinsum's
host-side per-slice
clone_dtodloop.gridDim.z65535 cap (fix(cuda): chunk batched tropical GEMM over the gridDim.z 65535 cap #61); reject batchedstrides past
i32::MAXinstead of truncating.TROPICAL_GEMM_BODY(was copy-pasted 8x across {f32,f64,i32,i64} × {single,batched}); reuse
grid_dims_f32instead of a hardcoded64; collapse a dead grid-tuple param.Net
tropical_gemm.cu−505 lines.Validation (HKUST-GZ A800, sm_80, CUDA 12.1)
cargo test -p tropical-gemm-cuda --lib --release: 57/57 passed on thisexact tree (rebased onto current main).
kernels byte-identical; 12 batched kernels have identical instruction
histograms and register/smem usage — no correctness or performance regression.
🤖 Generated with Claude Code