[CodeGen][RegAllocFast] Add MRI delegate callback to notify VReg spill

[cdevadas]

Summary

Add a callback to notify targets whenever a virtual register is
spilled/restored during fast regalloc. The physical register is
directly spilled in case of RegAllocFast and targets sometimes
need to track the spills for certain virtual registers involved
in special operations.

Also, used cloneVirtualRegister in LiveIntervals while cloning
a virtual register.

[cdevadas]

Testcase added with D134951
llvm/test/CodeGen/AMDGPU/wwm-register-spill-during-regalloc.ll

[cdevadas]

Ping.

[qcolombet]

Hi,

The LiveIntervals and LiveRangeEdit changes make sense to me (though in practice they are NFC.)

However, I'm not a fan of the MachineRegisterInfo change. Since it sounds like you need that field only for fast regalloc, could you do something specifically in that pass/with the related target hooks in that pass?

Cheers,
-Quentin

[cdevadas]

However, I'm not a fan of the MachineRegisterInfo change. Since it sounds like you need that field only for fast regalloc, could you do something specifically in that pass/with the related target hooks in that pass?

The alternative would be to pass the flags as an argument to the target hooks storeRegToStackSlot & loadRegFromStackSlot which are common to both fast regalloc and greedy.

[cdevadas]

Added a new delegate callback to notify targets whenever a virtual register is spilled/restored during RegAllocFast.

[cdevadas]

Removed the enum class DelegateReceiver and changed the code to notify all delegates when an event occurs.

[cdevadas]

Rebase after recent changes in D134950.

[qcolombet]

LGTM.

llvm/lib/CodeGen/RegAllocFast.cpp

494

Nit: we may want to have two different events for spill and reload. But I'm fine keeping this as is for now if that is enough for your needs.

[arsenm]

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

189

I assume TheDelegates are pointers, so const X *TheDelegate?

llvm/lib/CodeGen/LiveRangeEdit.cpp

36

Is the clone still strictly necessary for this patch?

[cdevadas]

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

189

Yes, they are pointers, will use *

llvm/lib/CodeGen/LiveRangeEdit.cpp

36

Not necessarily in the same patch. They are basically needed to notify the targets when the RegAllocFast clone registers (for AMDGPU they are needed to propagate the vreg flags). They are included in the earlier patch where the VRegFlags are added directly in the MachineRegisterInfo file.

[cdevadas]

Used auto *Delegate in the range-based loop.

[cdevadas]

D138656 handled it better and hence abanding it.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

public:
   public:
     virtual ~Delegate() = default;
 
     virtual void MRI_NoteNewVirtualRegister(Register Reg) = 0;
     virtual void MRI_NotecloneVirtualRegister(Register NewReg,
                                               Register SrcReg) {
       MRI_NoteNewVirtualRegister(NewReg);
     }
+    virtual void MRI_NoteVirtualRegisterSpill(Register Reg) {}
   };
 
 private:
   MachineFunction *MF;
   SmallPtrSet<Delegate *, 1> TheDelegates;
 
   /// True if subregister liveness is tracked.
   const bool TracksSubRegLiveness;
for (auto *TheDelegate : TheDelegates)
       TheDelegate->MRI_NoteNewVirtualRegister(Reg);
   }
 
   void noteCloneVirtualRegister(Register NewReg, Register SrcReg) {
     for (auto *TheDelegate : TheDelegates)
       TheDelegate->MRI_NotecloneVirtualRegister(NewReg, SrcReg);
   }
 
+  /// Callback to notify a virtual register spill/reload.
+  void noteVirtualRegisterSpill(Register Reg) {
+    for (auto *TheDelegate : TheDelegates)
+      TheDelegate->MRI_NoteVirtualRegisterSpill(Reg);
+  }
+
   //===--------------------------------------------------------------------===//
   // Function State
   //===--------------------------------------------------------------------===//
 
   // isSSA - Returns true when the machine function is in SSA form. Early
   // passes require the machine function to be in SSA form where every virtual
   // register has a single defining instruction.
   //

llvm/lib/CodeGen/LiveIntervals.cpp

 void LiveIntervals::splitSeparateComponents(LiveInterval &LI,
     SmallVectorImpl<LiveInterval*> &SplitLIs) {
   ConnectedVNInfoEqClasses ConEQ(*this);
   unsigned NumComp = ConEQ.Classify(LI);
   if (NumComp <= 1)
     return;
   LLVM_DEBUG(dbgs() << "  Split " << NumComp << " components: " << LI << '\n');
   Register Reg = LI.reg();
-  const TargetRegisterClass *RegClass = MRI->getRegClass(Reg);
   for (unsigned I = 1; I < NumComp; ++I) {
-    Register NewVReg = MRI->createVirtualRegister(RegClass);
+    Register NewVReg = MRI->cloneVirtualRegister(Reg);
     LiveInterval &NewLI = createEmptyInterval(NewVReg);
     SplitLIs.push_back(&NewLI);
   }
   ConEQ.Distribute(LI, SplitLIs.data(), *MRI);
 }
 
 void LiveIntervals::constructMainRangeFromSubranges(LiveInterval &LI) {
   assert(LICalc && "LICalc not initialized.");
   LICalc->reset(MF, getSlotIndexes(), DomTree, &getVNInfoAllocator());
   LICalc->constructMainRangeFromSubranges(LI);
 }

llvm/lib/CodeGen/LiveRangeEdit.cpp

 STATISTIC(NumDCEFoldedLoads,    "Number of single use loads folded after DCE");
 STATISTIC(NumFracRanges,        "Number of live ranges fractured by DCE");
 STATISTIC(NumReMaterialization, "Number of instructions rematerialized");
 
 void LiveRangeEdit::Delegate::anchor() { }
 
 LiveInterval &LiveRangeEdit::createEmptyIntervalFrom(Register OldReg,
                                                      bool createSubRanges) {
-  Register VReg = MRI.createVirtualRegister(MRI.getRegClass(OldReg));
+  Register VReg = MRI.cloneVirtualRegister(OldReg);
   if (VRM)
     VRM->setIsSplitFromReg(VReg, VRM->getOriginal(OldReg));
 
   LiveInterval &LI = LIS.createEmptyInterval(VReg);
   if (Parent && !Parent->isSpillable())
     LI.markNotSpillable();
   if (createSubRanges) {
     // Create empty subranges if the OldReg's interval has them. Do not create
     // the main range here---it will be constructed later after the subranges
     // have been finalized.
     LiveInterval &OldLI = LIS.getInterval(OldReg);
     VNInfo::Allocator &Alloc = LIS.getVNInfoAllocator();
     for (LiveInterval::SubRange &S : OldLI.subranges())
       LI.createSubRange(Alloc, S.LaneMask);
   }
   return LI;
 }
 
 Register LiveRangeEdit::createFrom(Register OldReg) {
-  Register VReg = MRI.createVirtualRegister(MRI.getRegClass(OldReg));
+  Register VReg = MRI.cloneVirtualRegister(OldReg);
   if (VRM) {
     VRM->setIsSplitFromReg(VReg, VRM->getOriginal(OldReg));
   }
   // FIXME: Getting the interval here actually computes it.
   // In theory, this may not be what we want, but in practice
   // the createEmptyIntervalFrom API is used when this is not
   // the case. Generally speaking we just want to annotate the
   // LiveInterval when it gets created but we cannot do that at

llvm/lib/CodeGen/RegAllocFast.cpp

 /// Insert spill instruction for \p AssignedReg before \p Before. Update
 /// DBG_VALUEs with \p VirtReg operands with the stack slot.
 void RegAllocFast::spill(MachineBasicBlock::iterator Before, Register VirtReg,
                          MCPhysReg AssignedReg, bool Kill, bool LiveOut) {
   LLVM_DEBUG(dbgs() << "Spilling " << printReg(VirtReg, TRI)
                     << " in " << printReg(AssignedReg, TRI));
   int FI = getStackSpaceFor(VirtReg);
   LLVM_DEBUG(dbgs() << " to stack slot #" << FI << '\n');
+  MRI->noteVirtualRegisterSpill(VirtReg);
 
   const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
   TII->storeRegToStackSlot(*MBB, Before, AssignedReg, Kill, FI, &RC, TRI);
   ++NumStores;
 
   MachineBasicBlock::iterator FirstTerm = MBB->getFirstTerminator();
 
   // When we spill a virtual register, we will have spill instructions behind
void RegAllocFast::spill(MachineBasicBlock::iterator Before, Register VirtReg,
   LRIDbgOperands.clear();
 }
 
 /// Insert reload instruction for \p PhysReg before \p Before.
 void RegAllocFast::reload(MachineBasicBlock::iterator Before, Register VirtReg,
                           MCPhysReg PhysReg) {
   LLVM_DEBUG(dbgs() << "Reloading " << printReg(VirtReg, TRI) << " into "
                     << printReg(PhysReg, TRI) << '\n');
+  MRI->noteVirtualRegisterSpill(VirtReg);
+
   int FI = getStackSpaceFor(VirtReg);
   const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
   TII->loadRegFromStackSlot(*MBB, Before, PhysReg, FI, &RC, TRI);
   ++NumLoads;
 }
 
 /// Get basic block begin insertion point.
 /// This is not just MBB.begin() because surprisingly we have EH_LABEL
 /// instructions marking the begin of a basic block. This means we must insert
 /// new instructions after such labels...