Add call branch annotation for ICP promoted direct call in SamplePGO mode.

[danielcdh]

Summary

SamplePGO uses branch_weight annotation to represent callsite hotness. When ICP promotes an indirect call to direct call, we need to make sure the direct call is annotated with branch_weight in SamplePGO mode, so that downstream function inliner can use hot callsite heuristic.

[xur]

Not familiar with samplePGO. Are the counters always annotated in the call sites (or the sampled BB)?
The call counts cannot be inferred from entry counts and the branch-weight?
If this the way samplePGO works, this looks fine to me.

[danielcdh]

Yes, SamplePGO annotate all hot callsites with absolute branch weight to prevent BFI from messing up with inaccurate branch probability.

[davidxl]

lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

526

document this is needed for sample PGO callsite hotness.

[danielcdh]

update

[davidxl]

lgtm

include/llvm/Passes/PassBuilder.h

 class AAManager;
 class TargetMachine;
 
 /// A struct capturing PGO tunables.
 struct PGOOptions {
   std::string ProfileGenFile = "";
   std::string ProfileUseFile = "";
   bool RunProfileGen = false;
+  bool SamplePGO = false;
 };
 
 /// \brief This class provides access to building LLVM's passes.
 ///
 /// It's members provide the baseline state available to passes during their
 /// construction. The \c PassRegistry.def file specifies how to construct all
 /// of the built-in passes, and those may reference these members during
 /// construction.

include/llvm/Transforms/Instrumentation.h

 
 ModulePass *createGCOVProfilerPass(const GCOVOptions &Options =
                                    GCOVOptions::getDefault());
 
 // PGO Instrumention
 ModulePass *createPGOInstrumentationGenLegacyPass();
 ModulePass *
 createPGOInstrumentationUseLegacyPass(StringRef Filename = StringRef(""));
-ModulePass *createPGOIndirectCallPromotionLegacyPass(bool InLTO = false);
+ModulePass *createPGOIndirectCallPromotionLegacyPass(bool InLTO = false,
+                                                     bool SamplePGO = false);
 
 // Helper function to check if it is legal to promote indirect call \p Inst
 // to a direct call of function \p F. Stores the reason in \p Reason.
 bool isLegalToPromote(Instruction *Inst, Function *F, const char **Reason);
 
 // Helper function that transforms Inst (either an indirect-call instruction, or
 // an invoke instruction , to a conditional call to F. This is like:
 //     if (Inst.CalledValue == F)
 //        F(...);
 //     else
 //        Inst(...);
 //     end
 // TotalCount is the profile count value that the instruction executes.
 // Count is the profile count value that F is the target function.
 // These two values are used to update the branch weight.
+// If \p AttachProfToDirectCall is true, a prof metadata is attached to the
+// new direct call to contain \p Count.
 // Returns the promoted direct call instruction.
 Instruction *promoteIndirectCall(Instruction *Inst, Function *F, uint64_t Count,
-                                 uint64_t TotalCount);
+                                 uint64_t TotalCount,
+                                 bool AttachProfToDirectCall);
 
 /// Options for the frontend instrumentation based profiling pass.
 struct InstrProfOptions {
   // Add the 'noredzone' attribute to added runtime library calls.
   bool NoRedZone = false;
 
   // Name of the profile file to use as output
   std::string InstrProfileOutput;

include/llvm/Transforms/PGOInstrumentation.h

 
 private:
   std::string ProfileFileName;
 };
 
 /// The indirect function call promotion pass.
 class PGOIndirectCallPromotion : public PassInfoMixin<PGOIndirectCallPromotion> {
 public:
-  PGOIndirectCallPromotion(bool IsInLTO = false) : InLTO(IsInLTO) {}
+  PGOIndirectCallPromotion(bool IsInLTO = false, bool SamplePGO = false)
+      : InLTO(IsInLTO), SamplePGO(SamplePGO) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+
 private:
   bool InLTO;
+  bool SamplePGO;
 };
 
 } // End llvm namespace
 #endif

lib/Passes/PassBuilder.cpp

PassBuilder::buildPerModuleDefaultPipeline(OptimizationLevel Level,
   // optimizations.
   FunctionPassManager GlobalCleanupPM(DebugLogging);
   GlobalCleanupPM.addPass(InstCombinePass());
   GlobalCleanupPM.addPass(SimplifyCFGPass());
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(GlobalCleanupPM)));
 
   // Add all the requested passes for PGO Instrumentation, if requested.
   if (PGOOpt) {
-    assert(PGOOpt->RunProfileGen || !PGOOpt->ProfileUseFile.empty());
+    assert(PGOOpt->RunProfileGen || PGOOpt->SamplePGO ||
+           !PGOOpt->ProfileUseFile.empty());
     addPGOInstrPasses(MPM, DebugLogging, Level, PGOOpt->RunProfileGen,
                       PGOOpt->ProfileGenFile, PGOOpt->ProfileUseFile);
   }
 
   // Indirect call promotion that promotes intra-module targes only.
-  MPM.addPass(PGOIndirectCallPromotion());
+  MPM.addPass(PGOIndirectCallPromotion(false, PGOOpt && PGOOpt->SamplePGO));
 
   // Require the GlobalsAA analysis for the module so we can query it within
   // the CGSCC pipeline.
   MPM.addPass(RequireAnalysisPass<GlobalsAA, Module>());
 
   // Now begin the main postorder CGSCC pipeline.
   // FIXME: The current CGSCC pipeline has its origins in the legacy pass
   // manager and trying to emulate its precise behavior. Much of this doesn't
ModulePassManager PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   // libraries and other oracles.
   MPM.addPass(InferFunctionAttrsPass());
 
   if (Level > 1) {
     // Indirect call promotion. This should promote all the targets that are
     // left by the earlier promotion pass that promotes intra-module targets.
     // This two-step promotion is to save the compile time. For LTO, it should
     // produce the same result as if we only do promotion here.
-    MPM.addPass(PGOIndirectCallPromotion(true /* InLTO */));
+    MPM.addPass(PGOIndirectCallPromotion(true /* InLTO */,
+                                         PGOOpt && PGOOpt->SamplePGO));
 
     // Propagate constants at call sites into the functions they call.  This
     // opens opportunities for globalopt (and inlining) by substituting function
     // pointers passed as arguments to direct uses of functions.
    MPM.addPass(IPSCCPPass());
   }
 
   // Now deduce any function attributes based in the current code.

lib/Transforms/IPO/PassManagerBuilder.cpp

void PassManagerBuilder::populateModulePassManager(
   // For ThinLTO there are two passes of indirect call promotion. The
   // first is during the compile phase when PerformThinLTO=false and
   // intra-module indirect call targets are promoted. The second is during
   // the ThinLTO backend when PerformThinLTO=true, when we promote imported
   // inter-module indirect calls. For that we perform indirect call promotion
   // earlier in the pass pipeline, here before globalopt. Otherwise imported
   // available_externally functions look unreferenced and are removed.
   if (PerformThinLTO)
-    MPM.add(createPGOIndirectCallPromotionLegacyPass(/*InLTO = */ true));
+    MPM.add(createPGOIndirectCallPromotionLegacyPass(/*InLTO = */ true,
+                                                     !PGOSampleUse.empty()));
 
   if (!DisableUnitAtATime) {
     // Infer attributes about declarations if possible.
     MPM.add(createInferFunctionAttrsLegacyPass());
 
     addExtensionsToPM(EP_ModuleOptimizerEarly, MPM);
 
     MPM.add(createIPSCCPPass());          // IP SCCP
void PassManagerBuilder::populateModulePassManager(
 
   if (!PerformThinLTO) {
     /// PGO instrumentation is added during the compile phase for ThinLTO, do
     /// not run it a second time
     addPGOInstrPasses(MPM);
     // Indirect call promotion that promotes intra-module targets only.
     // For ThinLTO this is done earlier due to interactions with globalopt
     // for imported functions.
-    MPM.add(createPGOIndirectCallPromotionLegacyPass());
+    MPM.add(
+        createPGOIndirectCallPromotionLegacyPass(false, !PGOSampleUse.empty()));
   }
 
   if (EnableNonLTOGlobalsModRef)
     // We add a module alias analysis pass here. In part due to bugs in the
     // analysis infrastructure this "works" in that the analysis stays alive
     // for the entire SCC pass run below.
     MPM.add(createGlobalsAAWrapperPass());
 
void PassManagerBuilder::addLTOOptimizationPasses(legacy::PassManagerBase &PM) {
   // Infer attributes about declarations if possible.
   PM.add(createInferFunctionAttrsLegacyPass());
 
   if (OptLevel > 1) {
     // Indirect call promotion. This should promote all the targets that are
     // left by the earlier promotion pass that promotes intra-module targets.
     // This two-step promotion is to save the compile time. For LTO, it should
     // produce the same result as if we only do promotion here.
-    PM.add(createPGOIndirectCallPromotionLegacyPass(true));
+    PM.add(
+        createPGOIndirectCallPromotionLegacyPass(true, !PGOSampleUse.empty()));
 
     // Propagate constants at call sites into the functions they call.  This
     // opens opportunities for globalopt (and inlining) by substituting function
     // pointers passed as arguments to direct uses of functions.
     PM.add(createIPSCCPPass());
   }
 
   // Infer attributes about definitions. The readnone attribute in particular is

lib/Transforms/IPO/SampleProfile.cpp

for (auto I : CIS) {
         const char *Reason = "Callee function not available";
         CalledFunction = F.getParent()->getFunction(CalleeFunctionName);
         if (CalledFunction && isLegalToPromote(I, CalledFunction, &Reason)) {
           // The indirect target was promoted and inlined in the profile, as a
           // result, we do not have profile info for the branch probability.
           // We set the probability to 80% taken to indicate that the static
           // call is likely taken.
           DI = dyn_cast<Instruction>(
-              promoteIndirectCall(I, CalledFunction, 80, 100)
+              promoteIndirectCall(I, CalledFunction, 80, 100, false)
                   ->stripPointerCasts());
           PromotedInsns.insert(I);
         } else {
           DEBUG(dbgs() << "\nFailed to promote indirect call to "
                        << CalleeFunctionName << " because " << Reason << "\n");
           continue;
         }
       }

lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

 
 // Set if the pass is called in LTO optimization. The difference for LTO mode
 // is the pass won't prefix the source module name to the internal linkage
 // symbols.
 static cl::opt<bool> ICPLTOMode("icp-lto", cl::init(false), cl::Hidden,
                                 cl::desc("Run indirect-call promotion in LTO "
                                          "mode"));
 
+// Set if the pass is called in SamplePGO mode. The difference for SamplePGO
+// mode is it will add prof metadatato the created direct call.
+static cl::opt<bool>
+    ICPSamplePGOMode("icp-samplepgo", cl::init(false), cl::Hidden,
+                     cl::desc("Run indirect-call promotion in SamplePGO mode"));
+
 // If the option is set to true, only call instructions will be considered for
 // transformation -- invoke instructions will be ignored.
 static cl::opt<bool>
     ICPCallOnly("icp-call-only", cl::init(false), cl::Hidden,
                 cl::desc("Run indirect-call promotion for call instructions "
                          "only"));
 
 // If the option is set to true, only invoke instructions will be considered for
static cl::opt<bool>
     ICPDUMPAFTER("icp-dumpafter", cl::init(false), cl::Hidden,
                  cl::desc("Dump IR after transformation happens"));
 
 namespace {
 class PGOIndirectCallPromotionLegacyPass : public ModulePass {
 public:
   static char ID;
 
-  PGOIndirectCallPromotionLegacyPass(bool InLTO = false)
-      : ModulePass(ID), InLTO(InLTO) {
+  PGOIndirectCallPromotionLegacyPass(bool InLTO = false, bool SamplePGO = false)
+      : ModulePass(ID), InLTO(InLTO), SamplePGO(SamplePGO) {
     initializePGOIndirectCallPromotionLegacyPassPass(
         *PassRegistry::getPassRegistry());
   }
 
   StringRef getPassName() const override { return "PGOIndirectCallPromotion"; }
 
 private:
   bool runOnModule(Module &M) override;
 
   // If this pass is called in LTO. We need to special handling the PGOFuncName
   // for the static variables due to LTO's internalization.
   bool InLTO;
+
+  // If this pass is called in SamplePGO. We need to add the prof metadata to
+  // the promoted direct call.
+  bool SamplePGO;
 };
 } // end anonymous namespace
 
 char PGOIndirectCallPromotionLegacyPass::ID = 0;
 INITIALIZE_PASS(PGOIndirectCallPromotionLegacyPass, "pgo-icall-prom",
                 "Use PGO instrumentation profile to promote indirect calls to "
                 "direct calls.",
                 false, false)
 
-ModulePass *llvm::createPGOIndirectCallPromotionLegacyPass(bool InLTO) {
-  return new PGOIndirectCallPromotionLegacyPass(InLTO);
+ModulePass *llvm::createPGOIndirectCallPromotionLegacyPass(bool InLTO,
+                                                           bool SamplePGO) {
+  return new PGOIndirectCallPromotionLegacyPass(InLTO, SamplePGO);
 }
 
 namespace {
 // The class for main data structure to promote indirect calls to conditional
 // direct calls.
 class ICallPromotionFunc {
 private:
   Function &F;
   Module *M;
 
   // Symtab that maps indirect call profile values to function names and
   // defines.
   InstrProfSymtab *Symtab;
 
+  bool SamplePGO;
+
   // Test if we can legally promote this direct-call of Target.
   bool isPromotionLegal(Instruction *Inst, uint64_t Target, Function *&F,
                         const char **Reason = nullptr);
 
   // A struct that records the direct target and it's call count.
   struct PromotionCandidate {
     Function *TargetFunction;
     uint64_t Count;
uint32_t tryToPromote(Instruction *Inst,
                         const std::vector<PromotionCandidate> &Candidates,
                         uint64_t &TotalCount);
 
   // Noncopyable
   ICallPromotionFunc(const ICallPromotionFunc &other) = delete;
   ICallPromotionFunc &operator=(const ICallPromotionFunc &other) = delete;
 
 public:
-  ICallPromotionFunc(Function &Func, Module *Modu, InstrProfSymtab *Symtab)
-      : F(Func), M(Modu), Symtab(Symtab) {
-  }
+  ICallPromotionFunc(Function &Func, Module *Modu, InstrProfSymtab *Symtab,
+                     bool SamplePGO)
+      : F(Func), M(Modu), Symtab(Symtab), SamplePGO(SamplePGO) {}
 
   bool processFunction();
 };
 } // end anonymous namespace
 
 bool llvm::isLegalToPromote(Instruction *Inst, Function *F,
                             const char **Reason) {
   // Check the return type.

 // It transforms to:
 //     if (Foo == DirectCallee)
 //        Ret1 = DirectCallee(Args);
 //     else
 //        Ret2 = (*Foo)(Args);
 //     Ret = phi(Ret1, Ret2);
 // It adds type casts for the args do not match the parameters and the return
 // value. Branch weights metadata also updated.
+// If \p AttachProfToDirectCall is true, a prof metadata is attached to the
+// new direct call to contain \p Count.
 // Returns the promoted direct call instruction.
 Instruction *llvm::promoteIndirectCall(Instruction *Inst,
                                        Function *DirectCallee, uint64_t Count,
-                                       uint64_t TotalCount) {
+                                       uint64_t TotalCount,
+                                       bool AttachProfToDirectCall) {
   assert(DirectCallee != nullptr);
   BasicBlock *BB = Inst->getParent();
   // Just to suppress the non-debug build warning.
   (void)BB;
   DEBUG(dbgs() << "\n\n== Basic Block Before ==\n");
   DEBUG(dbgs() << *BB << "\n");
 
   BasicBlock *DirectCallBB, *IndirectCallBB, *MergeBB;
   createIfThenElse(Inst, DirectCallee, Count, TotalCount, &DirectCallBB,
                    &IndirectCallBB, &MergeBB);
 
   Instruction *NewInst =
       createDirectCallInst(Inst, DirectCallee, DirectCallBB, MergeBB);
 
+  if (AttachProfToDirectCall) {
+    SmallVector<uint32_t, 1> Weights;
+    Weights.push_back(Count);
+    MDBuilder MDB(NewInst->getContext());
+    dyn_cast<Instruction>(NewInst->stripPointerCasts())
+        ->setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
+  }
+
   // Move Inst from MergeBB to IndirectCallBB.
   Inst->removeFromParent();
   IndirectCallBB->getInstList().insert(IndirectCallBB->getFirstInsertionPt(),
                                        Inst);
 
   if (InvokeInst *II = dyn_cast<InvokeInst>(Inst)) {
     // At this point, the original indirect invoke instruction has the original
     // UnwindDest and NormalDest. For the direct invoke instruction, the

 // Promote indirect-call to conditional direct-call for one callsite.
 uint32_t ICallPromotionFunc::tryToPromote(
     Instruction *Inst, const std::vector<PromotionCandidate> &Candidates,
     uint64_t &TotalCount) {
   uint32_t NumPromoted = 0;
 
   for (auto &C : Candidates) {
     uint64_t Count = C.Count;
-    promoteIndirectCall(Inst, C.TargetFunction, Count, TotalCount);
+    promoteIndirectCall(Inst, C.TargetFunction, Count, TotalCount, SamplePGO);
     assert(TotalCount >= Count);
     TotalCount -= Count;
     NumOfPGOICallPromotion++;
     NumPromoted++;
   }
   return NumPromoted;
 }
 
for (auto &I : findIndirectCallSites(F)) {
     // Otherwise we need update with the un-promoted records back.
     annotateValueSite(*M, *I, ICallProfDataRef.slice(NumPromoted), TotalCount,
                       IPVK_IndirectCallTarget, NumCandidates);
   }
   return Changed;
 }
 
 // A wrapper function that does the actual work.
-static bool promoteIndirectCalls(Module &M, bool InLTO) {
+static bool promoteIndirectCalls(Module &M, bool InLTO, bool SamplePGO) {
   if (DisableICP)
     return false;
   InstrProfSymtab Symtab;
   Symtab.create(M, InLTO);
   bool Changed = false;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     if (F.hasFnAttribute(Attribute::OptimizeNone))
       continue;
-    ICallPromotionFunc ICallPromotion(F, &M, &Symtab);
+    ICallPromotionFunc ICallPromotion(F, &M, &Symtab, SamplePGO);
     bool FuncChanged = ICallPromotion.processFunction();
     if (ICPDUMPAFTER && FuncChanged) {
       DEBUG(dbgs() << "\n== IR Dump After =="; F.print(dbgs()));
       DEBUG(dbgs() << "\n");
     }
     Changed |= FuncChanged;
     if (ICPCutOff != 0 && NumOfPGOICallPromotion >= ICPCutOff) {
       DEBUG(dbgs() << " Stop: Cutoff reached.\n");
       break;
     }
   }
   return Changed;
 }
 
 bool PGOIndirectCallPromotionLegacyPass::runOnModule(Module &M) {
   // Command-line option has the priority for InLTO.
-  return promoteIndirectCalls(M, InLTO | ICPLTOMode);
+  return promoteIndirectCalls(M, InLTO | ICPLTOMode,
+                              SamplePGO | ICPSamplePGOMode);
 }
 
-PreservedAnalyses PGOIndirectCallPromotion::run(Module &M, ModuleAnalysisManager &AM) {
-  if (!promoteIndirectCalls(M, InLTO | ICPLTOMode))
+PreservedAnalyses PGOIndirectCallPromotion::run(Module &M,
+                                                ModuleAnalysisManager &AM) {
+  if (!promoteIndirectCalls(M, InLTO | ICPLTOMode,
+                            SamplePGO | ICPSamplePGOMode))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }

test/Transforms/PGOProfile/indirect_call_promotion.ll

 ; RUN: opt < %s -pgo-icall-prom -S | FileCheck %s --check-prefix=ICALL-PROM
+; RUN: opt < %s -pgo-icall-prom -S -icp-samplepgo | FileCheck %s --check-prefix=ICALL-PROM
+; RUN: opt < %s -pgo-icall-prom -S -icp-samplepgo | FileCheck %s --check-prefix=ICALL-PROM-SAMPLEPGO
 ; RUN: opt < %s -passes=pgo-icall-prom -S | FileCheck %s --check-prefix=ICALL-PROM
 ; RUN: opt < %s -pgo-icall-prom -S -pass-remarks=pgo-icall-prom -icp-count-threshold=0 -icp-percent-threshold=0 -icp-max-prom=4 2>&1 | FileCheck %s --check-prefix=PASS-REMARK
 ; RUN: opt < %s -passes=pgo-icall-prom -S -pass-remarks=pgo-icall-prom -icp-count-threshold=0 -icp-percent-threshold=0 -icp-max-prom=4 2>&1 | FileCheck %s --check-prefix=PASS-REMARK
 ; PASS-REMARK: remark: <unknown>:0:0: Promote indirect call to func4 with count 1030 out of 1600
 ; PASS-REMARK: remark: <unknown>:0:0: Promote indirect call to func2 with count 410 out of 570
 ; PASS-REMARK: remark: <unknown>:0:0: Promote indirect call to func3 with count 150 out of 160
 ; PASS-REMARK: remark: <unknown>:0:0: Promote indirect call to func1 with count 10 out of 10
 

 define i32 @bar() {
 entry:
   %tmp = load i32 ()*, i32 ()** @foo, align 8
 ; ICALL-PROM:   [[BITCAST:%[0-9]+]] = bitcast i32 ()* %tmp to i8*
 ; ICALL-PROM:   [[CMP:%[0-9]+]] = icmp eq i8* [[BITCAST]], bitcast (i32 ()* @func4 to i8*)
 ; ICALL-PROM:   br i1 [[CMP]], label %if.true.direct_targ, label %if.false.orig_indirect, !prof [[BRANCH_WEIGHT:![0-9]+]]
 ; ICALL-PROM: if.true.direct_targ:
 ; ICALL-PROM:   [[DIRCALL_RET:%[0-9]+]] = call i32 @func4()
+; ICALL-PROM-SAMPLEPGO: call i32 @func4(), !prof [[CALL_METADATA:![0-9]+]]
 ; ICALL-PROM:   br label %if.end.icp
   %call = call i32 %tmp(), !prof !1
 ; ICALL-PROM: if.false.orig_indirect:
 ; ICALL-PROM:   %call = call i32 %tmp(), !prof [[NEW_VP_METADATA:![0-9]+]]
   ret i32 %call
 ; ICALL-PROM: if.end.icp:
 ; ICALL-PROM:   [[PHI_RET:%[0-9]+]] = phi i32 [ %call, %if.false.orig_indirect ], [ [[DIRCALL_RET]], %if.true.direct_targ ]
 ; ICALL-PROM:   ret i32 [[PHI_RET]]
 }
 
 !1 = !{!"VP", i32 0, i64 1600, i64 7651369219802541373, i64 1030, i64 -4377547752858689819, i64 410, i64 -6929281286627296573, i64 150, i64 -2545542355363006406, i64 10}
 
 ; ICALL-PROM: [[BRANCH_WEIGHT]] = !{!"branch_weights", i32 1030, i32 570}
 ; ICALL-PROM: [[NEW_VP_METADATA]] = !{!"VP", i32 0, i64 570, i64 -4377547752858689819, i64 410}
+; ICALL-PROM-SAMPLEPGO: [[CALL_METADATA]] = !{!"branch_weights", i32 1030}