bolt/lib/Passes/ShrinkWrapping.cpp

   1 //===- bolt/Passes/ShrinkWrapping.cpp -------------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements the ShrinkWrapping class.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "bolt/Passes/ShrinkWrapping.h"
  14 #include "bolt/Passes/DataflowInfoManager.h"
  15 #include "bolt/Passes/MCF.h"
  16 #include "bolt/Utils/CommandLineOpts.h"
  17 #include <numeric>
  18 #include <optional>
  19 #include <stack>
  20
  21 #define DEBUG_TYPE "shrinkwrapping"
  22
  23 using namespace llvm;
  24
  25 namespace opts {
  26
  27 extern cl::opt<bool> TimeOpts;
  28 extern cl::OptionCategory BoltOptCategory;
  29
  30 static cl::opt<unsigned> ShrinkWrappingThreshold(
  31     "shrink-wrapping-threshold",
  32     cl::desc("Percentage of prologue execution count to use as threshold when"
  33              " evaluating whether a block is cold enough to be profitable to"
  34              " move eligible spills there"),
  35     cl::init(30), cl::ZeroOrMore, cl::cat(BoltOptCategory));
  36 } // namespace opts
  37
  38 namespace llvm {
  39 namespace bolt {
  40
  41 void CalleeSavedAnalysis::analyzeSaves() {
  42   ReachingDefOrUse</*Def=*/true> &RD = Info.getReachingDefs();
  43   StackReachingUses &SRU = Info.getStackReachingUses();
  44   auto &InsnToBB = Info.getInsnToBBMap();
  45   BitVector BlacklistedRegs(BC.MRI->getNumRegs(), false);
  46
  47   LLVM_DEBUG(dbgs() << "Checking spill locations\n");
  48   for (BinaryBasicBlock &BB : BF) {
  49     LLVM_DEBUG(dbgs() << "\tNow at BB " << BB.getName() << "\n");
  50     const MCInst *Prev = nullptr;
  51     for (MCInst &Inst : BB) {
  52       if (ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(Inst)) {
  53         // Blacklist weird stores we don't understand
  54         if ((!FIE->IsSimple || FIE->StackOffset >= 0) && FIE->IsStore &&
  55             FIE->IsStoreFromReg) {
  56           BlacklistedRegs.set(FIE->RegOrImm);
  57           CalleeSaved.reset(FIE->RegOrImm);
  58           Prev = &Inst;
  59           continue;
  60         }
  61
  62         if (!FIE->IsStore || !FIE->IsStoreFromReg ||
  63             BlacklistedRegs[FIE->RegOrImm]) {
  64           Prev = &Inst;
  65           continue;
  66         }
  67
  68         // If this reg is defined locally, it is not a callee-saved reg
  69         if (RD.isReachedBy(FIE->RegOrImm,
  70                            Prev ? RD.expr_begin(*Prev) : RD.expr_begin(BB))) {
  71           BlacklistedRegs.set(FIE->RegOrImm);
  72           CalleeSaved.reset(FIE->RegOrImm);
  73           Prev = &Inst;
  74           continue;
  75         }
  76
  77         // If this stack position is accessed in another function, we are
  78         // probably dealing with a parameter passed in a stack -- do not mess
  79         // with it
  80         if (SRU.isStoreUsed(*FIE,
  81                             Prev ? SRU.expr_begin(*Prev) : SRU.expr_begin(BB)),
  82             /*IncludeLocalAccesses=*/false) {
  83           BlacklistedRegs.set(FIE->RegOrImm);
  84           CalleeSaved.reset(FIE->RegOrImm);
  85           Prev = &Inst;
  86           continue;
  87         }
  88
  89         // If this stack position is loaded elsewhere in another reg, we can't
  90         // update it, so blacklist it.
  91         if (SRU.isLoadedInDifferentReg(*FIE, Prev ? SRU.expr_begin(*Prev)
  92                                                   : SRU.expr_begin(BB))) {
  93           BlacklistedRegs.set(FIE->RegOrImm);
  94           CalleeSaved.reset(FIE->RegOrImm);
  95           Prev = &Inst;
  96           continue;
  97         }
  98
  99         // Ignore regs with multiple saves
 100         if (CalleeSaved[FIE->RegOrImm]) {
 101           BlacklistedRegs.set(FIE->RegOrImm);
 102           CalleeSaved.reset(FIE->RegOrImm);
 103           Prev = &Inst;
 104           continue;
 105         }
 106
 107         CalleeSaved.set(FIE->RegOrImm);
 108         SaveFIEByReg[FIE->RegOrImm] = &*FIE;
 109         SavingCost[FIE->RegOrImm] += InsnToBB[&Inst]->getKnownExecutionCount();
 110         BC.MIB->addAnnotation(Inst, getSaveTag(), FIE->RegOrImm, AllocatorId);
 111         OffsetsByReg[FIE->RegOrImm] = FIE->StackOffset;
 112         LLVM_DEBUG(dbgs() << "Logging new candidate for Callee-Saved Reg: "
 113                           << FIE->RegOrImm << "\n");
 114       }
 115       Prev = &Inst;
 116     }
 117   }
 118 }
 119
 120 void CalleeSavedAnalysis::analyzeRestores() {
 121   ReachingDefOrUse</*Def=*/false> &RU = Info.getReachingUses();
 122
 123   // Now compute all restores of these callee-saved regs
 124   for (BinaryBasicBlock &BB : BF) {
 125     const MCInst *Prev = nullptr;
 126     for (MCInst &Inst : llvm::reverse(BB)) {
 127       if (ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(Inst)) {
 128         if (!FIE->IsLoad || !CalleeSaved[FIE->RegOrImm]) {
 129           Prev = &Inst;
 130           continue;
 131         }
 132
 133         // If this reg is used locally after a restore, then we are probably
 134         // not dealing with a callee-saved reg. Except if this use is by
 135         // another store, but we don't cover this case yet.
 136         // Also not callee-saved if this load accesses caller stack or isn't
 137         // simple.
 138         if (!FIE->IsSimple || FIE->StackOffset >= 0 ||
 139             RU.isReachedBy(FIE->RegOrImm,
 140                            Prev ? RU.expr_begin(*Prev) : RU.expr_begin(BB))) {
 141           CalleeSaved.reset(FIE->RegOrImm);
 142           Prev = &Inst;
 143           continue;
 144         }
 145         // If stack offsets between saves/store don't agree with each other,
 146         // we don't completely understand what's happening here
 147         if (FIE->StackOffset != OffsetsByReg[FIE->RegOrImm]) {
 148           CalleeSaved.reset(FIE->RegOrImm);
 149           LLVM_DEBUG(dbgs() << "Dismissing Callee-Saved Reg because we found a "
 150                                "mismatching restore: "
 151                             << FIE->RegOrImm << "\n");
 152           Prev = &Inst;
 153           continue;
 154         }
 155
 156         LLVM_DEBUG(dbgs() << "Adding matching restore for: " << FIE->RegOrImm
 157                           << "\n");
 158         if (LoadFIEByReg[FIE->RegOrImm] == nullptr)
 159           LoadFIEByReg[FIE->RegOrImm] = &*FIE;
 160         BC.MIB->addAnnotation(Inst, getRestoreTag(), FIE->RegOrImm,
 161                               AllocatorId);
 162         HasRestores.set(FIE->RegOrImm);
 163       }
 164       Prev = &Inst;
 165     }
 166   }
 167 }
 168
 169 std::vector<MCInst *> CalleeSavedAnalysis::getSavesByReg(uint16_t Reg) {
 170   std::vector<MCInst *> Results;
 171   for (BinaryBasicBlock &BB : BF)
 172     for (MCInst &Inst : BB)
 173       if (getSavedReg(Inst) == Reg)
 174         Results.push_back(&Inst);
 175   return Results;
 176 }
 177
 178 std::vector<MCInst *> CalleeSavedAnalysis::getRestoresByReg(uint16_t Reg) {
 179   std::vector<MCInst *> Results;
 180   for (BinaryBasicBlock &BB : BF)
 181     for (MCInst &Inst : BB)
 182       if (getRestoredReg(Inst) == Reg)
 183         Results.push_back(&Inst);
 184   return Results;
 185 }
 186
 187 CalleeSavedAnalysis::~CalleeSavedAnalysis() {
 188   for (BinaryBasicBlock &BB : BF) {
 189     for (MCInst &Inst : BB) {
 190       BC.MIB->removeAnnotation(Inst, getSaveTag());
 191       BC.MIB->removeAnnotation(Inst, getRestoreTag());
 192     }
 193   }
 194 }
 195
 196 void StackLayoutModifier::blacklistRegion(int64_t Offset, int64_t Size) {
 197   if (BlacklistedRegions[Offset] < Size)
 198     BlacklistedRegions[Offset] = Size;
 199 }
 200
 201 bool StackLayoutModifier::isRegionBlacklisted(int64_t Offset, int64_t Size) {
 202   for (std::pair<const int64_t, int64_t> Elem : BlacklistedRegions)
 203     if (Offset + Size > Elem.first && Offset < Elem.first + Elem.second)
 204       return true;
 205   return false;
 206 }
 207
 208 bool StackLayoutModifier::blacklistAllInConflictWith(int64_t Offset,
 209                                                      int64_t Size) {
 210   bool HasConflict = false;
 211   for (auto Iter = AvailableRegions.begin(); Iter != AvailableRegions.end();) {
 212     std::pair<const int64_t, int64_t> &Elem = *Iter;
 213     if (Offset + Size > Elem.first && Offset < Elem.first + Elem.second &&
 214         (Offset != Elem.first || Size != Elem.second)) {
 215       Iter = AvailableRegions.erase(Iter);
 216       HasConflict = true;
 217       continue;
 218     }
 219     ++Iter;
 220   }
 221   if (HasConflict) {
 222     blacklistRegion(Offset, Size);
 223     return true;
 224   }
 225   return false;
 226 }
 227
 228 void StackLayoutModifier::checkFramePointerInitialization(MCInst &Point) {
 229   StackPointerTracking &SPT = Info.getStackPointerTracking();
 230   if (!BC.MII->get(Point.getOpcode())
 231            .hasDefOfPhysReg(Point, BC.MIB->getFramePointer(), *BC.MRI))
 232     return;
 233
 234   int SPVal, FPVal;
 235   std::tie(SPVal, FPVal) = *SPT.getStateBefore(Point);
 236   std::pair<MCPhysReg, int64_t> FP;
 237
 238   if (FPVal != SPT.EMPTY && FPVal != SPT.SUPERPOSITION)
 239     FP = std::make_pair(BC.MIB->getFramePointer(), FPVal);
 240   else
 241     FP = std::make_pair(0, 0);
 242   std::pair<MCPhysReg, int64_t> SP;
 243
 244   if (SPVal != SPT.EMPTY && SPVal != SPT.SUPERPOSITION)
 245     SP = std::make_pair(BC.MIB->getStackPointer(), SPVal);
 246   else
 247     SP = std::make_pair(0, 0);
 248
 249   int64_t Output;
 250   if (!BC.MIB->evaluateStackOffsetExpr(Point, Output, SP, FP))
 251     return;
 252
 253   // Not your regular frame pointer initialization... bail
 254   if (Output != SPVal)
 255     blacklistRegion(0, 0);
 256 }
 257
 258 void StackLayoutModifier::checkStackPointerRestore(MCInst &Point) {
 259   StackPointerTracking &SPT = Info.getStackPointerTracking();
 260   if (!BC.MII->get(Point.getOpcode())
 261            .hasDefOfPhysReg(Point, BC.MIB->getStackPointer(), *BC.MRI))
 262     return;
 263   // Check if the definition of SP comes from FP -- in this case, this
 264   // value may need to be updated depending on our stack layout changes
 265   bool UsesFP = llvm::any_of(BC.MIB->useOperands(Point), [&](MCOperand &Op) {
 266     return Op.isReg() && Op.getReg() == BC.MIB->getFramePointer();
 267   });
 268   if (!UsesFP)
 269     return;
 270
 271   // Setting up evaluation
 272   int SPVal, FPVal;
 273   std::tie(SPVal, FPVal) = *SPT.getStateBefore(Point);
 274   std::pair<MCPhysReg, int64_t> FP;
 275
 276   if (FPVal != SPT.EMPTY && FPVal != SPT.SUPERPOSITION)
 277     FP = std::make_pair(BC.MIB->getFramePointer(), FPVal);
 278   else
 279     FP = std::make_pair(0, 0);
 280   std::pair<MCPhysReg, int64_t> SP;
 281
 282   if (SPVal != SPT.EMPTY && SPVal != SPT.SUPERPOSITION)
 283     SP = std::make_pair(BC.MIB->getStackPointer(), SPVal);
 284   else
 285     SP = std::make_pair(0, 0);
 286
 287   int64_t Output;
 288   if (!BC.MIB->evaluateStackOffsetExpr(Point, Output, SP, FP))
 289     return;
 290
 291   // If the value is the same of FP, no need to adjust it
 292   if (Output == FPVal)
 293     return;
 294
 295   // If an allocation happened through FP, bail
 296   if (Output <= SPVal) {
 297     blacklistRegion(0, 0);
 298     return;
 299   }
 300
 301   // We are restoring SP to an old value based on FP. Mark it as a stack
 302   // access to be fixed later.
 303   BC.MIB->addAnnotation(Point, getSlotTag(), Output, AllocatorId);
 304 }
 305
 306 void StackLayoutModifier::classifyStackAccesses() {
 307   // Understand when stack slots are being used non-locally
 308   StackReachingUses &SRU = Info.getStackReachingUses();
 309
 310   for (BinaryBasicBlock &BB : BF) {
 311     const MCInst *Prev = nullptr;
 312     for (MCInst &Inst : llvm::reverse(BB)) {
 313       checkFramePointerInitialization(Inst);
 314       checkStackPointerRestore(Inst);
 315       ErrorOr<const FrameIndexEntry &> FIEX = FA.getFIEFor(Inst);
 316       if (!FIEX) {
 317         Prev = &Inst;
 318         continue;
 319       }
 320       if (!FIEX->IsSimple || (FIEX->IsStore && !FIEX->IsStoreFromReg)) {
 321         blacklistRegion(FIEX->StackOffset, FIEX->Size);
 322         Prev = &Inst;
 323         continue;
 324       }
 325       // If this stack position is accessed in another function, we are
 326       // probably dealing with a parameter passed in a stack -- do not mess
 327       // with it
 328       if (SRU.isStoreUsed(*FIEX,
 329                           Prev ? SRU.expr_begin(*Prev) : SRU.expr_begin(BB),
 330                           /*IncludeLocalAccesses=*/false)) {
 331         blacklistRegion(FIEX->StackOffset, FIEX->Size);
 332         Prev = &Inst;
 333         continue;
 334       }
 335       // Now we have a clear stack slot access. Check if its blacklisted or if
 336       // it conflicts with another chunk.
 337       if (isRegionBlacklisted(FIEX->StackOffset, FIEX->Size) ||
 338           blacklistAllInConflictWith(FIEX->StackOffset, FIEX->Size)) {
 339         Prev = &Inst;
 340         continue;
 341       }
 342       // We are free to go. Add it as available stack slot which we know how
 343       // to move it.
 344       AvailableRegions[FIEX->StackOffset] = FIEX->Size;
 345       BC.MIB->addAnnotation(Inst, getSlotTag(), FIEX->StackOffset, AllocatorId);
 346       RegionToRegMap[FIEX->StackOffset].insert(FIEX->RegOrImm);
 347       RegToRegionMap[FIEX->RegOrImm].insert(FIEX->StackOffset);
 348       LLVM_DEBUG(dbgs() << "Adding region " << FIEX->StackOffset << " size "
 349                         << (int)FIEX->Size << "\n");
 350     }
 351   }
 352 }
 353
 354 void StackLayoutModifier::classifyCFIs() {
 355   std::stack<std::pair<int64_t, uint16_t>> CFIStack;
 356   int64_t CfaOffset = -8;
 357   uint16_t CfaReg = 7;
 358
 359   auto recordAccess = [&](MCInst *Inst, int64_t Offset) {
 360     const uint16_t Reg = *BC.MRI->getLLVMRegNum(CfaReg, /*isEH=*/false);
 361     if (Reg == BC.MIB->getStackPointer() || Reg == BC.MIB->getFramePointer()) {
 362       BC.MIB->addAnnotation(*Inst, getSlotTag(), Offset, AllocatorId);
 363       LLVM_DEBUG(dbgs() << "Recording CFI " << Offset << "\n");
 364     } else {
 365       IsSimple = false;
 366       return;
 367     }
 368   };
 369
 370   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
 371     for (MCInst &Inst : *BB) {
 372       if (!BC.MIB->isCFI(Inst))
 373         continue;
 374       const MCCFIInstruction *CFI = BF.getCFIFor(Inst);
 375       switch (CFI->getOperation()) {
 376       case MCCFIInstruction::OpDefCfa:
 377         CfaOffset = -CFI->getOffset();
 378         recordAccess(&Inst, CfaOffset);
 379         [[fallthrough]];
 380       case MCCFIInstruction::OpDefCfaRegister:
 381         CfaReg = CFI->getRegister();
 382         break;
 383       case MCCFIInstruction::OpDefCfaOffset:
 384         CfaOffset = -CFI->getOffset();
 385         recordAccess(&Inst, CfaOffset);
 386         break;
 387       case MCCFIInstruction::OpOffset:
 388         recordAccess(&Inst, CFI->getOffset());
 389         BC.MIB->addAnnotation(Inst, getOffsetCFIRegTag(),
 390                               BC.MRI->getLLVMRegNum(CFI->getRegister(),
 391                                                     /*isEH=*/false),
 392                               AllocatorId);
 393         break;
 394       case MCCFIInstruction::OpSameValue:
 395         BC.MIB->addAnnotation(Inst, getOffsetCFIRegTag(),
 396                               BC.MRI->getLLVMRegNum(CFI->getRegister(),
 397                                                     /*isEH=*/false),
 398                               AllocatorId);
 399         break;
 400       case MCCFIInstruction::OpRememberState:
 401         CFIStack.push(std::make_pair(CfaOffset, CfaReg));
 402         break;
 403       case MCCFIInstruction::OpRestoreState: {
 404         assert(!CFIStack.empty() && "Corrupt CFI stack");
 405         std::pair<int64_t, uint16_t> &Elem = CFIStack.top();
 406         CFIStack.pop();
 407         CfaOffset = Elem.first;
 408         CfaReg = Elem.second;
 409         break;
 410       }
 411       case MCCFIInstruction::OpRelOffset:
 412       case MCCFIInstruction::OpAdjustCfaOffset:
 413         llvm_unreachable("Unhandled AdjustCfaOffset");
 414         break;
 415       default:
 416         break;
 417       }
 418     }
 419   }
 420 }
 421
 422 void StackLayoutModifier::scheduleChange(
 423     MCInst &Inst, StackLayoutModifier::WorklistItem Item) {
 424   auto &WList = BC.MIB->getOrCreateAnnotationAs<std::vector<WorklistItem>>(
 425       Inst, getTodoTag(), AllocatorId);
 426   WList.push_back(Item);
 427 }
 428
 429 bool StackLayoutModifier::canCollapseRegion(MCInst *DeletedPush) {
 430   if (!IsSimple || !BC.MIB->isPush(*DeletedPush))
 431     return false;
 432
 433   ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(*DeletedPush);
 434   if (!FIE)
 435     return false;
 436
 437   return canCollapseRegion(FIE->StackOffset);
 438 }
 439
 440 bool StackLayoutModifier::canCollapseRegion(int64_t RegionAddr) {
 441   if (!IsInitialized)
 442     initialize();
 443   if (!IsSimple)
 444     return false;
 445
 446   if (CollapsedRegions.count(RegionAddr))
 447     return true;
 448
 449   // Check if it is possible to readjust all accesses below RegionAddr
 450   if (!BlacklistedRegions.empty())
 451     return false;
 452
 453   return true;
 454 }
 455
 456 bool StackLayoutModifier::collapseRegion(MCInst *DeletedPush) {
 457   ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(*DeletedPush);
 458   if (!FIE)
 459     return false;
 460   int64_t RegionAddr = FIE->StackOffset;
 461   int64_t RegionSz = FIE->Size;
 462   return collapseRegion(DeletedPush, RegionAddr, RegionSz);
 463 }
 464
 465 bool StackLayoutModifier::collapseRegion(MCInst *Alloc, int64_t RegionAddr,
 466                                          int64_t RegionSz) {
 467   if (!canCollapseRegion(RegionAddr))
 468     return false;
 469
 470   assert(IsInitialized);
 471   StackAllocationAnalysis &SAA = Info.getStackAllocationAnalysis();
 472
 473   for (BinaryBasicBlock &BB : BF) {
 474     for (MCInst &Inst : BB) {
 475       if (!BC.MIB->hasAnnotation(Inst, getSlotTag()))
 476         continue;
 477       auto Slot =
 478           BC.MIB->getAnnotationAs<decltype(FrameIndexEntry::StackOffset)>(
 479               Inst, getSlotTag());
 480       if (!AvailableRegions.count(Slot))
 481         continue;
 482       // We need to ensure this access is affected by the deleted push
 483       if (!(*SAA.getStateBefore(Inst))[SAA.ExprToIdx[Alloc]])
 484         continue;
 485
 486       if (BC.MIB->isCFI(Inst)) {
 487         if (Slot > RegionAddr)
 488           continue;
 489         scheduleChange(Inst, WorklistItem(WorklistItem::AdjustCFI, RegionSz));
 490         continue;
 491       }
 492       ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(Inst);
 493       if (!FIE) {
 494         if (Slot > RegionAddr)
 495           continue;
 496         // SP update based on frame pointer
 497         scheduleChange(
 498             Inst, WorklistItem(WorklistItem::AdjustLoadStoreOffset, RegionSz));
 499         continue;
 500       }
 501
 502       if (Slot == RegionAddr) {
 503         BC.MIB->addAnnotation(Inst, "AccessesDeletedPos", 0U, AllocatorId);
 504         continue;
 505       }
 506       if (BC.MIB->isPush(Inst) || BC.MIB->isPop(Inst))
 507         continue;
 508
 509       if (FIE->StackPtrReg == BC.MIB->getStackPointer() && Slot < RegionAddr)
 510         continue;
 511
 512       if (FIE->StackPtrReg == BC.MIB->getFramePointer() && Slot > RegionAddr)
 513         continue;
 514
 515       scheduleChange(
 516           Inst, WorklistItem(WorklistItem::AdjustLoadStoreOffset, RegionSz));
 517     }
 518   }
 519
 520   CollapsedRegions.insert(RegionAddr);
 521   return true;
 522 }
 523
 524 void StackLayoutModifier::setOffsetForCollapsedAccesses(int64_t NewOffset) {
 525   for (BinaryBasicBlock &BB : BF) {
 526     for (MCInst &Inst : BB) {
 527       if (!BC.MIB->hasAnnotation(Inst, "AccessesDeletedPos"))
 528         continue;
 529       BC.MIB->removeAnnotation(Inst, "AccessesDeletedPos");
 530       scheduleChange(
 531           Inst, WorklistItem(WorklistItem::AdjustLoadStoreOffset, NewOffset));
 532     }
 533   }
 534 }
 535
 536 bool StackLayoutModifier::canInsertRegion(ProgramPoint P) {
 537   if (!IsInitialized)
 538     initialize();
 539   if (!IsSimple)
 540     return false;
 541
 542   StackPointerTracking &SPT = Info.getStackPointerTracking();
 543   int64_t RegionAddr = SPT.getStateBefore(P)->first;
 544   if (RegionAddr == SPT.SUPERPOSITION || RegionAddr == SPT.EMPTY)
 545     return false;
 546
 547   if (InsertedRegions.count(RegionAddr))
 548     return true;
 549
 550   // Check if we are going to screw up stack accesses at call sites that
 551   // pass parameters via stack
 552   if (!BlacklistedRegions.empty())
 553     return false;
 554
 555   return true;
 556 }
 557
 558 bool StackLayoutModifier::insertRegion(ProgramPoint P, int64_t RegionSz) {
 559   if (!canInsertRegion(P))
 560     return false;
 561
 562   assert(IsInitialized);
 563   StackPointerTracking &SPT = Info.getStackPointerTracking();
 564   // This RegionAddr is slightly different from the one seen in collapseRegion
 565   // This is the value of SP before the allocation the user wants to make.
 566   int64_t RegionAddr = SPT.getStateBefore(P)->first;
 567   if (RegionAddr == SPT.SUPERPOSITION || RegionAddr == SPT.EMPTY)
 568     return false;
 569
 570   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 571
 572   for (BinaryBasicBlock &BB : BF) {
 573     for (MCInst &Inst : BB) {
 574       if (!BC.MIB->hasAnnotation(Inst, getSlotTag()))
 575         continue;
 576       auto Slot =
 577           BC.MIB->getAnnotationAs<decltype(FrameIndexEntry::StackOffset)>(
 578               Inst, getSlotTag());
 579       if (!AvailableRegions.count(Slot))
 580         continue;
 581
 582       if (!(DA.doesADominateB(P, Inst)))
 583         continue;
 584
 585       if (BC.MIB->isCFI(Inst)) {
 586         if (Slot >= RegionAddr)
 587           continue;
 588         scheduleChange(Inst, WorklistItem(WorklistItem::AdjustCFI, -RegionSz));
 589         continue;
 590       }
 591       ErrorOr<const FrameIndexEntry &> FIE = FA.getFIEFor(Inst);
 592       if (!FIE) {
 593         if (Slot >= RegionAddr)
 594           continue;
 595         scheduleChange(
 596             Inst, WorklistItem(WorklistItem::AdjustLoadStoreOffset, -RegionSz));
 597         continue;
 598       }
 599
 600       if (FIE->StackPtrReg == BC.MIB->getStackPointer() && Slot < RegionAddr)
 601         continue;
 602       if (FIE->StackPtrReg == BC.MIB->getFramePointer() && Slot >= RegionAddr)
 603         continue;
 604       if (BC.MIB->isPush(Inst) || BC.MIB->isPop(Inst))
 605         continue;
 606       scheduleChange(
 607           Inst, WorklistItem(WorklistItem::AdjustLoadStoreOffset, -RegionSz));
 608     }
 609   }
 610
 611   InsertedRegions.insert(RegionAddr);
 612   return true;
 613 }
 614
 615 void StackLayoutModifier::performChanges() {
 616   std::set<uint32_t> ModifiedCFIIndices;
 617   for (BinaryBasicBlock &BB : BF) {
 618     for (MCInst &Inst : llvm::reverse(BB)) {
 619       if (BC.MIB->hasAnnotation(Inst, "AccessesDeletedPos")) {
 620         assert(BC.MIB->isPop(Inst) || BC.MIB->isPush(Inst));
 621         BC.MIB->removeAnnotation(Inst, "AccessesDeletedPos");
 622       }
 623       if (!BC.MIB->hasAnnotation(Inst, getTodoTag()))
 624         continue;
 625       auto &WList = BC.MIB->getAnnotationAs<std::vector<WorklistItem>>(
 626           Inst, getTodoTag());
 627       int64_t Adjustment = 0;
 628       WorklistItem::ActionType AdjustmentType = WorklistItem::None;
 629       for (WorklistItem &WI : WList) {
 630         if (WI.Action == WorklistItem::None)
 631           continue;
 632         assert(WI.Action == WorklistItem::AdjustLoadStoreOffset ||
 633                WI.Action == WorklistItem::AdjustCFI);
 634         assert((AdjustmentType == WorklistItem::None ||
 635                 AdjustmentType == WI.Action) &&
 636                "Conflicting actions requested at the same program point");
 637         AdjustmentType = WI.Action;
 638         Adjustment += WI.OffsetUpdate;
 639       }
 640       if (!Adjustment)
 641         continue;
 642       if (AdjustmentType != WorklistItem::AdjustLoadStoreOffset) {
 643         assert(BC.MIB->isCFI(Inst));
 644         uint32_t CFINum = Inst.getOperand(0).getImm();
 645         if (ModifiedCFIIndices.count(CFINum))
 646           continue;
 647         ModifiedCFIIndices.insert(CFINum);
 648         const MCCFIInstruction *CFI = BF.getCFIFor(Inst);
 649         const MCCFIInstruction::OpType Operation = CFI->getOperation();
 650         if (Operation == MCCFIInstruction::OpDefCfa ||
 651             Operation == MCCFIInstruction::OpDefCfaOffset)
 652           Adjustment = 0 - Adjustment;
 653         LLVM_DEBUG(dbgs() << "Changing CFI offset from " << CFI->getOffset()
 654                           << " to " << (CFI->getOffset() + Adjustment) << "\n");
 655         BF.mutateCFIOffsetFor(Inst, CFI->getOffset() + Adjustment);
 656         continue;
 657       }
 658       int32_t SrcImm = 0;
 659       MCPhysReg Reg = 0;
 660       MCPhysReg StackPtrReg = 0;
 661       int64_t StackOffset = 0;
 662       bool IsIndexed = false;
 663       bool IsLoad = false;
 664       bool IsStore = false;
 665       bool IsSimple = false;
 666       bool IsStoreFromReg = false;
 667       uint8_t Size = 0;
 668       bool Success = false;
 669       Success = BC.MIB->isStackAccess(Inst, IsLoad, IsStore, IsStoreFromReg,
 670                                       Reg, SrcImm, StackPtrReg, StackOffset,
 671                                       Size, IsSimple, IsIndexed);
 672       if (!Success) {
 673         // SP update based on FP value
 674         Success = BC.MIB->addToImm(Inst, Adjustment, &*BC.Ctx);
 675         assert(Success);
 676         continue;
 677       }
 678       assert(Success && IsSimple && !IsIndexed && (!IsStore || IsStoreFromReg));
 679       if (StackPtrReg != BC.MIB->getFramePointer())
 680         Adjustment = -Adjustment;
 681       if (IsLoad)
 682         BC.MIB->createRestoreFromStack(Inst, StackPtrReg,
 683                                        StackOffset + Adjustment, Reg, Size);
 684       else if (IsStore)
 685         BC.MIB->createSaveToStack(Inst, StackPtrReg, StackOffset + Adjustment,
 686                                   Reg, Size);
 687       LLVM_DEBUG({
 688         dbgs() << "Adjusted instruction: ";
 689         Inst.dump();
 690       });
 691     }
 692   }
 693 }
 694
 695 void StackLayoutModifier::initialize() {
 696   classifyStackAccesses();
 697   classifyCFIs();
 698   IsInitialized = true;
 699 }
 700
 701 std::atomic<std::uint64_t> ShrinkWrapping::SpillsMovedRegularMode{0};
 702 std::atomic<std::uint64_t> ShrinkWrapping::SpillsMovedPushPopMode{0};
 703 std::atomic<std::uint64_t> ShrinkWrapping::SpillsMovedDynamicCount{0};
 704 std::atomic<std::uint64_t> ShrinkWrapping::SpillsFailedDynamicCount{0};
 705 std::atomic<std::uint64_t> ShrinkWrapping::InstrDynamicCount{0};
 706 std::atomic<std::uint64_t> ShrinkWrapping::StoreDynamicCount{0};
 707
 708 using BBIterTy = BinaryBasicBlock::iterator;
 709
 710 void ShrinkWrapping::classifyCSRUses() {
 711   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 712   StackPointerTracking &SPT = Info.getStackPointerTracking();
 713   UsesByReg = std::vector<BitVector>(BC.MRI->getNumRegs(),
 714                                      BitVector(DA.NumInstrs, false));
 715
 716   const BitVector &FPAliases = BC.MIB->getAliases(BC.MIB->getFramePointer());
 717   for (BinaryBasicBlock &BB : BF) {
 718     for (MCInst &Inst : BB) {
 719       if (BC.MIB->isCFI(Inst))
 720         continue;
 721       BitVector BV = BitVector(BC.MRI->getNumRegs(), false);
 722       BC.MIB->getTouchedRegs(Inst, BV);
 723       BV &= CSA.CalleeSaved;
 724       for (int I : BV.set_bits()) {
 725         if (I == 0)
 726           continue;
 727         if (CSA.getSavedReg(Inst) != I && CSA.getRestoredReg(Inst) != I)
 728           UsesByReg[I].set(DA.ExprToIdx[&Inst]);
 729       }
 730       if (!SPT.HasFramePointer || !BC.MIB->isCall(Inst))
 731         continue;
 732       BV = CSA.CalleeSaved;
 733       BV &= FPAliases;
 734       for (int I : BV.set_bits())
 735         UsesByReg[I].set(DA.ExprToIdx[&Inst]);
 736     }
 737   }
 738 }
 739
 740 void ShrinkWrapping::pruneUnwantedCSRs() {
 741   BitVector ParamRegs = BC.MIB->getRegsUsedAsParams();
 742   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
 743     if (!CSA.CalleeSaved[I])
 744       continue;
 745     if (ParamRegs[I]) {
 746       CSA.CalleeSaved.reset(I);
 747       continue;
 748     }
 749     if (UsesByReg[I].empty()) {
 750       LLVM_DEBUG(
 751           dbgs()
 752           << "Dismissing Callee-Saved Reg because we found no uses of it:" << I
 753           << "\n");
 754       CSA.CalleeSaved.reset(I);
 755       continue;
 756     }
 757     if (!CSA.HasRestores[I]) {
 758       LLVM_DEBUG(
 759           dbgs() << "Dismissing Callee-Saved Reg because it does not have "
 760                     "restores:"
 761                  << I << "\n");
 762       CSA.CalleeSaved.reset(I);
 763     }
 764   }
 765 }
 766
 767 void ShrinkWrapping::computeSaveLocations() {
 768   BestSavePos = std::vector<std::vector<MCInst *>>(BC.MRI->getNumRegs());
 769   ReachingInsns<true> &RI = Info.getReachingInsnsBackwards();
 770   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 771   StackPointerTracking &SPT = Info.getStackPointerTracking();
 772
 773   LLVM_DEBUG(dbgs() << "Checking save/restore possibilities\n");
 774   for (BinaryBasicBlock &BB : BF) {
 775     LLVM_DEBUG(dbgs() << "\tNow at BB " << BB.getName() << "\n");
 776
 777     MCInst *First = BB.begin() != BB.end() ? &*BB.begin() : nullptr;
 778     if (!First)
 779       continue;
 780
 781     // Use reaching instructions to detect if we are inside a loop - if we
 782     // are, do not consider this BB as valid placement for saves.
 783     if (RI.isInLoop(BB))
 784       continue;
 785
 786     const std::pair<int, int> SPFP = *SPT.getStateBefore(*First);
 787     // If we don't know stack state at this point, bail
 788     if ((SPFP.first == SPT.SUPERPOSITION || SPFP.first == SPT.EMPTY) &&
 789         (SPFP.second == SPT.SUPERPOSITION || SPFP.second == SPT.EMPTY))
 790       continue;
 791
 792     for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
 793       if (!CSA.CalleeSaved[I])
 794         continue;
 795
 796       BitVector BBDominatedUses = BitVector(DA.NumInstrs, false);
 797       for (int J : UsesByReg[I].set_bits())
 798         if (DA.doesADominateB(*First, J))
 799           BBDominatedUses.set(J);
 800       LLVM_DEBUG(dbgs() << "\t\tBB " << BB.getName() << " dominates "
 801                         << BBDominatedUses.count() << " uses for reg " << I
 802                         << ". Total uses for reg is " << UsesByReg[I].count()
 803                         << "\n");
 804       BBDominatedUses &= UsesByReg[I];
 805       if (BBDominatedUses == UsesByReg[I]) {
 806         LLVM_DEBUG(dbgs() << "\t\t\tAdded " << BB.getName()
 807                           << " as a save pos for " << I << "\n");
 808         BestSavePos[I].push_back(First);
 809         LLVM_DEBUG({
 810           dbgs() << "Dominated uses are:\n";
 811           for (int J : UsesByReg[I].set_bits()) {
 812             dbgs() << "Idx " << J << ": ";
 813             BC.printInstruction(dbgs(), *DA.Expressions[J]);
 814             DA.Expressions[J]->dump();
 815           }
 816         });
 817       }
 818     }
 819   }
 820
 821   BestSaveCount = std::vector<std::vector<uint64_t>>(BC.MRI->getNumRegs());
 822
 823   auto &InsnToBB = Info.getInsnToBBMap();
 824   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
 825     if (!CSA.CalleeSaved[I])
 826       continue;
 827
 828     std::stable_sort(BestSavePos[I].begin(), BestSavePos[I].end(),
 829                      [&](const MCInst *A, const MCInst *B) {
 830                        const BinaryBasicBlock *BBA = InsnToBB[A];
 831                        const BinaryBasicBlock *BBB = InsnToBB[B];
 832                        const uint64_t CountA = BBA->getKnownExecutionCount();
 833                        const uint64_t CountB = BBB->getKnownExecutionCount();
 834                        return CountB < CountA;
 835                      });
 836
 837     for (MCInst *Pos : BestSavePos[I]) {
 838       const BinaryBasicBlock *BB = InsnToBB[Pos];
 839       const uint64_t Count = BB->getKnownExecutionCount();
 840       BestSaveCount[I].push_back(Count);
 841     }
 842   }
 843 }
 844
 845 void ShrinkWrapping::computeDomOrder() {
 846   DomOrder = std::vector<MCPhysReg>(BC.MRI->getNumRegs(), 0);
 847   std::vector<MCPhysReg> Order;
 848   for (MCPhysReg I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
 849     Order.push_back(I);
 850   }
 851
 852   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 853   auto &InsnToBB = Info.getInsnToBBMap();
 854   llvm::sort(Order, [&](const MCPhysReg &A, const MCPhysReg &B) {
 855     BinaryBasicBlock *BBA =
 856         BestSavePos[A].size() ? InsnToBB[BestSavePos[A].back()] : nullptr;
 857     BinaryBasicBlock *BBB =
 858         BestSavePos[B].size() ? InsnToBB[BestSavePos[B].back()] : nullptr;
 859     if (BBA == BBB)
 860       return A < B;
 861     if (!BBA && BBB)
 862       return false;
 863     if (BBA && !BBB)
 864       return true;
 865     if (DA.doesADominateB(*BestSavePos[A].back(), *BestSavePos[B].back()))
 866       return true;
 867     if (DA.doesADominateB(*BestSavePos[B].back(), *BestSavePos[A].back()))
 868       return false;
 869     return A < B;
 870   });
 871
 872   for (MCPhysReg I = 0, E = BC.MRI->getNumRegs(); I != E; ++I)
 873     DomOrder[Order[I]] = I;
 874 }
 875
 876 bool ShrinkWrapping::isBestSavePosCold(unsigned CSR, MCInst *&BestPosSave,
 877                                        uint64_t &TotalEstimatedWin) {
 878   const uint64_t CurSavingCost = CSA.SavingCost[CSR];
 879   if (!CSA.CalleeSaved[CSR])
 880     return false;
 881
 882   assert(BestSaveCount[CSR].size() == BestSavePos[CSR].size() &&
 883          "save position vectors out of sync");
 884   if (BestSaveCount[CSR].empty())
 885     return false;
 886
 887   const uint64_t BestCount = BestSaveCount[CSR].back();
 888   BestPosSave = BestSavePos[CSR].back();
 889   if (BestCount >= (opts::ShrinkWrappingThreshold / 100.0) * CurSavingCost)
 890     return false;
 891
 892   LLVM_DEBUG({
 893     auto &InsnToBB = Info.getInsnToBBMap();
 894     dbgs() << "Better position for saves found in func " << BF.getPrintName()
 895            << " count << " << BF.getKnownExecutionCount() << "\n";
 896     dbgs() << "Reg: " << CSR << "; New BB: " << InsnToBB[BestPosSave]->getName()
 897            << " Freq reduction: " << (CurSavingCost - BestCount) << "\n";
 898   });
 899
 900   TotalEstimatedWin = CurSavingCost - BestCount;
 901   return true;
 902 }
 903
 904 /// Auxiliary function used to create basic blocks for critical edges and update
 905 /// the dominance frontier with these new locations
 906 void ShrinkWrapping::splitFrontierCritEdges(
 907     BinaryFunction *Func, SmallVector<ProgramPoint, 4> &Frontier,
 908     const SmallVector<bool, 4> &IsCritEdge,
 909     const SmallVector<BinaryBasicBlock *, 4> &From,
 910     const SmallVector<SmallVector<BinaryBasicBlock *, 4>, 4> &To) {
 911   LLVM_DEBUG(dbgs() << "splitFrontierCritEdges: Now handling func "
 912                     << BF.getPrintName() << "\n");
 913   // For every FromBB, there might be one or more critical edges, with
 914   // To[I] containing destination BBs. It's important to memorize
 915   // the original size of the Frontier as we may append to it while splitting
 916   // critical edges originating with blocks with multiple destinations.
 917   for (size_t I = 0, IE = Frontier.size(); I < IE; ++I) {
 918     if (!IsCritEdge[I])
 919       continue;
 920     if (To[I].empty())
 921       continue;
 922     BinaryBasicBlock *FromBB = From[I];
 923     LLVM_DEBUG(dbgs() << " - Now handling FrontierBB " << FromBB->getName()
 924                       << "\n");
 925     // Split edge for every DestinationBBs
 926     for (size_t DI = 0, DIE = To[I].size(); DI < DIE; ++DI) {
 927       BinaryBasicBlock *DestinationBB = To[I][DI];
 928       LLVM_DEBUG(dbgs() << "   - Dest : " << DestinationBB->getName() << "\n");
 929       BinaryBasicBlock *NewBB = Func->splitEdge(FromBB, DestinationBB);
 930       // Insert dummy instruction so this BB is never empty (we need this for
 931       // PredictiveStackPointerTracking to work, since it annotates instructions
 932       // and not BBs).
 933       if (NewBB->empty()) {
 934         MCInst NewInst;
 935         BC.MIB->createNoop(NewInst);
 936         NewBB->addInstruction(std::move(NewInst));
 937         scheduleChange(&*NewBB->begin(), WorklistItem(WorklistItem::Erase, 0));
 938       }
 939
 940       // Update frontier
 941       ProgramPoint NewFrontierPP = ProgramPoint::getLastPointAt(*NewBB);
 942       if (DI == 0) {
 943         // Update frontier inplace
 944         Frontier[I] = NewFrontierPP;
 945         LLVM_DEBUG(dbgs() << "   - Update frontier with " << NewBB->getName()
 946                           << '\n');
 947       } else {
 948         // Append new frontier to the end of the list
 949         Frontier.push_back(NewFrontierPP);
 950         LLVM_DEBUG(dbgs() << "   - Append frontier " << NewBB->getName()
 951                           << '\n');
 952       }
 953     }
 954   }
 955 }
 956
 957 SmallVector<ProgramPoint, 4>
 958 ShrinkWrapping::doRestorePlacement(MCInst *BestPosSave, unsigned CSR,
 959                                    uint64_t TotalEstimatedWin) {
 960   SmallVector<ProgramPoint, 4> Frontier;
 961   SmallVector<bool, 4> IsCritEdge;
 962   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
 963
 964   SmallVector<BinaryBasicBlock *, 4> CritEdgesFrom;
 965   SmallVector<SmallVector<BinaryBasicBlock *, 4>, 4> CritEdgesTo;
 966   // In case of a critical edge, we need to create extra BBs to host restores
 967   // into edges transitioning to the dominance frontier, otherwise we pull these
 968   // restores to inside the dominated area.
 969   Frontier = DA.getDominanceFrontierFor(*BestPosSave).takeVector();
 970   LLVM_DEBUG({
 971     dbgs() << "Dumping dominance frontier for ";
 972     BC.printInstruction(dbgs(), *BestPosSave);
 973     for (ProgramPoint &PP : Frontier)
 974       if (PP.isInst())
 975         BC.printInstruction(dbgs(), *PP.getInst());
 976       else
 977         dbgs() << PP.getBB()->getName() << "\n";
 978   });
 979   for (ProgramPoint &PP : Frontier) {
 980     bool HasCritEdges = false;
 981     if (PP.isInst() && BC.MIB->isTerminator(*PP.getInst()) &&
 982         doesInstUsesCSR(*PP.getInst(), CSR)) {
 983       Frontier.clear();
 984       return Frontier;
 985     }
 986     BinaryBasicBlock *FrontierBB = Info.getParentBB(PP);
 987     CritEdgesFrom.emplace_back(FrontierBB);
 988     CritEdgesTo.emplace_back(0);
 989     SmallVector<BinaryBasicBlock *, 4> &Dests = CritEdgesTo.back();
 990     // Check for invoke instructions at the dominance frontier, which indicates
 991     // the landing pad is not dominated.
 992     if (PP.isInst() && BC.MIB->isInvoke(*PP.getInst())) {
 993       Frontier.clear();
 994       return Frontier;
 995     }
 996     doForAllSuccs(*FrontierBB, [&](ProgramPoint P) {
 997       if (!DA.doesADominateB(*BestPosSave, P)) {
 998         Dests.emplace_back(Info.getParentBB(P));
 999         return;
1000       }
1001       HasCritEdges = true;
1002     });
1003     IsCritEdge.push_back(HasCritEdges);
1004   }
1005   // Restores cannot be placed in empty BBs because we have a dataflow
1006   // analysis that depends on insertions happening before real instructions
1007   // (PredictiveStackPointerTracking). Detect now for empty BBs and add a
1008   // dummy nop that is scheduled to be removed later.
1009   bool InvalidateRequired = false;
1010   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
1011     if (BB->size() != 0)
1012       continue;
1013     MCInst NewInst;
1014     BC.MIB->createNoop(NewInst);
1015     auto II = BB->addInstruction(std::move(NewInst));
1016     scheduleChange(&*II, WorklistItem(WorklistItem::Erase, 0));
1017     InvalidateRequired = true;
1018   }
1019   if (std::accumulate(IsCritEdge.begin(), IsCritEdge.end(), 0)) {
1020     LLVM_DEBUG({
1021       dbgs() << "Now detected critical edges in the following frontier:\n";
1022       for (ProgramPoint &PP : Frontier) {
1023         if (PP.isBB()) {
1024           dbgs() << "  BB: " << PP.getBB()->getName() << "\n";
1025         } else {
1026           dbgs() << "  Inst: ";
1027           PP.getInst()->dump();
1028         }
1029       }
1030     });
1031     splitFrontierCritEdges(&BF, Frontier, IsCritEdge, CritEdgesFrom,
1032                            CritEdgesTo);
1033     InvalidateRequired = true;
1034   }
1035   if (InvalidateRequired) {
1036     // BitVectors that represent all insns of the function are invalid now
1037     // since we changed BBs/Insts. Re-run steps that depend on pointers being
1038     // valid
1039     Info.invalidateAll();
1040     classifyCSRUses();
1041   }
1042   return Frontier;
1043 }
1044
1045 bool ShrinkWrapping::validatePushPopsMode(unsigned CSR, MCInst *BestPosSave,
1046                                           int64_t SaveOffset) {
1047   if (FA.requiresAlignment(BF)) {
1048     LLVM_DEBUG({
1049       dbgs() << "Reg " << CSR
1050              << " is not using push/pops due to function "
1051                 "alignment requirements.\n";
1052     });
1053     return false;
1054   }
1055   if (FA.hasStackArithmetic(BF)) {
1056     LLVM_DEBUG({
1057       dbgs() << "Reg " << CSR
1058              << " is not using push/pops due to function "
1059                 "taking the address of a stack position.\n";
1060     });
1061     return false;
1062   }
1063   for (MCInst *Save : CSA.getSavesByReg(CSR)) {
1064     if (!SLM.canCollapseRegion(Save)) {
1065       LLVM_DEBUG(dbgs() << "Reg " << CSR << " cannot collapse region.\n");
1066       return false;
1067     }
1068   }
1069   // Abort if one of the restores for this CSR is not a POP.
1070   for (MCInst *Load : CSA.getRestoresByReg(CSR)) {
1071     if (!BC.MIB->isPop(*Load)) {
1072       LLVM_DEBUG(dbgs() << "Reg " << CSR << " has a mismatching restore.\n");
1073       return false;
1074     }
1075   }
1076
1077   StackPointerTracking &SPT = Info.getStackPointerTracking();
1078   // Abort if we are inserting a push into an entry BB (offset -8) and this
1079   // func sets up a frame pointer.
1080   if (!SLM.canInsertRegion(BestPosSave) || SaveOffset == SPT.SUPERPOSITION ||
1081       SaveOffset == SPT.EMPTY || (SaveOffset == -8 && SPT.HasFramePointer)) {
1082     LLVM_DEBUG({
1083       dbgs() << "Reg " << CSR
1084              << " cannot insert region or we are "
1085                 "trying to insert a push into entry bb.\n";
1086     });
1087     return false;
1088   }
1089   return true;
1090 }
1091
1092 SmallVector<ProgramPoint, 4> ShrinkWrapping::fixPopsPlacements(
1093     const SmallVector<ProgramPoint, 4> &RestorePoints, int64_t SaveOffset,
1094     unsigned CSR) {
1095   SmallVector<ProgramPoint, 4> FixedRestorePoints = RestorePoints;
1096   // Moving pop locations to the correct sp offset
1097   ReachingInsns<true> &RI = Info.getReachingInsnsBackwards();
1098   StackPointerTracking &SPT = Info.getStackPointerTracking();
1099   for (ProgramPoint &PP : FixedRestorePoints) {
1100     BinaryBasicBlock *BB = Info.getParentBB(PP);
1101     bool Found = false;
1102     if (SPT.getStateAt(ProgramPoint::getLastPointAt(*BB))->first ==
1103         SaveOffset) {
1104       BitVector BV = *RI.getStateAt(ProgramPoint::getLastPointAt(*BB));
1105       BV &= UsesByReg[CSR];
1106       if (!BV.any()) {
1107         Found = true;
1108         PP = BB;
1109         continue;
1110       }
1111     }
1112     for (MCInst &Inst : llvm::reverse(*BB)) {
1113       if (SPT.getStateBefore(Inst)->first == SaveOffset) {
1114         BitVector BV = *RI.getStateAt(Inst);
1115         BV &= UsesByReg[CSR];
1116         if (!BV.any()) {
1117           Found = true;
1118           PP = &Inst;
1119           break;
1120         }
1121       }
1122     }
1123     if (!Found) {
1124       LLVM_DEBUG({
1125         dbgs() << "Could not find restore insertion point for " << CSR
1126                << ", falling back to load/store mode\n";
1127       });
1128       FixedRestorePoints.clear();
1129       return FixedRestorePoints;
1130     }
1131   }
1132   return FixedRestorePoints;
1133 }
1134
1135 void ShrinkWrapping::scheduleOldSaveRestoresRemoval(unsigned CSR,
1136                                                     bool UsePushPops) {
1137
1138   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
1139     std::vector<MCInst *> CFIs;
1140     for (MCInst &Inst : llvm::reverse(*BB)) {
1141       if (BC.MIB->isCFI(Inst)) {
1142         // Delete all offset CFIs related to this CSR
1143         if (SLM.getOffsetCFIReg(Inst) == CSR) {
1144           HasDeletedOffsetCFIs[CSR] = true;
1145           scheduleChange(&Inst, WorklistItem(WorklistItem::Erase, CSR));
1146           continue;
1147         }
1148         CFIs.push_back(&Inst);
1149         continue;
1150       }
1151
1152       uint16_t SavedReg = CSA.getSavedReg(Inst);
1153       uint16_t RestoredReg = CSA.getRestoredReg(Inst);
1154       if (SavedReg != CSR && RestoredReg != CSR) {
1155         CFIs.clear();
1156         continue;
1157       }
1158
1159       scheduleChange(&Inst, WorklistItem(UsePushPops
1160                                              ? WorklistItem::Erase
1161                                              : WorklistItem::ChangeToAdjustment,
1162                                          CSR));
1163
1164       // Delete associated CFIs
1165       const bool RecordDeletedPushCFIs =
1166           SavedReg == CSR && DeletedPushCFIs[CSR].empty();
1167       const bool RecordDeletedPopCFIs =
1168           RestoredReg == CSR && DeletedPopCFIs[CSR].empty();
1169       for (MCInst *CFI : CFIs) {
1170         const MCCFIInstruction *MCCFI = BF.getCFIFor(*CFI);
1171         // Do not touch these...
1172         if (MCCFI->getOperation() == MCCFIInstruction::OpRestoreState ||
1173             MCCFI->getOperation() == MCCFIInstruction::OpRememberState)
1174           continue;
1175         scheduleChange(CFI, WorklistItem(WorklistItem::Erase, CSR));
1176         if (RecordDeletedPushCFIs) {
1177           // Do not record this to be replayed later because we are going to
1178           // rebuild it.
1179           if (MCCFI->getOperation() == MCCFIInstruction::OpDefCfaOffset)
1180             continue;
1181           DeletedPushCFIs[CSR].push_back(CFI->getOperand(0).getImm());
1182         }
1183         if (RecordDeletedPopCFIs) {
1184           if (MCCFI->getOperation() == MCCFIInstruction::OpDefCfaOffset)
1185             continue;
1186           DeletedPopCFIs[CSR].push_back(CFI->getOperand(0).getImm());
1187         }
1188       }
1189       CFIs.clear();
1190     }
1191   }
1192 }
1193
1194 bool ShrinkWrapping::doesInstUsesCSR(const MCInst &Inst, uint16_t CSR) {
1195   if (BC.MIB->isCFI(Inst) || CSA.getSavedReg(Inst) == CSR ||
1196       CSA.getRestoredReg(Inst) == CSR)
1197     return false;
1198   BitVector BV = BitVector(BC.MRI->getNumRegs(), false);
1199   BC.MIB->getTouchedRegs(Inst, BV);
1200   return BV[CSR];
1201 }
1202
1203 void ShrinkWrapping::scheduleSaveRestoreInsertions(
1204     unsigned CSR, MCInst *BestPosSave,
1205     SmallVector<ProgramPoint, 4> &RestorePoints, bool UsePushPops) {
1206   auto &InsnToBB = Info.getInsnToBBMap();
1207   const FrameIndexEntry *FIESave = CSA.SaveFIEByReg[CSR];
1208   const FrameIndexEntry *FIELoad = CSA.LoadFIEByReg[CSR];
1209   assert(FIESave && FIELoad && "Invalid CSR");
1210
1211   LLVM_DEBUG({
1212     dbgs() << "Scheduling save insertion at: ";
1213     BestPosSave->dump();
1214   });
1215
1216   scheduleChange(BestPosSave,
1217                  UsePushPops ? WorklistItem::InsertPushOrPop
1218                              : WorklistItem::InsertLoadOrStore,
1219                  *FIESave, CSR);
1220
1221   for (ProgramPoint &PP : RestorePoints) {
1222     BinaryBasicBlock *FrontierBB = Info.getParentBB(PP);
1223     LLVM_DEBUG({
1224       dbgs() << "Scheduling restore insertion at: ";
1225       if (PP.isInst())
1226         PP.getInst()->dump();
1227       else
1228         dbgs() << PP.getBB()->getName() << "\n";
1229     });
1230     MCInst *Term =
1231         FrontierBB->getTerminatorBefore(PP.isInst() ? PP.getInst() : nullptr);
1232     if (Term)
1233       PP = Term;
1234     bool PrecededByPrefix = false;
1235     if (PP.isInst()) {
1236       auto Iter = FrontierBB->findInstruction(PP.getInst());
1237       if (Iter != FrontierBB->end() && Iter != FrontierBB->begin()) {
1238         --Iter;
1239         PrecededByPrefix = BC.MIB->isPrefix(*Iter);
1240       }
1241     }
1242     if (PP.isInst() &&
1243         (doesInstUsesCSR(*PP.getInst(), CSR) || PrecededByPrefix)) {
1244       assert(!InsnToBB[PP.getInst()]->hasTerminatorAfter(PP.getInst()) &&
1245              "cannot move to end of bb");
1246       scheduleChange(InsnToBB[PP.getInst()],
1247                      UsePushPops ? WorklistItem::InsertPushOrPop
1248                                  : WorklistItem::InsertLoadOrStore,
1249                      *FIELoad, CSR);
1250       continue;
1251     }
1252     scheduleChange(PP,
1253                    UsePushPops ? WorklistItem::InsertPushOrPop
1254                                : WorklistItem::InsertLoadOrStore,
1255                    *FIELoad, CSR);
1256   }
1257 }
1258
1259 void ShrinkWrapping::moveSaveRestores() {
1260   bool DisablePushPopMode = false;
1261   bool UsedPushPopMode = false;
1262   // Keeps info about successfully moved regs: reg index, save position and
1263   // save size
1264   std::vector<std::tuple<unsigned, MCInst *, size_t>> MovedRegs;
1265   uint64_t TotalEstimatedWin = 0;
1266
1267   computeDomOrder();
1268   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
1269     MCInst *BestPosSave = nullptr;
1270     uint64_t EstimatedWin = 0;
1271     SmallVector<ProgramPoint, 4> RestorePoints;
1272     while (RestorePoints.empty() &&
1273            isBestSavePosCold(I, BestPosSave, EstimatedWin)) {
1274       RestorePoints = doRestorePlacement(BestPosSave, I, EstimatedWin);
1275       if (RestorePoints.empty()) {
1276         LLVM_DEBUG({
1277           dbgs() << "Dropping opportunity because restore placement failed"
1278                     " -- total est. freq reduc: "
1279                  << EstimatedWin << ". Will try "
1280                  << (BestSaveCount[I].size() - 1) << " more times.\n";
1281         });
1282         BestSaveCount[I].pop_back();
1283         BestSavePos[I].pop_back();
1284         computeDomOrder();
1285       }
1286     }
1287     if (RestorePoints.empty()) {
1288       SpillsFailedDynamicCount += EstimatedWin;
1289       continue;
1290     }
1291
1292     const FrameIndexEntry *FIESave = CSA.SaveFIEByReg[I];
1293     const FrameIndexEntry *FIELoad = CSA.LoadFIEByReg[I];
1294     (void)FIELoad;
1295     assert(FIESave && FIELoad);
1296     StackPointerTracking &SPT = Info.getStackPointerTracking();
1297     const std::pair<int, int> SPFP = *SPT.getStateBefore(*BestPosSave);
1298     int SaveOffset = SPFP.first;
1299     uint8_t SaveSize = FIESave->Size;
1300
1301     // If we don't know stack state at this point, bail
1302     if ((SPFP.first == SPT.SUPERPOSITION || SPFP.first == SPT.EMPTY) &&
1303         (SPFP.second == SPT.SUPERPOSITION || SPFP.second == SPT.EMPTY)) {
1304       SpillsFailedDynamicCount += EstimatedWin;
1305       continue;
1306     }
1307
1308     // Operation mode: if true, will insert push/pops instead of loads/restores
1309     bool UsePushPops = validatePushPopsMode(I, BestPosSave, SaveOffset);
1310
1311     if (UsePushPops) {
1312       SmallVector<ProgramPoint, 4> FixedRestorePoints =
1313           fixPopsPlacements(RestorePoints, SaveOffset, I);
1314       if (FixedRestorePoints.empty())
1315         UsePushPops = false;
1316       else
1317         RestorePoints = FixedRestorePoints;
1318     }
1319
1320     // Disable push-pop mode for all CSRs in this function
1321     if (!UsePushPops)
1322       DisablePushPopMode = true;
1323     else
1324       UsedPushPopMode = true;
1325
1326     scheduleOldSaveRestoresRemoval(I, UsePushPops);
1327     scheduleSaveRestoreInsertions(I, BestPosSave, RestorePoints, UsePushPops);
1328     MovedRegs.emplace_back(std::make_tuple(I, BestPosSave, SaveSize));
1329     TotalEstimatedWin += EstimatedWin;
1330   }
1331
1332   // Revert push-pop mode if it failed for a single CSR
1333   if (DisablePushPopMode && UsedPushPopMode) {
1334     UsedPushPopMode = false;
1335     for (BinaryBasicBlock &BB : BF) {
1336       auto WRI = Todo.find(&BB);
1337       if (WRI != Todo.end()) {
1338         std::vector<WorklistItem> &TodoList = WRI->second;
1339         for (WorklistItem &Item : TodoList)
1340           if (Item.Action == WorklistItem::InsertPushOrPop)
1341             Item.Action = WorklistItem::InsertLoadOrStore;
1342       }
1343       for (MCInst &Inst : llvm::reverse(BB)) {
1344         auto TodoList = BC.MIB->tryGetAnnotationAs<std::vector<WorklistItem>>(
1345             Inst, getAnnotationIndex());
1346         if (!TodoList)
1347           continue;
1348         bool isCFI = BC.MIB->isCFI(Inst);
1349         for (WorklistItem &Item : *TodoList) {
1350           if (Item.Action == WorklistItem::InsertPushOrPop)
1351             Item.Action = WorklistItem::InsertLoadOrStore;
1352           if (!isCFI && Item.Action == WorklistItem::Erase)
1353             Item.Action = WorklistItem::ChangeToAdjustment;
1354         }
1355       }
1356     }
1357   }
1358   SpillsMovedDynamicCount += TotalEstimatedWin;
1359
1360   // Update statistics
1361   if (!UsedPushPopMode) {
1362     SpillsMovedRegularMode += MovedRegs.size();
1363     return;
1364   }
1365
1366   // Schedule modifications to stack-accessing instructions via
1367   // StackLayoutModifier.
1368   SpillsMovedPushPopMode += MovedRegs.size();
1369   for (std::tuple<unsigned, MCInst *, size_t> &I : MovedRegs) {
1370     unsigned RegNdx;
1371     MCInst *SavePos;
1372     size_t SaveSize;
1373     std::tie(RegNdx, SavePos, SaveSize) = I;
1374     for (MCInst *Save : CSA.getSavesByReg(RegNdx))
1375       SLM.collapseRegion(Save);
1376     SLM.insertRegion(SavePos, SaveSize);
1377   }
1378 }
1379
1380 namespace {
1381 /// Helper function to identify whether two basic blocks created by splitting
1382 /// a critical edge have the same contents.
1383 bool isIdenticalSplitEdgeBB(const BinaryContext &BC, const BinaryBasicBlock &A,
1384                             const BinaryBasicBlock &B) {
1385   if (A.succ_size() != B.succ_size())
1386     return false;
1387   if (A.succ_size() != 1)
1388     return false;
1389
1390   if (*A.succ_begin() != *B.succ_begin())
1391     return false;
1392
1393   if (A.size() != B.size())
1394     return false;
1395
1396   // Compare instructions
1397   auto I = A.begin(), E = A.end();
1398   auto OtherI = B.begin(), OtherE = B.end();
1399   while (I != E && OtherI != OtherE) {
1400     if (I->getOpcode() != OtherI->getOpcode())
1401       return false;
1402     if (!BC.MIB->equals(*I, *OtherI, [](const MCSymbol *A, const MCSymbol *B) {
1403           return true;
1404         }))
1405       return false;
1406     ++I;
1407     ++OtherI;
1408   }
1409   return true;
1410 }
1411 } // namespace
1412
1413 bool ShrinkWrapping::foldIdenticalSplitEdges() {
1414   bool Changed = false;
1415   for (auto Iter = BF.begin(); Iter != BF.end(); ++Iter) {
1416     BinaryBasicBlock &BB = *Iter;
1417     if (!BB.getName().starts_with(".LSplitEdge"))
1418       continue;
1419     for (BinaryBasicBlock &RBB : llvm::reverse(BF)) {
1420       if (&RBB == &BB)
1421         break;
1422       if (!RBB.getName().starts_with(".LSplitEdge") || !RBB.isValid() ||
1423           !isIdenticalSplitEdgeBB(BC, *Iter, RBB))
1424         continue;
1425       assert(RBB.pred_size() == 1 && "Invalid split edge BB");
1426       BinaryBasicBlock *Pred = *RBB.pred_begin();
1427       uint64_t OrigCount = Pred->branch_info_begin()->Count;
1428       uint64_t OrigMispreds = Pred->branch_info_begin()->MispredictedCount;
1429       BF.replaceJumpTableEntryIn(Pred, &RBB, &BB);
1430       Pred->replaceSuccessor(&RBB, &BB, OrigCount, OrigMispreds);
1431       Changed = true;
1432       // Remove the block from CFG
1433       RBB.markValid(false);
1434     }
1435   }
1436
1437   return Changed;
1438 }
1439
1440 namespace {
1441
1442 // A special StackPointerTracking that compensates for our future plans
1443 // in removing/adding insn.
1444 class PredictiveStackPointerTracking
1445     : public StackPointerTrackingBase<PredictiveStackPointerTracking> {
1446   friend class DataflowAnalysis<PredictiveStackPointerTracking,
1447                                 std::pair<int, int>>;
1448   decltype(ShrinkWrapping::Todo) &TodoMap;
1449   DataflowInfoManager &Info;
1450
1451   std::optional<unsigned> AnnotationIndex;
1452
1453 protected:
1454   void compNextAux(const MCInst &Point,
1455                    const std::vector<ShrinkWrapping::WorklistItem> &TodoItems,
1456                    std::pair<int, int> &Res) {
1457     for (const ShrinkWrapping::WorklistItem &Item : TodoItems) {
1458       if (Item.Action == ShrinkWrapping::WorklistItem::Erase &&
1459           BC.MIB->isPush(Point)) {
1460         Res.first += BC.MIB->getPushSize(Point);
1461         continue;
1462       }
1463       if (Item.Action == ShrinkWrapping::WorklistItem::Erase &&
1464           BC.MIB->isPop(Point)) {
1465         Res.first -= BC.MIB->getPopSize(Point);
1466         continue;
1467       }
1468       if (Item.Action == ShrinkWrapping::WorklistItem::InsertPushOrPop &&
1469           Item.FIEToInsert.IsStore) {
1470         Res.first -= Item.FIEToInsert.Size;
1471         continue;
1472       }
1473       if (Item.Action == ShrinkWrapping::WorklistItem::InsertPushOrPop &&
1474           Item.FIEToInsert.IsLoad) {
1475         Res.first += Item.FIEToInsert.Size;
1476         continue;
1477       }
1478     }
1479   }
1480
1481   std::pair<int, int> computeNext(const MCInst &Point,
1482                                   const std::pair<int, int> &Cur) {
1483     std::pair<int, int> Res =
1484         StackPointerTrackingBase<PredictiveStackPointerTracking>::computeNext(
1485             Point, Cur);
1486     if (Res.first == StackPointerTracking::SUPERPOSITION ||
1487         Res.first == StackPointerTracking::EMPTY)
1488       return Res;
1489     auto TodoItems =
1490         BC.MIB->tryGetAnnotationAs<std::vector<ShrinkWrapping::WorklistItem>>(
1491             Point, ShrinkWrapping::getAnnotationName());
1492     if (TodoItems)
1493       compNextAux(Point, *TodoItems, Res);
1494     auto &InsnToBBMap = Info.getInsnToBBMap();
1495     if (&*InsnToBBMap[&Point]->rbegin() != &Point)
1496       return Res;
1497     auto WRI = TodoMap.find(InsnToBBMap[&Point]);
1498     if (WRI == TodoMap.end())
1499       return Res;
1500     compNextAux(Point, WRI->second, Res);
1501     return Res;
1502   }
1503
1504   StringRef getAnnotationName() const {
1505     return StringRef("PredictiveStackPointerTracking");
1506   }
1507
1508 public:
1509   PredictiveStackPointerTracking(BinaryFunction &BF,
1510                                  decltype(ShrinkWrapping::Todo) &TodoMap,
1511                                  DataflowInfoManager &Info,
1512                                  MCPlusBuilder::AllocatorIdTy AllocatorId = 0)
1513       : StackPointerTrackingBase<PredictiveStackPointerTracking>(BF,
1514                                                                  AllocatorId),
1515         TodoMap(TodoMap), Info(Info) {}
1516
1517   void run() {
1518     StackPointerTrackingBase<PredictiveStackPointerTracking>::run();
1519   }
1520 };
1521
1522 } // end anonymous namespace
1523
1524 void ShrinkWrapping::insertUpdatedCFI(unsigned CSR, int SPValPush,
1525                                       int SPValPop) {
1526   MCInst *SavePoint = nullptr;
1527   for (BinaryBasicBlock &BB : BF) {
1528     for (MCInst &Inst : llvm::reverse(BB)) {
1529       int32_t SrcImm = 0;
1530       MCPhysReg Reg = 0;
1531       MCPhysReg StackPtrReg = 0;
1532       int64_t StackOffset = 0;
1533       bool IsIndexed = false;
1534       bool IsLoad = false;
1535       bool IsStore = false;
1536       bool IsSimple = false;
1537       bool IsStoreFromReg = false;
1538       uint8_t Size = 0;
1539       if (!BC.MIB->isStackAccess(Inst, IsLoad, IsStore, IsStoreFromReg, Reg,
1540                                  SrcImm, StackPtrReg, StackOffset, Size,
1541                                  IsSimple, IsIndexed))
1542         continue;
1543       if (Reg != CSR || !IsStore || !IsSimple)
1544         continue;
1545       SavePoint = &Inst;
1546       break;
1547     }
1548     if (SavePoint)
1549       break;
1550   }
1551   assert(SavePoint);
1552   LLVM_DEBUG({
1553     dbgs() << "Now using as save point for reg " << CSR << " :";
1554     SavePoint->dump();
1555   });
1556   bool PrevAffectedZone = false;
1557   BinaryBasicBlock *PrevBB = nullptr;
1558   DominatorAnalysis<false> &DA = Info.getDominatorAnalysis();
1559   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
1560     if (BB->size() == 0)
1561       continue;
1562     const bool InAffectedZoneAtEnd = DA.count(*BB->rbegin(), *SavePoint);
1563     const bool InAffectedZoneAtBegin =
1564         (*DA.getStateBefore(*BB->begin()))[DA.ExprToIdx[SavePoint]];
1565     bool InAffectedZone = InAffectedZoneAtBegin;
1566     for (auto InstIter = BB->begin(); InstIter != BB->end(); ++InstIter) {
1567       const bool CurZone = DA.count(*InstIter, *SavePoint);
1568       if (InAffectedZone != CurZone) {
1569         auto InsertionIter = InstIter;
1570         ++InsertionIter;
1571         InAffectedZone = CurZone;
1572         if (InAffectedZone)
1573           InstIter = insertCFIsForPushOrPop(*BB, InsertionIter, CSR, true, 0,
1574                                             SPValPop);
1575         else
1576           InstIter = insertCFIsForPushOrPop(*BB, InsertionIter, CSR, false, 0,
1577                                             SPValPush);
1578         --InstIter;
1579       }
1580     }
1581     // Are we at the first basic block or hot-cold split point?
1582     if (!PrevBB || (BF.isSplit() && BB->isCold() != PrevBB->isCold())) {
1583       if (InAffectedZoneAtBegin)
1584         insertCFIsForPushOrPop(*BB, BB->begin(), CSR, true, 0, SPValPush);
1585     } else if (InAffectedZoneAtBegin != PrevAffectedZone) {
1586       if (InAffectedZoneAtBegin)
1587         insertCFIsForPushOrPop(*PrevBB, PrevBB->end(), CSR, true, 0, SPValPush);
1588       else
1589         insertCFIsForPushOrPop(*PrevBB, PrevBB->end(), CSR, false, 0, SPValPop);
1590     }
1591     PrevAffectedZone = InAffectedZoneAtEnd;
1592     PrevBB = BB;
1593   }
1594 }
1595
1596 void ShrinkWrapping::rebuildCFIForSP() {
1597   for (BinaryBasicBlock &BB : BF) {
1598     for (MCInst &Inst : BB) {
1599       if (!BC.MIB->isCFI(Inst))
1600         continue;
1601       const MCCFIInstruction *CFI = BF.getCFIFor(Inst);
1602       if (CFI->getOperation() == MCCFIInstruction::OpDefCfaOffset)
1603         BC.MIB->addAnnotation(Inst, "DeleteMe", 0U, AllocatorId);
1604     }
1605   }
1606
1607   int PrevSPVal = -8;
1608   BinaryBasicBlock *PrevBB = nullptr;
1609   StackPointerTracking &SPT = Info.getStackPointerTracking();
1610   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
1611     if (BB->size() == 0)
1612       continue;
1613     const int SPValAtEnd = SPT.getStateAt(*BB->rbegin())->first;
1614     const int SPValAtBegin = SPT.getStateBefore(*BB->begin())->first;
1615     int SPVal = SPValAtBegin;
1616     for (auto Iter = BB->begin(); Iter != BB->end(); ++Iter) {
1617       const int CurVal = SPT.getStateAt(*Iter)->first;
1618       if (SPVal != CurVal) {
1619         auto InsertionIter = Iter;
1620         ++InsertionIter;
1621         Iter = BF.addCFIInstruction(
1622             BB, InsertionIter,
1623             MCCFIInstruction::cfiDefCfaOffset(nullptr, -CurVal));
1624         SPVal = CurVal;
1625       }
1626     }
1627     if (BF.isSplit() && PrevBB && BB->isCold() != PrevBB->isCold())
1628       BF.addCFIInstruction(
1629           BB, BB->begin(),
1630           MCCFIInstruction::cfiDefCfaOffset(nullptr, -SPValAtBegin));
1631     else if (SPValAtBegin != PrevSPVal)
1632       BF.addCFIInstruction(
1633           PrevBB, PrevBB->end(),
1634           MCCFIInstruction::cfiDefCfaOffset(nullptr, -SPValAtBegin));
1635     PrevSPVal = SPValAtEnd;
1636     PrevBB = BB;
1637   }
1638
1639   for (BinaryBasicBlock &BB : BF)
1640     for (auto I = BB.begin(); I != BB.end();)
1641       if (BC.MIB->hasAnnotation(*I, "DeleteMe"))
1642         I = BB.eraseInstruction(I);
1643       else
1644         ++I;
1645 }
1646
1647 Expected<MCInst> ShrinkWrapping::createStackAccess(int SPVal, int FPVal,
1648                                                    const FrameIndexEntry &FIE,
1649                                                    bool CreatePushOrPop) {
1650   MCInst NewInst;
1651   if (SPVal != StackPointerTracking::SUPERPOSITION &&
1652       SPVal != StackPointerTracking::EMPTY) {
1653     if (FIE.IsLoad) {
1654       BC.MIB->createRestoreFromStack(NewInst, BC.MIB->getStackPointer(),
1655                                      FIE.StackOffset - SPVal, FIE.RegOrImm,
1656                                      FIE.Size);
1657     } else {
1658       BC.MIB->createSaveToStack(NewInst, BC.MIB->getStackPointer(),
1659                                 FIE.StackOffset - SPVal, FIE.RegOrImm,
1660                                 FIE.Size);
1661     }
1662     if (CreatePushOrPop)
1663       BC.MIB->changeToPushOrPop(NewInst);
1664     return NewInst;
1665   }
1666   assert(FPVal != StackPointerTracking::SUPERPOSITION &&
1667          FPVal != StackPointerTracking::EMPTY);
1668
1669   if (FIE.IsLoad) {
1670     BC.MIB->createRestoreFromStack(NewInst, BC.MIB->getFramePointer(),
1671                                    FIE.StackOffset - FPVal, FIE.RegOrImm,
1672                                    FIE.Size);
1673   } else {
1674     BC.MIB->createSaveToStack(NewInst, BC.MIB->getFramePointer(),
1675                               FIE.StackOffset - FPVal, FIE.RegOrImm, FIE.Size);
1676   }
1677   return NewInst;
1678 }
1679
1680 void ShrinkWrapping::updateCFIInstOffset(MCInst &Inst, int64_t NewOffset) {
1681   const MCCFIInstruction *CFI = BF.getCFIFor(Inst);
1682   if (UpdatedCFIs.count(CFI))
1683     return;
1684
1685   switch (CFI->getOperation()) {
1686   case MCCFIInstruction::OpDefCfa:
1687   case MCCFIInstruction::OpDefCfaRegister:
1688   case MCCFIInstruction::OpDefCfaOffset:
1689     CFI = BF.mutateCFIOffsetFor(Inst, -NewOffset);
1690     break;
1691   case MCCFIInstruction::OpOffset:
1692   default:
1693     break;
1694   }
1695
1696   UpdatedCFIs.insert(CFI);
1697 }
1698
1699 BBIterTy ShrinkWrapping::insertCFIsForPushOrPop(BinaryBasicBlock &BB,
1700                                                 BBIterTy Pos, unsigned Reg,
1701                                                 bool isPush, int Sz,
1702                                                 int64_t NewOffset) {
1703   if (isPush) {
1704     for (uint32_t Idx : DeletedPushCFIs[Reg]) {
1705       Pos = BF.addCFIPseudo(&BB, Pos, Idx);
1706       updateCFIInstOffset(*Pos++, NewOffset);
1707     }
1708     if (HasDeletedOffsetCFIs[Reg]) {
1709       Pos = BF.addCFIInstruction(
1710           &BB, Pos,
1711           MCCFIInstruction::createOffset(
1712               nullptr, BC.MRI->getDwarfRegNum(Reg, false), NewOffset));
1713       ++Pos;
1714     }
1715   } else {
1716     for (uint32_t Idx : DeletedPopCFIs[Reg]) {
1717       Pos = BF.addCFIPseudo(&BB, Pos, Idx);
1718       updateCFIInstOffset(*Pos++, NewOffset);
1719     }
1720     if (HasDeletedOffsetCFIs[Reg]) {
1721       Pos = BF.addCFIInstruction(
1722           &BB, Pos,
1723           MCCFIInstruction::createSameValue(
1724               nullptr, BC.MRI->getDwarfRegNum(Reg, false)));
1725       ++Pos;
1726     }
1727   }
1728   return Pos;
1729 }
1730
1731 Expected<BBIterTy> ShrinkWrapping::processInsertion(BBIterTy InsertionPoint,
1732                                                     BinaryBasicBlock *CurBB,
1733                                                     const WorklistItem &Item,
1734                                                     int64_t SPVal,
1735                                                     int64_t FPVal) {
1736   // Trigger CFI reconstruction for this CSR if necessary - writing to
1737   // PushOffsetByReg/PopOffsetByReg *will* trigger CFI update
1738   if ((Item.FIEToInsert.IsStore &&
1739        !DeletedPushCFIs[Item.AffectedReg].empty()) ||
1740       (Item.FIEToInsert.IsLoad && !DeletedPopCFIs[Item.AffectedReg].empty()) ||
1741       HasDeletedOffsetCFIs[Item.AffectedReg]) {
1742     if (Item.Action == WorklistItem::InsertPushOrPop) {
1743       if (Item.FIEToInsert.IsStore)
1744         PushOffsetByReg[Item.AffectedReg] = SPVal - Item.FIEToInsert.Size;
1745       else
1746         PopOffsetByReg[Item.AffectedReg] = SPVal;
1747     } else {
1748       if (Item.FIEToInsert.IsStore)
1749         PushOffsetByReg[Item.AffectedReg] = Item.FIEToInsert.StackOffset;
1750       else
1751         PopOffsetByReg[Item.AffectedReg] = Item.FIEToInsert.StackOffset;
1752     }
1753   }
1754
1755   LLVM_DEBUG({
1756     dbgs() << "Creating stack access with SPVal = " << SPVal
1757            << "; stack offset = " << Item.FIEToInsert.StackOffset
1758            << " Is push = " << (Item.Action == WorklistItem::InsertPushOrPop)
1759            << "\n";
1760   });
1761   Expected<MCInst> NewInstOrErr =
1762       createStackAccess(SPVal, FPVal, Item.FIEToInsert,
1763                         Item.Action == WorklistItem::InsertPushOrPop);
1764   if (auto E = NewInstOrErr.takeError())
1765     return Error(std::move(E));
1766   MCInst &NewInst = *NewInstOrErr;
1767   if (InsertionPoint != CurBB->end()) {
1768     LLVM_DEBUG({
1769       dbgs() << "Adding before Inst: ";
1770       InsertionPoint->dump();
1771       dbgs() << "the following inst: ";
1772       NewInst.dump();
1773     });
1774     BBIterTy Iter =
1775         CurBB->insertInstruction(InsertionPoint, std::move(NewInst));
1776     return ++Iter;
1777   }
1778   CurBB->addInstruction(std::move(NewInst));
1779   LLVM_DEBUG(dbgs() << "Adding to BB!\n");
1780   return CurBB->end();
1781 }
1782
1783 Expected<BBIterTy> ShrinkWrapping::processInsertionsList(
1784     BBIterTy InsertionPoint, BinaryBasicBlock *CurBB,
1785     std::vector<WorklistItem> &TodoList, int64_t SPVal, int64_t FPVal) {
1786   bool HasInsertions = llvm::any_of(TodoList, [&](WorklistItem &Item) {
1787     return Item.Action == WorklistItem::InsertLoadOrStore ||
1788            Item.Action == WorklistItem::InsertPushOrPop;
1789   });
1790
1791   if (!HasInsertions)
1792     return InsertionPoint;
1793
1794   assert(((SPVal != StackPointerTracking::SUPERPOSITION &&
1795            SPVal != StackPointerTracking::EMPTY) ||
1796           (FPVal != StackPointerTracking::SUPERPOSITION &&
1797            FPVal != StackPointerTracking::EMPTY)) &&
1798          "Cannot insert if we have no idea of the stack state here");
1799
1800   // Revert the effect of PSPT for this location, we want SP Value before
1801   // insertions
1802   if (InsertionPoint == CurBB->end()) {
1803     for (WorklistItem &Item : TodoList) {
1804       if (Item.Action != WorklistItem::InsertPushOrPop)
1805         continue;
1806       if (Item.FIEToInsert.IsStore)
1807         SPVal += Item.FIEToInsert.Size;
1808       if (Item.FIEToInsert.IsLoad)
1809         SPVal -= Item.FIEToInsert.Size;
1810     }
1811   }
1812
1813   // Reorder POPs to obey the correct dominance relation between them
1814   llvm::stable_sort(TodoList, [&](const WorklistItem &A,
1815                                   const WorklistItem &B) {
1816     if ((A.Action != WorklistItem::InsertPushOrPop || !A.FIEToInsert.IsLoad) &&
1817         (B.Action != WorklistItem::InsertPushOrPop || !B.FIEToInsert.IsLoad))
1818       return false;
1819     if ((A.Action != WorklistItem::InsertPushOrPop || !A.FIEToInsert.IsLoad))
1820       return true;
1821     if ((B.Action != WorklistItem::InsertPushOrPop || !B.FIEToInsert.IsLoad))
1822       return false;
1823     return DomOrder[B.AffectedReg] < DomOrder[A.AffectedReg];
1824   });
1825
1826   // Process insertions
1827   for (WorklistItem &Item : TodoList) {
1828     if (Item.Action == WorklistItem::Erase ||
1829         Item.Action == WorklistItem::ChangeToAdjustment)
1830       continue;
1831
1832     auto InsertionPointOrErr =
1833         processInsertion(InsertionPoint, CurBB, Item, SPVal, FPVal);
1834     if (auto E = InsertionPointOrErr.takeError())
1835       return Error(std::move(E));
1836     InsertionPoint = *InsertionPointOrErr;
1837     if (Item.Action == WorklistItem::InsertPushOrPop &&
1838         Item.FIEToInsert.IsStore)
1839       SPVal -= Item.FIEToInsert.Size;
1840     if (Item.Action == WorklistItem::InsertPushOrPop &&
1841         Item.FIEToInsert.IsLoad)
1842       SPVal += Item.FIEToInsert.Size;
1843   }
1844   return InsertionPoint;
1845 }
1846
1847 Expected<bool> ShrinkWrapping::processInsertions() {
1848   PredictiveStackPointerTracking PSPT(BF, Todo, Info, AllocatorId);
1849   PSPT.run();
1850
1851   bool Changes = false;
1852   for (BinaryBasicBlock &BB : BF) {
1853     // Process insertions before some inst.
1854     for (auto I = BB.begin(); I != BB.end(); ++I) {
1855       MCInst &Inst = *I;
1856       auto TodoList = BC.MIB->tryGetAnnotationAs<std::vector<WorklistItem>>(
1857           Inst, getAnnotationIndex());
1858       if (!TodoList)
1859         continue;
1860       Changes = true;
1861       std::vector<WorklistItem> List = *TodoList;
1862       LLVM_DEBUG({
1863         dbgs() << "Now processing insertions in " << BB.getName()
1864                << " before inst: ";
1865         Inst.dump();
1866       });
1867       auto Iter = I;
1868       std::pair<int, int> SPTState =
1869           *PSPT.getStateAt(Iter == BB.begin() ? (ProgramPoint)&BB : &*(--Iter));
1870       auto IterOrErr =
1871           processInsertionsList(I, &BB, List, SPTState.first, SPTState.second);
1872       if (auto E = IterOrErr.takeError())
1873         return Error(std::move(E));
1874       I = *IterOrErr;
1875     }
1876     // Process insertions at the end of bb
1877     auto WRI = Todo.find(&BB);
1878     if (WRI != Todo.end()) {
1879       std::pair<int, int> SPTState = *PSPT.getStateAt(*BB.rbegin());
1880       if (auto E = processInsertionsList(BB.end(), &BB, WRI->second,
1881                                          SPTState.first, SPTState.second)
1882                        .takeError())
1883         return Error(std::move(E));
1884       Changes = true;
1885     }
1886   }
1887   return Changes;
1888 }
1889
1890 void ShrinkWrapping::processDeletions() {
1891   LivenessAnalysis &LA = Info.getLivenessAnalysis();
1892   for (BinaryBasicBlock &BB : BF) {
1893     for (auto II = BB.begin(); II != BB.end(); ++II) {
1894       MCInst &Inst = *II;
1895       auto TodoList = BC.MIB->tryGetAnnotationAs<std::vector<WorklistItem>>(
1896           Inst, getAnnotationIndex());
1897       if (!TodoList)
1898         continue;
1899       // Process all deletions
1900       for (WorklistItem &Item : *TodoList) {
1901         if (Item.Action != WorklistItem::Erase &&
1902             Item.Action != WorklistItem::ChangeToAdjustment)
1903           continue;
1904
1905         if (Item.Action == WorklistItem::ChangeToAdjustment) {
1906           // Is flag reg alive across this func?
1907           bool DontClobberFlags = LA.isAlive(&Inst, BC.MIB->getFlagsReg());
1908           if (int Sz = BC.MIB->getPushSize(Inst)) {
1909             BC.MIB->createStackPointerIncrement(Inst, Sz, DontClobberFlags);
1910             continue;
1911           }
1912           if (int Sz = BC.MIB->getPopSize(Inst)) {
1913             BC.MIB->createStackPointerDecrement(Inst, Sz, DontClobberFlags);
1914             continue;
1915           }
1916         }
1917
1918         LLVM_DEBUG({
1919           dbgs() << "Erasing: ";
1920           BC.printInstruction(dbgs(), Inst);
1921         });
1922         II = std::prev(BB.eraseInstruction(II));
1923         break;
1924       }
1925     }
1926   }
1927 }
1928
1929 void ShrinkWrapping::rebuildCFI() {
1930   const bool FP = Info.getStackPointerTracking().HasFramePointer;
1931   Info.invalidateAll();
1932   if (!FP) {
1933     rebuildCFIForSP();
1934     Info.invalidateAll();
1935   }
1936   for (unsigned I = 0, E = BC.MRI->getNumRegs(); I != E; ++I) {
1937     if (PushOffsetByReg[I] == 0 || PopOffsetByReg[I] == 0)
1938       continue;
1939     const int64_t SPValPush = PushOffsetByReg[I];
1940     const int64_t SPValPop = PopOffsetByReg[I];
1941     insertUpdatedCFI(I, SPValPush, SPValPop);
1942     Info.invalidateAll();
1943   }
1944 }
1945
1946 Expected<bool> ShrinkWrapping::perform(bool HotOnly) {
1947   HasDeletedOffsetCFIs = BitVector(BC.MRI->getNumRegs(), false);
1948   PushOffsetByReg = std::vector<int64_t>(BC.MRI->getNumRegs(), 0LL);
1949   PopOffsetByReg = std::vector<int64_t>(BC.MRI->getNumRegs(), 0LL);
1950
1951   // Update pass statistics
1952   uint64_t TotalInstrs = 0ULL;
1953   uint64_t TotalStoreInstrs = 0ULL;
1954   for (BinaryBasicBlock *BB : BF.getLayout().blocks()) {
1955     uint64_t BBExecCount = BB->getExecutionCount();
1956     if (!BBExecCount || BBExecCount == BinaryBasicBlock::COUNT_NO_PROFILE)
1957       continue;
1958     for (const auto &Instr : *BB) {
1959       if (BC.MIB->isPseudo(Instr))
1960         continue;
1961       if (BC.MIB->mayStore(Instr))
1962         TotalStoreInstrs += BBExecCount;
1963       TotalInstrs += BBExecCount;
1964     }
1965   }
1966   InstrDynamicCount += TotalInstrs;
1967   StoreDynamicCount += TotalStoreInstrs;
1968
1969   if (!FA.hasFrameInfo(BF))
1970     return false;
1971
1972   if (HotOnly && (BF.getKnownExecutionCount() < BC.getHotThreshold()))
1973     return false;
1974
1975   if (opts::EqualizeBBCounts)
1976     equalizeBBCounts(Info, BF);
1977
1978   if (BF.checkForAmbiguousJumpTables()) {
1979     LLVM_DEBUG(dbgs() << "BOLT-DEBUG: ambiguous JTs in " << BF.getPrintName()
1980                       << ".\n");
1981     // We could call disambiguateJumpTables here, but it is probably not worth
1982     // the cost (of duplicating potentially large jump tables that could regress
1983     // dcache misses). Moreover, ambiguous JTs are rare and coming from code
1984     // written in assembly language. Just bail.
1985     return false;
1986   }
1987   SLM.initialize();
1988   CSA.compute();
1989   classifyCSRUses();
1990   pruneUnwantedCSRs();
1991   computeSaveLocations();
1992   moveSaveRestores();
1993   LLVM_DEBUG({
1994     dbgs() << "Func before shrink-wrapping: \n";
1995     BF.dump();
1996   });
1997   SLM.performChanges();
1998   // Early exit if processInsertions doesn't detect any todo items
1999   auto ModifiedOrErr = processInsertions();
2000   if (auto E = ModifiedOrErr.takeError())
2001     return Error(std::move(E));
2002   const bool Modified = *ModifiedOrErr;
2003   if (!Modified)
2004     return false;
2005   processDeletions();
2006   if (foldIdenticalSplitEdges()) {
2007     const std::pair<unsigned, uint64_t> Stats = BF.eraseInvalidBBs();
2008     (void)Stats;
2009     LLVM_DEBUG(dbgs() << "Deleted " << Stats.first
2010                       << " redundant split edge BBs (" << Stats.second
2011                       << " bytes) for " << BF.getPrintName() << "\n");
2012   }
2013   rebuildCFI();
2014   // We may have split edges, creating BBs that need correct branching
2015   BF.fixBranches();
2016   LLVM_DEBUG({
2017     dbgs() << "Func after shrink-wrapping: \n";
2018     BF.dump();
2019   });
2020   return true;
2021 }
2022
2023 void ShrinkWrapping::printStats(BinaryContext &BC) {
2024   BC.outs() << "BOLT-INFO: Shrink wrapping moved " << SpillsMovedRegularMode
2025             << " spills inserting load/stores and " << SpillsMovedPushPopMode
2026             << " spills inserting push/pops\n";
2027   if (!InstrDynamicCount || !StoreDynamicCount)
2028     return;
2029   BC.outs() << "BOLT-INFO: Shrink wrapping reduced " << SpillsMovedDynamicCount
2030             << " store executions ("
2031             << format("%.1lf%%",
2032                       (100.0 * SpillsMovedDynamicCount / InstrDynamicCount))
2033             << " total instructions executed, "
2034             << format("%.1lf%%",
2035                       (100.0 * SpillsMovedDynamicCount / StoreDynamicCount))
2036             << " store instructions)\n";
2037   BC.outs() << "BOLT-INFO: Shrink wrapping failed at reducing "
2038             << SpillsFailedDynamicCount << " store executions ("
2039             << format("%.1lf%%",
2040                       (100.0 * SpillsFailedDynamicCount / InstrDynamicCount))
2041             << " total instructions executed, "
2042             << format("%.1lf%%",
2043                       (100.0 * SpillsFailedDynamicCount / StoreDynamicCount))
2044             << " store instructions)\n";
2045 }
2046
2047 // Operators necessary as a result of using MCAnnotation
2048 raw_ostream &operator<<(raw_ostream &OS,
2049                         const std::vector<ShrinkWrapping::WorklistItem> &Vec) {
2050   OS << "SWTodo[";
2051   const char *Sep = "";
2052   for (const ShrinkWrapping::WorklistItem &Item : Vec) {
2053     OS << Sep;
2054     switch (Item.Action) {
2055     case ShrinkWrapping::WorklistItem::Erase:
2056       OS << "Erase";
2057       break;
2058     case ShrinkWrapping::WorklistItem::ChangeToAdjustment:
2059       OS << "ChangeToAdjustment";
2060       break;
2061     case ShrinkWrapping::WorklistItem::InsertLoadOrStore:
2062       OS << "InsertLoadOrStore";
2063       break;
2064     case ShrinkWrapping::WorklistItem::InsertPushOrPop:
2065       OS << "InsertPushOrPop";
2066       break;
2067     }
2068     Sep = ", ";
2069   }
2070   OS << "]";
2071   return OS;
2072 }
2073
2074 raw_ostream &
2075 operator<<(raw_ostream &OS,
2076            const std::vector<StackLayoutModifier::WorklistItem> &Vec) {
2077   OS << "SLMTodo[";
2078   const char *Sep = "";
2079   for (const StackLayoutModifier::WorklistItem &Item : Vec) {
2080     OS << Sep;
2081     switch (Item.Action) {
2082     case StackLayoutModifier::WorklistItem::None:
2083       OS << "None";
2084       break;
2085     case StackLayoutModifier::WorklistItem::AdjustLoadStoreOffset:
2086       OS << "AdjustLoadStoreOffset";
2087       break;
2088     case StackLayoutModifier::WorklistItem::AdjustCFI:
2089       OS << "AdjustCFI";
2090       break;
2091     }
2092     Sep = ", ";
2093   }
2094   OS << "]";
2095   return OS;
2096 }
2097
2098 bool operator==(const ShrinkWrapping::WorklistItem &A,
2099                 const ShrinkWrapping::WorklistItem &B) {
2100   return (A.Action == B.Action && A.AffectedReg == B.AffectedReg &&
2101           A.Adjustment == B.Adjustment &&
2102           A.FIEToInsert.IsLoad == B.FIEToInsert.IsLoad &&
2103           A.FIEToInsert.IsStore == B.FIEToInsert.IsStore &&
2104           A.FIEToInsert.RegOrImm == B.FIEToInsert.RegOrImm &&
2105           A.FIEToInsert.Size == B.FIEToInsert.Size &&
2106           A.FIEToInsert.IsSimple == B.FIEToInsert.IsSimple &&
2107           A.FIEToInsert.StackOffset == B.FIEToInsert.StackOffset);
2108 }
2109
2110 bool operator==(const StackLayoutModifier::WorklistItem &A,
2111                 const StackLayoutModifier::WorklistItem &B) {
2112   return (A.Action == B.Action && A.OffsetUpdate == B.OffsetUpdate);
2113 }
2114
2115 } // end namespace bolt
2116 } // end namespace llvm