[PartialInlining] Add support in partial inliner to use cost analysis in inliner

[davidxl]

Summary

Added all the hooks such that getInlineCost interface can be used in partial inliner. It is important to make sure the partial inliner share the cost model as the primary one.
(partial ininlining is still not enabled)

[davide]

lib/Transforms/IPO/PartialInlining.cpp

111

Why don't we preserve TTI? (not that it matters for performances, but)

302–306

Can you add a test that triggers this?

309–323

and these as well, maybe

[davidxl]

Address review feedback -- add one more test case.

[eraman]

LGTM with a few minor comments.

lib/Transforms/IPO/PartialInlining.cpp

80

I know this is how it is in InlineCost.cpp, but still: why not use function_ref everywhere?

302

This is fine, but perhaps you can filter out functions with always_inline attribute earlier and not do the analysis at all. That can be done in a separate patch. Similarly for noinline as well.

475

This should also be moved into the loop above (and the description of the statistic changed to indicate that it is the number of callsites)

lib/Transforms/IPO/PartialInlining.cpp

 // that surrounds the body of the function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationDiagnosticInfo.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/Cloning.h"
struct FunctionOutliningInfo {
   BasicBlock *ReturnBlock;
   // The dominating block of the region ot be outlined.
   BasicBlock *NonReturnBlock;
   // The set of blocks in Entries that that are predecessors to ReturnBlock
   SmallVector<BasicBlock *, 4> ReturnBlockPreds;
 };
 
 struct PartialInlinerImpl {
-  PartialInlinerImpl(InlineFunctionInfo IFI) : IFI(std::move(IFI)) {}
+  PartialInlinerImpl(
+      std::function<AssumptionCache &(Function &)> *GetAC,
+      std::function<TargetTransformInfo &(Function &)> *GTTI,
+      Optional<function_ref<BlockFrequencyInfo &(Function &)>> GBFI,
+      ProfileSummaryInfo *ProfSI)
+      : GetAssumptionCache(GetAC), GetTTI(GTTI), GetBFI(GBFI), PSI(ProfSI) {}
   bool run(Module &M);
   Function *unswitchFunction(Function *F);
 
   std::unique_ptr<FunctionOutliningInfo> computeOutliningInfo(Function *F);
 
 private:
-  InlineFunctionInfo IFI;
   int NumPartialInlining = 0;
+  std::function<AssumptionCache &(Function &)> *GetAssumptionCache;
+  std::function<TargetTransformInfo &(Function &)> *GetTTI;
+  Optional<function_ref<BlockFrequencyInfo &(Function &)>> GetBFI;
+  ProfileSummaryInfo *PSI;
 
+  bool shouldPartialInline(CallSite CS, OptimizationRemarkEmitter &ORE);
   bool IsLimitReached() {
     return (MaxNumPartialInlining != -1 &&
             NumPartialInlining >= MaxNumPartialInlining);
   }
 };
 
 struct PartialInlinerLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   PartialInlinerLegacyPass() : ModulePass(ID) {
     initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
+    AU.addRequired<ProfileSummaryInfoWrapperPass>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
   }
   bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
 
     AssumptionCacheTracker *ACT = &getAnalysis<AssumptionCacheTracker>();
+    TargetTransformInfoWrapperPass *TTIWP =
+        &getAnalysis<TargetTransformInfoWrapperPass>();
+    ProfileSummaryInfo *PSI =
+        getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
+
     std::function<AssumptionCache &(Function &)> GetAssumptionCache =
         [&ACT](Function &F) -> AssumptionCache & {
       return ACT->getAssumptionCache(F);
     };
-    InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
-    return PartialInlinerImpl(IFI).run(M);
+
+    std::function<TargetTransformInfo &(Function &)> GetTTI =
+        [&TTIWP](Function &F) -> TargetTransformInfo & {
+      return TTIWP->getTTI(F);
+    };
+
+    return PartialInlinerImpl(&GetAssumptionCache, &GetTTI, None, PSI).run(M);
   }
 };
 }
 
 std::unique_ptr<FunctionOutliningInfo>
 PartialInlinerImpl::computeOutliningInfo(Function *F) {
   BasicBlock *EntryBlock = &F->front();
   BranchInst *BR = dyn_cast<BranchInst>(EntryBlock->getTerminator());
while (OutliningInfo->GetNumInlinedBlocks() < MaxNumInlineBlocks) {
     OutliningInfo->NonReturnBlock = NonReturnBlock;
     OutliningInfo->ReturnBlockPreds.push_back(Cand);
     Entries.insert(Cand);
   }
 
   return OutliningInfo;
 }
 
+bool PartialInlinerImpl::shouldPartialInline(CallSite CS,
+                                             OptimizationRemarkEmitter &ORE) {
+  // TODO : more sharing with shouldInline in Inliner.cpp
+  using namespace ore;
+  Instruction *Call = CS.getInstruction();
+  Function *Callee = CS.getCalledFunction();
+  Function *Caller = CS.getCaller();
+  auto &CalleeTTI = (*GetTTI)(*Callee);
+  InlineCost IC = getInlineCost(CS, getInlineParams(), CalleeTTI,
+                                *GetAssumptionCache, GetBFI, PSI);
+
+  if (IC.isAlways()) {
+    ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "AlwaysInline", Call)
+             << NV("Callee", Callee)
+             << " should always be fully inlined, not partially");
+    return false;
+  }
+
+  if (IC.isNever()) {
+    ORE.emit(OptimizationRemarkMissed(DEBUG_TYPE, "NeverInline", Call)
+             << NV("Callee", Callee) << " not partially inlined into "
+             << NV("Caller", Caller)
+             << " because it should never be inlined (cost=never)");
+    return false;
+  }
+
+  if (!IC) {
+    ORE.emit(OptimizationRemarkMissed(DEBUG_TYPE, "TooCostly", Call)
+             << NV("Callee", Callee) << " not partially inlined into "
+             << NV("Caller", Caller) << " because too costly to inline (cost="
+             << NV("Cost", IC.getCost()) << ", threshold="
+             << NV("Threshold", IC.getCostDelta() + IC.getCost()) << ")");
+    return false;
+  }
+
+  ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "CanBePartiallyInlined", Call)
+           << NV("Callee", Callee) << " can be partially inlined into "
+           << NV("Caller", Caller) << " with cost=" << NV("Cost", IC.getCost())
+           << " (threshold="
+           << NV("Threshold", IC.getCostDelta() + IC.getCost()) << ")");
+  return true;
+}
+
 Function *PartialInlinerImpl::unswitchFunction(Function *F) {
 
   if (F->hasAddressTaken())
     return nullptr;
 
   std::unique_ptr<FunctionOutliningInfo> OutliningInfo =
       computeOutliningInfo(F);
 
   if (!OutliningInfo)
     return nullptr;
 
   // Clone the function, so that we can hack away on it.
   ValueToValueMapTy VMap;
   Function *DuplicateFunction = CloneFunction(F, VMap);
-  DuplicateFunction->setLinkage(GlobalValue::InternalLinkage);
   BasicBlock *NewReturnBlock =
       cast<BasicBlock>(VMap[OutliningInfo->ReturnBlock]);
   BasicBlock *NewNonReturnBlock =
       cast<BasicBlock>(VMap[OutliningInfo->NonReturnBlock]);
   DenseSet<BasicBlock *> NewEntries;
   for (BasicBlock *BB : OutliningInfo->Entries) {
     NewEntries.insert(cast<BasicBlock>(VMap[BB]));
   }
for (User *User : Users) {
     else if (InvokeInst *II = dyn_cast<InvokeInst>(User))
       CS = CallSite(II);
     else
       llvm_unreachable("All uses must be calls");
 
     if (IsLimitReached())
       continue;
 
-    NumPartialInlining++;
-
     OptimizationRemarkEmitter ORE(CS.getCaller());
+    if (!shouldPartialInline(CS, ORE))
+      continue;
+
     DebugLoc DLoc = CS.getInstruction()->getDebugLoc();
     BasicBlock *Block = CS.getParent();
     ORE.emit(OptimizationRemark(DEBUG_TYPE, "PartiallyInlined", DLoc, Block)
              << ore::NV("Callee", F) << " partially inlined into "
              << ore::NV("Caller", CS.getCaller()));
 
+    InlineFunctionInfo IFI(nullptr, GetAssumptionCache);
     InlineFunction(CS, IFI);
+    NumPartialInlining++;
   }
 
   // Ditch the duplicate, since we're done with it, and rewrite all remaining
   // users (function pointers, etc.) back to the original function.
   DuplicateFunction->replaceAllUsesWith(F);
   DuplicateFunction->eraseFromParent();
 
   ++NumPartialInlined;
 
   return ExtractedFunction;
 }
 
 bool PartialInlinerImpl::run(Module &M) {
   if (DisablePartialInlining)
     return false;
 
bool PartialInlinerImpl::run(Module &M) {
 
   return Changed;
 }
 
 char PartialInlinerLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(PartialInlinerLegacyPass, "partial-inliner",
                       "Partial Inliner", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(PartialInlinerLegacyPass, "partial-inliner",
                     "Partial Inliner", false, false)
 
 ModulePass *llvm::createPartialInliningPass() {
   return new PartialInlinerLegacyPass();
 }
 
 PreservedAnalyses PartialInlinerPass::run(Module &M,
                                           ModuleAnalysisManager &AM) {
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+
   std::function<AssumptionCache &(Function &)> GetAssumptionCache =
       [&FAM](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
-  InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
-  if (PartialInlinerImpl(IFI).run(M))
+
+  std::function<BlockFrequencyInfo &(Function &)> GetBFI =
+      [&FAM](Function &F) -> BlockFrequencyInfo & {
+    return FAM.getResult<BlockFrequencyAnalysis>(F);
+  };
+
+  std::function<TargetTransformInfo &(Function &)> GetTTI =
+      [&FAM](Function &F) -> TargetTransformInfo & {
+    return FAM.getResult<TargetIRAnalysis>(F);
+  };
+
+  ProfileSummaryInfo *PSI = &AM.getResult<ProfileSummaryAnalysis>(M);
+
+  if (PartialInlinerImpl(&GetAssumptionCache, &GetTTI, {GetBFI}, PSI).run(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
 }

test/Transforms/CodeExtractor/PartialInlineOptRemark.ll

 ; RUN: opt -S -partial-inliner -pass-remarks=partial-inlining  -disable-output < %s 2>&1 | FileCheck %s
 ; RUN: opt -S -passes=partial-inliner  -pass-remarks=partial-inlining -disable-output < %s 2>&1 | FileCheck %s
 ; RUN: opt -S -partial-inliner -pass-remarks=partial-inlining  -disable-output -max-partial-inlining=1 < %s 2>&1 | FileCheck %s
 ; RUN: opt -S -passes=partial-inliner  -pass-remarks=partial-inlining -disable-output -max-partial-inlining=1 < %s 2>&1 | FileCheck %s
 
 ; RUN: opt -S -partial-inliner -pass-remarks=partial-inlining  -disable-partial-inlining < %s 2>&1 | FileCheck --check-prefix=LIMIT %s
 ; RUN: opt -S -passes=partial-inliner  -pass-remarks=partial-inlining  --disable-partial-inlining < %s 2>&1 | FileCheck  --check-prefix=LIMIT %s
 ; RUN: opt -S -partial-inliner -pass-remarks=partial-inlining   -max-partial-inlining=0 < %s 2>&1 | FileCheck --check-prefix=LIMIT  %s
 ; RUN: opt -S -passes=partial-inliner  -pass-remarks=partial-inlining  -max-partial-inlining=0 < %s 2>&1 | FileCheck --check-prefix=LIMIT  %s
+; RUN: opt -S -partial-inliner -pass-remarks=partial-inlining   -inline-threshold=0 < %s 2>&1 | FileCheck --check-prefix=LIMIT  %s
+; RUN: opt -S -passes=partial-inliner  -pass-remarks=partial-inlining  -inline-threshold=0 < %s 2>&1 | FileCheck --check-prefix=LIMIT  %s
 
 define i32 @bar(i32 %arg) local_unnamed_addr #0 !dbg !5 {
 bb:
   %tmp = icmp slt i32 %arg, 0, !dbg !7
   br i1 %tmp, label %bb1, label %bb2, !dbg !8
 
 bb1:                                              ; preds = %bb
   tail call void (...) @foo() #0, !dbg !9
   tail call void (...) @foo() #0, !dbg !10
   tail call void (...) @foo() #0, !dbg !11
   tail call void (...) @foo() #0, !dbg !12
   tail call void (...) @foo() #0, !dbg !13
   tail call void (...) @foo() #0, !dbg !14
   tail call void (...) @foo() #0, !dbg !15
   tail call void (...) @foo() #0, !dbg !16
   tail call void (...) @foo() #0, !dbg !17
   br label %bb2, !dbg !18
 
 bb2:                                              ; preds = %bb1, %bb
   %tmp3 = phi i32 [ 0, %bb1 ], [ 1, %bb ]
   ret i32 %tmp3, !dbg !19
 }
 
+define i32 @bar_noinline(i32 %arg) local_unnamed_addr #1 !dbg !5 {
+bb:
+  %tmp = icmp slt i32 %arg, 0, !dbg !7
+  br i1 %tmp, label %bb1, label %bb2, !dbg !8
+
+bb1:                                              ; preds = %bb
+  tail call void (...) @foo() #0, !dbg !9
+  tail call void (...) @foo() #0, !dbg !10
+  tail call void (...) @foo() #0, !dbg !11
+  br label %bb2, !dbg !18
+
+bb2:                                              ; preds = %bb1, %bb
+  %tmp3 = phi i32 [ 0, %bb1 ], [ 1, %bb ]
+  ret i32 %tmp3, !dbg !19
+}
+
+define i32 @bar_alwaysinline(i32 %arg) local_unnamed_addr #2 !dbg !5 {
+bb:
+  %tmp = icmp slt i32 %arg, 0, !dbg !7
+  br i1 %tmp, label %bb1, label %bb2, !dbg !8
+
+bb1:                                              ; preds = %bb
+  tail call void (...) @foo() #0, !dbg !9
+  tail call void (...) @foo() #0, !dbg !10
+  tail call void (...) @foo() #0, !dbg !11
+  br label %bb2, !dbg !18
+
+bb2:                                              ; preds = %bb1, %bb
+  %tmp3 = phi i32 [ 0, %bb1 ], [ 1, %bb ]
+  ret i32 %tmp3, !dbg !19
+}
+
 ; Function Attrs: nounwind
 declare void @foo(...) local_unnamed_addr #0
 
 ; Function Attrs: nounwind
 define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 !dbg !20 {
 bb:
 ; CHECK:remark{{.*}}bar partially inlined into dummy_caller
+; CHECK-NOT:remark{{.*}}bar_noinline partially inlined into dummy_caller
+; CHECK-NOT:remark{{.*}}bar_alwaysinline partially inlined into dummy_caller
 ; LIMIT-NOT:remark{{.*}}bar partially inlined into dummy_caller
   %tmp = tail call i32 @bar(i32 %arg), !dbg !21
+  %tmp2 = tail call i32 @bar_noinline(i32 %arg), !dbg !21
+  %tmp3 = tail call i32 @bar_alwaysinline(i32 %arg), !dbg !21
   ret i32 %tmp, !dbg !22
 }
 
 attributes #0 = { nounwind }
+attributes #1 = { noinline nounwind }
+attributes #2 = { alwaysinline nounwind }
 
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!3}
 !llvm.ident = !{!4}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang", isOptimized: true, runtimeVersion: 0, emissionKind: NoDebug, enums: !2)
 !1 = !DIFile(filename: "t.c", directory: "/tmp")
 !2 = !{}

test/Transforms/CodeExtractor/PartialInlineOr.ll

 ; RUN: opt < %s -partial-inliner -S | FileCheck %s
 ; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
 ; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck --check-prefix=LIMIT %s
 ; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=LIMIT %s
 
-; Function Attrs: noinline nounwind uwtable
+; Function Attrs: nounwind uwtable
 define i32 @bar(i32 %arg) local_unnamed_addr #0 {
 bb:
   %tmp = icmp slt i32 %arg, 0
   br i1 %tmp, label %bb4, label %bb1
 
 bb1:                                              ; preds = %bb
   %tmp2 = tail call i32 (...) @channels() #1
   %tmp3 = icmp slt i32 %tmp2, %arg
bb5:                                              ; preds = %bb4, %bb1
   %.0 = phi i32 [ 0, %bb4 ], [ 1, %bb1 ]
   ret i32 %.0
 }
 
 declare i32 @channels(...) local_unnamed_addr
 
 declare void @foo(...) local_unnamed_addr
 
-; Function Attrs: noinline nounwind uwtable
+; Function Attrs: nounwind uwtable
 define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
 bb:
 ; CHECK-LABEL: @dummy_caller
 ; CHECK: br i1
 ; CHECK: br i1
 ; CHECK: call void @bar.2_
 ; LIMIT-LABEL: @dummy_caller
 ; LIMIT-NOT: br

 define i32 @dummy_caller2(i32 %arg) local_unnamed_addr #0 {
 ; CHECK: br i1
 ; CHECK: br i1
 ; CHECK: call {{.*}} @bar_multi_ret.1_
   %tmp = tail call i32 @bar_multi_ret(i32 %arg)
   ret i32 %tmp
 }
 
-attributes #0 = { noinline nounwind uwtable }
+attributes #0 = { nounwind uwtable }
 attributes #1 = { nounwind }
 
 !llvm.ident = !{!0}
 
 !0 = !{!"clang version 5.0.0 (trunk 300576)"}