[BOLT] Increase coverage of shrink wrapping [4/5]

[rafauler]

Summary

Change shrink-wrapping to try a priority list of save
positions, instead of trying the best one and giving up if it doesn't
work. This also increases coverage.

[Amir]

To reiterate:

1. This diff changes the way save location is handled: instead of computing and keeping a single place for a given register, store the list of potential locations.
2. This is done in order to work around the issue when the single best location might be discarded down the stream, cancelling the move of a spill-fill for a given register.
3. In order to support iterating over save locations, computeDomOrder needs to be updated multiple times.

As a side effect of this change, we also fix an issue where the critical edge is split for a given reg placement, but the placement is canceled.

bolt/include/bolt/Passes/ShrinkWrapping.h

class ShrinkWrapping {
   DenseMap<unsigned, std::vector<uint32_t>> DeletedPopCFIs;
   BitVector HasDeletedOffsetCFIs;
   SmallPtrSet<const MCCFIInstruction *, 16> UpdatedCFIs;
   std::vector<BitVector> UsesByReg;
   std::vector<int64_t> PushOffsetByReg;
   std::vector<int64_t> PopOffsetByReg;
   std::vector<MCPhysReg> DomOrder;
   CalleeSavedAnalysis CSA;
-  std::vector<SmallSetVector<MCInst *, 4>> SavePos;
-  std::vector<uint64_t> BestSaveCount;
-  std::vector<MCInst *> BestSavePos;
+  std::vector<std::vector<uint64_t>> BestSaveCount;
+  std::vector<std::vector<MCInst *>> BestSavePos;
 
   /// Pass stats
   static std::atomic_uint64_t SpillsMovedRegularMode;
   static std::atomic_uint64_t SpillsMovedPushPopMode;
   static std::atomic_uint64_t SpillsMovedDynamicCount;
   static std::atomic_uint64_t SpillsFailedDynamicCount;
   static std::atomic_uint64_t InstrDynamicCount;
   static std::atomic_uint64_t StoreDynamicCount;

bolt/lib/Passes/ShrinkWrapping.cpp

if (!CSA.HasRestores[I]) {
                     "restores:"
                  << I << "\n");
       CSA.CalleeSaved.reset(I);
     }
   }
 }
 
 void ShrinkWrapping::computeSaveLocations() {
-  SavePos = std::vector<SmallSetVector<MCInst *, 4>>(BC.MRI->getNumRegs());
+  BestSavePos = std::vector<std::vector<MCInst *>>(BC.MRI->getNumRegs());
   ReachingInsns<true> &RI = Info.getReachingInsnsBackwards();
   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
   StackPointerTracking &SPT = Info.getStackPointerTracking();
 
   LLVM_DEBUG(dbgs() << "Checking save/restore possibilities\n");
   for (BinaryBasicBlock &BB : BF) {
     LLVM_DEBUG(dbgs() << "\tNow at BB " << BB.getName() << "\n");
 
for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
       LLVM_DEBUG(dbgs() << "\t\tBB " << BB.getName() << " dominates "
                         << BBDominatedUses.count() << " uses for reg " << I
                         << ". Total uses for reg is " << UsesByReg[I].count()
                         << "\n");
       BBDominatedUses &= UsesByReg[I];
       if (BBDominatedUses == UsesByReg[I]) {
         LLVM_DEBUG(dbgs() << "\t\t\tAdded " << BB.getName()
                           << " as a save pos for " << I << "\n");
-        SavePos[I].insert(First);
+        BestSavePos[I].push_back(First);
         LLVM_DEBUG({
           dbgs() << "Dominated uses are:\n";
           for (int J : UsesByReg[I].set_bits()) {
             dbgs() << "Idx " << J << ": ";
+            BC.printInstruction(dbgs(), *DA.Expressions[J]);
             DA.Expressions[J]->dump();
           }
         });
       }
     }
   }
 
-  BestSaveCount = std::vector<uint64_t>(BC.MRI->getNumRegs(),
-                                        std::numeric_limits<uint64_t>::max());
-  BestSavePos = std::vector<MCInst *>(BC.MRI->getNumRegs(), nullptr);
+  BestSaveCount = std::vector<std::vector<uint64_t>>(BC.MRI->getNumRegs());
+
   auto &InsnToBB = Info.getInsnToBBMap();
   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
     if (!CSA.CalleeSaved[I])
       continue;
 
-    for (MCInst *Pos : SavePos[I]) {
-      BinaryBasicBlock *BB = InsnToBB[Pos];
-      uint64_t Count = BB->getExecutionCount();
-      if (Count != BinaryBasicBlock::COUNT_NO_PROFILE &&
-          Count < BestSaveCount[I]) {
-        BestSavePos[I] = Pos;
-        BestSaveCount[I] = Count;
-      }
+    std::stable_sort(BestSavePos[I].begin(), BestSavePos[I].end(),
+                     [&](const MCInst *A, const MCInst *B) {
+                       const BinaryBasicBlock *BBA = InsnToBB[A];
+                       const BinaryBasicBlock *BBB = InsnToBB[B];
+                       const uint64_t CountA = BBA->getKnownExecutionCount();
+                       const uint64_t CountB = BBB->getKnownExecutionCount();
+                       return CountB < CountA;
+                     });
+
+    for (MCInst *Pos : BestSavePos[I]) {
+      const BinaryBasicBlock *BB = InsnToBB[Pos];
+      const uint64_t Count = BB->getKnownExecutionCount();
+      BestSaveCount[I].push_back(Count);
     }
   }
 }
 
 void ShrinkWrapping::computeDomOrder() {
+  DomOrder = std::vector<MCPhysReg>(BC.MRI->getNumRegs(), 0);
   std::vector<MCPhysReg> Order;
   for (MCPhysReg I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
     Order.push_back(I);
   }
 
   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
   auto &InsnToBB = Info.getInsnToBBMap();
   llvm::sort(Order, [&](const MCPhysReg &A, const MCPhysReg &B) {
-    BinaryBasicBlock *BBA = BestSavePos[A] ? InsnToBB[BestSavePos[A]] : nullptr;
-    BinaryBasicBlock *BBB = BestSavePos[B] ? InsnToBB[BestSavePos[B]] : nullptr;
+    BinaryBasicBlock *BBA =
+        BestSavePos[A].size() ? InsnToBB[BestSavePos[A].back()] : nullptr;
+    BinaryBasicBlock *BBB =
+        BestSavePos[B].size() ? InsnToBB[BestSavePos[B].back()] : nullptr;
     if (BBA == BBB)
       return A < B;
     if (!BBA && BBB)
       return false;
     if (BBA && !BBB)
       return true;
-    if (DA.doesADominateB(*BestSavePos[A], *BestSavePos[B]))
+    if (DA.doesADominateB(*BestSavePos[A].back(), *BestSavePos[B].back()))
       return true;
-    if (DA.doesADominateB(*BestSavePos[B], *BestSavePos[A]))
+    if (DA.doesADominateB(*BestSavePos[B].back(), *BestSavePos[A].back()))
       return false;
     return A < B;
   });
 
   for (MCPhysReg I = 0, E = BC.MRI->getNumRegs(); I != E; ++I)
     DomOrder[Order[I]] = I;
 }
 
 bool ShrinkWrapping::isBestSavePosCold(unsigned CSR, MCInst *&BestPosSave,
                                        uint64_t &TotalEstimatedWin) {
   const uint64_t CurSavingCost = CSA.SavingCost[CSR];
   if (!CSA.CalleeSaved[CSR])
     return false;
 
-  uint64_t BestCount = BestSaveCount[CSR];
-  BestPosSave = BestSavePos[CSR];
-  bool ShouldMove = false;
-  if (BestCount != std::numeric_limits<uint64_t>::max() &&
-      BestCount < (opts::ShrinkWrappingThreshold / 100.0) * CurSavingCost) {
-    LLVM_DEBUG({
-      auto &InsnToBB = Info.getInsnToBBMap();
-      dbgs() << "Better position for saves found in func " << BF.getPrintName()
-             << " count << " << BF.getKnownExecutionCount() << "\n";
-      dbgs() << "Reg: " << CSR
-             << "; New BB: " << InsnToBB[BestPosSave]->getName()
-             << " Freq reduction: " << (CurSavingCost - BestCount) << "\n";
-    });
-    TotalEstimatedWin += CurSavingCost - BestCount;
-    ShouldMove = true;
-  }
+  assert(BestSaveCount[CSR].size() == BestSavePos[CSR].size() &&
+         "save position vectors out of sync");
+  if (BestSaveCount[CSR].empty())
+    return false;
+
+  const uint64_t BestCount = BestSaveCount[CSR].back();
+  BestPosSave = BestSavePos[CSR].back();
+  if (BestCount >= (opts::ShrinkWrappingThreshold / 100.0) * CurSavingCost)
+    return false;
+
+    dbgs() << "Reg: " << CSR << "; New BB: " << InsnToBB[BestPosSave]->getName()
 
-  if (!ShouldMove)
-    return false;
-  if (!BestPosSave) {
-    LLVM_DEBUG({
-      dbgs() << "Dropping opportunity because we don't know where to put "
-                "stores -- total est. freq reduc: "
-             << TotalEstimatedWin << "\n";
-    });
-    return false;
-  }
+  TotalEstimatedWin = CurSavingCost - BestCount;
   return true;
 }
 
 /// Auxiliar function used to create basic blocks for critical edges and update
 /// the dominance frontier with these new locations
 void ShrinkWrapping::splitFrontierCritEdges(
     BinaryFunction *Func, SmallVector<ProgramPoint, 4> &Frontier,
     const SmallVector<bool, 4> &IsCritEdge,
void ShrinkWrapping::splitFrontierCritEdges(
   }
 }
 
 SmallVector<ProgramPoint, 4>
 ShrinkWrapping::doRestorePlacement(MCInst *BestPosSave, unsigned CSR,
                                    uint64_t TotalEstimatedWin) {
   SmallVector<ProgramPoint, 4> Frontier;
   SmallVector<bool, 4> IsCritEdge;
-  bool CannotPlace = false;
   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 
   SmallVector<BinaryBasicBlock *, 4> CritEdgesFrom;
   SmallVector<SmallVector<BinaryBasicBlock *, 4>, 4> CritEdgesTo;
   // In case of a critical edge, we need to create extra BBs to host restores
   // into edges transitioning to the dominance frontier, otherwise we pull these
   // restores to inside the dominated area.
   Frontier = DA.getDominanceFrontierFor(*BestPosSave).takeVector();
   LLVM_DEBUG({
     dbgs() << "Dumping dominance frontier for ";
     BC.printInstruction(dbgs(), *BestPosSave);
     for (ProgramPoint &PP : Frontier)
       if (PP.isInst())
         BC.printInstruction(dbgs(), *PP.getInst());
       else
         dbgs() << PP.getBB()->getName() << "\n";
   });
   for (ProgramPoint &PP : Frontier) {
     bool HasCritEdges = false;
     if (PP.isInst() && BC.MIB->isTerminator(*PP.getInst()) &&
-        doesInstUsesCSR(*PP.getInst(), CSR))
-      CannotPlace = true;
+        doesInstUsesCSR(*PP.getInst(), CSR)) {
+      Frontier.clear();
+      return Frontier;
+    }
     BinaryBasicBlock *FrontierBB = Info.getParentBB(PP);
     CritEdgesFrom.emplace_back(FrontierBB);
     CritEdgesTo.emplace_back(0);
     SmallVector<BinaryBasicBlock *, 4> &Dests = CritEdgesTo.back();
     // Check for invoke instructions at the dominance frontier, which indicates
     // the landing pad is not dominated.
     if (PP.isInst() && BC.MIB->isInvoke(*PP.getInst())) {
-      LLVM_DEBUG(
-          dbgs() << "Bailing on restore placement to avoid LP splitting\n");
       Frontier.clear();
       return Frontier;
     }
     doForAllSuccs(*FrontierBB, [&](ProgramPoint P) {
       if (!DA.doesADominateB(*BestPosSave, P)) {
         Dests.emplace_back(Info.getParentBB(P));
         return;
       }
ShrinkWrapping::doRestorePlacement(MCInst *BestPosSave, unsigned CSR,
   }
   if (InvalidateRequired) {
     // BitVectors that represent all insns of the function are invalid now
     // since we changed BBs/Insts. Re-run steps that depend on pointers being
     // valid
     Info.invalidateAll();
     classifyCSRUses();
   }
-  if (CannotPlace) {
-    LLVM_DEBUG({
-      dbgs() << "Dropping opportunity because restore placement failed"
-                " -- total est. freq reduc: "
-             << TotalEstimatedWin << "\n";
-    });
-    Frontier.clear();
-    return Frontier;
-  }
   return Frontier;
 }
 
 bool ShrinkWrapping::validatePushPopsMode(unsigned CSR, MCInst *BestPosSave,
                                           int64_t SaveOffset) {
   if (FA.requiresAlignment(BF)) {
     LLVM_DEBUG({
       dbgs() << "Reg " << CSR

 void ShrinkWrapping::moveSaveRestores() {
   bool DisablePushPopMode = false;
   bool UsedPushPopMode = false;
   // Keeps info about successfully moved regs: reg index, save position and
   // save size
   std::vector<std::tuple<unsigned, MCInst *, size_t>> MovedRegs;
   uint64_t TotalEstimatedWin = 0;
 
+  computeDomOrder();
   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
     MCInst *BestPosSave = nullptr;
     uint64_t EstimatedWin = 0;
-    if (!isBestSavePosCold(I, BestPosSave, EstimatedWin))
-      continue;
-    SmallVector<ProgramPoint, 4> RestorePoints =
-        doRestorePlacement(BestPosSave, I, EstimatedWin);
+    SmallVector<ProgramPoint, 4> RestorePoints;
+    while (RestorePoints.empty() &&
+           isBestSavePosCold(I, BestPosSave, EstimatedWin)) {
+      RestorePoints = doRestorePlacement(BestPosSave, I, EstimatedWin);
+      if (RestorePoints.empty()) {
+        LLVM_DEBUG({
+          dbgs() << "Dropping opportunity because restore placement failed"
+                    " -- total est. freq reduc: "
+                 << EstimatedWin << ". Will try "
+                 << (BestSaveCount[I].size() - 1) << " more times.\n";
+        });
+        BestSaveCount[I].pop_back();
+        BestSavePos[I].pop_back();
+        computeDomOrder();
+      }
+    }
     if (RestorePoints.empty()) {
       SpillsFailedDynamicCount += EstimatedWin;
       continue;
     }
 
     const FrameIndexEntry *FIESave = CSA.SaveFIEByReg[I];
     const FrameIndexEntry *FIELoad = CSA.LoadFIEByReg[I];
     (void)FIELoad;
for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
     Info.invalidateAll();
   }
 }
 
 bool ShrinkWrapping::perform(bool HotOnly) {
   HasDeletedOffsetCFIs = BitVector(BC.MRI->getNumRegs(), false);
   PushOffsetByReg = std::vector<int64_t>(BC.MRI->getNumRegs(), 0LL);
   PopOffsetByReg = std::vector<int64_t>(BC.MRI->getNumRegs(), 0LL);
-  DomOrder = std::vector<MCPhysReg>(BC.MRI->getNumRegs(), 0);
 
   // Update pass statistics
   uint64_t TotalInstrs = 0ULL;
   uint64_t TotalStoreInstrs = 0ULL;
   for (BinaryBasicBlock *BB : BF.layout()) {
     uint64_t BBExecCount = BB->getExecutionCount();
     if (!BBExecCount || BBExecCount == BinaryBasicBlock::COUNT_NO_PROFILE)
       continue;
if (BF.checkForAmbiguousJumpTables()) {
     // written in assembly language. Just bail.
     return false;
   }
   SLM.initialize();
   CSA.compute();
   classifyCSRUses();
   pruneUnwantedCSRs();
   computeSaveLocations();
-  computeDomOrder();
   moveSaveRestores();
   LLVM_DEBUG({
     dbgs() << "Func before shrink-wrapping: \n";
     BF.dump();
   });
   SLM.performChanges();
   // Early exit if processInsertions doesn't detect any todo items
   if (!processInsertions())