[water] provenance-based propagation heuristic#1206
Merged
Conversation
Water Code Coverage |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates Wave index-expression propagation to track and propagate an additional “source vector shape” (the originating op’s vector shape) in the lattice, in order to preserve Python-prototype propagation/stopper behavior that can’t be reconstructed from the per-value vector shapes alone.
Changes:
- Extend
IndexExprsLatticeStoragewithsourceVectorShapeandsourceVectorShapePriority, and join them using priority semantics. - Update the propagation-stopper heuristic (
shouldPropagateIndexExprs) and key ops (MMA init, Write sideways propagation, Permute stride permutation) to preserve/use the new source vector shape. - Adjust and expand MLIR tests and design notes to reflect the new behavior.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| water/lib/Dialect/Wave/IR/WaveOps.cpp | Updates propagation heuristic and ensures MMA/Write/Permute flows preserve sourceVectorShape. |
| water/lib/Dialect/Wave/IR/WaveInterfaces.cpp | Implements storage, printing, equality, and join semantics for sourceVectorShape. |
| water/include/water/Dialect/Wave/IR/WaveInterfaces.h | Documents new lattice fields and updated propagation heuristic behavior. |
| water/test/Dialect/Wave/infer-index-exprs.mlir | Adjusts an existing test expectation around backward propagation into a register. |
| water/test/Dialect/Wave/infer-index-exprs-lattice.mlir | Updates skip tests and adds new tests covering sourceVectorShape join behavior and conflicts. |
| docs/wave/ir_design_notes.rst | Updates design notes to describe the new source-vector-shape-based heuristic. |
Comment on lines
+1696
to
+1702
| // Equal priorities: identical values join cleanly, different go to top. | ||
| if (lhsSVS == rhsSVS) { | ||
| joinedSourceVectorShape = lhsSVS; | ||
| joinedSourceVectorShapePriority = lhsPri; | ||
| } else { | ||
| return IndexExprsLatticeStorage::top(); | ||
| } |
There was a problem hiding this comment.
IndexExprsLatticeStorage::join can now return top() due to a sourceVectorShape conflict (equal sourceVectorShapePriority but different sourceVectorShape). Current diagnostics in propagation treat such cases as "index expression" conflicts (since getJoinedVectorShape can still succeed), which is misleading and makes failures hard to debug. Consider adding an explicit source-vector-shape join/check helper so callers can classify and report this as a distinct "source vector shape" conflict (and include both source shapes/priorities in the notes).
Contributor
There was a problem hiding this comment.
valid concern. Maybe in a follow-up or maybe not as we do not intend on keeping this forever
Contributor
Author
There was a problem hiding this comment.
This should never happen outside of synthetic cases. I just didn't want to assert in the core lattice method.
Contributor
martin-luecke
left a comment
martin-luecke
left a comment
There was a problem hiding this comment.
Not 100% done reviewing yet, but leaving some comments to address
Comment on lines
+1696
to
+1702
| // Equal priorities: identical values join cleanly, different go to top. | ||
| if (lhsSVS == rhsSVS) { | ||
| joinedSourceVectorShape = lhsSVS; | ||
| joinedSourceVectorShapePriority = lhsPri; | ||
| } else { | ||
| return IndexExprsLatticeStorage::top(); | ||
| } |
Contributor
There was a problem hiding this comment.
valid concern. Maybe in a follow-up or maybe not as we do not intend on keeping this forever
Comment on lines
+2035
to
+2046
| // CHECK-DAG: M : <[#wave.index_symbol<T0>] -> (T0 * 32, 1, 1)> | ||
| // CHECK-DAG: N : <[#wave.index_symbol<T0>] -> (T0 * 10, 1, 1)> | ||
| %result = wave.add %a, %b {wave_test.override_operand_index = [ | ||
| [{M = 5 : i32, N = 5 : i32}, { | ||
| M = #wave.index_mapping<[#wave.index_symbol<T0>] -> (T0 * 32, 1, 1)>, | ||
| N = #wave.index_mapping<[#wave.index_symbol<T0>] -> (T0 * 10, 1, 1)> | ||
| }, {M = 16 : i64, N = 16 : i64}], | ||
| [{M = 3 : i32, N = 3 : i32}, { | ||
| M = #wave.index_mapping<[#wave.index_symbol<T0>] -> (T0 * 32, 1, 1)>, | ||
| N = #wave.index_mapping<[#wave.index_symbol<T0>] -> (T0 * 10, 1, 1)> | ||
| }, {M = 16 : i64, N = 0 : i64}] | ||
| ]} |
Contributor
There was a problem hiding this comment.
We don't actually check here that {M=16,N=16} wins in propagation. Having a consuming op where this is propagated to would enable checking that this works as expected. Same for the other non-failure tests
Contributor
Author
There was a problem hiding this comment.
So the actual problem here is we don't even print source vector shape anywhere, do we can't check it directly, only by proxy.. I'll factor this out into a separate file.
martin-luecke
approved these changes
Apr 1, 2026
Contributor
martin-luecke
left a comment
martin-luecke
left a comment
There was a problem hiding this comment.
LGTM, with a few comments and questions to address
ftynse
force-pushed
the
users/ftynse/heuristic
branch
from
2dbb998 to
1d4587e
Compare
It is unfortunately necessary to introduce this additional complexity to reflect the naive non-dataflow approach in the python counterpart: keep track of the vector shape of the operation from which the lattice originally propagates. This is required by the propagation-stopper heuristic that cannot infer this value from per-value vector shapes that water correctly cleans up to correspond to the shape. The shape is stored as a separate lattice field that is simply carried over based on priority. Since all Mma operations now have different priorities by construction, this never leads to lattice tops outside of synthetic test cases. For Write operations, the other source, the priority is equal but so are vector shapes since those are simply taken from the hardware constraint object. This is highly dubious due to non-local propagation rules and blind overwriting of vector shapes in the python code, but necessary to match the current behavior. It must be revised as soon as possible to avoid priorities and non-local behavior altogether. Signed-off-by: Alex Zinenko <git@ozinenko.com>
Signed-off-by: Alex Zinenko <git@ozinenko.com>
Signed-off-by: Alex Zinenko <git@ozinenko.com>
ftynse
force-pushed
the
users/ftynse/heuristic
branch
from
1d4587e to
c70921d
Compare
tgymnich
approved these changes
Apr 2, 2026
Contributor
tgymnich
left a comment
tgymnich
left a comment
There was a problem hiding this comment.
LGTM. maybe add ScaledMMA
Comment on lines
+2162
to
+2163
| SmallVector<wave::IndexExprsLatticeStorage> slots = llvm::map_to_vector( | ||
| valuesForIndexExpr, [&](Value v) { return getLatticeValue(v); }); |
Contributor
There was a problem hiding this comment.
Suggested change
| SmallVector<wave::IndexExprsLatticeStorage> slots = llvm::map_to_vector( | |
| valuesForIndexExpr, [&](Value v) { return getLatticeValue(v); }); | |
| SmallVector<wave::IndexExprsLatticeStorage> slots = | |
| llvm::map_to_vector(valuesForIndexExpr, getLatticeValue); |
Contributor
Author
|
I'll update scaled mma separately to avoid issues with an interferring refactor |
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.
It is unfortunately necessary to introduce this additional complexity to reflect the naive non-dataflow approach in the python counterpart: keep track of the vector shape of the operation from which the lattice originally propagates. This is required by the propagation-stopper heuristic that cannot infer this value from per-value vector shapes that water correctly cleans up to correspond to the shape. The shape is stored as a separate lattice field that is simply carried over based on priority. Since all Mma operations now have different priorities by construction, this never leads to lattice tops outside of synthetic test cases. For Write operations, the other source, the priority is equal but so are vector shapes since those are simply taken from the hardware constraint object.
This is highly dubious due to non-local propagation rules and blind overwriting of vector shapes in the python code, but necessary to match the current behavior. It must be revised as soon as possible to avoid priorities and non-local behavior altogether.