[IRSim] Reduce IROutliner Compile time by reducing the number of time the CodeExtractorCache is rebuilt.

[AndrewLitteken]

Summary

Currently, in the IROutliner, we rebuild the CodeExtractorCache for every function when it needs to be analyzed for cost, or when it is outlined. For large pieces of code, this takes a long time.

We can instead build the caches at the beginning of the run and then adapt the cache as necessary. This required making the ache accessible from the CodeExtractor, and then adapting the IROutliner to build everything at once, and then reference it when necessary, removing instructions that were no longer in the original functions once they are outlined.

This reduces compile time significantly.

[AndrewLitteken]

Fix diff order

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

357

comment?

368

we probably want comments for all calls to this explaining why we're doing this

2849

For a follow-up patch: maybe Candidates should have a getFunction function?

llvm/lib/Transforms/Utils/CodeExtractor.cpp

322

comments in this function?

335

comments here too?

[AndrewLitteken]

llvm/lib/Transforms/IPO/IROutliner.cpp

2849

It turns out that they do, and I just forgot to use it.

llvm/include/llvm/Transforms/IPO/IROutliner.h

struct OutlinableRegion {
   /// The global value numbers that are used as outputs for this section. Once
   /// extracted, each output will be stored to an output register.  This
   /// documents the global value numbers that are used in this pattern.
   SmallVector<unsigned, 4> GVNStores;
 
   /// Used to create an outlined function.
   CodeExtractor *CE = nullptr;
 
+  /// Used to create an outlined function.
+  CodeExtractorAnalysisCache *CEAC = nullptr;
+
   /// The call site of the extracted region.
   CallInst *Call = nullptr;
 
   /// The function for the extracted region.
   Function *ExtractedFunction = nullptr;
 
   /// Flag for whether we have split out the IRSimilarityCanidate. That is,
   /// make the region contained the IRSimilarityCandidate its own BasicBlock.
private:
   /// TargetTransformInfo lambda for target specific information.
   function_ref<TargetTransformInfo &(Function &)> getTTI;
 
   /// A mapping from newly created reloaded output values to the original value.
   /// If an value is replace by an output from an outlined region, this maps
   /// that Value, back to its original Value.
   DenseMap<Value *, Value *> OutputMappings;
 
+  /// A mapping from functions in the original modules to CodeExtractorAnalyis
+  /// Caches.  These are updated as functions are outlined in order to be
+  /// consistent with vairbale sthat may have been removed, or moved.
+  DenseMap<Function *, CodeExtractorAnalysisCache *> CEACs;
+
   /// IRSimilarityIdentifier lambda to retrieve IRSimilarityIdentifier.
   function_ref<IRSimilarityIdentifier &(Module &)> getIRSI;
 
   /// The optimization remark emitter for the pass.
   function_ref<OptimizationRemarkEmitter &(Function &)> getORE;
 
   /// The memory allocator used to allocate the CodeExtractors.
   SpecificBumpPtrAllocator<CodeExtractor> ExtractorAllocator;
 
+  /// The memory allocator used to allocate the CodeExtractors.
+  SpecificBumpPtrAllocator<CodeExtractorAnalysisCache> CEACAllocator;
+
   /// The memory allocator used to allocate the OutlinableRegions.
   SpecificBumpPtrAllocator<OutlinableRegion> RegionAllocator;
 
   /// The memory allocator used to allocate new IRInstructionData.
   SpecificBumpPtrAllocator<IRInstructionData> InstDataAllocator;
 
   /// Custom InstVisitor to classify different instructions for whether it can
   /// be analyzed for similarity.  This is needed as there may be instruction we

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

class CodeExtractorAnalysisCache {
   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;
   
+  /// The set of Instructions that have side effects, and the BasicBlock they
+  /// were initially in.
+  DenseMap<BasicBlock *, Instruction *> SideEffectingInstructions;
+
   void findSideEffectInfoForBlock(BasicBlock &BB);
 
 public:
   CodeExtractorAnalysisCache(Function &F);
 
+  /// For a \p NewBlock that has been created from \p OldBlock, such that \p
+  /// OldBlock is a predecessor of \p NewBlock, check if \p NewBlock constains
+  /// any side effecting instructions such as the corresponding allocations for
+  /// loads, and stores, and record them.
+  ///
+  /// \param OldBlock - The original BasicBlock that was in the function.
+  /// \param NewBlock - The new BasicBlock with \p OldBlock as a predecessor.
+  void TestAndAddSideEffectingBlock(BasicBlock *OldBlock, BasicBlock *NewBlock);
+
+  /// Remove enteries in the CodeExtractorAnalysisCache reference \p RemoveBlock
+  /// when it is no longer included in the Function that the cache was created
+  /// from.
+  ///
+  /// \param RemoveBlock - The BasicBlock to remove references to in the cache.
+  void TestAndRemoveSideEffectingBlock(BasicBlock *RemovedBlock);
+
   /// 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;

llvm/lib/Transforms/IPO/IROutliner.cpp

void OutlinableRegion::splitCandidate() {
 
   StartBB = PrevBB->splitBasicBlock(StartInst, OriginalName + "_to_outline");
   PrevBB->replaceSuccessorsPhiUsesWith(PrevBB, StartBB);
   // If there was a PHINode with an incoming block outside the region,
   // make sure is correctly updated in the newly split block.
   if (PHIPredBlock)
     PrevBB->replaceSuccessorsPhiUsesWith(PHIPredBlock, PrevBB);
 
+  // Add StartBB to the cache, and record its corresponding side effects so that
+  // inputs and outputs will be correctly analyzed.
+  CEAC->TestAndAddSideEffectingBlock(PrevBB, StartBB);
+
   CandidateSplit = true;
   if (!BackInst->isTerminator()) {
     EndBB = EndInst->getParent();
     FollowBB = EndBB->splitBasicBlock(EndInst, OriginalName + "_after_outline");
     EndBB->replaceSuccessorsPhiUsesWith(EndBB, FollowBB);
     FollowBB->replaceSuccessorsPhiUsesWith(PrevBB, FollowBB);
+
+    if (EndBB == StartBB)
+      // Add FollowBB to the cache, and record its corresponding side effects so
+      // that inputs and outputs will be correctly analyzed.
+      CEAC->TestAndAddSideEffectingBlock(PrevBB, FollowBB);
+    else
+      // Add FollowBB to the cache, and record its corresponding side effects so
+      // that inputs and outputs will be correctly analyzed. In this case we had
+      // multiple BasicBlocks, and EndBB is the parent block.
+      CEAC->TestAndAddSideEffectingBlock(EndBB, FollowBB);
   } else {
     EndBB = BackInst->getParent();
     EndsInBranch = true;
     FollowBB = nullptr;
   }
 
   // Refind the basic block set.
   BBSet.clear();
if (!ExtractedFunction) {
     Candidate->getBasicBlocks(BBSet);
 
     replaceTargetsFromPHINode(StartBB, StartBB, PrevBB, BBSet);
     if (!EndsInBranch)
       replaceTargetsFromPHINode(FollowBB, FollowBB, EndBB, BBSet);
   }
 
   moveBBContents(*StartBB, *PrevBB);
+  // As StartBB no longer exists, we remove references to it from the
+  // CodeExtractorAnalysis cache.
+  CEAC->TestAndRemoveSideEffectingBlock(StartBB);
 
   BasicBlock *PlacementBB = PrevBB;
   if (StartBB != EndBB)
     PlacementBB = EndBB;
   if (!EndsInBranch && PlacementBB->getUniqueSuccessor() != nullptr) {
     assert(FollowBB != nullptr && "FollowBB for Candidate is not defined!");
     assert(PlacementBB->getTerminator() && "Terminator removed from EndBB!");
     PlacementBB->getTerminator()->eraseFromParent();
     moveBBContents(*FollowBB, *PlacementBB);
+    // As FollowBB no longer exists, we remove references to it from the
+    // CodeExtractorAnalysis cache.
+    CEAC->TestAndRemoveSideEffectingBlock(FollowBB);
     PlacementBB->replaceSuccessorsPhiUsesWith(FollowBB, PlacementBB);
     FollowBB->eraseFromParent();
   }
 
   PrevBB->replaceSuccessorsPhiUsesWith(StartBB, PrevBB);
   StartBB->eraseFromParent();
 
   // Make sure to save changes back to the StartBB.
static void getCodeExtractorArguments(
   // the correct set of extracted outputs will be in the final Outputs ValueSet.
   SetVector<Value *> OverallInputs, PremappedInputs, SinkCands, HoistCands,
       DummyOutputs;
 
   // Use the code extractor to get the inputs and outputs, without sunken
   // allocas or removing llvm.assumes.
   CodeExtractor *CE = Region.CE;
   CE->findInputsOutputs(OverallInputs, DummyOutputs, SinkCands);
-  assert(Region.StartBB && "Region must have a start BasicBlock!");
-  Function *OrigF = Region.StartBB->getParent();
-  CodeExtractorAnalysisCache CEAC(*OrigF);
-  BasicBlock *Dummy = nullptr;
-
   // The region may be ineligible due to VarArgs in the parent function. In this
   // case we ignore the region.
   if (!CE->isEligible()) {
     Region.IgnoreRegion = true;
     return;
   }
 
+  assert(Region.StartBB && "Region must have a start BasicBlock!");
+  BasicBlock *Dummy = nullptr;
+
   // Find if any values are going to be sunk into the function when extracted
-  CE->findAllocas(CEAC, SinkCands, HoistCands, Dummy);
+  CE->findAllocas(*Region.CEAC, SinkCands, HoistCands, Dummy);
   CE->findInputsOutputs(PremappedInputs, Outputs, SinkCands);
 
   // TODO: Support regions with sunken allocas: values whose lifetimes are
   // contained completely within the outlined region.  These are not guaranteed
   // to be the same in every region, so we must elevate them all to arguments
   // when they appear.  If these values are not equal, it means there is some
   // Input in OverallInputs that was removed for ArgInputs.
   if (OverallInputs.size() != PremappedInputs.size()) {
if (OutputMappings.find(Outputs[OutputIdx.value()]) == OutputMappings.end()) {
     OutputMappings.insert(std::make_pair(LI, Orig));
   }
 }
 
 bool IROutliner::extractSection(OutlinableRegion &Region) {
   SetVector<Value *> ArgInputs, Outputs, SinkCands;
   assert(Region.StartBB && "StartBB for the OutlinableRegion is nullptr!");
   BasicBlock *InitialStart = Region.StartBB;
-  Function *OrigF = Region.StartBB->getParent();
-  CodeExtractorAnalysisCache CEAC(*OrigF);
   Region.ExtractedFunction =
-      Region.CE->extractCodeRegion(CEAC, ArgInputs, Outputs);
+      Region.CE->extractCodeRegion(*Region.CEAC, ArgInputs, Outputs);
 
   // If the extraction was successful, find the BasicBlock, and reassign the
   // OutlinableRegion blocks
   if (!Region.ExtractedFunction) {
     LLVM_DEBUG(dbgs() << "CodeExtractor failed to outline " << Region.StartBB
                       << "\n");
     Region.reattachCandidate();
     return false;
for (SimilarityGroup &CandidateVec : SimilarityCandidates) {
     if (CurrentGroup.IgnoreGroup)
       continue;
 
     // Create a CodeExtractor for each outlinable region. Identify inputs and
     // outputs for each section using the code extractor and create the argument
     // types for the Aggregate Outlining Function.
     OutlinedRegions.clear();
     for (OutlinableRegion *OS : CurrentGroup.Regions) {
+      Function *OrigF = OS->Candidate->getFunction();
+      DenseMap<Function *, CodeExtractorAnalysisCache *>::iterator CEACIt =
+          CEACs.find(OrigF);
+      if (CEACIt == CEACs.end()) {
+        bool Inserted = false;
+        CodeExtractorAnalysisCache *CEAC =
+          new (CEACAllocator.Allocate()) CodeExtractorAnalysisCache(*OrigF);
+        std::tie(CEACIt, Inserted) = CEACs.insert(std::make_pair(OrigF, CEAC));
+      }
+      assert(CEACIt != CEACs.end() &&
+             "Could not find a CodeExtractorAnalysisCache for function!");
+      OS->CEAC = CEACIt->second;
+  
       // Break the outlinable region out of its parent BasicBlock into its own
       // BasicBlocks (see function implementation).
       OS->splitCandidate();
 
       // There's a chance that when the region is split, extra instructions are
       // added to the region. This makes the region no longer viable
       // to be split, so we ignore it for outlining.
       if (!OS->CandidateSplit)

llvm/lib/Transforms/Utils/CodeExtractor.cpp

for (BasicBlock &BB : F) {
     for (Instruction &II : BB.instructionsWithoutDebug())
       if (auto *AI = dyn_cast<AllocaInst>(&II))
         Allocas.push_back(AI);
 
     findSideEffectInfoForBlock(BB);
   }
 }
 
+void CodeExtractorAnalysisCache::TestAndAddSideEffectingBlock(
+    BasicBlock *OldBlock, BasicBlock *NewBlock) {
+  if (BaseMemAddrs.find(OldBlock) != BaseMemAddrs.end()) {
+    findSideEffectInfoForBlock(*NewBlock);
+    return;
+  }
+
+  if (SideEffectingBlocks.find(OldBlock) == SideEffectingBlocks.end())
+    return;
+  
+  findSideEffectInfoForBlock(*NewBlock);
+}
+
+void CodeExtractorAnalysisCache::TestAndRemoveSideEffectingBlock(
+    BasicBlock *RemovedBlock) {
+  if (BaseMemAddrs.find(RemovedBlock) != BaseMemAddrs.end())
+    BaseMemAddrs.erase(RemovedBlock);
+
+  if (SideEffectingBlocks.find(RemovedBlock) == SideEffectingBlocks.end())
+    return;
+  
+  SideEffectingBlocks.erase(RemovedBlock);
+}
+
 void CodeExtractorAnalysisCache::findSideEffectInfoForBlock(BasicBlock &BB) {
   for (Instruction &II : BB.instructionsWithoutDebug()) {
     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 (isa<Constant>(MemAddr))
         break;
       Value *Base = MemAddr->stripInBoundsConstantOffsets();
       if (!isa<AllocaInst>(Base)) {
         SideEffectingBlocks.insert(&BB);
+        SideEffectingInstructions.insert(std::make_pair(&BB, &II));
         return;
       }
       BaseMemAddrs[&BB].insert(Base);
       break;
     }
     default: {
       IntrinsicInst *IntrInst = dyn_cast<IntrinsicInst>(&II);
       if (IntrInst) {
         if (IntrInst->isLifetimeStartOrEnd())
           break;
         SideEffectingBlocks.insert(&BB);
+        SideEffectingInstructions.insert(std::make_pair(&BB, &II));
         return;
       }
       // Treat all the other cases conservatively if it has side effects.
       if (II.mayHaveSideEffects()) {
         SideEffectingBlocks.insert(&BB);
+        SideEffectingInstructions.insert(std::make_pair(&BB, &II));
         return;
       }
     }
     }
   }
 }
 
 bool CodeExtractorAnalysisCache::doesBlockContainClobberOfAddr(