[llvm] [CodeGen] Ensure clearRegisterKills clears inside bundles. (PR #149177)

[Ricardo Jesus via llvm-commits]

https://github.com/rj-jesus updated https://github.com/llvm/llvm-project/pull/149177

>From dfa67db235787d763db1afa13023928e759a76dd Mon Sep 17 00:00:00 2001
From: Ricardo Jesus <rjj at nvidia.com>
Date: Wed, 16 Jul 2025 10:20:09 -0700
Subject: [PATCH 1/2] [CodeGen] Ensure clearRegisterKills works inside bundles.

---
 llvm/lib/CodeGen/MachineInstr.cpp             |  8 ++++----
 .../test/CodeGen/AArch64/sve-vls-ldst-opt.mir | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineInstr.cpp b/llvm/lib/CodeGen/MachineInstr.cpp
index da3665b3b6a0b..a966add78a97e 100644
--- a/llvm/lib/CodeGen/MachineInstr.cpp
+++ b/llvm/lib/CodeGen/MachineInstr.cpp
@@ -2177,12 +2177,12 @@ void MachineInstr::clearRegisterKills(Register Reg,
                                       const TargetRegisterInfo *RegInfo) {
   if (!Reg.isPhysical())
     RegInfo = nullptr;
-  for (MachineOperand &MO : operands()) {
-    if (!MO.isReg() || !MO.isUse() || !MO.isKill())
+  for (MIBundleOperands O(*this); O.isValid(); ++O) {
+    if (!O->isReg() || !O->isUse() || !O->isKill())
       continue;
-    Register OpReg = MO.getReg();
+    Register OpReg = O->getReg();
     if ((RegInfo && RegInfo->regsOverlap(Reg, OpReg)) || Reg == OpReg)
-      MO.setIsKill(false);
+      O->setIsKill(false);
   }
 }
 
diff --git a/llvm/test/CodeGen/AArch64/sve-vls-ldst-opt.mir b/llvm/test/CodeGen/AArch64/sve-vls-ldst-opt.mir
index 49453bc178914..28c5220642ec5 100644
--- a/llvm/test/CodeGen/AArch64/sve-vls-ldst-opt.mir
+++ b/llvm/test/CodeGen/AArch64/sve-vls-ldst-opt.mir
@@ -72,3 +72,23 @@ body: |
 # CHECK: STURQi killed renamable $q1, renamable $x1, 16 :: (store (s128))
 # CHECK: STURQi killed renamable $q2, renamable $x1, 48 :: (store (s128))
 # CHECK: STR_ZXI killed renamable $z3, renamable $x1, 4 :: (store (<vscale x 1 x s128>))
+---
+name: clear-kill-in-bundle
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $x0, $z0
+
+    STR_ZXI $z0, $x0, 0 :: (store (<vscale x 1 x s128>))
+    BUNDLE implicit-def $z1, implicit-def $q1, implicit killed $z0 {
+      $z1 = ADD_ZZZ_D $z0, killed $z0
+    }
+    STR_ZXI renamable $z1, $x0, 1 :: (store (<vscale x 1 x s128>))
+
+    RET_ReallyLR
+...
+# CHECK-LABEL: name: clear-kill-in-bundle
+# CHECK: BUNDLE implicit-def $z1, implicit-def $q1, implicit $z0 {
+# CHECK:   $z1 = ADD_ZZZ_D $z0, $z0
+# CHECK: }
+# CHECK: STPQi $q0, $q1, $x0, 0 :: (store (<vscale x 1 x s128>))

>From a3632c764ae5213aa609e557e51d867d49054d61 Mon Sep 17 00:00:00 2001
From: Ricardo Jesus <rjj at nvidia.com>
Date: Fri, 18 Jul 2025 08:37:36 -0700
Subject: [PATCH 2/2] Make clearRegisterKills bundle aware.

---
 llvm/include/llvm/CodeGen/MachineInstr.h |  9 +++++----
 llvm/lib/CodeGen/MachineInstr.cpp        | 25 ++++++++++++++++--------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineInstr.h b/llvm/include/llvm/CodeGen/MachineInstr.h
index 94d04b82666be..68eba13fef454 100644
--- a/llvm/include/llvm/CodeGen/MachineInstr.h
+++ b/llvm/include/llvm/CodeGen/MachineInstr.h
@@ -880,9 +880,9 @@ class MachineInstr
   /// queries but they are bundle aware.
 
   enum QueryType {
-    IgnoreBundle,    // Ignore bundles
-    AnyInBundle,     // Return true if any instruction in bundle has property
-    AllInBundle      // Return true if all instructions in bundle have property
+    IgnoreBundle, // Ignore bundles
+    AnyInBundle,  // Check/update property for any instruction in bundle
+    AllInBundle   // Check/update property for all instructions in bundle
   };
 
   /// Return true if the instruction (or in the case of a bundle,
@@ -1700,7 +1700,8 @@ class MachineInstr
   /// Clear all kill flags affecting Reg.  If RegInfo is provided, this includes
   /// all aliasing registers.
   LLVM_ABI void clearRegisterKills(Register Reg,
-                                   const TargetRegisterInfo *RegInfo);
+                                   const TargetRegisterInfo *RegInfo,
+                                   QueryType Type = AllInBundle);
 
   /// We have determined MI defined a register without a use.
   /// Look for the operand that defines it and mark it as IsDead. If
diff --git a/llvm/lib/CodeGen/MachineInstr.cpp b/llvm/lib/CodeGen/MachineInstr.cpp
index a966add78a97e..1672e65a70111 100644
--- a/llvm/lib/CodeGen/MachineInstr.cpp
+++ b/llvm/lib/CodeGen/MachineInstr.cpp
@@ -2174,16 +2174,25 @@ bool MachineInstr::addRegisterKilled(Register IncomingReg,
 }
 
 void MachineInstr::clearRegisterKills(Register Reg,
-                                      const TargetRegisterInfo *RegInfo) {
+                                      const TargetRegisterInfo *RegInfo,
+                                      QueryType Type) {
   if (!Reg.isPhysical())
     RegInfo = nullptr;
-  for (MIBundleOperands O(*this); O.isValid(); ++O) {
-    if (!O->isReg() || !O->isUse() || !O->isKill())
-      continue;
-    Register OpReg = O->getReg();
-    if ((RegInfo && RegInfo->regsOverlap(Reg, OpReg)) || Reg == OpReg)
-      O->setIsKill(false);
-  }
+
+  auto clearKills = [&](auto Operands) {
+    for (MachineOperand &MO : Operands) {
+      if (!MO.isReg() || !MO.isUse() || !MO.isKill())
+        continue;
+      Register OpReg = MO.getReg();
+      if ((RegInfo && RegInfo->regsOverlap(Reg, OpReg)) || Reg == OpReg)
+        MO.setIsKill(false);
+    }
+  };
+
+  if (Type != AllInBundle || !isBundled() || isBundledWithPred())
+    clearKills(operands());
+  else
+    clearKills(mi_bundle_ops(*this));
 }
 
 bool MachineInstr::addRegisterDead(Register Reg,