Handle Invoke during sample profiler annotation: make it inlinable.

[danielcdh]

Summary

Previously we reline on inst-combine to remove inlinable invoke instructions. This causes trouble because a few extra optimizations are schedule early that could introduce too much CFG change (e.g. simplifycfg removes too much control flow). This patch handles invoke instruction in-place during sample profile annotation, so that we do not rely on instcombine to remove those invoke instructions.

[dnovillo]

LGTM.

Apologies for the delay.

lib/Transforms/IPO/SampleProfile.cpp

public:
   void dump() { Reader->dump(); }
 
 protected:
   bool runOnFunction(Function &F);
   unsigned getFunctionLoc(Function &F);
   bool emitAnnotations(Function &F);
   ErrorOr<uint64_t> getInstWeight(const Instruction &I) const;
   ErrorOr<uint64_t> getBlockWeight(const BasicBlock *BB) const;
-  const FunctionSamples *findCalleeFunctionSamples(const CallInst &I) const;
+  const FunctionSamples *findCalleeFunctionSamples(const Instruction &I) const;
   const FunctionSamples *findFunctionSamples(const Instruction &I) const;
   bool inlineHotFunctions(Function &F);
   void printEdgeWeight(raw_ostream &OS, Edge E);
   void printBlockWeight(raw_ostream &OS, const BasicBlock *BB) const;
   void printBlockEquivalence(raw_ostream &OS, const BasicBlock *BB);
   bool computeBlockWeights(Function &F);
   void findEquivalenceClasses(Function &F);
   void findEquivalencesFor(BasicBlock *BB1, ArrayRef<BasicBlock *> Descendants,
bool doInitialization(Module &M) override {
     return SampleLoader.doInitialization(M);
   }
   const char *getPassName() const override { return "Sample profile pass"; }
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
   }
+
 private:
   SampleProfileLoader SampleLoader;
 };
 
 class SampleCoverageTracker {
 public:
   SampleCoverageTracker() : SampleCoverage(), TotalUsedSamples(0) {}
 
if (!FS)
     return std::error_code();
 
   // Ignore all intrinsics and branch instructions.
   // Branch instruction usually contains debug info from sources outside of
   // the residing basic block, thus we ignore them during annotation.
   if (isa<BranchInst>(Inst) || isa<IntrinsicInst>(Inst))
     return std::error_code();
 
-  // If a call instruction is inlined in profile, but not inlined here,
+  // If a call/invoke instruction is inlined in profile, but not inlined here,
   // it means that the inlined callsite has no sample, thus the call
   // instruction should have 0 count.
-  const CallInst *CI = dyn_cast<CallInst>(&Inst);
-  if (CI && findCalleeFunctionSamples(*CI))
-    return 0; 
+  bool IsCall = isa<CallInst>(Inst) || isa<InvokeInst>(Inst);
+  if (IsCall && findCalleeFunctionSamples(Inst))
+    return 0;
 
   const DILocation *DIL = DLoc;
   unsigned Lineno = DLoc.getLine();
   unsigned HeaderLineno = DIL->getScope()->getSubprogram()->getLine();
 
   uint32_t LineOffset = getOffset(Lineno, HeaderLineno);
   uint32_t Discriminator = DIL->getDiscriminator();
   ErrorOr<uint64_t> R = FS->findSamplesAt(LineOffset, Discriminator);

 /// \param BB The basic block to query.
 ///
 /// \returns the weight for \p BB.
 ErrorOr<uint64_t>
 SampleProfileLoader::getBlockWeight(const BasicBlock *BB) const {
   DenseMap<uint64_t, uint64_t> CM;
   for (auto &I : BB->getInstList()) {
     const ErrorOr<uint64_t> &R = getInstWeight(I);
-    if (R) CM[R.get()]++;
+    if (R)
+      CM[R.get()]++;
   }
-  if (CM.size() == 0) return std::error_code();
+  if (CM.size() == 0)
+    return std::error_code();
   uint64_t W = 0, C = 0;
   for (const auto &C_W : CM) {
     if (C_W.second == W) {
       C = std::max(C, C_W.first);
     } else if (C_W.second > W) {
       C = C_W.first;
       W = C_W.second;
     }
for (const auto &BB : F) {
     DEBUG(printBlockWeight(dbgs(), &BB));
   }
 
   return Changed;
 }
 
 /// \brief Get the FunctionSamples for a call instruction.
 ///
-/// The FunctionSamples of a call instruction \p Inst is the inlined
+/// The FunctionSamples of a call/invoke instruction \p Inst is the inlined
 /// instance in which that call instruction is calling to. It contains
 /// all samples that resides in the inlined instance. We first find the
 /// inlined instance in which the call instruction is from, then we
 /// traverse its children to find the callsite with the matching
-/// location and callee function name.
+/// location.
 ///
-/// \param Inst Call instruction to query.
+/// \param Inst Call/Invoke instruction to query.
 ///
 /// \returns The FunctionSamples pointer to the inlined instance.
 const FunctionSamples *
-SampleProfileLoader::findCalleeFunctionSamples(const CallInst &Inst) const {
+SampleProfileLoader::findCalleeFunctionSamples(const Instruction &Inst) const {
   const DILocation *DIL = Inst.getDebugLoc();
   if (!DIL) {
     return nullptr;
   }
   DISubprogram *SP = DIL->getScope()->getSubprogram();
   if (!SP)
     return nullptr;
 
if (S.size() == 0)
     return Samples;
   const FunctionSamples *FS = Samples;
   for (int i = S.size() - 1; i >= 0 && FS != nullptr; i--) {
     FS = FS->findFunctionSamplesAt(S[i]);
   }
   return FS;
 }
 
-
 /// \brief Iteratively inline hot callsites of a function.
 ///
 /// Iteratively traverse all callsites of the function \p F, and find if
 /// the corresponding inlined instance exists and is hot in profile. If
 /// it is hot enough, inline the callsites and adds new callsites of the
 /// callee into the caller.
 ///
 /// TODO: investigate the possibility of not invoking InlineFunction directly.
 ///
 /// \param F function to perform iterative inlining.
 ///
 /// \returns True if there is any inline happened.
 bool SampleProfileLoader::inlineHotFunctions(Function &F) {
   bool Changed = false;
   LLVMContext &Ctx = F.getContext();
   std::function<AssumptionCache &(Function &)> GetAssumptionCache = [&](
       Function &F) -> AssumptionCache & { return ACT->getAssumptionCache(F); };
   while (true) {
     bool LocalChanged = false;
-    SmallVector<CallInst *, 10> CIS;
+    SmallVector<Instruction *, 10> CIS;
     for (auto &BB : F) {
       for (auto &I : BB.getInstList()) {
-        CallInst *CI = dyn_cast<CallInst>(&I);
-        if (CI && callsiteIsHot(Samples, findCalleeFunctionSamples(*CI)))
-          CIS.push_back(CI);
+        const FunctionSamples *FS = nullptr;
+        if ((isa<CallInst>(I) || isa<InvokeInst>(I)) &&
+            (FS = findCalleeFunctionSamples(I))) {
+
+          if (callsiteIsHot(Samples, FS))
+            CIS.push_back(&I);
+        }
       }
     }
-    for (auto CI : CIS) {
+    for (auto I : CIS) {
       InlineFunctionInfo IFI(nullptr, ACT ? &GetAssumptionCache : nullptr);
-      Function *CalledFunction = CI->getCalledFunction();
-      DebugLoc DLoc = CI->getDebugLoc();
-      uint64_t NumSamples = findCalleeFunctionSamples(*CI)->getTotalSamples();
-      if (InlineFunction(CI, IFI)) {
+      CallInst *CI = dyn_cast<CallInst>(I);
+      InvokeInst *II = dyn_cast<InvokeInst>(I);
+      Function *CalledFunction =
+          (CI == nullptr ? II->getCalledFunction() : CI->getCalledFunction());
+      DebugLoc DLoc = I->getDebugLoc();
+      uint64_t NumSamples = findCalleeFunctionSamples(*I)->getTotalSamples();
+      if ((CI && InlineFunction(CI, IFI)) || (II && InlineFunction(II, IFI))) {
         LocalChanged = true;
         emitOptimizationRemark(Ctx, DEBUG_TYPE, F, DLoc,
                                Twine("inlined hot callee '") +
                                    CalledFunction->getName() + "' with " +
                                    Twine(NumSamples) + " samples into '" +
                                    F.getName() + "'");
       }
     }
for (auto &BI : F) {
     BasicBlock *BB = &BI;
 
     if (BlockWeights[BB]) {
       for (auto &I : BB->getInstList()) {
         if (CallInst *CI = dyn_cast<CallInst>(&I)) {
           if (!dyn_cast<IntrinsicInst>(&I)) {
             SmallVector<uint32_t, 1> Weights;
             Weights.push_back(BlockWeights[BB]);
-            CI->setMetadata(LLVMContext::MD_prof, 
+            CI->setMetadata(LLVMContext::MD_prof,
                             MDB.createBranchWeights(Weights));
           }
         }
       }
     }
     TerminatorInst *TI = BB->getTerminator();
     if (TI->getNumSuccessors() == 1)
       continue;
if (!F.isDeclaration()) {
       clearFunctionData();
       retval |= runOnFunction(F);
     }
   M.setProfileSummary(Reader->getSummary().getMD(M.getContext()));
   return retval;
 }
 
 bool SampleProfileLoaderLegacyPass::runOnModule(Module &M) {
-  // FIXME: pass in AssumptionCache correctly for the new pass manager. 
+  // FIXME: pass in AssumptionCache correctly for the new pass manager.
   SampleLoader.setACT(&getAnalysis<AssumptionCacheTracker>());
   return SampleLoader.runOnModule(M);
 }
 
 bool SampleProfileLoader::runOnFunction(Function &F) {
   F.setEntryCount(0);
   Samples = Reader->getSamplesFor(F);
   if (!Samples->empty())

test/Transforms/SampleProfile/Inputs/einline.prof

+_Z3foov:200:100
+ 1: _Z3barv:100
+ 3: _Z3barv:100

test/Transforms/SampleProfile/early-inline.ll

+; RUN: opt < %s -instcombine -sample-profile -sample-profile-file=%S/Inputs/einline.prof | FileCheck %s
+
+; Checks if both call and invoke can be inlined early if their inlined
+; instances are hot in profile.
+
+target triple = "x86_64-unknown-linux-gnu"
+
+@_ZTIi = external constant i8*
+
+; Function Attrs: uwtable
+define void @_Z3foov() #0 personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) !dbg !6 {
+  %1 = alloca i8*
+  %2 = alloca i32
+  %3 = alloca i32, align 4
+; CHECK-NOT: call
+  call void @_ZL3barv(), !dbg !9
+; CHECK-NOT: invoke
+  invoke void @_ZL3barv()
+          to label %4 unwind label %5, !dbg !10
+
+; <label>:4:
+  ret void
+
+; <label>:5:
+  %6 = landingpad { i8*, i32 }
+          catch i8* bitcast (i8** @_ZTIi to i8*)
+  ret void
+}
+
+; Function Attrs: nounwind uwtable
+define internal void @_ZL3barv() #1 {
+  ret void
+}
+
+declare i32 @__gxx_personality_v0(...)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: !1)
+!1 = !DIFile(filename: "a", directory: "b/")
+!3 = !{i32 2, !"Dwarf Version", i32 4}
+!4 = !{i32 2, !"Debug Info Version", i32 3}
+!6 = distinct !DISubprogram(linkageName: "_Z3foov", scope: !1, line: 5, scopeLine: 5, unit: !0)
+!9 = !DILocation(line: 6, column: 3, scope: !6)
+!10 = !DILocation(line: 8, column: 5, scope: !11)
+!11 = distinct !DILexicalBlock(scope: !6, file: !1, line: 7, column: 7)