VirtRegMap: Support partially allocated virtual registers

[arsenm]

Summary

Don't assert if there are unassigned virtual registers.
Maintain LiveIntervals by removing the RegUnits for allocated
registers, since they should not longer be necessary.

One part I find somewhat questionable is the special
handling necessary for handleIdentityCopy. The LiveIntervals
for the relevant regunits needs to be removed.

[arsenm]

ping

[arsenm]

ping

[qcolombet]

The LiveIntervals for the relevant regunits needs to be removed.

Could you elaborate on why we need to remove these?

lib/CodeGen/VirtRegMap.cpp

624 ↗

(On Diff #176668)

Although that's unlikely to happen, we could imagine that another physical reg that share regunits with the one we are clearing is still relevant for subsequence allocation.

Put differently, why is this code valid?

In theory we may clear RegA's regunits, while they overlap with RegB's regunits and RegB's live interval needs to be valid.

[arsenm]

lib/CodeGen/VirtRegMap.cpp

624 ↗

(On Diff #176668)

For the real use case I have, that can't happen. If these aren't removed, there will be a huge amount of work to maintain these accurately on the off chance this happens. I think this should just be documented as a restriction for the register class filter functions. If one regunit is disabled, any register class that has a register overlapping that regunit also needs to be disabled.

[qcolombet]

lib/CodeGen/VirtRegMap.cpp

624 ↗

(On Diff #176668)

I was thinking along the same line.
Could you think of some assert that we could use to make sure this doesn't happen?

[arsenm]

Rebase

[arsenm]

55a29c6b71c9b80353ccb17c0dd15dde5c9940b3

lib/CodeGen/VirtRegMap.cpp

624 ↗

(On Diff #176668)

I'm not sure what the assert would look like, since you can't generically ask what the register class for a physical register is. In practice you would hit verifier errors after this point. If the filter was a list of register classes rather than a callback, I might be able to do something with that.

llvm/lib/CodeGen/VirtRegMap.cpp

 class VirtRegRewriter : public MachineFunctionPass {
   MachineFunction *MF;
   const TargetRegisterInfo *TRI;
   const TargetInstrInfo *TII;
   MachineRegisterInfo *MRI;
   SlotIndexes *Indexes;
   LiveIntervals *LIS;
   VirtRegMap *VRM;
+  DenseSet<Register> RewriteRegs;
   bool ClearVirtRegs;
 
   void rewrite();
   void addMBBLiveIns();
   bool readsUndefSubreg(const MachineOperand &MO) const;
-  void addLiveInsForSubRanges(const LiveInterval &LI, Register PhysReg) const;
-  void handleIdentityCopy(MachineInstr &MI) const;
+  void addLiveInsForSubRanges(const LiveInterval &LI, MCRegister PhysReg) const;
+  void handleIdentityCopy(MachineInstr &MI);
   void expandCopyBundle(MachineInstr &MI) const;
   bool subRegLiveThrough(const MachineInstr &MI, MCRegister SuperPhysReg) const;
 
 public:
   static char ID;
   VirtRegRewriter(bool ClearVirtRegs_ = true) :
     MachineFunctionPass(ID),
     ClearVirtRegs(ClearVirtRegs_) {}

 INITIALIZE_PASS_DEPENDENCY(LiveStacks)
 INITIALIZE_PASS_DEPENDENCY(VirtRegMap)
 INITIALIZE_PASS_END(VirtRegRewriter, "virtregrewriter",
                     "Virtual Register Rewriter", false, false)
 
 void VirtRegRewriter::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addRequired<LiveIntervals>();
+  AU.addPreserved<LiveIntervals>();
   AU.addRequired<SlotIndexes>();
   AU.addPreserved<SlotIndexes>();
   AU.addRequired<LiveDebugVariables>();
   AU.addRequired<LiveStacks>();
   AU.addPreserved<LiveStacks>();
   AU.addRequired<VirtRegMap>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
if (ClearVirtRegs) {
     VRM->clearAllVirt();
     MRI->clearVirtRegs();
   }
 
   return true;
 }
 
 void VirtRegRewriter::addLiveInsForSubRanges(const LiveInterval &LI,
-                                             Register PhysReg) const {
+                                             MCRegister PhysReg) const {
   assert(!LI.empty());
   assert(LI.hasSubRanges());
 
   using SubRangeIteratorPair =
       std::pair<const LiveInterval::SubRange *, LiveInterval::const_iterator>;
 
   SmallVector<SubRangeIteratorPair, 4> SubRanges;
   SlotIndex First;
for (unsigned Idx = 0, IdxE = MRI->getNumVirtRegs(); Idx != IdxE; ++Idx) {
     if (MRI->reg_nodbg_empty(VirtReg))
       continue;
     LiveInterval &LI = LIS->getInterval(VirtReg);
     if (LI.empty() || LIS->intervalIsInOneMBB(LI))
       continue;
     // This is a virtual register that is live across basic blocks. Its
     // assigned PhysReg must be marked as live-in to those blocks.
     Register PhysReg = VRM->getPhys(VirtReg);
-    assert(PhysReg != VirtRegMap::NO_PHYS_REG && "Unmapped virtual register.");
+    if (PhysReg == VirtRegMap::NO_PHYS_REG) {
+      // There may be no physical register assigned if only some register
+      // classes were already allocated.
+      assert(!ClearVirtRegs && "Unmapped virtual register");
+      continue;
+    }
 
     if (LI.hasSubRanges()) {
       addLiveInsForSubRanges(LI, PhysReg);
     } else {
       // Go over MBB begin positions and see if we have segments covering them.
       // The following works because segments and the MBBIndex list are both
       // sorted by slot indexes.
       SlotIndexes::MBBIndexIterator I = Indexes->MBBIndexBegin();
bool VirtRegRewriter::readsUndefSubreg(const MachineOperand &MO) const {
   // See if any of the relevant subregister liveranges is defined at this point.
   for (const LiveInterval::SubRange &SR : LI.subranges()) {
     if ((SR.LaneMask & UseMask).any() && SR.liveAt(BaseIndex))
       return false;
   }
   return true;
 }
 
-void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) const {
+void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) {
   if (!MI.isIdentityCopy())
     return;
   LLVM_DEBUG(dbgs() << "Identity copy: " << MI);
   ++NumIdCopies;
 
+  Register DstReg = MI.getOperand(0).getReg();
+
+  // We may have deferred allocation of the virtual register, and the rewrite
+  // regs code doesn't handle the liveness update.
+  if (DstReg.isVirtual())
+    return;
+
+  RewriteRegs.insert(DstReg);
+
   // Copies like:
   //    %r0 = COPY undef %r0
   //    %al = COPY %al, implicit-def %eax
   // give us additional liveness information: The target (super-)register
   // must not be valid before this point. Replace the COPY with a KILL
   // instruction to maintain this information.
   if (MI.getOperand(1).isUndef() || MI.getNumOperands() > 2) {
     MI.setDesc(TII->get(TargetOpcode::KILL));
for (MachineBasicBlock::instr_iterator
         // Make sure MRI knows about registers clobbered by regmasks.
         if (MO.isRegMask())
           MRI->addPhysRegsUsedFromRegMask(MO.getRegMask());
 
         if (!MO.isReg() || !MO.getReg().isVirtual())
           continue;
         Register VirtReg = MO.getReg();
         MCRegister PhysReg = VRM->getPhys(VirtReg);
-        assert(PhysReg != VirtRegMap::NO_PHYS_REG &&
-               "Instruction uses unmapped VirtReg");
+        if (PhysReg == VirtRegMap::NO_PHYS_REG)
+          continue;
+
+        assert(Register(PhysReg).isPhysical());
+
+        RewriteRegs.insert(PhysReg);
         assert(!MRI->isReserved(PhysReg) && "Reserved register assignment");
 
         // Preserve semantics of sub-register operands.
         unsigned SubReg = MO.getSubReg();
         if (SubReg != 0) {
           if (NoSubRegLiveness || !MRI->shouldTrackSubRegLiveness(VirtReg)) {
             // A virtual register kill refers to the whole register, so we may
             // have to add implicit killed operands for the super-register.  A
for (MachineBasicBlock::instr_iterator
       LLVM_DEBUG(dbgs() << "> " << *MI);
 
       expandCopyBundle(*MI);
 
       // We can remove identity copies right now.
       handleIdentityCopy(*MI);
     }
   }
+
+  if (LIS) {
+    // Don't bother maintaining accurate LiveIntervals for registers which were
+    // already allocated.
+    for (Register PhysReg : RewriteRegs) {
+      for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid();
+           ++Units) {
+        LIS->removeRegUnit(*Units);
+      }
+    }
+  }
+
+  RewriteRegs.clear();
 }
 
 FunctionPass *llvm::createVirtRegRewriter(bool ClearVirtRegs) {
   return new VirtRegRewriter(ClearVirtRegs);
 }