[Codegen][LLVMGPU] Add layout support for ArgCompare operations#23693
Merged
bangtianliu merged 2 commits intoiree-org:mainfrom Mar 30, 2026
Merged
[Codegen][LLVMGPU] Add layout support for ArgCompare operations#23693bangtianliu merged 2 commits intoiree-org:mainfrom
bangtianliu merged 2 commits intoiree-org:mainfrom
Conversation
bangtianliu
requested review from
Groverkss,
Max191,
krzysz00,
kuhar,
nirvedhmeshram and
qedawkins
as code owners
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
4 times, most recently
from
95a1651 to
09f915e
Compare
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
from
09f915e to
1b3345e
Compare
11 tasks
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
3 times, most recently
from
1974732 to
fd40274
Compare
Max191
approved these changes
Mar 27, 2026
Contributor
Max191
left a comment
Max191
left a comment
There was a problem hiding this comment.
LGTM, but maybe wait for someone who is more well-versed with layout analysis as well.
Comment on lines
+291
to
+300
| if (argCompare.getInitValue() == val) { | ||
| LDBG() << " Propagating init_value layout to result\n"; | ||
| addCandidate(argCompare.getResultValue(), layout); | ||
| continue; | ||
| } | ||
| if (argCompare.getInitIndex() == val) { | ||
| LDBG() << " Propagating init_index layout to result\n"; | ||
| addCandidate(argCompare.getResultIndex(), layout); | ||
| continue; | ||
| } |
Contributor
There was a problem hiding this comment.
Weird edge case: Couldn't the init value and init index technically be the same tensor? I highly doubt there would ever be an arg_compare like that, but just for completeness, should we remove the continues in these init cases and put the continue at the end of the containing if (auto argCompare = dyn_cast<ArgCompareOp>(user)) { instead?
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
from
fd40274 to
fd0e107
Compare
Contributor
Author
|
cc @Groverkss @sommerlukas for review |
sommerlukas
reviewed
Mar 30, 2026
Contributor
sommerlukas
left a comment
sommerlukas
left a comment
There was a problem hiding this comment.
Just a few small things on the tests.
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Outdated
Show resolved
Hide resolved
sommerlukas
reviewed
Mar 30, 2026
Contributor
sommerlukas
left a comment
sommerlukas
left a comment
There was a problem hiding this comment.
Thanks for addressing the comments, but I think we should have the check lines actually match the expected value.
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
Outdated
Show resolved
Hide resolved
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
from
a60ed3c to
d51a71f
Compare
sommerlukas
approved these changes
Mar 30, 2026
Groverkss
approved these changes
Mar 30, 2026
compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
Outdated
Show resolved
Hide resolved
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
bangtianliu
force-pushed
the
argcompare-layout-propagation
branch
from
d51a71f to
371a77e
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
4 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 PR adds tensor layout propagation for iree_linalg_ext.arg_compare operations in the LLVMGPU backend, enabling arg_compare to participate in the vector distribution pipeline.
Issue: #23005
Assisted-by: Claude Code