[llvm] [LLVM][CodeGen] Fix register lane liveness tracking in RegisterPressure (PR #88892)

[Krzysztof Parzyszek via llvm-commits]

https://github.com/kparzysz updated https://github.com/llvm/llvm-project/pull/88892

>From c7205234d8b96a1db05d5f8e75d49c0f21bf1716 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 12 Apr 2024 12:47:17 -0500
Subject: [PATCH 1/2] [LLVM][CodeGen] Fix register lane liveness tracking in
 RegisterPressure

Re-enable an old assertion in `decreaseSetPressure`.
---
 llvm/lib/CodeGen/RegisterPressure.cpp | 39 ++++++++++++++-------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/CodeGen/RegisterPressure.cpp b/llvm/lib/CodeGen/RegisterPressure.cpp
index f86aa3a167202f..1b89cd57b9f2c3 100644
--- a/llvm/lib/CodeGen/RegisterPressure.cpp
+++ b/llvm/lib/CodeGen/RegisterPressure.cpp
@@ -64,7 +64,7 @@ static void increaseSetPressure(std::vector<unsigned> &CurrSetPressure,
 static void decreaseSetPressure(std::vector<unsigned> &CurrSetPressure,
                                 const MachineRegisterInfo &MRI, Register Reg,
                                 LaneBitmask PrevMask, LaneBitmask NewMask) {
-  //assert((NewMask & !PrevMask) == 0 && "Must not add bits");
+  assert((NewMask & ~PrevMask).none() && "Must not add bits");
   if (NewMask.any() || PrevMask.none())
     return;
 
@@ -617,17 +617,11 @@ void RegisterOperands::adjustLaneLiveness(const LiveIntervals &LIS,
       ++I;
     }
   }
-  for (auto *I = Uses.begin(); I != Uses.end();) {
-    LaneBitmask LiveBefore = getLiveLanesAt(LIS, MRI, true, I->RegUnit,
-                                            Pos.getBaseIndex());
-    LaneBitmask LaneMask = I->LaneMask & LiveBefore;
-    if (LaneMask.none()) {
-      I = Uses.erase(I);
-    } else {
-      I->LaneMask = LaneMask;
-      ++I;
-    }
-  }
+
+  // For uses just copy the copy the information from LIS.
+  for (auto &[RegUnit, LaneMask] : Uses)
+    LaneMask = getLiveLanesAt(LIS, MRI, true, RegUnit, Pos.getBaseIndex());
+
   if (AddFlagsMI != nullptr) {
     for (const RegisterMaskPair &P : DeadDefs) {
       Register RegUnit = P.RegUnit;
@@ -1060,18 +1054,27 @@ void RegPressureTracker::bumpUpwardPressure(const MachineInstr *MI) {
   // Kill liveness at live defs.
   for (const RegisterMaskPair &P : RegOpers.Defs) {
     Register Reg = P.RegUnit;
-    LaneBitmask LiveLanes = LiveRegs.contains(Reg);
+    LaneBitmask LiveAfter = LiveRegs.contains(Reg);
     LaneBitmask UseLanes = getRegLanes(RegOpers.Uses, Reg);
     LaneBitmask DefLanes = P.LaneMask;
-    LaneBitmask LiveAfter = (LiveLanes & ~DefLanes) | UseLanes;
-    decreaseRegPressure(Reg, LiveLanes, LiveAfter);
+    LaneBitmask LiveBefore = (LiveAfter & ~DefLanes) | UseLanes;
+
+    // There may be parts of the register that were dead before the
+    // instruction, but became live afterwards. Similarly, some parts
+    // may have been killed in this instruction.
+    decreaseRegPressure(Reg, LiveAfter, LiveAfter & LiveBefore);
+    increaseRegPressure(Reg, LiveAfter, ~LiveAfter & LiveBefore);
   }
   // Generate liveness for uses.
   for (const RegisterMaskPair &P : RegOpers.Uses) {
     Register Reg = P.RegUnit;
-    LaneBitmask LiveLanes = LiveRegs.contains(Reg);
-    LaneBitmask LiveAfter = LiveLanes | P.LaneMask;
-    increaseRegPressure(Reg, LiveLanes, LiveAfter);
+    // If this register was also in a def operand, we've handled it
+    // with defs.
+    if (getRegLanes(RegOpers.Defs, Reg).any())
+      continue;
+    LaneBitmask LiveAfter = LiveRegs.contains(Reg);
+    LaneBitmask LiveBefore = LiveAfter | P.LaneMask;
+    increaseRegPressure(Reg, LiveAfter, LiveBefore);
   }
 }
 

>From 14c05f6577aa10404513ed9bb37524d22dd537ce Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Tue, 16 Apr 2024 09:08:29 -0500
Subject: [PATCH 2/2] Fix comment

---
 llvm/lib/CodeGen/RegisterPressure.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/RegisterPressure.cpp b/llvm/lib/CodeGen/RegisterPressure.cpp
index 1b89cd57b9f2c3..3fa22447f416b7 100644
--- a/llvm/lib/CodeGen/RegisterPressure.cpp
+++ b/llvm/lib/CodeGen/RegisterPressure.cpp
@@ -618,7 +618,7 @@ void RegisterOperands::adjustLaneLiveness(const LiveIntervals &LIS,
     }
   }
 
-  // For uses just copy the copy the information from LIS.
+  // For uses just copy the information from LIS.
   for (auto &[RegUnit, LaneMask] : Uses)
     LaneMask = getLiveLanesAt(LIS, MRI, true, RegUnit, Pos.getBaseIndex());