[AMDGPU] Mark WMMA machine instructions as convergent #165602
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
Reduced MIR test.
|
@llvm/pr-subscribers-backend-amdgpu Author: Syadus Sefat (mssefat) Changes[AMDGPU] Mark WMMA machine instructions as convergent The WMMA MI(s) are missing the isConvergent flag. This causes incorrect behavior in passes like machine-sink, where WMMA instructions get sunk into divergent branches. This patch fixes the issue by setting the isConvergent flag to 1 in the VOP3PInstructions.td file. Full diff: https://github.com/llvm/llvm-project/pull/165602.diff 2 Files Affected:
diff --git a/llvm/lib/Target/AMDGPU/VOP3PInstructions.td b/llvm/lib/Target/AMDGPU/VOP3PInstructions.td
index 4ae2c1ed04dae..31d8bce4d0c87 100644
--- a/llvm/lib/Target/AMDGPU/VOP3PInstructions.td
+++ b/llvm/lib/Target/AMDGPU/VOP3PInstructions.td
@@ -1707,7 +1707,7 @@ multiclass WMMAInstGFX12<string Instr, VOP3PWMMA_Profile WMMAProfile, string Pse
defvar WMMAConstraints2Addr = !if(DiffVdstSrc2, "@earlyclobber $vdst", "@earlyclobber $vdst,$vdst = $src2");
defvar WMMAConstraints3Addr = "@earlyclobber $vdst";
- let Mnemonic = Instr, mayRaiseFPException = 0, ReadsModeReg = 0 in {
+ let Mnemonic = Instr, mayRaiseFPException = 0, ReadsModeReg = 0, isConvergent = 1 in {
let Constraints = WMMAConstraints2Addr, isConvertibleToThreeAddress = 1 in
def _twoaddr : VOP3P_Pseudo<Instr, WMMAProfile>, WMMAInstInfo {
let PseudoInstr = Instr#PseudoInstrSuffix;
@@ -1734,7 +1734,7 @@ multiclass SWMMACInstGFX12<string Instr, VOP3PWMMA_Profile WMMAProfile, string P
let mayRaiseFPException = 0;
let ReadsModeReg = 0;
let AsmMatchConverter = "cvtSWMMAC";
-
+ let isConvergent = 1;
let Constraints = "@earlyclobber $vdst,$vdst = $srcTiedDef";
}
}
@@ -1906,8 +1906,10 @@ defm V_WMMA_SCALE_F32_32X16X128_F4_w32 : WMMAInstGFX12<"v_wmma_scale_f32_32x16
defm V_WMMA_SCALE16_F32_32X16X128_F4_w32 : WMMAInstGFX12<"v_wmma_scale16_f32_32x16x128_f4", F32_32X16X128_F4_SCALE16_w32, "_w32">;
} // End is_wmma_xdl = 1.
-defm V_WMMA_LD_SCALE_PAIRED_B32 : VOP3PInst<"v_wmma_ld_scale_paired_b32", VOP_WMMA_LD_SCALE<i32, VCSrc_b32_Lo256>>;
-defm V_WMMA_LD_SCALE16_PAIRED_B64 : VOP3PInst<"v_wmma_ld_scale16_paired_b64", VOP_WMMA_LD_SCALE<i64, VCSrc_b64_Lo256>>;
+let isConvergent = 1 in {
+ defm V_WMMA_LD_SCALE_PAIRED_B32 : VOP3PInst<"v_wmma_ld_scale_paired_b32", VOP_WMMA_LD_SCALE<i32, VCSrc_b32_Lo256>>;
+ defm V_WMMA_LD_SCALE16_PAIRED_B64 : VOP3PInst<"v_wmma_ld_scale16_paired_b64", VOP_WMMA_LD_SCALE<i64, VCSrc_b64_Lo256>>;
+}
} // End SubtargetPredicate = isGFX125xOnly
} // End WaveSizePredicate = isWave32
diff --git a/llvm/test/CodeGen/AMDGPU/wmma-gfx12-convergent.mir b/llvm/test/CodeGen/AMDGPU/wmma-gfx12-convergent.mir
new file mode 100644
index 0000000000000..eef36674dba35
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/wmma-gfx12-convergent.mir
@@ -0,0 +1,78 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx12-generic -run-pass=machine-sink %s -o - | FileCheck %s
+
+--- |
+ ; ModuleID = 'test-wmma-convergent'
+ target triple = "amdgcn-amd-amdhsa"
+
+ define void @wmma_test() {
+ entry:
+ br label %if.then
+
+ if.then:
+ br label %if.end
+
+ if.end:
+ ret void
+ }
+
+...
+---
+name: wmma_test
+alignment: 1
+tracksRegLiveness: true
+body: |
+ ; CHECK-LABEL: name: wmma_test
+ ; CHECK: bb.0.entry:
+ ; CHECK-NEXT: successors: %bb.1(0x40000000), %bb.2(0x40000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[DEF:%[0-9]+]]:vreg_128 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF1:%[0-9]+]]:vreg_128 = IMPLICIT_DEF
+ ; CHECK-NEXT: [[DEF2:%[0-9]+]]:sreg_32 = IMPLICIT_DEF
+ ; CHECK-NEXT: early-clobber %3:vreg_256 = V_WMMA_F32_16X16X16_F16_w32_threeaddr 8, [[DEF]], 8, [[DEF1]], 8, 0, 0, 0, implicit $exec
+ ; CHECK-NEXT: [[SI_IF:%[0-9]+]]:sreg_32 = SI_IF [[DEF2]], %bb.2, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+ ; CHECK-NEXT: S_BRANCH %bb.1
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.1.if.then:
+ ; CHECK-NEXT: successors: %bb.2(0x80000000)
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: [[COPY:%[0-9]+]]:vgpr_32 = COPY %3.sub1
+ ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vgpr_32 = COPY %3.sub3
+ ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vgpr_32 = COPY %3.sub5
+ ; CHECK-NEXT: [[COPY3:%[0-9]+]]:vgpr_32 = COPY %3.sub7
+ ; CHECK-NEXT: [[COPY4:%[0-9]+]]:vgpr_32 = COPY %3.sub6
+ ; CHECK-NEXT: [[COPY5:%[0-9]+]]:vgpr_32 = COPY %3.sub4
+ ; CHECK-NEXT: [[COPY6:%[0-9]+]]:vgpr_32 = COPY %3.sub2
+ ; CHECK-NEXT: [[COPY7:%[0-9]+]]:vgpr_32 = COPY %3.sub0
+ ; CHECK-NEXT: {{ $}}
+ ; CHECK-NEXT: bb.2.if.end:
+ ; CHECK-NEXT: SI_END_CF [[SI_IF]], implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+ ; CHECK-NEXT: S_ENDPGM 0
+
+ bb.0.entry:
+ successors: %bb.1, %bb.2
+
+ %0:vreg_128 = IMPLICIT_DEF
+ %1:vreg_128 = IMPLICIT_DEF
+ %2:sreg_32 = IMPLICIT_DEF
+ early-clobber %3:vreg_256 = V_WMMA_F32_16X16X16_F16_w32_threeaddr 8, %0:vreg_128, 8, %1:vreg_128, 8, 0, 0, 0, implicit $exec
+ %4:sreg_32 = SI_IF %2:sreg_32, %bb.2, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+ S_BRANCH %bb.1
+
+ bb.1.if.then:
+ successors: %bb.2
+
+ %5:vgpr_32 = COPY %3.sub1:vreg_256
+ %6:vgpr_32 = COPY %3.sub3:vreg_256
+ %7:vgpr_32 = COPY %3.sub5:vreg_256
+ %8:vgpr_32 = COPY %3.sub7:vreg_256
+ %9:vgpr_32 = COPY %3.sub6:vreg_256
+ %10:vgpr_32 = COPY %3.sub4:vreg_256
+ %11:vgpr_32 = COPY %3.sub2:vreg_256
+ %12:vgpr_32 = COPY %3.sub0:vreg_256
+
+ bb.2.if.end:
+ SI_END_CF %4:sreg_32, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+ S_ENDPGM 0
+
+...
|
|
@arsenm Hi Matt, could you please review? Thank you. |
bcahoon
reviewed
Oct 30, 2025
|
|
||
| defm V_WMMA_LD_SCALE_PAIRED_B32 : VOP3PInst<"v_wmma_ld_scale_paired_b32", VOP_WMMA_LD_SCALE<i32, VCSrc_b32_Lo256>>; | ||
| defm V_WMMA_LD_SCALE16_PAIRED_B64 : VOP3PInst<"v_wmma_ld_scale16_paired_b64", VOP_WMMA_LD_SCALE<i64, VCSrc_b64_Lo256>>; | ||
| let isConvergent = 1 in { |
There was a problem hiding this comment.
Should we move isConvergent VOP_WMMA_LD_SCALE?
There was a problem hiding this comment.
I think we should avoid moving isConvergent = 1 directly into VOP_WMMA_LD_SCALE since it extends a profile class and isConvergent is an instruction property.
Instead we can create VOP3PInst_WMMA_LD that extends VOP3PInst specifically for WMMA_LD_SCALE_PAIRED instructions:
multiclass VOP3PInst_WMMA_LD<string OpName, VOPProfile P, SDPatternOperator node = null_frag> {
let isConvergent = 1 in {
defm NAME : VOP3PInst<OpName, P, node>;
}
}
Then we can use it for WMMA LD_SCALE_PAIRED instructions:
defm V_WMMA_LD_SCALE_PAIRED_B32 : VOP3PInst_WMMA_LD<"v_wmma_ld_scale_paired_b32", VOP_WMMA_LD_SCALE<i32, VCSrc_b32_Lo256>>;
defm V_WMMA_LD_SCALE16_PAIRED_B64 : VOP3PInst_WMMA_LD<"v_wmma_ld_scale16_paired_b64", VOP_WMMA_LD_SCALE16<i64, VCSrc_b64_Lo256>>;
Please let me know if this looks better to you or you suggest me to stick with the existing one.
arsenm
reviewed
Oct 31, 2025
Comment on lines
+4
to
+19
| --- | | ||
| ; ModuleID = 'test-wmma-convergent' | ||
| target triple = "amdgcn-amd-amdhsa" | ||
|
|
||
| define void @wmma_test() { | ||
| entry: | ||
| br label %if.then | ||
|
|
||
| if.then: | ||
| br label %if.end | ||
|
|
||
| if.end: | ||
| ret void | ||
| } | ||
|
|
||
| ... |
There was a problem hiding this comment.
Suggested change
| --- | | |
| ; ModuleID = 'test-wmma-convergent' | |
| target triple = "amdgcn-amd-amdhsa" | |
| define void @wmma_test() { | |
| entry: | |
| br label %if.then | |
| if.then: | |
| br label %if.end | |
| if.end: | |
| ret void | |
| } | |
| ... |
| ... | ||
| --- | ||
| name: wmma_test | ||
| alignment: 1 |
There was a problem hiding this comment.
Suggested change
| alignment: 1 |
Comment on lines
+65
to
+72
| %5:vgpr_32 = COPY %3.sub1:vreg_256 | ||
| %6:vgpr_32 = COPY %3.sub3:vreg_256 | ||
| %7:vgpr_32 = COPY %3.sub5:vreg_256 | ||
| %8:vgpr_32 = COPY %3.sub7:vreg_256 | ||
| %9:vgpr_32 = COPY %3.sub6:vreg_256 | ||
| %10:vgpr_32 = COPY %3.sub4:vreg_256 | ||
| %11:vgpr_32 = COPY %3.sub2:vreg_256 | ||
| %12:vgpr_32 = COPY %3.sub0:vreg_256 |
There was a problem hiding this comment.
You can probably simplify this to a simple full copy
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.
[AMDGPU] Mark WMMA machine instructions as convergent
The WMMA MI(s) are missing the isConvergent flag. This causes incorrect behavior in passes like machine-sink, where WMMA instructions get sunk into divergent branches.
This patch fixes the issue by setting the isConvergent flag to 1 in the VOP3PInstructions.td file.