VirtRegMap: Improve addMBBLiveIns() using SlotIndex::MBBIndexIterator; NFC

[MatzeB]

Summary

Now that we have an explicit iterator over the idx2MBBMap in SlotIndices
we can use the fact that segments and the idx2MBBMap are sorted by
SlotIndex position. We can advance both simultaneously instead of
starting from the beginning for each segment.

This complicates the code for the subregister case somewhat because we have to
keep a set of iterators for each subrange in sync. The result should however be
more efficient and has the advantage that we get the complete lanemask for each
block immediately instead of incrementally building it. This is important for a
subsequent change.

Removes the now unused SlotIndexes::findMBBLiveIns function.

[qcolombet]

Hi Matthias,

LGTM with maybe an additional comment and an assert.
See the inlined comments.

Thanks,
-Quentin

lib/CodeGen/VirtRegMap.cpp

309

Reiterate the assumption here so that we know why we do not need to start over for each segment.

311

I do not remember if we check that somewhere, like the machine verifier, but I guess it won’t hurt to have an assert to check that the list is indeed sorted.

What do you think?

[MatzeB]

Thanks for the review. I created a separate commit that adds a check that the MBBIdxMap is sorted to the MachineVerifier (r247166).

include/llvm/CodeGen/LiveIntervalAnalysis.h

public:
     void RemoveMachineInstrFromMaps(MachineInstr *MI) {
       Indexes->removeMachineInstrFromMaps(MI);
     }
 
     void ReplaceMachineInstrInMaps(MachineInstr *MI, MachineInstr *NewMI) {
       Indexes->replaceMachineInstrInMaps(MI, NewMI);
     }
 
-    bool findLiveInMBBs(SlotIndex Start, SlotIndex End,
-                        SmallVectorImpl<MachineBasicBlock*> &MBBs) const {
-      return Indexes->findLiveInMBBs(Start, End, MBBs);
-    }
-
     VNInfo::Allocator& getVNInfoAllocator() { return VNInfoAllocator; }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override;
     void releaseMemory() override;
 
     /// runOnMachineFunction - pass entry point
     bool runOnMachineFunction(MachineFunction&) override;
 

include/llvm/CodeGen/SlotIndexes.h

MachineBasicBlock* getMBBFromIndex(SlotIndex index) const {
          (I == MBBIndexEnd() && !idx2MBBMap.empty())) ? std::prev(I) : I;
 
       assert(J != MBBIndexEnd() && J->first <= index &&
              index < getMBBEndIdx(J->second) &&
              "index does not correspond to an MBB");
       return J->second;
     }
 
-    bool findLiveInMBBs(SlotIndex start, SlotIndex end,
-                        SmallVectorImpl<MachineBasicBlock*> &mbbs) const {
-      bool resVal = false;
-      for (MBBIndexIterator itr = findMBBIndex(start);
-           itr != MBBIndexEnd() && itr->first < end; ++itr) {
-        mbbs.push_back(itr->second);
-        resVal = true;
-      }
-      return resVal;
-    }
-
     /// Returns the MBB covering the given range, or null if the range covers
     /// more than one basic block.
     MachineBasicBlock* getMBBCoveringRange(SlotIndex start, SlotIndex end) const {
 
       assert(start < end && "Backwards ranges not allowed.");
       MBBIndexIterator itr = findMBBIndex(start);
       if (itr == MBBIndexEnd()) {
         itr = std::prev(itr);

lib/CodeGen/VirtRegMap.cpp

class VirtRegRewriter : public MachineFunctionPass {
   MachineRegisterInfo *MRI;
   SlotIndexes *Indexes;
   LiveIntervals *LIS;
   VirtRegMap *VRM;
 
   void rewrite();
   void addMBBLiveIns();
   bool readsUndefSubreg(const MachineOperand &MO) const;
+  void addLiveInsForSubRanges(const LiveInterval &LI, unsigned PhysReg) const;
+
 public:
   static char ID;
   VirtRegRewriter() : MachineFunctionPass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   bool runOnMachineFunction(MachineFunction&) override;
 };
bool VirtRegRewriter::runOnMachineFunction(MachineFunction &fn) {
 
   // All machine operands and other references to virtual registers have been
   // replaced. Remove the virtual registers and release all the transient data.
   VRM->clearAllVirt();
   MRI->clearVirtRegs();
   return true;
 }
 
+void VirtRegRewriter::addLiveInsForSubRanges(const LiveInterval &LI,
+                                             unsigned PhysReg) const {
+  assert(!LI.empty());
+  assert(LI.hasSubRanges());
+
+  typedef std::pair<const LiveInterval::SubRange *,
+                    LiveInterval::const_iterator> SubRangeIteratorPair;
+  SmallVector<SubRangeIteratorPair, 4> SubRanges;
+  SlotIndex First;
+  SlotIndex Last;
+  for (const LiveInterval::SubRange &SR : LI.subranges()) {
+    SubRanges.push_back(std::make_pair(&SR, SR.begin()));
+    if (!First.isValid() || SR.segments.front().start < First)
+      First = SR.segments.front().start;
+    if (!Last.isValid() || SR.segments.back().end > Last)
+      Last = SR.segments.back().end;
+  }
+
+  // Check all mbb start positions between First and Last while
+  // simulatenously advancing an iterator for each subrange.
+  for (SlotIndexes::MBBIndexIterator MBBI = Indexes->findMBBIndex(First);
+       MBBI != Indexes->MBBIndexEnd() && MBBI->first <= Last; ++MBBI) {
+    SlotIndex MBBBegin = MBBI->first;
+    // Advance all subrange iterators so that their end position is just
+    // behind MBBBegin (or the iterator is at the end).
+    unsigned LaneMask = 0;
+    for (auto &RangeIterPair : SubRanges) {
+      const LiveInterval::SubRange *SR = RangeIterPair.first;
+      LiveInterval::const_iterator &SRI = RangeIterPair.second;
+      while (SRI != SR->end() && SRI->end <= MBBBegin)
+        ++SRI;
+      if (SRI == SR->end())
+        continue;
+      if (SRI->start <= MBBBegin)
+        LaneMask |= SR->LaneMask;
+    }
+    if (LaneMask == 0)
+      continue;
+    MachineBasicBlock *MBB = MBBI->second;
+    for (MCSubRegIndexIterator SR(PhysReg, TRI); SR.isValid(); ++SR) {
+      unsigned SubReg = SR.getSubReg();
+      unsigned SubRegIndex = SR.getSubRegIndex();
+      unsigned SubRegLaneMask = TRI->getSubRegIndexLaneMask(SubRegIndex);
+      if ((SubRegLaneMask & LaneMask) != 0)
+        MBB->addLiveIn(SubReg);
+    }
+  }
+}
+
 // Compute MBB live-in lists from virtual register live ranges and their
 // assignments.
 void VirtRegRewriter::addMBBLiveIns() {
-  SmallVector<MachineBasicBlock*, 16> LiveIn;
   for (unsigned Idx = 0, IdxE = MRI->getNumVirtRegs(); Idx != IdxE; ++Idx) {
     unsigned VirtReg = TargetRegisterInfo::index2VirtReg(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 (LI.hasSubRanges()) {
-      for (LiveInterval::SubRange &S : LI.subranges()) {
-        for (const auto &Seg : S.segments) {
-          if (!Indexes->findLiveInMBBs(Seg.start, Seg.end, LiveIn))
-            continue;
-          for (MCSubRegIndexIterator SR(PhysReg, TRI); SR.isValid(); ++SR) {
-            unsigned SubReg = SR.getSubReg();
-            unsigned SubRegIndex = SR.getSubRegIndex();
-            unsigned SubRegLaneMask = TRI->getSubRegIndexLaneMask(SubRegIndex);
-            if ((SubRegLaneMask & S.LaneMask) == 0)
-              continue;
-            for (unsigned i = 0, e = LiveIn.size(); i != e; ++i) {
-              LiveIn[i]->addLiveIn(SubReg);
-            }
-          }
-          LiveIn.clear();
-        }
-      }
+      addLiveInsForSubRanges(LI, PhysReg);
     } else {
-      // Scan the segments of LI.
-      for (const auto &Seg : LI.segments) {
-        if (!Indexes->findLiveInMBBs(Seg.start, Seg.end, LiveIn))
-          continue;
-        for (unsigned i = 0, e = LiveIn.size(); i != e; ++i)
-          LiveIn[i]->addLiveIn(PhysReg);
-        LiveIn.clear();
+      // Go over MBB begin positions and see if we have segments covering them.
+      SlotIndexes::MBBIndexIterator I = Indexes->MBBIndexBegin();
+      for (const auto &Seg : LI) {
+        I = Indexes->advanceMBBIndex(I, Seg.start);
+        for (; I != Indexes->MBBIndexEnd() && I->first < Seg.end; ++I) {
+          MachineBasicBlock *MBB = I->second;
+          MBB->addLiveIn(PhysReg);
+        }
       }
     }
   }
 
   // Sort and unique MBB LiveIns as we've not checked if SubReg/PhysReg were in
   // each MBB's LiveIns set before calling addLiveIn on them.
   for (MachineBasicBlock &MBB : *MF)
     MBB.sortUniqueLiveIns();