From 55a29c6b71c9b80353ccb17c0dd15dde5c9940b3 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Thu, 25 Oct 2018 14:47:57 -0700
Subject: [PATCH] VirtRegMap: Support partially allocated virtual registers

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.
---
 llvm/lib/CodeGen/VirtRegMap.cpp | 47 +++++++++++++++++++++++++++++++++++------
 1 file changed, 40 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/CodeGen/VirtRegMap.cpp b/llvm/lib/CodeGen/VirtRegMap.cpp
index 1882233..607f18c 100644
--- a/llvm/lib/CodeGen/VirtRegMap.cpp
+++ b/llvm/lib/CodeGen/VirtRegMap.cpp
@@ -181,13 +181,14 @@ class VirtRegRewriter : public MachineFunctionPass {
   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;
 
@@ -230,6 +231,7 @@ INITIALIZE_PASS_END(VirtRegRewriter, "virtregrewriter",
 void VirtRegRewriter::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesCFG();
   AU.addRequired<LiveIntervals>();
+  AU.addPreserved<LiveIntervals>();
   AU.addRequired<SlotIndexes>();
   AU.addPreserved<SlotIndexes>();
   AU.addRequired<LiveDebugVariables>();
@@ -274,7 +276,7 @@ bool VirtRegRewriter::runOnMachineFunction(MachineFunction &fn) {
 }
 
 void VirtRegRewriter::addLiveInsForSubRanges(const LiveInterval &LI,
-                                             Register PhysReg) const {
+                                             MCRegister PhysReg) const {
   assert(!LI.empty());
   assert(LI.hasSubRanges());
 
@@ -330,7 +332,12 @@ void VirtRegRewriter::addMBBLiveIns() {
     // 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);
@@ -381,12 +388,21 @@ bool VirtRegRewriter::readsUndefSubreg(const MachineOperand &MO) const {
   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
@@ -526,8 +542,12 @@ void VirtRegRewriter::rewrite() {
           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.
@@ -599,6 +619,19 @@ void VirtRegRewriter::rewrite() {
       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) {
-- 
2.7.4