RegAllocFast: Fix verifier errors after assigning to reserved registers (#128281)

Author: arsenm

llvm/lib/CodeGen/RegAllocFast.cpp

@@ -206,6 +206,7 @@ class RegAllocFastImpl {
     bool Error = false;              ///< Could not allocate.
 
     explicit LiveReg(Register VirtReg) : VirtReg(VirtReg) {}
+    explicit LiveReg() {}
 
     unsigned getSparseSetIndex() const { return VirtReg.virtRegIndex(); }
   };
@@ -216,7 +217,7 @@ class RegAllocFastImpl {
   LiveRegMap LiveVirtRegs;
 
   /// Stores assigned virtual registers present in the bundle MI.
-  DenseMap<Register, MCPhysReg> BundleVirtRegsMap;
+  DenseMap<Register, LiveReg> BundleVirtRegsMap;
 
   DenseMap<unsigned, SmallVector<MachineOperand *, 2>> LiveDbgValueMap;
   /// List of DBG_VALUE that we encountered without the vreg being assigned
@@ -374,7 +375,8 @@ class RegAllocFastImpl {
                             SmallSet<Register, 2> &PrologLiveIns) const;
 
   void reloadAtBegin(MachineBasicBlock &MBB);
-  bool setPhysReg(MachineInstr &MI, MachineOperand &MO, MCPhysReg PhysReg);
+  bool setPhysReg(MachineInstr &MI, MachineOperand &MO,
+                  const LiveReg &Assignment);
 
   Register traceCopies(Register VirtReg) const;
   Register traceCopyChain(Register Reg) const;
@@ -1005,7 +1007,7 @@ void RegAllocFastImpl::allocVirtRegUndef(MachineOperand &MO) {
     MO.setSubReg(0);
   }
   MO.setReg(PhysReg);
-  MO.setIsRenamable(true);
+  MO.setIsRenamable(!LRI->Error);
 }
 
 /// Variation of defineVirtReg() with special handling for livethrough regs
@@ -1109,10 +1111,10 @@ bool RegAllocFastImpl::defineVirtReg(MachineInstr &MI, unsigned OpNum,
     LRI->Reloaded = false;
   }
   if (MI.getOpcode() == TargetOpcode::BUNDLE) {
-    BundleVirtRegsMap[VirtReg] = PhysReg;
+    BundleVirtRegsMap[VirtReg] = *LRI;
   }
   markRegUsedInInstr(PhysReg);
-  return setPhysReg(MI, MO, PhysReg);
+  return setPhysReg(MI, MO, *LRI);
 }
 
 /// Allocates a register for a VirtReg use.
@@ -1158,10 +1160,10 @@ bool RegAllocFastImpl::useVirtReg(MachineInstr &MI, MachineOperand &MO,
   LRI->LastUse = &MI;
 
   if (MI.getOpcode() == TargetOpcode::BUNDLE) {
-    BundleVirtRegsMap[VirtReg] = LRI->PhysReg;
+    BundleVirtRegsMap[VirtReg] = *LRI;
   }
   markRegUsedInInstr(LRI->PhysReg);
-  return setPhysReg(MI, MO, LRI->PhysReg);
+  return setPhysReg(MI, MO, *LRI);
 }
 
 /// Query a physical register to use as a filler in contexts where the
@@ -1215,16 +1217,27 @@ MCPhysReg RegAllocFastImpl::getErrorAssignment(const LiveReg &LR,
 /// Changes operand OpNum in MI the refer the PhysReg, considering subregs.
 /// \return true if MI's MachineOperands were re-arranged/invalidated.
 bool RegAllocFastImpl::setPhysReg(MachineInstr &MI, MachineOperand &MO,
-                                  MCPhysReg PhysReg) {
+                                  const LiveReg &Assignment) {
+  MCPhysReg PhysReg = Assignment.PhysReg;
+  assert(PhysReg && "assignments should always be to a valid physreg");
+
+  if (LLVM_UNLIKELY(Assignment.Error)) {
+    // Make sure we don't set renamable in error scenarios, as we may have
+    // assigned to a reserved register.
+    if (MO.isUse())
+      MO.setIsUndef(true);
+  }
+
   if (!MO.getSubReg()) {
     MO.setReg(PhysReg);
-    MO.setIsRenamable(true);
+    MO.setIsRenamable(!Assignment.Error);
     return false;
   }
 
   // Handle subregister index.
-  MO.setReg(PhysReg ? TRI->getSubReg(PhysReg, MO.getSubReg()) : MCRegister());
-  MO.setIsRenamable(true);
+  MO.setReg(TRI->getSubReg(PhysReg, MO.getSubReg()));
+  MO.setIsRenamable(!Assignment.Error);
+
   // Note: We leave the subreg number around a little longer in case of defs.
   // This is so that the register freeing logic in allocateInstruction can still
   // recognize this as subregister defs. The code there will clear the number.
@@ -1706,7 +1719,7 @@ void RegAllocFastImpl::handleDebugValue(MachineInstr &MI) {
     if (LRI != LiveVirtRegs.end() && LRI->PhysReg) {
       // Update every use of Reg within MI.
       for (auto &RegMO : DbgOps)
-        setPhysReg(MI, *RegMO, LRI->PhysReg);
+        setPhysReg(MI, *RegMO, *LRI);
     } else {
       DanglingDbgValues[Reg].push_back(&MI);
     }
@@ -1729,8 +1742,7 @@ void RegAllocFastImpl::handleBundle(MachineInstr &MI) {
       if (!Reg.isVirtual() || !shouldAllocateRegister(Reg))
         continue;
 
-      DenseMap<Register, MCPhysReg>::iterator DI;
-      DI = BundleVirtRegsMap.find(Reg);
+      DenseMap<Register, LiveReg>::iterator DI = BundleVirtRegsMap.find(Reg);
       assert(DI != BundleVirtRegsMap.end() && "Unassigned virtual register");
 
       setPhysReg(MI, MO, DI->second);

llvm/test/CodeGen/AMDGPU/ran-out-of-registers-error-all-regs-reserved.ll

@@ -1,7 +1,6 @@
 ; RUN: not llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 -vgpr-regalloc=greedy -verify-machineinstrs -filetype=null %s 2>&1 | FileCheck -implicit-check-not=error %s
 ; RUN: not llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 -vgpr-regalloc=basic -verify-machineinstrs -filetype=null %s 2>&1 | FileCheck -implicit-check-not=error %s
-; TODO: Fix verifier error after fast
-; TODO: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 -vgpr-regalloc=fast -verify-machineinstrs -filetype=null %s 2>&1 | FileCheck -implicit-check-not=error %s
+; RUN: not llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx908 -vgpr-regalloc=fast -verify-machineinstrs -filetype=null %s 2>&1 | FileCheck -implicit-check-not=error %s
 
 declare <32 x i32> @llvm.amdgcn.mfma.i32.32x32x4i8(i32, i32, <32 x i32>, i32 immarg, i32 immarg, i32 immarg)