[CodeGen][RegAllocFast] Map PhysReg to its current VirtReg

[cdevadas]

Summary

It is not possible to have the one-to-one mapping for a
virtual register and the physical register allocated for
it during fast regalloc. Due to the in-place allocation,
a virtual register, while having multiple live ranges,
might get allocated to different physical registers. This
makes it impossible to map the physical register back to
its current live range inside the target hooks, like the
spill interface.

This patch introduces the physical register to the current
virtual register mapping whenever a physical register is
actually allocated.

[cdevadas]

This should replace the patch D134951 which used the delegate method for a specific purpose.
This is a more generic approach for targets to help identify the exact live range (virtual register) a physical register was currently assigned for.
While the allocation is in progress, this would serve the purpose of mapping a physical register to its virtual register.

[cdevadas]

PhysToCurrentVirtReg should also be updated while reloading a physical register from the spill location. It happens when a RegMask (for instance, a function call) clobbers a register that is used in an instruction later in the BB. Also, at a BB entry, the LiveIn registers are reloaded. The reload point begins a new live range and it should be properly reflected in the mapping.

[cdevadas]

See the Usecase of PhysToCurrentVirtReg in D124196 (look for getPhysToCurrentVirtReg)

[arsenm]

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

781

Why pollute MRI with this, this is state local to RegAllocFast

[cdevadas]

The idea of mapping a physical register into the current virtual register in RegAllocFast would be likely a risky choice with a global variable. The same was true for the initial patch I posted D134951 which does the same thing specifically around the spill and reload functions via. the delegate.
The effort was to identify the physical register allocated for some special machine operands (whole-wave vector registers for AMDGPU), especially inside storeRegToStackSlot & loadRegFromStackSlot spill interfaces called from fast regalloc that directly spills the physRegs.
The better approach would be by passing an additional parameter (the virtual register) to the spiller interface. The interface is common to Greedy, PEI, and to some other places that have no use for this additional parameter. Still, I think it is the right way to handle the problem.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

781

It should be accessible for targets and hence I thought MRI would be a feasible place.

[cdevadas]

Posted D138656 to pass the virtual register as an additional argument to the spiller interface. This review can be abandoned once D138656 gets approved.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

private:
   /// vector is the frozen set of reserved registers when register allocation
   /// started.
   BitVector ReservedRegs;
 
   using VRegToTypeMap = IndexedMap<LLT, VirtReg2IndexFunctor>;
   /// Map generic virtual registers to their low-level type.
   VRegToTypeMap VRegToType;
 
+  /// Map a physical register into the virtual register it currently represents.
+  /// Useful for some target-hooks during the RegAllocFast to map a physical
+  /// register back to its virtual register. For instance, AMDGPU needs this
+  /// mapping inside the target spill hooks to identify the spills for certain
+  /// machine operands. This field is relevant only during RegAllocFast and will
+  /// be cleared elsewhere.
+  Register PhysToCurrentVirtReg;
+
   /// Keep track of the physical registers that are live in to the function.
   /// Live in values are typically arguments in registers.  LiveIn values are
   /// allowed to have virtual registers associated with them, stored in the
   /// second element.
   std::vector<std::pair<MCRegister, Register>> LiveIns;
 
 public:
   explicit MachineRegisterInfo(MachineFunction *MF);
public:
 
   /// Set the low-level type of \p VReg to \p Ty.
   void setType(Register VReg, LLT Ty);
 
   /// Create and return a new generic virtual register with low-level
   /// type \p Ty.
   Register createGenericVirtualRegister(LLT Ty, StringRef Name = "");
 
+  void setPhysToCurrentVirtReg(Register Reg) { PhysToCurrentVirtReg = Reg; }
+  Register getPhysToCurrentVirtReg() const { return PhysToCurrentVirtReg; }
+
   /// Remove all types associated to virtual registers (after instruction
   /// selection and constraining of all generic virtual registers).
   void clearVirtRegTypes();
 
   /// Creates a new virtual register that has no register class, register bank
   /// or size assigned yet. This is only allowed to be used
   /// temporarily while constructing machine instructions. Most operations are
   /// undefined on an incomplete register until one of setRegClass(),

llvm/lib/CodeGen/RegAllocFast.cpp

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->setPhysToCurrentVirtReg(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
void RegAllocFast::assignVirtToPhysReg(MachineInstr &AtMI, LiveReg &LR,
                                        MCPhysReg PhysReg) {
   Register VirtReg = LR.VirtReg;
   LLVM_DEBUG(dbgs() << "Assigning " << printReg(VirtReg, TRI) << " to "
                     << printReg(PhysReg, TRI) << '\n');
   assert(LR.PhysReg == 0 && "Already assigned a physreg");
   assert(PhysReg != 0 && "Trying to assign no register");
   LR.PhysReg = PhysReg;
   setPhysRegState(PhysReg, VirtReg);
+  MRI->setPhysToCurrentVirtReg(VirtReg);
 
   assignDanglingDebugValues(AtMI, VirtReg, PhysReg);
 }
 
 static bool isCoalescable(const MachineInstr &MI) {
   return MI.isFullCopy();
 }
 
for (MachineBasicBlock &MBB : MF)
     allocateBasicBlock(MBB);
 
   if (ClearVirtRegs) {
     // All machine operands and other references to virtual registers have been
     // replaced. Remove the virtual registers.
     MRI->clearVirtRegs();
   }
 
+  // Clear PhysToCurrentVirtReg.
+  MRI->setPhysToCurrentVirtReg(Register());
+
   StackSlotForVirtReg.clear();
   LiveDbgValueMap.clear();
   return true;
 }
 
 FunctionPass *llvm::createFastRegisterAllocator() {
   return new RegAllocFast();
 }
 
 FunctionPass *llvm::createFastRegisterAllocator(RegClassFilterFunc Ftor,
                                                 bool ClearVirtRegs) {
   return new RegAllocFast(Ftor, ClearVirtRegs);
 }