Repository navigation
feat(rocke): add tensor_transforms -- coordinate graphs, required bound checks and split addresses - #13876
Merged
Conversation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
CongMa13
approved these changes
Oct 16, 2026
cgmillette
force-pushed
the
users/cgmillette/rocke/tensor-transforms
branch
2 times, most recently
from
October 20, 2026 01:59
107a3a1 to
a108157
Compare
cgmillette
marked this pull request as ready for review
October 20, 2026 01:59
cgmillette
enabled auto-merge (squash)
October 20, 2026 15:32
cgmillette
force-pushed
the
users/cgmillette/rocke/tensor-transforms
branch
from
October 20, 2026 17:57
a108157 to
67491c5
Compare
cgmillette
force-pushed
the
users/cgmillette/rocke/tensor-transforms
branch
from
October 20, 2026 19:44
67491c5 to
41780cd
Compare
bartekxk
approved these changes
Oct 20, 2026
Add the tensor_transforms package with its first module, graph.py: a frozen TransformGraph that maps named input coordinates to named output coordinates through eight node kinds (embed, merge, unmerge, offset, slice_window, freeze, xor, broadcast), assembled by make_transform_graph from inputs(...), the steps, and outputs(...). The graph holds no tensor and no bound. A domain is declared only where a kind needs one (an input's extent, unmerge pieces, the product a merge writes, the power-of-two extent an xor carries); offset, embed and slice outputs have none. There is no PAD kind: a halo is a negative offset and the tensor's own length masks it. Assembly checks single-writer names, repeated steps, undeclared reads, cycles, restated domains (unmerge product, merge lengths), xor's power-of-two domain, unwritten outputs and unread inputs, each with an actionable message. Nodes are stored in topological order with ties broken by (kind, read, write), so equality and hash do not depend on the order the steps were listed; input order is part of the graph. evaluate and resolve take the coordinates as a tuple in inputs(...) order (a mapping is refused) and compute in pure Python what every lane computes, valid or not: an unmerge read past its domain drives the outer piece past its length. resolve returns every intermediate; edge(name) reports a coordinate's domain and its writer. Tests check each kind against numpy over its whole domain (square and rectangular, power-of-two and not), the conv input graph against a numpy im2col index, the positional call rule, canonical equality, and every builder, assembly and call error text. Co-Authored-By: Claude <noreply@anthropic.com>
…nsforms Add ranges.py, the one home of range arithmetic and of the rule that decides which side of which coordinate needs a run-time check. Interval and Congruence hold one value's range and residue. Their methods (plus, shifted, scaled, floor_div, floor_mod, minimum, maximum, overlaps, intersect, snapped; join for congruences) are pure Python and use floor semantics for every sign. An end is an int, None (unbounded) or a RuntimeValue: a runtime integer with its stated envelope and multiple-of fact, which stays a RuntimeValue only under a shift or a unit scale and otherwise becomes its numeric end. minimum, maximum, overlaps and intersect are symmetric; intersect keeps a runtime upper end over an int one so a runtime length can be proven by identity, and raises exactly when overlaps says the two ranges share no value. Two RuntimeValues over one IR value that state different facts are refused. value_ranges reports every coordinate's unconditioned range (what every lane computes, valid or not). required_checks returns one BoundCheck per coordinate whose range can leave its input domain, a domain a step requires of it (an unmerge read past its pieces: a K tail), or its tensor bound, with each side reported separately and a proven side left None. Checks are AND-ed, so each coordinate is narrowed by its own domains before its readers compute their ranges. graph_injective says whether distinct in-domain inputs always reach distinct outputs. The toolbox forms take and return plain (lo, hi) pairs, one per graph input in inputs(...) order, and take the bounds from the tensor's lengths; internal forms take Intervals and explicit AxisBounds. graph.py gains the internal identity_graph, the graph of a plain window. The package door exports RuntimeValue, value_ranges, required_checks, graph_injective, BoundCheck and CheckReason. Tests check every operation against exhaustive enumeration, the runtime end rules and their symmetry (with random pairs against substitution), the conv checks, the stride-2 bound taken from the tensor, the whole tile proven inside a runtime length, the plain-window decisions against enumeration over every origin, that no elided check hides an out-of-domain value, injectivity, that public values carry no internal type, and every error text. Co-Authored-By: Claude <noreply@anthropic.com>
Add tensor_transforms/address.py with the pure-Python golden every emitted address is tested against: - address_of(tensor, position): sum(position[a] * strides[a]) for one int per tensor axis, in axis order; - transform_address_of(graph, tensor, coords): the same address for the graph's outputs at coords, one int per graph input in inputs(...) order. Both are positional (a mapping or any non-tuple is refused), compute in Python integers, and do not bound the position, so they give the address a lane computes whether or not it is valid. A negative stride is refused with the structured-address message, and a position or graph that does not match the tensor rank raises its rank message. TensorDesc is imported for type checking only. The positional coordinate check and the graph-outputs-vs-tensor-rank check move into graph.py helpers so evaluate, resolve, required_checks and the address functions share one text each. Both functions are exported on the package door. Tests compare the identity graph with the plain strided address, the conv graph with numpy NHWC offsets, and fire every error with its text. Co-Authored-By: Claude <noreply@anthropic.com>
…from the IR Add walk.py, the IR range walk of tensor_transforms. It reads how a plain i32 IR value was computed and emits nothing. value_facts derives the value's range and residue over a closed rule set: constants, thread and block ids, readfirstlane, loop counters, range markers, add / sub, mul / div / mod by a constant, shl by a constant, and with a low mask, smin / smax. Every rule maps onto a range operation of ranges.py. An op outside the set, or a rule whose precondition fails, contributes an unknown range and residue and the rows combine as usual, so a partial fact survives (x * 128 is a multiple of 128 whatever x is). shl keeps its scaled range only when that range is proven inside i32 (shl wraps), and keeps the residue in every case. A loop counter gets the tight range of its trips for constant bounds, a one-value range for a loop that never runs, and an end over the runtime length itself when the stop is a length plus constants, so a loop over full tiles proves counter + tile <= length by identity. A range marker intersects its stated range with the derived one and snaps it to the residue; a stated fact the IR contradicts is the intersect error. runtime_value_of rebuilds the RuntimeValue an envelope marker states from the marker's stated facts, never from the derived range, so the rebuilt length equals the one its producer returned. split_coordinate turns a value into an AffineCoord: a constant, wave-uniform parts and one per-lane part, each with its own facts. Only block ids, kernel parameters, readfirstlane results, counters of loops with uniform bounds, markers on uniform values and pure arithmetic on them are uniform; loads, intrinsics, thread and lane ids are per-lane. AffineCoord and UniformPart are added to ranges.py with of_int and an emit-free +, beside a small Interval.within query. leaf_varying_within_workgroup returns the first leaf that may differ between the waves of a workgroup, accepting a scalar load only through a parameter declared readonly and noalias (or constant memory) at a workgroup-uniform index. The first walk that reaches a thread id requires an explicit max_workgroup_size and records it in the kernel attributes; a later walk refuses a kernel whose size changed since. A coordinate that is not i32 is refused. Tests build small kernels and check every rule row, the worked examples, soundness against i32 evaluation (including shl and the low-mask and with a negative control), the loop-counter cases, the full-tile proof, runtime ends, the split, the workgroup-uniform walk, the workgroup-size record, every error text, and that the walk emits nothing. Co-Authored-By: Claude <noreply@anthropic.com>
Add emit.py to tensor_transforms with two toolbox functions: emit_coords gives every named coordinate of a graph as one emitted value, and emit_addresses gives one ElementAddress per access: the element address split into wave-uniform terms (UniformTerm, one per set of graph inputs), one i32 per-lane part with its range, and a Python-int constant, plus the element size in bits. Each coordinate is a plain i32 value or an int; the IR range walk splits it into a constant, uniform parts and a per-lane part. All accesses of a call share one table, so a sub-expression they have in common is emitted once, in first-use order. A per-lane part whose range starts below 0 is rebased once per call and its lower end moves into the constant. Affine steps keep the parts apart and scale uniform terms per input set. unmerge splits the uniform side, the per-lane part and the constant separately and recombines them with exact 0/1 carries, innermost piece first: one division of the uniform side per level, at most one shared per-lane division per level, and one compare + select per access; a uniform side aligned to the inner length needs no remainder, and a uniform side not proven in [0, 2^31) falls back to splitting the whole coordinate or raises an error that names the tools that bound an origin. xor adds its reads' parts first and masks with and. Widths follow the unconditioned ranges: uniform terms are i64 unless proven inside i32, and a per-lane part outside 32 bits is an error. Every entry checks its accesses before emitting (tuple, coordinate type, count, tensor rank, non-negative strides). AffineCoord gains its combined range and residue; address.py shares its stride check. Tests run the emitted IR on the host for every lane against transform_address_of and fail on any overflowing op, pin where the divisions, compares and selects land, pin the IR text with reviewed golden strings, compile each case for gfx90a, gfx942 and gfx950, and on gfx90a compare each lane's parts with the golden, including an address past 2^31 and a convolution input of more than 2^31 elements. Co-Authored-By: Claude <noreply@anthropic.com>
…cesses Add the access entries of tensor_transforms and the one bound predicate they are built on. - emit_bound_check(b, value, lo, hi): lo <= value < hi for one plain i32 value, split by the IR range walk. The per-wave side folds into two ends saturated into [0, L] and one unsigned compare decides; a side the walked range proves emits nothing. The fold is 32-bit when every range it touches is proven inside i32, else each end is pre-clamped once in 64 bits per distinct (per-wave side, end) of a call and finished per access constant in 32 bits, so no subtraction can wrap. A lone i32 per-wave value against [0, hi) with hi below 2^31 is a single compare. _emit_bound_check takes a coordinate the tiling layer already holds split. - emit_accesses(b, graph, tensor, accesses, vector_axis=, vector_width=) returns one Access per access: its split element address and its validity against the tensor's own lengths and the graph's domains, over the same table the address uses. A vector along one graph input is addressed by its first element; its validity is one answer when every check has one for the whole vector, else the first and last element's checks plus valid_span, the valid elements, from the same folded ends. - emit_accesses_except_buffer_checked leaves named sides out of valid for the tiling layer's buffer delivery; it is not on the package door. - emit_access_validity gives each access's validity alone. - runs.py: address_contiguous_run, validity_uniform_over_run and checks_monotone_over_run, from each coordinate's step along the vector and its start residue. Tests cover each fold width and the single compare (both checked against each other on a gfx90a over every value they admit), the per-lane rebase in validity, the entry and vector errors, the door, the callers allowed to state a range marker, the run queries against exhaustive enumeration at every start residue, the gfx90a ISA of the folds, and the validity and valid_span of convolution and ragged vector accesses on a gfx90a. Co-Authored-By: Claude <noreply@anthropic.com>
Add the two author docs of the tensor_transforms package: - tensor_transforms_api.md: the one home of the package's facts -- the units and their tiers, the eight kinds (with the flip note against CK and V4, and why a halo is an offset plus the tensor's own length rather than a PAD kind), the index spaces and the positional rule, the author-facing results, the signatures by tier, the validity rule, the IR range walk's split and range rules, the structured address and its integer emit rules, the bound-check fold, the vector rules, the internal values, and every error message verbatim. - tensor_transforms_guide.md: how to declare a graph, check it with evaluate, add a halo, split K, build a swizzle, get an address, see the required checks, how to write coordinates the layer can split, and the CK pad transform in rocke with the coordinate-side reasons rocke differs. Every example is runnable code. Point to them from the tiling API surface and contract, the label-flow doc, the package README and docstring, and the two architecture docs on coordinate transforms. Add the coordinate-graph step to the tiling-API design skill, the matching checklist items and domain keywords to the tiling expert, and rebuild the dispatch registry rows from the experts' frontmatter. Co-Authored-By: Claude <noreply@anthropic.com>
…rs of two Every attention config that reached the no-unsigned-wrap magic-division sum divided by a power of two. For those divisors the magic multiplier is 1, the multiply-high term is 0 and the sum is just the dividend, so the emitted add reads the same whatever its flags. Each affected family now has one config whose divisor gives a real multiply-high term: - fmha_splitkv_decode: 40 query heads over 8 KV heads (head to KV head divides by 5). - sage_attention: Q and K scale blocks of 48 positions. - sparse_attention: jenga blocks of 3 query rows and 48 keys (both block divisions are by 3). - the scalar unified attention kernels, 2D and split 3D: 40 query heads over 8 KV heads, added to the representative-IR golden for all three LLVM flavors. Only the two new cases are added; no existing hash changes. The Python and C parity emitters and the binding parity configs carry the same new config, so the byte-identity gate compares both engines on it. Co-Authored-By: Claude <noreply@anthropic.com>
The scalar unified attention kernel is arch-neutral and compiles for the device it launches on, but its numeric test only ran on gfx950. It now runs on any GPU; the tiled-kernel tests in the same file keep their gfx950 gate. The test also gains a 40 query head over 8 KV head case. That ratio of 5 gives the kernel's head to KV-head magic division a multiplier other than 1, which the existing 4 over 2 case never reaches. Co-Authored-By: Claude <noreply@anthropic.com>
The per-lane address, bound-check, access-validity and run tests skipped unless the device was gfx90a, although the IR they run is arch-neutral. They now run on whatever GPU is present and compile the kernel for that GPU's arch, through the test helper current_device_arch that the other numeric tests use; test_emit drops its own copy of the gfx90a gate. The ISA text assertions stay gfx90a-specific: they compile for gfx90a offline and need no GPU. Co-Authored-By: Claude <noreply@anthropic.com>
… RDNA and gfx1250 The golden cases compiled only for CDNA targets, although the ops they emit (the unsigned compare, the no-unsigned-wrap add and the range marker) lower for RDNA and gfx1250 too. Each case now covers gfx90a, gfx942, gfx950, gfx1151, gfx1201 and gfx1250: - a new test lowers it with both engines at every LLVM flavor and requires the same IR text from each; - the compile test builds a code object for each target the loaded comgr can target, and skips naming the target when it cannot (an older ROCm does not know gfx1250). Co-Authored-By: Claude <noreply@anthropic.com>
… graph Attention is one of the workloads the coordinate graph serves, but no example or test built an attention-shaped graph. The guide gains a section on mapping query heads onto KV heads: unmerge the query head into (KV head, group) and leave the group unread. A test in test_emit emits the K-cache address of that graph for 40 query heads over 8 KV heads, checks it against transform_address_of for every lane, query head and a sample of token tiles, and pins that the division by 5 is one magic division on the per-wave side with no per-lane divide. Co-Authored-By: Claude <noreply@anthropic.com>
CongMa13
approved these changes
Oct 20, 2026
CongMa13
left a comment
Contributor
There was a problem hiding this comment.
No new P0 or P1 issues were identified during this review. The changes are approved.
cgmillette
force-pushed
the
users/cgmillette/rocke/tensor-transforms
branch
from
October 20, 2026 22:37
90e53fb to
ebcb208
Compare
Contributor
TheRock Submodule Bump ActivityNewest first
|
cgmillette
added a commit
that referenced
this pull request
Oct 26, 2026
The differential golden had gone stale, because nothing ran the check: 675 of its 1116 LLVM IR entries (compared at llvm20) and 608 of its 1089 canonical IR entries no longer match the tree. Recorded afresh in the new format, with the Python reference emitters, from a tree whose two-engine gate is green at llvm20, llvm22 and llvm23: - ll at llvm20 and llvm22: 82 families, 1474 configs per flavor. At llvm23: 83 families, 1578 configs, where gfx1250_lowering has configs that need LLVM 23 and are refused below it, and every conv_direct_dgrad_wmma config needs LLVM 23. - ir_canonical: 82 families, 1561 configs (conv_abi's emitter has no IR mode). Canonical IR was checked to be identical at llvm20 and llvm23. What the re-bless absorbs: every mismatch that could be traced is an intended code generation change that did not re-bless this file -- the shared LDS pool (#8844, about 425 configs), the ds.swizzle declaration (#9583, about 177), renamed values in the attention transpose reads (#9133, 57), the smaller LDS pool (#10112, about 39), the flat work group size of one-wave wave32 kernels (#12710, about 39), #9220 (6) and the conv changes of #13109 (6). 441 entries were already stale at the start of the history available here and cannot be attributed; no unexplained change was found after it. The two repeated keys (conv_direct_grouped, conv_wgrad_workspace_reduce), which hid their newer copies, are gone. It also takes in the code generation changes merged since then, each traced by recording the file at the commit and at its parent; the count is per section unless a flavor is named: - the magic division sum emitted without signed-overflow poison (#13802): 215 configs in 19 families -- attention_unified 6, attention_unified_stress 32, conv_direct_grouped 16, conv_implicit_gemm 26, conv_implicit_gemm_dgrad 31, conv_implicit_gemm_wgrad 42, deep_fused_conv_pool 2, gfx1201_deep_fused_conv_pool 5, gfx950_deep_fused_conv_pool 5, dynamic_helpers 2, fmha_splitkv_decode 6, img2col 5, moe_fused_mega 5, moe_gemm_fused 6, moe_sorting 6, pooling 4, sage_attention 4, sparse_attention 6, transpose 6; - attention decoding on gfx942 and gfx950 (#13492): 23 changed and 25 new configs in attention_strided_kv, gfx942_attention_tiled_3d and gfx950_attention_tiled_3d; - the f64 follow-ups (#13887): 3 changed and 1 new elementwise config, 5 changed reduce configs; - MX scale and fp6 decoding (#13885): 6 mx_gemm configs; - the gfx950 conv dgrad speed-up (#13075): 36 new conv_direct_dgrad_ext configs; 22 changed, 15 new and 2 removed conv_implicit_gemm_dgrad configs; 2 changed and 3 new conv_direct_grouped configs; 2 new dynamic_helpers and 3 new target_intrinsics configs; - removing the ds.read.tr16.b128 declaration (#14054): gfx1250_lowering, 100 changed and 2 removed configs at llvm23 and in canonical IR, 5 changed, 2 new and 4 removed at llvm20 and llvm22; - optional LSE output on unified split-KV attention (#13955): 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs per LLVM flavor, canonical IR unchanged; - the running softmax max in 3D split-KV (#14276): 14 attention_strided_kv, 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs; - depthwise 7x7 and grouped wgrad on batched MFMA (#14157): 7 changed, 10 new and 1 removed conv_direct_grouped configs; - the retirement of wmma_fmha_fwd (#13897): the gfx1151_wmma_fmha_fwd (157 configs) and gfx1201_wmma_fmha_fwd (12) families are gone; - new configs, every existing hash unchanged: the barrier_fence family (#13575, 10), one conv_implicit_gemm_dgrad config (#13782), 14 conv_direct_grouped configs for the split-C and grouped LDS-halo kernels (#13993), one config each in fmha_splitkv_decode, sage_attention and sparse_attention (#13876), 23 conv_direct_grouped configs for the direct wgrad kernels (#13990), 2 conv_implicit_gemm_dgrad configs for patch-embedding dgrad (#13988), 6 f64_mma configs (#13794), 5 conv_direct_grouped and 7 conv_implicit_gemm configs for depthwise and grouped conv forward (#13987), and the conv_direct_dgrad_wmma family (7 configs, llvm23 and canonical IR only, #14156). #13885, #13955 and #14276 updated the old file's entries for their changes; the new llvm20 and canonical IR sections reproduce every one of those hashes. Co-Authored-By: Claude <noreply@anthropic.com>
cgmillette
added a commit
that referenced
this pull request
Oct 26, 2026
The differential golden had gone stale, because nothing ran the check: 675 of its 1116 LLVM IR entries (compared at llvm20) and 608 of its 1089 canonical IR entries no longer match the tree. Recorded afresh in the new format, with the Python reference emitters, from a tree whose two-engine gate is green at llvm20, llvm22 and llvm23: - ll at llvm20 and llvm22: 82 families, 1474 configs per flavor. At llvm23: 83 families, 1578 configs, where gfx1250_lowering has configs that need LLVM 23 and are refused below it, and every conv_direct_dgrad_wmma config needs LLVM 23. - ir_canonical: 82 families, 1561 configs (conv_abi's emitter has no IR mode). Canonical IR was checked to be identical at llvm20 and llvm23. What the re-bless absorbs: every mismatch that could be traced is an intended code generation change that did not re-bless this file -- the shared LDS pool (#8844, about 425 configs), the ds.swizzle declaration (#9583, about 177), renamed values in the attention transpose reads (#9133, 57), the smaller LDS pool (#10112, about 39), the flat work group size of one-wave wave32 kernels (#12710, about 39), #9220 (6) and the conv changes of #13109 (6). 441 entries were already stale at the start of the history available here and cannot be attributed; no unexplained change was found after it. The two repeated keys (conv_direct_grouped, conv_wgrad_workspace_reduce), which hid their newer copies, are gone. It also takes in the code generation changes merged since then, each traced by recording the file at the commit and at its parent; the count is per section unless a flavor is named: - the magic division sum emitted without signed-overflow poison (#13802): 215 configs in 19 families -- attention_unified 6, attention_unified_stress 32, conv_direct_grouped 16, conv_implicit_gemm 26, conv_implicit_gemm_dgrad 31, conv_implicit_gemm_wgrad 42, deep_fused_conv_pool 2, gfx1201_deep_fused_conv_pool 5, gfx950_deep_fused_conv_pool 5, dynamic_helpers 2, fmha_splitkv_decode 6, img2col 5, moe_fused_mega 5, moe_gemm_fused 6, moe_sorting 6, pooling 4, sage_attention 4, sparse_attention 6, transpose 6; - attention decoding on gfx942 and gfx950 (#13492): 23 changed and 25 new configs in attention_strided_kv, gfx942_attention_tiled_3d and gfx950_attention_tiled_3d; - the f64 follow-ups (#13887): 3 changed and 1 new elementwise config, 5 changed reduce configs; - MX scale and fp6 decoding (#13885): 6 mx_gemm configs; - the gfx950 conv dgrad speed-up (#13075): 36 new conv_direct_dgrad_ext configs; 22 changed, 15 new and 2 removed conv_implicit_gemm_dgrad configs; 2 changed and 3 new conv_direct_grouped configs; 2 new dynamic_helpers and 3 new target_intrinsics configs; - removing the ds.read.tr16.b128 declaration (#14054): gfx1250_lowering, 100 changed and 2 removed configs at llvm23 and in canonical IR, 5 changed, 2 new and 4 removed at llvm20 and llvm22; - optional LSE output on unified split-KV attention (#13955): 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs per LLVM flavor, canonical IR unchanged; - the running softmax max in 3D split-KV (#14276): 14 attention_strided_kv, 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs; - depthwise 7x7 and grouped wgrad on batched MFMA (#14157): 7 changed, 10 new and 1 removed conv_direct_grouped configs; - the retirement of wmma_fmha_fwd (#13897): the gfx1151_wmma_fmha_fwd (157 configs) and gfx1201_wmma_fmha_fwd (12) families are gone; - new configs, every existing hash unchanged: the barrier_fence family (#13575, 10), one conv_implicit_gemm_dgrad config (#13782), 14 conv_direct_grouped configs for the split-C and grouped LDS-halo kernels (#13993), one config each in fmha_splitkv_decode, sage_attention and sparse_attention (#13876), 23 conv_direct_grouped configs for the direct wgrad kernels (#13990), 2 conv_implicit_gemm_dgrad configs for patch-embedding dgrad (#13988), 6 f64_mma configs (#13794), 5 conv_direct_grouped and 7 conv_implicit_gemm configs for depthwise and grouped conv forward (#13987), and the conv_direct_dgrad_wmma family (7 configs, llvm23 and canonical IR only, #14156). #13885, #13955 and #14276 updated the old file's entries for their changes; the new llvm20 and canonical IR sections reproduce every one of those hashes. Co-Authored-By: Claude <noreply@anthropic.com>
This was referenced Oct 29, 2026
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.
Motivation
Kernels need one way to say "these loop coordinates become these tensor coordinates": implicit-GEMM convolution, halos, K splitting and swizzles. They also need bound checks emitted only where a coordinate can actually leave its tensor, and addresses split so the per-wave part is computed once. Today each kernel does this by hand, and the bound checks are either missing or emitted everywhere.
This PR adds
helpers/tiling/tensor_transforms, a coordinate-graph package for any workload (copy, prefetch, reduction, attention, convolution, GEMM). It is new files only; no existing kernel or API changes.Stacked on #13802 (it uses that PR's
cmp_ult,add_nuwand range marker); it will be retargeted todevelopwhen #13802 merges.ISSUE ID : AICK-2578
Technical Details
The graph (
graph.py)make_transform_graph(inputs(...), steps..., outputs(...))with eight kinds:embed,merge(N -> 1),unmerge(1 -> N),offset,slice_window,freeze,xor,broadcast.TensorDesc. A halo is anoffsetplus the tensor's own length, not a pad kind.evaluate/resolvetake a tuple ininputs(...)order; graph input i is window coordinate i. Strings name only the coordinates the author declares.Ranges and required checks (
ranges.py)Interval,Congruenceand their range arithmetic;RuntimeValue, a runtime integer carrying its stated range and multiple-of fact.value_ranges,required_checks(which side of which coordinate can leave its domain or tensor bound, with the reason) andgraph_injective.The IR range walk (
walk.py)value_factsderives a value's range and residue from how it was computed (constants, thread / block ids, loop counters, range markers, add / sub / mul / div / mod / shifts / masks / min / max); anything else is unknown and partial facts survive.split_coordinatesplits a plain i32 value into a constant, wave-uniform parts and one per-lane part.Emission (
emit.py,runs.py,address.py)emit_accessesreturns, per access, the element address split into wave-uniform terms, one per-lane part and a constant, and its validity against the tensor's own lengths and the graph's domains, all over one shared table.emit_bound_check: the one predicate,lo <= value < hi, folded so a side the walk proves emits nothing and the rest is one unsigned compare; widths are chosen so no subtraction can wrap.address_of/transform_address_of: the pure-Python address every emitted address is tested against.Docs
tensor_transforms_api.md(the package's facts, every error message verbatim) andtensor_transforms_guide.md(runnable examples: halo, K split, swizzle, addresses, required checks, and CK's pad transform expressed here).Test Plan
transform_address_of, failing on any overflowing op;golden IR text; every golden case lowered by both engines to the same IR at llvm20, llvm22 and llvm23 and
compiled for gfx90a, gfx942, gfx950, gfx1151, gfx1201 and gfx1250 (gfx1250 only where the loaded comgr can target it).
transform_address_of.and validity of convolution, halo and ragged vector accesses.
sage and sparse attention, and two scalar unified attention golden cases; the scalar unified attention mask
test now runs on any GPU and adds a 40 / 8 head case.
Test Result
compile cases skip because that host's comgr cannot target gfx1250.
once on gfx942 and gfx90a with non-power-of-two divisors against torch references (all within 5e-3).
add nuwcommit from fix(rocke): emit magic division's sum without signed-overflow poison #13802 (201 of its 700 hashes change: conv,conv dgrad / wgrad, direct conv, the fused conv-pool families, MoE sorting and the 5 gfx950 scalar
attention_unifiedcases, each at three LLVM flavors). The tensor_transforms commits do not change it. Thereview follow-up adds two cases (the scalar unified 2D and 3D kernels with 40 query heads over 8 KV heads) at
three flavors; no existing hash changes.
Submission Checklist
🤖 Generated with Claude Code