Separate wave control flow lowering into dedicated pass #646
+399
−217
Conversation
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
Creates a new LowerWaveControlFlowPass that specifically handles wave.iterate -> scf.for conversion, avoiding the issue where mixed lowering was creating unlegalizable IR. The problem is that if LowerWaveToMLIRPass tries to do both: 1. Control flow lowering (wave.iterate -> scf.for) 2. Operation lowering (wave.read, wave.mma, etc. -> standard ops) This will cause ` wave.iterate` bodies to be cloned with unconverted wave operations inside `scf.for` loops, creating IR that couldn't be legalized. The solution separates the pipeline into two phases: 1. LowerWaveControlFlowPass: wave.iterate -> scf.for (leaves other wave ops unchanged) 2. LowerWaveToMLIRPass: Remaining wave ops -> standard MLIR ops This enables a clean two-phase lowering pipeline where control flow is handled first, then remaining wave operations are lowered in the standard `scf.for` context. Signed-off-by: tyb0807 <sontuan.vu@amd.com>
ftynse
requested changes
Dec 29, 2025
|
|
||
| def LowerWaveToMLIRPass : Pass<"lower-wave-to-mlir"> { | ||
| let summary = "Lower Wave dialect to upstream MLIR dialects"; | ||
| let summary = "Lower Wave dialect operations to upstream MLIR dialects"; |
Contributor
There was a problem hiding this comment.
You know control flow operations are also operations?
| let description = [{ | ||
| This pass lowers operations from the Wave dialect to upstream MLIR | ||
| dialects, notably arith and vector. | ||
| This pass lowers remaining Wave dialect operations to upstream MLIR |
Contributor
There was a problem hiding this comment.
It is unclear what "remaining" means here.
| target.addLegalDialect< | ||
| // clang-format off | ||
| arith::ArithDialect, | ||
| func::FuncDialect, |
Comment on lines
+182
to
+185
| void getDependentDialects(DialectRegistry ®istry) const override { | ||
| registry.insert<arith::ArithDialect, scf::SCFDialect>(); | ||
| } | ||
|
|
Contributor
There was a problem hiding this comment.
You already listed a different set of dialects in the tablegen definition. This is a huge footgun!
| @@ -0,0 +1,148 @@ | |||
| // RUN: water-opt %s -allow-unregistered-dialect -lower-wave-control-flow --mlir-print-local-scope --split-input-file --verify-diagnostics | FileCheck %s | |||
Contributor
There was a problem hiding this comment.
Do we need unregistered dialect? -verify-diagnostics? Tests should be minimal.
|
|
||
| return | ||
| } | ||
| } No newline at end of file |
Contributor
There was a problem hiding this comment.
Trailing newline. If you run pre-commit, these should be fixed for you.
Comment on lines
+231
to
+234
| // Create register values - elements_per_thread determined by MMA backward propagation. | ||
| %lhs = wave.register %lhs_init : !wave.tensor<[@M, @K] of f16, <register>> | ||
| %rhs = wave.register %rhs_init : !wave.tensor<[@N, @K] of f16, <register>> | ||
| %acc = wave.register %acc_init : !wave.tensor<[@M, @N] of f32, <register>> |
Contributor
There was a problem hiding this comment.
Why is this test changed here?
Comment on lines
+257
to
+259
| // Create register values - elements_per_thread determined by MMA backward propagation. | ||
| %lhs = wave.register %lhs_init : !wave.tensor<[@M, @K] of f16, <register>> | ||
| %rhs = wave.register %rhs_init : !wave.tensor<[@N, @K] of f16, <register>> |
Contributor
|
I don't quite understand the problem from the description. We clone the body with operations that weren't converted, why are they not converted later on in the same pass? Is the infra somehow not revisiting them? |
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.
Creates a new LowerWaveControlFlowPass that specifically handles wave.iterate -> scf.for conversion, avoiding the issue where mixed lowering was creating unlegalizable IR.
The problem is that if LowerWaveToMLIRPass tries to do both:
This will cause
wave.iteratebodies to be cloned with unconverted wave operations insidescf.forloops, creating IR that couldn't be legalized.The solution separates the pipeline into two phases:
This enables a clean two-phase lowering pipeline where control flow is handled first, then remaining wave operations are lowered in the standard
scf.forcontext.