[TEMP] Fix LiveRegMatrix dangling pointers in postOptimization

Author: michaelselehov

llvm/include/llvm/CodeGen/LiveRangeEdit.h

@@ -20,6 +20,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
@@ -58,6 +59,11 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
     /// Called before shrinking the live range of a virtual register.
     virtual void LRE_WillShrinkVirtReg(Register) {}
 
+    /// Called when a virtual register's LiveInterval is about to become empty.
+    /// This happens when removeVRegDefAt removes the last definition.
+    /// Implementations should unassign from LiveRegMatrix before the interval is cleared.
+    virtual void LRE_WillClearVirtReg(Register, LiveInterval &) {}
+
     /// Called after cloning a virtual register.
     /// This is used for new registers representing connected components of Old.
     virtual void LRE_DidCloneVirtReg(Register New, Register Old) {}
@@ -75,6 +81,11 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
   /// FirstNew - Index of the first register added to NewRegs.
   const unsigned FirstNew;
 
+  /// Track which virtual registers are new (created during this edit).
+  /// Used to avoid calling Matrix->unassign on registers that were never
+  /// added to LiveRegMatrix.
+  SmallSet<Register, 8> NewVirtRegs;
+
   /// DeadRemats - The saved instructions which have already been dead after
   /// rematerialization but not deleted yet -- to be done in postOptimization.
   SmallPtrSet<MachineInstr *, 32> *DeadRemats;
@@ -142,6 +153,12 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
   bool empty() const { return size() == 0; }
   Register get(unsigned idx) const { return NewRegs[idx + FirstNew]; }
 
+  /// Check if a virtual register was created during this edit.
+  /// This is used to identify registers that were never added to LiveRegMatrix.
+  bool isNewVirtualRegister(Register VReg) const {
+    return NewVirtRegs.contains(VReg);
+  }
+
   /// pop_back - It allows LiveRangeEdit users to drop new registers.
   /// The context is when an original def instruction of a register is
   /// dead after rematerialization, we still want to keep it for following

llvm/lib/CodeGen/InlineSpiller.cpp

@@ -86,6 +86,7 @@ class HoistSpillHelper : private LiveRangeEdit::Delegate {
   const TargetInstrInfo &TII;
   const TargetRegisterInfo &TRI;
   const MachineBlockFrequencyInfo &MBFI;
+  LiveRegMatrix *Matrix;
 
   InsertPointAnalysis IPA;
 
@@ -129,16 +130,18 @@ class HoistSpillHelper : private LiveRangeEdit::Delegate {
 
 public:
   HoistSpillHelper(const Spiller::RequiredAnalyses &Analyses,
-                   MachineFunction &mf, VirtRegMap &vrm)
+                   MachineFunction &mf, VirtRegMap &vrm,
+                   LiveRegMatrix *matrix = nullptr)
       : MF(mf), LIS(Analyses.LIS), LSS(Analyses.LSS), MDT(Analyses.MDT),
         VRM(vrm), MRI(mf.getRegInfo()), TII(*mf.getSubtarget().getInstrInfo()),
         TRI(*mf.getSubtarget().getRegisterInfo()), MBFI(Analyses.MBFI),
-        IPA(LIS, mf.getNumBlockIDs()) {}
+        Matrix(matrix), IPA(LIS, mf.getNumBlockIDs()) {}
 
   void addToMergeableSpills(MachineInstr &Spill, int StackSlot,
                             Register Original);
   bool rmFromMergeableSpills(MachineInstr &Spill, int StackSlot);
   void hoistAllSpills();
+  void LRE_WillClearVirtReg(Register, LiveInterval &) override;
   void LRE_DidCloneVirtReg(Register, Register) override;
 };
 
@@ -191,7 +194,7 @@ class InlineSpiller : public Spiller {
       : MF(MF), LIS(Analyses.LIS), LSS(Analyses.LSS), VRM(VRM),
         MRI(MF.getRegInfo()), TII(*MF.getSubtarget().getInstrInfo()),
         TRI(*MF.getSubtarget().getRegisterInfo()), Matrix(Matrix),
-        HSpiller(Analyses, MF, VRM), VRAI(VRAI) {}
+        HSpiller(Analyses, MF, VRM, Matrix), VRAI(VRAI) {}
 
   void spill(LiveRangeEdit &, AllocationOrder *Order = nullptr) override;
   ArrayRef<Register> getSpilledRegs() override { return RegsToSpill; }
@@ -1750,6 +1753,17 @@ void HoistSpillHelper::hoistAllSpills() {
   }
 }
 
+/// Called when a LiveInterval is about to be cleared by removeVRegDefAt.
+/// Unassign from LiveRegMatrix to prevent dangling pointers (fixes LLVM bug #48911).
+void HoistSpillHelper::LRE_WillClearVirtReg(Register VirtReg,
+                                             LiveInterval &LI) {
+  // If this virtual register is assigned to a physical register, unassign it
+  // from LiveRegMatrix before the interval is cleared. Otherwise, LiveIntervalUnion
+  // will contain dangling pointers.
+  if (Matrix && VRM.hasPhys(VirtReg))
+    Matrix->unassign(LI);
+}
+
 /// For VirtReg clone, the \p New register should have the same physreg or
 /// stackslot as the \p old register.
 void HoistSpillHelper::LRE_DidCloneVirtReg(Register New, Register Old) {

llvm/lib/CodeGen/LiveRangeEdit.cpp

@@ -268,6 +268,11 @@ void LiveRangeEdit::eliminateDeadDef(MachineInstr *MI, ToShrinkSet &ToShrink) {
     if (MO.isDef()) {
       if (TheDelegate && LI.getVNInfoAt(Idx) != nullptr)
         TheDelegate->LRE_WillShrinkVirtReg(LI.reg());
+      // Notify delegate BEFORE clearing if this will make the interval empty.
+      // This allows Matrix->unassign to work with non-empty LiveRange.
+      // Skip new registers (clones) - they were never added to LiveRegMatrix.
+      if (TheDelegate && LI.size() == 1 && !isNewVirtualRegister(Reg))
+        TheDelegate->LRE_WillClearVirtReg(Reg, LI);
       LIS.removeVRegDefAt(LI, Idx);
       if (LI.empty())
         RegsToErase.push_back(Reg);
@@ -398,6 +403,7 @@ LiveRangeEdit::MRI_NoteNewVirtualRegister(Register VReg) {
     VRM->grow();
 
   NewRegs.push_back(VReg);
+  NewVirtRegs.insert(VReg);
 }
 
 void LiveRangeEdit::calculateRegClassAndHint(MachineFunction &MF,