[water] Add ordered_syms attribute to preserve dimension ordering#714
Merged
tyb0807 merged 2 commits intoiree-org:mainfrom Jan 20, 2026
Merged
[water] Add ordered_syms attribute to preserve dimension ordering#714tyb0807 merged 2 commits intoiree-org:mainfrom
tyb0807 merged 2 commits intoiree-org:mainfrom
Conversation
tyb0807
force-pushed
the
ordered_syms_attr
branch
from
4e56df1 to
cab717d
Compare
Contributor
Author
|
After offline discussion, I'll try to take a shot at introducing a proper attribute for index expressions that lets us keep stable ordering across dimensions, required when lowering index expressions to affine maps, used primarily by memory operations. In the meantime, this can still be used to unblock end-to-end execution of matmul and FA kernels. |
ftynse
reviewed
Jan 13, 2026
water/lib/Dialect/Wave/Transforms/ResolveDistributedAllocations.cpp
Outdated
Show resolved
Hide resolved
water/lib/Dialect/Wave/Transforms/ResolveDistributedAllocations.cpp
Outdated
Show resolved
Hide resolved
tyb0807
force-pushed
the
ordered_syms_attr
branch
7 times, most recently
from
1c559d5 to
fd2c2aa
Compare
Contributor
Author
|
Still pushing on this because it might take some time for #730 to land (if that's actually the right solution to this) |
martin-luecke
approved these changes
Jan 20, 2026
Contributor
martin-luecke
left a comment
martin-luecke
left a comment
There was a problem hiding this comment.
With your addressing of previous reviews, this looks good to me.
Please fix the conflict, and we are good to merge.
MLIR's DictionaryAttr sorts entries alphabetically by key, which loses the semantic dimension ordering from the Python frontend when index expressions are stored. This caused incorrect code generation with transposed memory indices. This change adds an ordered_syms attribute to ReadOp and WriteOp that captures the dimension ordering from WaveTensorType before it is converted to MemRefType. The ResolveDistributedAllocations pass now sets this attribute, and a new OrderedSymsSpecified normal form bit ensures this happens before LowerWaveToMLIR. For ops with WaveTensorType memory operands, the ordering is still extracted from the type's shape. For MemRefType memory operands (after allocation resolution), the ordered_syms attribute is used. Signed-off-by: tyb0807 <sontuan.vu@amd.com>
tyb0807
force-pushed
the
ordered_syms_attr
branch
2 times, most recently
from
a3ddcd8 to
949fda8
Compare
Signed-off-by: tyb0807 <sontuan.vu@amd.com>
tyb0807
force-pushed
the
ordered_syms_attr
branch
from
949fda8 to
f78391d
Compare
tyb0807
added a commit
to tyb0807/wave
that referenced
this pull request
Jan 26, 2026
…ee-org#714) MLIR's DictionaryAttr sorts entries alphabetically by key, which loses the semantic dimension ordering. This caused incorrect code generation with transposed memory indices. This change adds an ordered_syms attribute to ReadOp and WriteOp that captures the dimension ordering from WaveTensorType before it is converted to MemRefType. The ResolveDistributedAllocations pass now sets this attribute, and a new OrderedSymsSpecified normal form bit ensures this happens before LowerWaveToMLIR. For ops with WaveTensorType memory operands, the ordering is extracted from the type's shape. For MemRefType memory operands (after allocation resolution), the ordered_syms attribute is used. --------- Signed-off-by: tyb0807 <sontuan.vu@amd.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
3 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.
MLIR's DictionaryAttr sorts entries alphabetically by key, which loses the semantic dimension ordering. This caused incorrect code generation with transposed memory indices.
This change adds an ordered_syms attribute to ReadOp and WriteOp that captures the dimension ordering from WaveTensorType before it is converted to MemRefType. The ResolveDistributedAllocations pass now sets this attribute, and a new OrderedSymsSpecified normal form bit ensures this happens before LowerWaveToMLIR.
For ops with WaveTensorType memory operands, the ordering is extracted from the type's shape. For MemRefType memory operands (after allocation resolution), the ordered_syms attribute is used.