[TEMP] AMDGPU: Fix LiveRangeEdit+LiveRegMatrix use-after-delete of LiveInterval

Author: macurtis-amd

llvm/include/llvm/CodeGen/LiveRegMatrix.h

@@ -135,6 +135,9 @@ class LiveRegMatrix {
   /// the assignment and updates VirtRegMap accordingly.
   void unassign(const LiveInterval &VirtReg);
 
+  /// Remove any assignments for \p VirtReg.
+  void purge(const LiveInterval &VirtReg);
+
   /// Returns true if the given \p PhysReg has any live intervals assigned.
   bool isPhysRegUsed(MCRegister PhysReg) const;
 

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();
+  bool LRE_CanEraseVirtReg(Register) 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,15 @@ void HoistSpillHelper::hoistAllSpills() {
   }
 }
 
+bool HoistSpillHelper::LRE_CanEraseVirtReg(Register VirtReg) {
+  if (!Matrix || !LIS.hasInterval(VirtReg) || !VRM.hasPhys(VirtReg))
+    return true;
+
+  LiveInterval &LI = LIS.getInterval(VirtReg);
+  Matrix->purge(LI);
+  return true;
+}
+
 /// 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/LiveRegMatrix.cpp

@@ -142,6 +142,20 @@ void LiveRegMatrix::unassign(const LiveInterval &VirtReg) {
   LLVM_DEBUG(dbgs() << '\n');
 }
 
+void LiveRegMatrix::purge(const LiveInterval &VirtReg) {
+  VRM->clearVirt(VirtReg.reg());
+
+  for (unsigned i = 0; i < Matrix.size(); ++i) {
+    LiveIntervalUnion &LIU = Matrix[static_cast<MCRegUnit>(i)];
+    for (LiveIntervalUnion::SegmentIter SI = LIU.begin(); SI.valid();) {
+      if (SI.value() == &VirtReg) {
+        SI.erase();
+      } else
+        ++SI;
+    }
+  }
+}
+
 bool LiveRegMatrix::isPhysRegUsed(MCRegister PhysReg) const {
   for (MCRegUnit Unit : TRI->regunits(PhysReg)) {
     if (!Matrix[Unit].empty())