[CodeExtractor] Factor out and reuse shrinkwrap analysis

[vsk]

Summary

Factor out CodeExtractor's analysis of allocas (for shrinkwrapping
purposes), and allow the analysis to be reused.

This resolves a quadratic compile-time bug observed when compiling
AMDGPUDisassembler.cpp.o.

Pre-patch (Release + LTO clang):

 ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
176.5278 ( 57.8%)   0.4915 ( 18.5%)  177.0192 ( 57.4%)  177.4112 ( 57.3%)  Hot Cold Splitting

Post-patch (ReleaseAsserts clang):

 ---User Time---   --System Time--   --User+System--   ---Wall Time---  --- Name ---
1.4051 (  3.3%)   0.0079 (  0.3%)   1.4129 (  3.2%)   1.4129 (  3.2%)  Hot Cold Splitting

Testing: check-llvm, and comparing the AMDGPUDisassembler.cpp.o binary
pre- vs. post-patch.

An alternate approach is to hide CodeExtractorAnalysisCache from clients
of CodeExtractor, and to recompute the analysis from scratch inside of
CodeExtractor::extractCodeRegion(). This eliminates some redundant work
in the shrinkwrapping legality check. However, some clients continue to
exhibit O(n^2) compile time behavior as computing the analysis is O(n).

rdar://55912966

[vsk]

• Add some documentation.

[davidxl]

lgtm

llvm/include/llvm/Transforms/IPO/HotColdSplitting.h

 
 class Module;
 class ProfileSummaryInfo;
 class BlockFrequencyInfo;
 class TargetTransformInfo;
 class OptimizationRemarkEmitter;
 class AssumptionCache;
 class DominatorTree;
+class CodeExtractorAnalysisCache;
 
 /// A sequence of basic blocks.
 ///
 /// A 0-sized SmallVector is slightly cheaper to move than a std::vector.
 using BlockSequence = SmallVector<BasicBlock *, 0>;
 
 class HotColdSplitting {
 public:
   HotColdSplitting(ProfileSummaryInfo *ProfSI,
                    function_ref<BlockFrequencyInfo *(Function &)> GBFI,
                    function_ref<TargetTransformInfo &(Function &)> GTTI,
                    std::function<OptimizationRemarkEmitter &(Function &)> *GORE,
                    function_ref<AssumptionCache *(Function &)> LAC)
       : PSI(ProfSI), GetBFI(GBFI), GetTTI(GTTI), GetORE(GORE), LookupAC(LAC) {}
   bool run(Module &M);
 
 private:
   bool isFunctionCold(const Function &F) const;
   bool shouldOutlineFrom(const Function &F) const;
   bool outlineColdRegions(Function &F, bool HasProfileSummary);
-  Function *extractColdRegion(const BlockSequence &Region, DominatorTree &DT,
-                              BlockFrequencyInfo *BFI, TargetTransformInfo &TTI,
+  Function *extractColdRegion(const BlockSequence &Region,
+                              const CodeExtractorAnalysisCache &CEAC,
+                              DominatorTree &DT, BlockFrequencyInfo *BFI,
+                              TargetTransformInfo &TTI,
                               OptimizationRemarkEmitter &ORE,
                               AssumptionCache *AC, unsigned Count);
   ProfileSummaryInfo *PSI;
   function_ref<BlockFrequencyInfo *(Function &)> GetBFI;
   function_ref<TargetTransformInfo &(Function &)> GetTTI;
   std::function<OptimizationRemarkEmitter &(Function &)> *GetORE;
   function_ref<AssumptionCache *(Function &)> LookupAC;
 };

llvm/include/llvm/Transforms/Utils/CodeExtractor.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include <limits>
 
 namespace llvm {
 
+class AllocaInst;
 class BasicBlock;
 class BlockFrequency;
 class BlockFrequencyInfo;
 class BranchProbabilityInfo;
 class AssumptionCache;
 class CallInst;
 class DominatorTree;
 class Function;
 class Instruction;
 class Loop;
 class Module;
 class Type;
 class Value;
 
+/// A cache for the CodeExtractor analysis. The operation \ref
+/// CodeExtractor::extractCodeRegion is guaranteed not to invalidate this
+/// object. This object should conservatively be considered invalid if any
+/// other mutating operations on the IR occur.
+///
+/// Constructing this object is O(n) in the size of the function.
+class CodeExtractorAnalysisCache {
+  /// The allocas in the function.
+  SmallVector<AllocaInst *, 16> Allocas;
+
+  /// Base memory addresses of load/store instructions, grouped by block.
+  DenseMap<BasicBlock *, DenseSet<Value *>> BaseMemAddrs;
+
+  /// Blocks which contain instructions which may have unknown side-effects
+  /// on memory.
+  DenseSet<BasicBlock *> SideEffectingBlocks;
+
+  void findSideEffectInfoForBlock(BasicBlock &BB);
+
+public:
+  CodeExtractorAnalysisCache(Function &F);
+
+  /// Get the allocas in the function at the time the analysis was created.
+  /// Note that some of these allocas may no longer be present in the function,
+  /// due to \ref CodeExtractor::extractCodeRegion.
+  ArrayRef<AllocaInst *> getAllocas() const { return Allocas; }
+
+  /// Check whether \p BB contains an instruction thought to load from, store
+  /// to, or otherwise clobber the alloca \p Addr.
+  bool doesBlockContainClobberOfAddr(BasicBlock &BB, AllocaInst *Addr) const;
+};
+
   /// Utility class for extracting code into a new function.
   ///
   /// This utility provides a simple interface for extracting some sequence of
   /// code into its own function, replacing it with a call to that function. It
   /// also provides various methods to query about the nature and result of
   /// such a transformation.
   ///
   /// The rough algorithm used is:
CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
                   BranchProbabilityInfo *BPI = nullptr,
                   AssumptionCache *AC = nullptr,
                   std::string Suffix = "");
 
     /// Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
-    Function *extractCodeRegion();
+    Function *extractCodeRegion(const CodeExtractorAnalysisCache &CEAC);
 
     /// Verify that assumption cache isn't stale after a region is extracted.
     /// Returns false when verifier finds errors. AssumptionCache is passed as
     /// parameter to make this function stateless.
     static bool verifyAssumptionCache(const Function& F, AssumptionCache *AC);
 
     /// Test whether this code extractor is eligible.
     ///
public:
     /// significant impact on the cost however.
     void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
                            const ValueSet &Allocas) const;
 
     /// Check if life time marker nodes can be hoisted/sunk into the outline
     /// region.
     ///
     /// Returns true if it is safe to do the code motion.
-    bool isLegalToShrinkwrapLifetimeMarkers(Instruction *AllocaAddr) const;
+    bool
+    isLegalToShrinkwrapLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
+                                       Instruction *AllocaAddr) const;
 
     /// Find the set of allocas whose life ranges are contained within the
     /// outlined region.
     ///
     /// Allocas which have life_time markers contained in the outlined region
     /// should be pushed to the outlined function. The address bitcasts that
     /// are used by the lifetime markers are also candidates for shrink-
     /// wrapping. The instructions that need to be sunk are collected in
     /// 'Allocas'.
-    void findAllocas(ValueSet &SinkCands, ValueSet &HoistCands,
+    void findAllocas(const CodeExtractorAnalysisCache &CEAC,
+                     ValueSet &SinkCands, ValueSet &HoistCands,
                      BasicBlock *&ExitBlock) const;
 
     /// Find or create a block within the outline region for placing hoisted
     /// code.
     ///
     /// CommonExitBlock is block outside the outline region. It is the common
     /// successor of blocks inside the region. If there exists a single block
     /// inside the region that is the predecessor of CommonExitBlock, that block
     /// will be returned. Otherwise CommonExitBlock will be split and the
     /// original block will be added to the outline region.
     BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
 
   private:
     struct LifetimeMarkerInfo {
       bool SinkLifeStart = false;
       bool HoistLifeEnd = false;
       Instruction *LifeStart = nullptr;
       Instruction *LifeEnd = nullptr;
     };
 
-    LifetimeMarkerInfo getLifetimeMarkers(Instruction *Addr,
-                                          BasicBlock *ExitBlock) const;
+    LifetimeMarkerInfo
+    getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
+                       Instruction *Addr, BasicBlock *ExitBlock) const;
 
     void severSplitPHINodesOfEntry(BasicBlock *&Header);
     void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
     void splitReturnBlocks();
 
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,

llvm/lib/Transforms/IPO/BlockExtractor.cpp

for (BasicBlock *BB : BBs) {
                         << BB->getParent()->getName() << ":" << BB->getName()
                         << "\n");
       BlocksToExtractVec.push_back(BB);
       if (const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
         BlocksToExtractVec.push_back(II->getUnwindDest());
       ++NumExtracted;
       Changed = true;
     }
-    Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion();
+    CodeExtractorAnalysisCache CEAC(*BBs[0]->getParent());
+    Function *F = CodeExtractor(BlocksToExtractVec).extractCodeRegion(CEAC);
     if (F)
       LLVM_DEBUG(dbgs() << "Extracted group '" << (*BBs.begin())->getName()
                         << "' in: " << F->getName() << '\n');
     else
       LLVM_DEBUG(dbgs() << "Failed to extract for group '"
                         << (*BBs.begin())->getName() << "'\n");
   }
 

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

if (!SuccsOutsideRegion.empty()) {
     LLVM_DEBUG(dbgs() << "Applying penalty for: " << SuccsOutsideRegion.size()
                       << " non-region successors\n");
     Penalty += (SuccsOutsideRegion.size() - 1) * TargetTransformInfo::TCC_Basic;
   }
 
   return Penalty;
 }
 
-Function *HotColdSplitting::extractColdRegion(const BlockSequence &Region,
-                                              DominatorTree &DT,
-                                              BlockFrequencyInfo *BFI,
-                                              TargetTransformInfo &TTI,
-                                              OptimizationRemarkEmitter &ORE,
-                                              AssumptionCache *AC,
-                                              unsigned Count) {
+Function *HotColdSplitting::extractColdRegion(
+    const BlockSequence &Region, const CodeExtractorAnalysisCache &CEAC,
+    DominatorTree &DT, BlockFrequencyInfo *BFI, TargetTransformInfo &TTI,
+    OptimizationRemarkEmitter &ORE, AssumptionCache *AC, unsigned Count) {
   assert(!Region.empty());
 
   // TODO: Pass BFI and BPI to update profile information.
   CodeExtractor CE(Region, &DT, /* AggregateArgs */ false, /* BFI */ nullptr,
                    /* BPI */ nullptr, AC, /* AllowVarArgs */ false,
                    /* AllowAlloca */ false,
                    /* Suffix */ "cold." + std::to_string(Count));
 
   // Perform a simple cost/benefit analysis to decide whether or not to permit
   // splitting.
   SetVector<Value *> Inputs, Outputs, Sinks;
   CE.findInputsOutputs(Inputs, Outputs, Sinks);
   int OutliningBenefit = getOutliningBenefit(Region, TTI);
   int OutliningPenalty =
       getOutliningPenalty(Region, Inputs.size(), Outputs.size());
   LLVM_DEBUG(dbgs() << "Split profitability: benefit = " << OutliningBenefit
                     << ", penalty = " << OutliningPenalty << "\n");
   if (OutliningBenefit <= OutliningPenalty)
     return nullptr;
 
   Function *OrigF = Region[0]->getParent();
-  if (Function *OutF = CE.extractCodeRegion()) {
+  if (Function *OutF = CE.extractCodeRegion(CEAC)) {
     User *U = *OutF->user_begin();
     CallInst *CI = cast<CallInst>(U);
     CallSite CS(CI);
     NumColdRegionsOutlined++;
     if (TTI.useColdCCForColdCall(*OutF)) {
       OutF->setCallingConv(CallingConv::Cold);
       CS.setCallingConv(CallingConv::Cold);
     }
for (OutliningRegion &Region : Regions) {
       if (RegionsOverlap)
         continue;
 
       OutliningWorklist.emplace_back(std::move(Region));
       ++NumColdRegionsFound;
     }
   }
 
+  if (OutliningWorklist.empty())
+    return Changed;
+
   // Outline single-entry cold regions, splitting up larger regions as needed.
   unsigned OutlinedFunctionID = 1;
-  while (!OutliningWorklist.empty()) {
+  // Cache and recycle the CodeExtractor analysis to avoid O(n^2) compile-time.
+  CodeExtractorAnalysisCache CEAC(F);
+  do {
     OutliningRegion Region = OutliningWorklist.pop_back_val();
     assert(!Region.empty() && "Empty outlining region in worklist");
     do {
       BlockSequence SubRegion = Region.takeSingleEntrySubRegion(*DT);
       LLVM_DEBUG({
         dbgs() << "Hot/cold splitting attempting to outline these blocks:\n";
         for (BasicBlock *BB : SubRegion)
           BB->dump();
       });
 
-      Function *Outlined = extractColdRegion(SubRegion, *DT, BFI, TTI, ORE, AC,
-                                             OutlinedFunctionID);
+      Function *Outlined = extractColdRegion(SubRegion, CEAC, *DT, BFI, TTI,
+                                             ORE, AC, OutlinedFunctionID);
       if (Outlined) {
         ++OutlinedFunctionID;
         Changed = true;
       }
     } while (!Region.empty());
-  }
+  } while (!OutliningWorklist.empty());
 
   return Changed;
 }
 
 bool HotColdSplitting::run(Module &M) {
   bool Changed = false;
   bool HasProfileSummary = (M.getProfileSummary(/* IsCS */ false) != nullptr);
   for (auto It = M.begin(), End = M.end(); It != End; ++It) {

llvm/lib/Transforms/IPO/LoopExtractor.cpp

for (unsigned i = 0, e = ExitBlocks.size(); i != e; ++i)
         break;
       }
   }
 
   if (ShouldExtractLoop) {
     if (NumLoops == 0) return Changed;
     --NumLoops;
     AssumptionCache *AC = nullptr;
+    Function &Func = *L->getHeader()->getParent();
     if (auto *ACT = getAnalysisIfAvailable<AssumptionCacheTracker>())
-      AC = ACT->lookupAssumptionCache(*L->getHeader()->getParent());
+      AC = ACT->lookupAssumptionCache(Func);
+    CodeExtractorAnalysisCache CEAC(Func);
     CodeExtractor Extractor(DT, *L, false, nullptr, nullptr, AC);
-    if (Extractor.extractCodeRegion() != nullptr) {
+    if (Extractor.extractCodeRegion(CEAC) != nullptr) {
       Changed = true;
       // After extraction, the loop is replaced by a function call, so
       // we shouldn't try to run any more loop passes on it.
       LPM.markLoopAsDeleted(*L);
       LI.erase(L);
     }
     ++NumExtracted;
   }

llvm/lib/Transforms/IPO/PartialInlining.cpp

bool PartialInlinerImpl::FunctionCloner::doMultiRegionFunctionOutlining() {
   DominatorTree DT;
   DT.recalculate(*ClonedFunc);
 
   // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
   LoopInfo LI(DT);
   BranchProbabilityInfo BPI(*ClonedFunc, LI);
   ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
 
+  // Cache and recycle the CodeExtractor analysis to avoid O(n^2) compile-time.
+  CodeExtractorAnalysisCache CEAC(*ClonedFunc);
+
   SetVector<Value *> Inputs, Outputs, Sinks;
   for (FunctionOutliningMultiRegionInfo::OutlineRegionInfo RegionInfo :
        ClonedOMRI->ORI) {
     int CurrentOutlinedRegionCost = ComputeRegionCost(RegionInfo.Region);
 
     CodeExtractor CE(RegionInfo.Region, &DT, /*AggregateArgs*/ false,
                      ClonedFuncBFI.get(), &BPI,
                      LookupAC(*RegionInfo.EntryBlock->getParent()),
if (TracePartialInlining) {
       for (Value *output : Outputs)
         dbgs() << "instr used in func: " << *output << "\n";
     }
 #endif
     // Do not extract regions that have live exit variables.
     if (Outputs.size() > 0 && !ForceLiveExit)
       continue;
 
-    Function *OutlinedFunc = CE.extractCodeRegion();
+    Function *OutlinedFunc = CE.extractCodeRegion(CEAC);
 
     if (OutlinedFunc) {
       CallSite OCS = PartialInlinerImpl::getOneCallSiteTo(OutlinedFunc);
       BasicBlock *OutliningCallBB = OCS.getInstruction()->getParent();
       assert(OutliningCallBB->getParent() == ClonedFunc);
       OutlinedFunctions.push_back(std::make_pair(OutlinedFunc,OutliningCallBB));
       NumColdRegionsOutlined++;
       OutlinedRegionCost += CurrentOutlinedRegionCost;
if (!ToBeInlined(&BB) && &BB != ClonedOI->NonReturnBlock) {
       // FIXME: the code extractor may hoist/sink more code
       // into the outlined function which may make the outlining
       // overhead (the difference of the outlined function cost
       // and OutliningRegionCost) look larger.
       OutlinedRegionCost += computeBBInlineCost(&BB);
     }
 
   // Extract the body of the if.
+  CodeExtractorAnalysisCache CEAC(*ClonedFunc);
   Function *OutlinedFunc =
       CodeExtractor(ToExtract, &DT, /*AggregateArgs*/ false,
                     ClonedFuncBFI.get(), &BPI, LookupAC(*ClonedFunc),
                     /* AllowVarargs */ true)
-          .extractCodeRegion();
+          .extractCodeRegion(CEAC);
 
   if (OutlinedFunc) {
     BasicBlock *OutliningCallBB =
         PartialInlinerImpl::getOneCallSiteTo(OutlinedFunc)
             .getInstruction()
             ->getParent();
     assert(OutliningCallBB->getParent() == ClonedFunc);
     OutlinedFunctions.push_back(std::make_pair(OutlinedFunc, OutliningCallBB));

llvm/lib/Transforms/Utils/CodeExtractor.cpp

static BasicBlock *getCommonExitBlock(const SetVector<BasicBlock *> &Blocks) {
   };
 
   if (any_of(Blocks, hasNonCommonExitSucc))
     return nullptr;
 
   return CommonExitBlock;
 }
 
-bool CodeExtractor::isLegalToShrinkwrapLifetimeMarkers(
-    Instruction *Addr) const {
-  AllocaInst *AI = cast<AllocaInst>(Addr->stripInBoundsConstantOffsets());
-  Function *Func = (*Blocks.begin())->getParent();
-  for (BasicBlock &BB : *Func) {
-    if (Blocks.count(&BB))
-      continue;
-    for (Instruction &II : BB) {
-      if (isa<DbgInfoIntrinsic>(II))
-        continue;
+CodeExtractorAnalysisCache::CodeExtractorAnalysisCache(Function &F) {
+  for (BasicBlock &BB : F) {
+    for (Instruction &II : BB.instructionsWithoutDebug())
+      if (auto *AI = dyn_cast<AllocaInst>(&II))
+        Allocas.push_back(AI);
 
-      unsigned Opcode = II.getOpcode();
-      Value *MemAddr = nullptr;
-      switch (Opcode) {
-      case Instruction::Store:
-      case Instruction::Load: {
-        if (Opcode == Instruction::Store) {
-          StoreInst *SI = cast<StoreInst>(&II);
-          MemAddr = SI->getPointerOperand();
-        } else {
-          LoadInst *LI = cast<LoadInst>(&II);
-          MemAddr = LI->getPointerOperand();
-        }
-        // Global variable can not be aliased with locals.
-        if (dyn_cast<Constant>(MemAddr))
-          break;
-        Value *Base = MemAddr->stripInBoundsConstantOffsets();
-        if (!isa<AllocaInst>(Base) || Base == AI)
-          return false;
-        break;
-      }
-      default: {
-        IntrinsicInst *IntrInst = dyn_cast<IntrinsicInst>(&II);
-        if (IntrInst) {
-          if (IntrInst->isLifetimeStartOrEnd())
-            break;
-          return false;
-        }
-        // Treat all the other cases conservatively if it has side effects.
-        if (II.mayHaveSideEffects())
-          return false;
-      }
-      }
-    }
-  }
+    findSideEffectInfoForBlock(BB);
+  }
+}
+
+void CodeExtractorAnalysisCache::findSideEffectInfoForBlock(BasicBlock &BB) {
+  for (Instruction &II : BB.instructionsWithoutDebug()) {
+      if (!isa<AllocaInst>(Base)) {
+        SideEffectingBlocks.insert(&BB);
+        return;
+      }
+      BaseMemAddrs[&BB].insert(Base);
+        SideEffectingBlocks.insert(&BB);
+        return;
+      if (II.mayHaveSideEffects()) {
+        SideEffectingBlocks.insert(&BB);
+        return;
+      }
 
+bool CodeExtractorAnalysisCache::doesBlockContainClobberOfAddr(
+    BasicBlock &BB, AllocaInst *Addr) const {
+  if (SideEffectingBlocks.count(&BB))
+    return true;
+  auto It = BaseMemAddrs.find(&BB);
+  if (It != BaseMemAddrs.end())
+    return It->second.count(Addr);
+  return false;
+}
+
+bool CodeExtractor::isLegalToShrinkwrapLifetimeMarkers(
+    const CodeExtractorAnalysisCache &CEAC, Instruction *Addr) const {
+  AllocaInst *AI = cast<AllocaInst>(Addr->stripInBoundsConstantOffsets());
+  Function *Func = (*Blocks.begin())->getParent();
+  for (BasicBlock &BB : *Func) {
+    if (Blocks.count(&BB))
+      continue;
+    if (CEAC.doesBlockContainClobberOfAddr(BB, AI))
+      return false;
+  }
   return true;
 }
 
 BasicBlock *
 CodeExtractor::findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock) {
   BasicBlock *SinglePredFromOutlineRegion = nullptr;
   assert(!Blocks.count(CommonExitBlock) &&
          "Expect a block outside the region!");
#endif
   return CommonExitBlock;
 }
 
 // Find the pair of life time markers for address 'Addr' that are either
 // defined inside the outline region or can legally be shrinkwrapped into the
 // outline region. If there are not other untracked uses of the address, return
 // the pair of markers if found; otherwise return a pair of nullptr.
 CodeExtractor::LifetimeMarkerInfo
-CodeExtractor::getLifetimeMarkers(Instruction *Addr,
+CodeExtractor::getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
+                                  Instruction *Addr,
                                   BasicBlock *ExitBlock) const {
   LifetimeMarkerInfo Info;
 
   for (User *U : Addr->users()) {
     IntrinsicInst *IntrInst = dyn_cast<IntrinsicInst>(U);
     if (IntrInst) {
       if (IntrInst->getIntrinsicID() == Intrinsic::lifetime_start) {
         // Do not handle the case where Addr has multiple start markers.
CodeExtractor::getLifetimeMarkers(const CodeExtractorAnalysisCache &CEAC,
 
   if (!Info.LifeStart || !Info.LifeEnd)
     return {};
 
   Info.SinkLifeStart = !definedInRegion(Blocks, Info.LifeStart);
   Info.HoistLifeEnd = !definedInRegion(Blocks, Info.LifeEnd);
   // Do legality check.
   if ((Info.SinkLifeStart || Info.HoistLifeEnd) &&
-      !isLegalToShrinkwrapLifetimeMarkers(Addr))
+      !isLegalToShrinkwrapLifetimeMarkers(CEAC, Addr))
     return {};
 
   // Check to see if we have a place to do hoisting, if not, bail.
   if (Info.HoistLifeEnd && !ExitBlock)
     return {};
 
   return Info;
 }
 
-void CodeExtractor::findAllocas(ValueSet &SinkCands, ValueSet &HoistCands,
+void CodeExtractor::findAllocas(const CodeExtractorAnalysisCache &CEAC,
+                                ValueSet &SinkCands, ValueSet &HoistCands,
                                 BasicBlock *&ExitBlock) const {
   Function *Func = (*Blocks.begin())->getParent();
   ExitBlock = getCommonExitBlock(Blocks);
 
   auto moveOrIgnoreLifetimeMarkers =
       [&](const LifetimeMarkerInfo &LMI) -> bool {
     if (!LMI.LifeStart)
       return false;
     if (LMI.SinkLifeStart) {
       LLVM_DEBUG(dbgs() << "Sinking lifetime.start: " << *LMI.LifeStart
                         << "\n");
       SinkCands.insert(LMI.LifeStart);
     }
     if (LMI.HoistLifeEnd) {
       LLVM_DEBUG(dbgs() << "Hoisting lifetime.end: " << *LMI.LifeEnd << "\n");
       HoistCands.insert(LMI.LifeEnd);
     }
     return true;
   };
 
-  for (BasicBlock &BB : *Func) {
-    if (Blocks.count(&BB))
+  // Look up allocas in the original function in CodeExtractorAnalysisCache, as
+  // this is much faster than walking all the instructions.
+  for (AllocaInst *AI : CEAC.getAllocas()) {
+    BasicBlock *BB = AI->getParent();
+    if (Blocks.count(BB))
       continue;
-    for (Instruction &II : BB) {
-      auto *AI = dyn_cast<AllocaInst>(&II);
-      if (!AI)
-        continue;
+
+    // As a prior call to extractCodeRegion() may have shrinkwrapped the alloca,
+    // check whether it is actually still in the original function.
+    Function *AIFunc = BB->getParent();
+    if (AIFunc != Func)
 
-      LifetimeMarkerInfo MarkerInfo = getLifetimeMarkers(AI, ExitBlock);
-      bool Moved = moveOrIgnoreLifetimeMarkers(MarkerInfo);
-      if (Moved) {
-        LLVM_DEBUG(dbgs() << "Sinking alloca: " << *AI << "\n");
-        SinkCands.insert(AI);
-        continue;
-      }
+    LifetimeMarkerInfo MarkerInfo = getLifetimeMarkers(CEAC, AI, ExitBlock);
 
-      // Follow any bitcasts.
-      SmallVector<Instruction *, 2> Bitcasts;
-      SmallVector<LifetimeMarkerInfo, 2> BitcastLifetimeInfo;
-      for (User *U : AI->users()) {
-        if (U->stripInBoundsConstantOffsets() == AI) {
-          Instruction *Bitcast = cast<Instruction>(U);
-          LifetimeMarkerInfo LMI = getLifetimeMarkers(Bitcast, ExitBlock);
-          if (LMI.LifeStart) {
-            Bitcasts.push_back(Bitcast);
-            BitcastLifetimeInfo.push_back(LMI);
-            continue;
-          }
-        }
+        LifetimeMarkerInfo LMI = getLifetimeMarkers(CEAC, Bitcast, ExitBlock);
 
-        // Found unknown use of AI.
-        if (!definedInRegion(Blocks, U)) {
-          Bitcasts.clear();
-          break;
-        }
-      }
 
-      // Either no bitcasts reference the alloca or there are unknown uses.
-      if (Bitcasts.empty())
-        continue;
 
-      LLVM_DEBUG(dbgs() << "Sinking alloca (via bitcast): " << *AI << "\n");
-      SinkCands.insert(AI);
-      for (unsigned I = 0, E = Bitcasts.size(); I != E; ++I) {
-        Instruction *BitcastAddr = Bitcasts[I];
-        const LifetimeMarkerInfo &LMI = BitcastLifetimeInfo[I];
-        assert(LMI.LifeStart &&
-               "Unsafe to sink bitcast without lifetime markers");
-        moveOrIgnoreLifetimeMarkers(LMI);
-        if (!definedInRegion(Blocks, BitcastAddr)) {
-          LLVM_DEBUG(dbgs() << "Sinking bitcast-of-alloca: " << *BitcastAddr
-                            << "\n");
-          SinkCands.insert(BitcastAddr);
-        }
-      }
-    }
-  }
-}
 
 bool CodeExtractor::isEligible() const {
   if (Blocks.empty())
     return false;
   BasicBlock *Header = *Blocks.begin();
   Function *F = Header->getParent();
 
   // For functions with varargs, check that varargs handling is only done in the
for (unsigned I = 0, E = BranchDist.Weights.size(); I < E; ++I) {
     BranchProbability BP(Weight.Amount, BranchDist.Total);
     BPI->setEdgeProbability(CodeReplacer, Weight.TargetNode.Index, BP);
   }
   TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
-Function *CodeExtractor::extractCodeRegion() {
+Function *
+CodeExtractor::extractCodeRegion(const CodeExtractorAnalysisCache &CEAC) {
   if (!isEligible())
     return nullptr;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
   Function *oldFunction = header->getParent();
 
any_of(Blocks, [&BranchI](const BasicBlock *BB) {
         return true;
       });
     });
   }
   newFuncRoot->getInstList().push_back(BranchI);
 
   ValueSet inputs, outputs, SinkingCands, HoistingCands;
   BasicBlock *CommonExit = nullptr;
-  findAllocas(SinkingCands, HoistingCands, CommonExit);
+  findAllocas(CEAC, SinkingCands, HoistingCands, CommonExit);
   assert(HoistingCands.empty() || CommonExit);
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs, SinkingCands);
 
   // Now sink all instructions which only have non-phi uses inside the region.
   // Group the allocas at the start of the block, so that any bitcast uses of
   // the allocas are well-defined.

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

TEST(CodeExtractor, ExitStub) {
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 3> Candidates{ getBlockByName(Func, "header"),
                                            getBlockByName(Func, "body1"),
                                            getBlockByName(Func, "body2") };
 
   CodeExtractor CE(Candidates);
   EXPECT_TRUE(CE.isEligible());
 
-  Function *Outlined = CE.extractCodeRegion();
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit = getBlockByName(Func, "notExtracted");
   BasicBlock *ExitSplit = getBlockByName(Outlined, "notExtracted.split");
   // Ensure that PHI in exit block has only one incoming value (from code
   // replacer block).
   EXPECT_TRUE(Exit && cast<PHINode>(Exit->front()).getNumIncomingValues() == 1);
   // Ensure that there is a PHI in outlined function with 2 incoming values.
   EXPECT_TRUE(ExitSplit &&
TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
   SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "extracted1"),
     getBlockByName(Func, "extracted2")
   };
 
   CodeExtractor CE(ExtractedBlocks);
   EXPECT_TRUE(CE.isEligible());
 
-  Function *Outlined = CE.extractCodeRegion();
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit1 = getBlockByName(Func, "exit1");
   BasicBlock *Exit2 = getBlockByName(Func, "exit2");
   // Ensure that PHIs in exits are not splitted (since that they have only one
   // incoming value from extracted region).
   EXPECT_TRUE(Exit1 &&
           cast<PHINode>(Exit1->front()).getNumIncomingValues() == 2);
   EXPECT_TRUE(Exit2 &&
SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "invoke.cont2"),
     getBlockByName(Func, "invoke.cont3"),
     getBlockByName(Func, "lpad2")
   };
 
   CodeExtractor CE(ExtractedBlocks);
   EXPECT_TRUE(CE.isEligible());
 
-  Function *Outlined = CE.extractCodeRegion();
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined, &errs()));
   EXPECT_FALSE(verifyFunction(*Func, &errs()));
 }
 
 TEST(CodeExtractor, StoreOutputInvokeResultInExitStub) {
   LLVMContext Ctx;
   SMDiagnostic Err;
TEST(CodeExtractor, StoreOutputInvokeResultInExitStub) {
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "entry"),
                                        getBlockByName(Func, "lpad") };
 
   CodeExtractor CE(Blocks);
   EXPECT_TRUE(CE.isEligible());
 
-  Function *Outlined = CE.extractCodeRegion();
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
 TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
   LLVMContext Ctx;
   SMDiagnostic Err;
TEST(CodeExtractor, ExtractAndInvalidateAssumptionCache) {
 
   assert(M && "Could not parse module?");
   Function *Func = M->getFunction("test");
   SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "if.else") };
   AssumptionCache AC(*Func);
   CodeExtractor CE(Blocks, nullptr, false, nullptr, nullptr, &AC);
   EXPECT_TRUE(CE.isEligible());
 
-  Function *Outlined = CE.extractCodeRegion();
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
   EXPECT_FALSE(CE.verifyAssumptionCache(*Func, &AC));
 }
 } // end anonymous namespace