[llvm] [AMDGPU] SIFoldOperands: update BUNDLE header implicit use when folding (PR #201872)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Petar Avramovic (petar-avramovic)

<details>
<summary>Changes</summary>

When folding an operand inside a BUNDLE, also rewrite the matching
implicit use on the bundle header. LiveVariables iterates a
MachineBasicBlock with the bundle-aware iterator and only inspects the
header, so without this update its kill flags go stale and a later
MachineVerifier run reports "Using a killed virtual register".

Co-Authored-By: Claude Opus 4 <noreply@<!-- -->anthropic.com>

---
Full diff: https://github.com/llvm/llvm-project/pull/201872.diff


3 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIFoldOperands.cpp (+15) 
- (modified) llvm/test/CodeGen/AMDGPU/ds_gws_align.ll (+11) 
- (added) llvm/test/CodeGen/AMDGPU/si-fold-operands-bundle.mir (+51) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
index b88100ddd1fbe..6c1db80c08e7d 100644
--- a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
+++ b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
@@ -18,6 +18,7 @@
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBundle.h"
 #include "llvm/CodeGen/MachineOperand.h"
 
 #define DEBUG_TYPE "si-fold-operands"
@@ -737,12 +738,26 @@ bool SIFoldOperandsImpl::updateOperand(FoldCandidate &Fold) const {
   // 16-bit SGPRs instead of 32-bit ones.
   if (Old.getSubReg() == AMDGPU::lo16 && TRI->isSGPRReg(*MRI, New->getReg()))
     Old.setSubReg(AMDGPU::NoSubRegister);
+  Register OldReg = Old.getReg();
   if (New->getReg().isPhysical()) {
     Old.substPhysReg(New->getReg(), *TRI);
   } else {
     Old.substVirtReg(New->getReg(), New->getSubReg(), *TRI);
     Old.setIsUndef(New->isUndef());
   }
+
+  // If MI is inside a BUNDLE, point the header's matching implicit use at
+  // NewReg too, so LiveVariables sees the fold.
+  if (MI->isBundledWithPred()) {
+    MachineInstr &Header = *getBundleStart(MI->getIterator());
+    for (MachineOperand &MO : Header.operands()) {
+      if (MO.isReg() && MO.isImplicit() && !MO.isDef() &&
+          MO.getReg() == OldReg) {
+        MO.setReg(New->getReg());
+        MO.setSubReg(New->getSubReg());
+      }
+    }
+  }
   return true;
 }
 
diff --git a/llvm/test/CodeGen/AMDGPU/ds_gws_align.ll b/llvm/test/CodeGen/AMDGPU/ds_gws_align.ll
index 53bca0c2f6c7e..d860602492f0c 100644
--- a/llvm/test/CodeGen/AMDGPU/ds_gws_align.ll
+++ b/llvm/test/CodeGen/AMDGPU/ds_gws_align.ll
@@ -53,6 +53,17 @@ bb:
   ret void
 }
 
+; GCN-LABEL: {{^}}gws_init_subreg_uses:
+; GCN-COUNT-2: ds_gws_init v{{[0-9]+}} gds
+define amdgpu_ps void @gws_init_subreg_uses(ptr addrspace(1) %p) {
+  %v = load <4 x i32>, ptr addrspace(1) %p
+  %a = extractelement <4 x i32> %v, i32 0
+  %b = extractelement <4 x i32> %v, i32 1
+  call void @llvm.amdgcn.ds.gws.init(i32 %a, i32 0)
+  call void @llvm.amdgcn.ds.gws.init(i32 %b, i32 0)
+  ret void
+}
+
 declare void @llvm.amdgcn.ds.gws.init(i32, i32)
 declare void @llvm.amdgcn.ds.gws.sema.br(i32, i32)
 declare void @llvm.amdgcn.ds.gws.barrier(i32, i32)
diff --git a/llvm/test/CodeGen/AMDGPU/si-fold-operands-bundle.mir b/llvm/test/CodeGen/AMDGPU/si-fold-operands-bundle.mir
new file mode 100644
index 0000000000000..33ba62290c9be
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/si-fold-operands-bundle.mir
@@ -0,0 +1,51 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=amdgcn -mcpu=gfx908 -run-pass=si-fold-operands -verify-machineinstrs -o - %s | FileCheck %s
+
+# When SIFoldOperands folds a COPY source into a use inside a BUNDLE, it
+# must also update the BUNDLE header's matching implicit operand. Otherwise
+# LiveVariables (which only inspects the bundle header) miscomputes kills
+# and a later MachineVerifier run reports "Using a killed virtual register".
+
+---
+name: fold_into_bundled_use
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $vgpr0, $vgpr1, $vgpr2, $vgpr3
+
+    ; CHECK-LABEL: name: fold_into_bundled_use
+    ; CHECK: liveins: $vgpr0, $vgpr1, $vgpr2, $vgpr3
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:vgpr_32 = COPY $vgpr1
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:vgpr_32 = COPY $vgpr2
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:vgpr_32 = COPY $vgpr3
+    ; CHECK-NEXT: [[REG_SEQUENCE:%[0-9]+]]:vreg_128 = REG_SEQUENCE [[COPY]], %subreg.sub0, [[COPY1]], %subreg.sub1, [[COPY2]], %subreg.sub2, [[COPY3]], %subreg.sub3
+    ; CHECK-NEXT: $m0 = S_MOV_B32 0
+    ; CHECK-NEXT: BUNDLE implicit [[REG_SEQUENCE]].sub0, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource") {
+    ; CHECK-NEXT:   DS_GWS_INIT [[REG_SEQUENCE]].sub0, 0, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource")
+    ; CHECK-NEXT:   S_WAITCNT .Vmcnt_0_Expcnt_0_Lgkmcnt_0
+    ; CHECK-NEXT: }
+    ; CHECK-NEXT: BUNDLE implicit [[REG_SEQUENCE]].sub1, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource") {
+    ; CHECK-NEXT:   DS_GWS_INIT [[REG_SEQUENCE]].sub1, 0, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource")
+    ; CHECK-NEXT:   S_WAITCNT .Vmcnt_0_Expcnt_0_Lgkmcnt_0
+    ; CHECK-NEXT: }
+    ; CHECK-NEXT: S_ENDPGM 0
+    %0:vgpr_32 = COPY $vgpr0
+    %1:vgpr_32 = COPY $vgpr1
+    %2:vgpr_32 = COPY $vgpr2
+    %3:vgpr_32 = COPY $vgpr3
+    %4:vreg_128 = REG_SEQUENCE %0:vgpr_32, %subreg.sub0, %1:vgpr_32, %subreg.sub1, %2:vgpr_32, %subreg.sub2, %3:vgpr_32, %subreg.sub3
+    %5:vgpr_32 = COPY %4.sub0:vreg_128
+    %6:vgpr_32 = COPY %4.sub1:vreg_128
+    $m0 = S_MOV_B32 0
+    BUNDLE implicit %5:vgpr_32, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource") {
+      DS_GWS_INIT %5:vgpr_32, 0, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource")
+      S_WAITCNT 0
+    }
+    BUNDLE implicit %6:vgpr_32, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource") {
+      DS_GWS_INIT %6:vgpr_32, 0, implicit $m0, implicit $exec :: (store (s32) into custom "GWSResource")
+      S_WAITCNT 0
+    }
+    S_ENDPGM 0
+...

``````````

</details>


https://github.com/llvm/llvm-project/pull/201872