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

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

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

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

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.

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<unsigned> RewriteRegs;
   bool ClearVirtRegs;
 
   void rewrite();
   void addMBBLiveIns();
   bool readsUndefSubreg(const MachineOperand &MO) const;
   void addLiveInsForSubRanges(const LiveInterval &LI, unsigned PhysReg) const;
-  void handleIdentityCopy(MachineInstr &MI) const;
+  void handleIdentityCopy(MachineInstr &MI);
   void expandCopyBundle(MachineInstr &MI) const;
   bool subRegLiveThrough(const MachineInstr &MI, unsigned 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);
 }
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.
     unsigned 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;
 
+  RewriteRegs.insert(MI.getOperand(0).getReg());
+
   // 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(0).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() || !TargetRegisterInfo::isVirtualRegister(MO.getReg()))
           continue;
         unsigned VirtReg = MO.getReg();
         unsigned PhysReg = VRM->getPhys(VirtReg);
-        assert(PhysReg != VirtRegMap::NO_PHYS_REG &&
-               "Instruction uses unmapped VirtReg");
+
+        if (PhysReg == VirtRegMap::NO_PHYS_REG)
+          continue;
+
+        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 (unsigned PhysReg : RewriteRegs) {
+      for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid();
+           ++Units) {
+        LIS->removeRegUnit(*Units);
+      }
+    }
+  }
+
+  RewriteRegs.clear();
 }
 
 FunctionPass *llvm::createVirtRegRewriter(bool ClearVirtRegs) {
   return new VirtRegRewriter(ClearVirtRegs);
 }