[AlwaysInliner] Update BFI when inlining

[davidxl]

llvm/test/Transforms/Inline/prof-update-sample.ll

10

Clone this function with always inline attribute instead of modifying the original function.

[aeubanks]

copy test with alwaysinline callee

[davidxl]

lgtm

llvm/lib/Transforms/IPO/AlwaysInliner.cpp

 // are marked as "always inline".
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/AlwaysInliner.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/Inliner.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "inline"
 
 PreservedAnalyses AlwaysInlinerPass::run(Module &M,
                                          ModuleAnalysisManager &MAM) {
   // Add inline assumptions during code generation.
   FunctionAnalysisManager &FAM =
       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
-  InlineFunctionInfo IFI(/*cg=*/nullptr, GetAssumptionCache);
+  auto &PSI = MAM.getResult<ProfileSummaryAnalysis>(M);
 
   SmallSetVector<CallBase *, 16> Calls;
   bool Changed = false;
   SmallVector<Function *, 16> InlinedFunctions;
   for (Function &F : M)
     if (!F.isDeclaration() && F.hasFnAttribute(Attribute::AlwaysInline) &&
         isInlineViable(F).isSuccess()) {
       Calls.clear();
if (!F.isDeclaration() && F.hasFnAttribute(Attribute::AlwaysInline) &&
             [&](CallBase &CB) {
               return InlineCost::getAlways("always inline attribute");
             },
             ORE);
         assert(OIC);
         emitInlinedInto(ORE, CB->getDebugLoc(), CB->getParent(), F, *Caller,
                         *OIC, false, DEBUG_TYPE);
 
+        InlineFunctionInfo IFI(
+            /*cg=*/nullptr, GetAssumptionCache, &PSI,
+            &FAM.getResult<BlockFrequencyAnalysis>(*(CB->getCaller())),
+            &FAM.getResult<BlockFrequencyAnalysis>(F));
+
         InlineResult Res =
             InlineFunction(*CB, IFI, /*CalleeAAR=*/nullptr, InsertLifetime);
         assert(Res.isSuccess() && "unexpected failure to inline");
         (void)Res;
         Changed = true;
       }
 
       // Remember to try and delete this function afterward. This both avoids

llvm/test/Transforms/Inline/prof-update-sample.ll

 ; RUN: opt < %s -inline -S | FileCheck %s
+; RUN: opt < %s -passes=always-inline -S | FileCheck %s
 ; Checks if inliner updates branch_weights annotation for call instructions.
 
 declare void @ext();
 declare void @ext1();
 @func = global void ()* null
 
-; CHECK: define void @callee(i32 %n) !prof ![[ENTRY_COUNT:[0-9]*]]
-define void  @callee(i32 %n) !prof !15 {
+; CHECK: define void @callee(i32 %n) #0 !prof ![[ENTRY_COUNT:[0-9]*]]
+define void  @callee(i32 %n) #0 !prof !15 {
   %cond = icmp sle i32 %n, 10
   br i1 %cond, label %cond_true, label %cond_false
 cond_true:
 ; ext1 is optimized away, thus not updated.
 ; CHECK: call void @ext1(), !prof ![[COUNT_CALLEE1:[0-9]*]]
   call void @ext1(), !prof !16
   ret void
 cond_false: