From 051f2c144e1eee5ed94cc478075a64ee2d0f2745 Mon Sep 17 00:00:00 2001
From: Wenlei He <aktoon@gmail.com>
Date: Sun, 7 Mar 2021 09:21:09 -0800
Subject: [PATCH] [SamplePGO] Skip inlinee profile scaling for sample loader
 inlining

For CGSCC inline, we need to scale down a function's branch weights and entry counts when thee it's inlined at a callsite. This is done through updateCallProfile. Additionally, we also scale the weigths for the inlined clone based on call site count in updateCallerBFI. Neither is needed for inlining during sample profile loader as it's using context profile that is separated from inlinee's own profile. This change skip the inlinee profile scaling for sample loader inlining.

Differential Revision: https://reviews.llvm.org/D98187
---
 llvm/include/llvm/Transforms/Utils/Cloning.h |  9 +++++++--
 llvm/lib/Transforms/IPO/SampleProfile.cpp    |  1 +
 llvm/lib/Transforms/Utils/InlineFunction.cpp | 20 +++++++++++++-------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Utils/Cloning.h b/llvm/include/llvm/Transforms/Utils/Cloning.h
index 3eb7485623679..929db306e4ba2 100644
--- a/llvm/include/llvm/Transforms/Utils/Cloning.h
+++ b/llvm/include/llvm/Transforms/Utils/Cloning.h
@@ -197,9 +197,10 @@ class InlineFunctionInfo {
       function_ref<AssumptionCache &(Function &)> GetAssumptionCache = nullptr,
       ProfileSummaryInfo *PSI = nullptr,
       BlockFrequencyInfo *CallerBFI = nullptr,
-      BlockFrequencyInfo *CalleeBFI = nullptr)
+      BlockFrequencyInfo *CalleeBFI = nullptr, bool UpdateProfile = true)
       : CG(cg), GetAssumptionCache(GetAssumptionCache), PSI(PSI),
-        CallerBFI(CallerBFI), CalleeBFI(CalleeBFI) {}
+        CallerBFI(CallerBFI), CalleeBFI(CalleeBFI),
+        UpdateProfile(UpdateProfile) {}
 
   /// If non-null, InlineFunction will update the callgraph to reflect the
   /// changes it makes.
@@ -223,6 +224,10 @@ class InlineFunctionInfo {
   /// `InlinedCalls` above is used.
   SmallVector<CallBase *, 8> InlinedCallSites;
 
+  /// Update profile for callee as well as cloned version. We need to do this
+  /// for regular inlining, but not for inlining from sample profile loader.
+  bool UpdateProfile;
+
   void reset() {
     StaticAllocas.clear();
     InlinedCalls.clear();
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 22fd951c5c9da..5e53e63947769 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -1093,6 +1093,7 @@ bool SampleProfileLoader::tryInlineCandidate(
     return false;
 
   InlineFunctionInfo IFI(nullptr, GetAC);
+  IFI.UpdateProfile = false;
   if (InlineFunction(CB, IFI).isSuccess()) {
     // The call to InlineFunction erases I, so we can't pass it here.
     emitInlinedInto(*ORE, DLoc, BB, *CalledFunction, *BB->getParent(), Cost,
diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
index 5f75ead1247bb..e754721597c2c 100644
--- a/llvm/lib/Transforms/Utils/InlineFunction.cpp
+++ b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1954,13 +1954,19 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
     if (objcarc::hasAttachedCallOpBundle(&CB))
       inlineRetainOrClaimRVCalls(CB, Returns);
 
-    if (IFI.CallerBFI != nullptr && IFI.CalleeBFI != nullptr)
-      // Update the BFI of blocks cloned into the caller.
-      updateCallerBFI(OrigBB, VMap, IFI.CallerBFI, IFI.CalleeBFI,
-                      CalledFunc->front());
-
-    updateCallProfile(CalledFunc, VMap, CalledFunc->getEntryCount(), CB,
-                      IFI.PSI, IFI.CallerBFI);
+    // Updated caller/callee profiles only when requested. For sample loader
+    // inlining, the context-sensitive inlinee profile doesn't need to be
+    // subtracted from callee profile, and the inlined clone also doesn't need
+    // to be scaled based on call site count.
+    if (IFI.UpdateProfile) {
+      if (IFI.CallerBFI != nullptr && IFI.CalleeBFI != nullptr)
+        // Update the BFI of blocks cloned into the caller.
+        updateCallerBFI(OrigBB, VMap, IFI.CallerBFI, IFI.CalleeBFI,
+                        CalledFunc->front());
+
+      updateCallProfile(CalledFunc, VMap, CalledFunc->getEntryCount(), CB,
+                        IFI.PSI, IFI.CallerBFI);
+    }
 
     // Inject byval arguments initialization.
     for (std::pair<Value*, Value*> &Init : ByValInit)