Split struct-element kernel scratch into per-field buffers - #2982
Open
wsmoses wants to merge 1 commit into
Open
Conversation
Reduce-with-location kernels (MFEM's ArgMin/ArgMax/MinMaxLoc reducers) keep their block partials in shared scratch of struct type -- memref<256 x struct<(f64, i32)>> -- which cannot become a tensor whole. Every access reaches the scratch through a flat primitive view whose index pins a fixed byte offset within the struct: affine maps with a constant residue, plain arith index chains with a derivable residue, whole-struct integer moves, and fixed-size memcpys between struct-strided geps. Split the array-of-structs into one primitive scratch per field and rewrite each access form onto its field buffer, with same-size type punning kept as bitcasts and struct-wide integer moves recomposed with shifts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
wsmoses
force-pushed
the
pb/split-struct-scratch
branch
from
August 26, 2026 16:12
fb7e24b to
cbec48a
Compare
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.
MFEM's reduce-with-location kernels (
ArgMin/ArgMax/MinMaxLocreducers, exercised by the GPU unit-test suite'sReduce *cases) keep their block partials in shared scratch of struct type —memref<256 x !llvm.struct<(f64, i32)>>— which the raiser rejects as a non-primitive alloca.Every access reaches such scratch through a flat primitive view with a statically determined byte offset inside the struct, so the array-of-structs splits into one primitive scratch per field:
affine.load/storethroughpointer2memrefviews: the map's byte expression has a constant residue mod the struct size, picking the field ([t*4+2]on an i32 view of an(f64,i32)pair → the i32 field of pairt).memref.load/storewhose index is an arith chain: a small residue analysis (constants, sums, products, shifts, casts) derives the same field pinning for the tree-reduction phase's computed indices.i64move of thef64half) becomearith.bitcasts.i64whole-pair move of an(i32,i32)pair) split into per-field accesses recomposed with shifts.llvm.intr.memcpybetween struct-strided geps (16-byte pair copies) becomes per-field moves, including geps folded to constant indices or to the scratch base.Validated end to end on MFEM:
general/test_reduction.cppfully raises and the whole reduction battery (Sum/Mult/BAnd/BOr/Min/Max/MinMax/ArgMin/ArgMax) passes on GPU through raised XLA kernels, exactly.🤖 Generated with Claude Code
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD