Refactor recognition of loads of multiramps - #9261
Open
abadams wants to merge 10 commits into
Open
Conversation
get_subtile partitions accesses to a tile-memory allocation between the distinct sub-tiles it holds, and is_load_of_multiramp digs a load out from under the casts, broadcasts and lane permutations that can wrap it. Neither is specific to AMX, so move them next to the MultiRamp machinery they are built on. This drops 126 lines from ExtractTileOperations. Supporting this, is_multiramp learns to see through a shuffle of a single vector when the shuffle is a reshaping rather than a gather: either a transpose, or any permutation of a one-dimensional multiramp whose lane indices are themselves a multiramp. MultiRamp::transpose applies the former to a multiramp directly, splitting a dim in two where the transpose falls inside one. Routing the AMX operands through is_load_of_multiramp lets that pass match loads it previously missed, and lets it check the operand types against the type actually loaded from memory rather than the type of the expression wrapping it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
is_load_of_multiramp stripped any number of casts off a load and returned only the Load underneath, so a caller had no way to see that the values had been cast. ExtractTileOperations then typed its tile registers from the load, and the signedness of those types picks which of the four integer tdpb instructions runs. A pipeline multiplying uint8 buffers reinterpreted as int8 compiled to tdpbuud - an unsigned multiply for a signed algorithm. Peel at most one cast, so the element type of the original Expr and the type of the returned Load together say whether the values were cast and to what, and have the AMX integer path reject a mismatch. The float path still reads the load's type, which is what it wants: the bf16 to f32 widening cast is expected there, and tile_load takes the bf16. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A transpose's shuffle mask is itself a multiramp of constants - transposing n lanes with c columns gives a mask of shape (n/c, c) and strides (c, 1) - so the general reshaping-shuffle case in is_multiramp already covers it. Drop the separate is_transpose case and MultiRamp::transpose, which existed only to serve it and had no other caller. That leaves one rule for shuffles of a single vector: it is a reshaping rather than a gather if the mask is a multiramp of constants, and we can say what it reshapes to when the shuffled vector is one-dimensional. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
alexreinking
self-requested a review
July 30, 2026 21:46
alexreinking
previously approved these changes
Jul 31, 2026
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #9261 +/- ##
==========================================
+ Coverage 70.32% 70.37% +0.04%
==========================================
Files 258 258
Lines 79160 79238 +78
Branches 18997 19032 +35
==========================================
+ Hits 55672 55762 +90
+ Misses 17802 17787 -15
- Partials 5686 5689 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
expect_user_error accepted any CompileError, so a scenario could pass on an unrelated one - a schedule Halide rejects before lowering looks the same as the lowering error the scenario is for. Give it the message to look for, and name the path each scenario exercises. That leaves one of the checks in get_subtile untested, so add a case for it: two matmuls into the same allocation whose tiles have the same rank but different extents. Also drop the returns after the user_errors there, which throw. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
is_multiramp treated a shuffle of a single vector as a reshaping only when the vector was one-dimensional, which is the shape flatten_nested_ramps leaves a strided load in. A tile load that gets transposed is not that shape, so add MultiRamp::shuffle, which handles any mask whose lanes are a permutation of [0, total_lanes()) - that is, whose strides are the prefix products of its lane counts in some order. It works by refining the mask's dims and the multiramp's against each other until both are groupings of a common shape, then reordering that shape the way the mask asks for. MultiRamp::add was already doing the same walk to line two multiramps up, and was also duplicating strides_for_shape to rescale the strides as it went, so share the walk and let strides_for_shape do the rest. That makes add a good deal shorter, and makes the failure paths that a common refinement rules out into internal errors rather than silent rejections. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Member
Author
|
I had to redo shuffles-of-multiramps to support more than the 1d case for future work. Please take another look at multiramp.cpp |
alexreinking
self-requested a review
August 3, 2026 21:08
alexreinking
requested changes
Aug 5, 2026
Test that one of the two lane counts divides the other, rather than that they are not coprime. Being coprime was never the condition: 12 and 8 have a common factor but no common refinement, because the prefix products of a refinement each divide the next and have to include the prefix products of both shapes. Mismatched total lane counts are a caller error rather than a shape that has no common refinement, so assert instead of returning false. Each step consumes the same number of lanes from both shapes, so the check is just that they run out together. Track the lanes of the current dim not yet accounted for, taking the next dim at the top of the loop when the current one is used up. This drops a placeholder value that was never read, and an increment that used to hide in the condition of an if. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
alexreinking
approved these changes
Aug 5, 2026
Load::make lost the default for is_streaming when the short forms were added, so the call that left it off no longer compiles. The short form is what it wanted anyway: an unpredicated load from an internal buffer. Co-Authored-By: Claude Opus 5 <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
2 participants
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.
This is a refactoring of some AMX matrix-load-recognition logic so that it can be shared with other targets that need to recognize the same basic patterns: multi-dimensional loads, and multi-dimensional loads of non-overlapping subtiles. Recognizing a multi-dimensional load is slightly more tricky that asking if a load index is a multiramp, because there may be shuffles and broadcasts outside the load too - shuffles and loads commute. The second feature is needed to partition one logical Halide allocation between multiple AMX tile registers (or in the future, WMMA fragments).