[ctxprof] Flatten indirect call info in pre-thinlink compilation

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -21,6 +21,10 @@ namespace llvm {
 
 class CtxProfAnalysis;
 
+using FlatIndirectTargets = DenseMap<GlobalValue::GUID, uint64_t>;
+using CtxProfFlatIndirectCallProfile =
+    DenseMap<GlobalValue::GUID, DenseMap<uint32_t, FlatIndirectTargets>>;
+
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
   friend class CtxProfAnalysis;
@@ -101,6 +105,7 @@ class PGOContextualProfile {
   void visit(ConstVisitor, const Function *F = nullptr) const;
 
   const CtxProfFlatProfile flatten() const;
+  const CtxProfFlatIndirectCallProfile flattenVirtCalls() const;
 
   bool invalidate(Module &, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &) {

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -334,6 +334,20 @@ const CtxProfFlatProfile PGOContextualProfile::flatten() const {
   return Flat;
 }
 
+const CtxProfFlatIndirectCallProfile
+PGOContextualProfile::flattenVirtCalls() const {
+  CtxProfFlatIndirectCallProfile Ret;
+  preorderVisit<const PGOCtxProfContext::CallTargetMapTy,
+                const PGOCtxProfContext>(
+      Profiles.Contexts, [&](const PGOCtxProfContext &Ctx) {
+        auto &Targets = Ret[Ctx.guid()];
+        for (const auto &[ID, SubctxSet] : Ctx.callsites())
+          for (const auto &Subctx : SubctxSet)
+            Targets[ID][Subctx.first] += Subctx.second.getEntrycount();
+      });
+  return Ret;
+}
+
 void CtxProfAnalysis::collectIndirectCallPromotionList(
     CallBase &IC, Result &Profile,
     SetVector<std::pair<CallBase *, Function *>> &Candidates) {

llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp

@@ -36,9 +36,12 @@
 #include "llvm/Transforms/Scalar/DCE.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <deque>
+#include <functional>
 
 using namespace llvm;
 
+#define DEBUG_TYPE "ctx_prof_flatten"
+
 namespace {
 
 class ProfileAnnotator final {
@@ -414,6 +417,58 @@ void removeInstrumentation(Function &F) {
         I.eraseFromParent();
 }
 
+void annotateIndirectCall(
+    Module &M, CallBase &CB,
+    const DenseMap<uint32_t, FlatIndirectTargets> &FlatProf,
+    const InstrProfCallsite &Ins) {
+  auto Idx = Ins.getIndex()->getZExtValue();
+  auto FIt = FlatProf.find(Idx);
+  if (FIt == FlatProf.end())
+    return;
+  const auto &Targets = FIt->second;
+  SmallVector<InstrProfValueData, 2> Data;
+  uint64_t Sum = 0;
+  for (auto &[Guid, Count] : Targets) {
+    Data.push_back({/*.Value=*/Guid, /*.Count=*/Count});
+    Sum += Count;
+  }
+  struct InstrProfValueDataGTComparer {
+    bool operator()(const InstrProfValueData &A, const InstrProfValueData &B) {
+      return A.Count > B.Count;
+    }
+  };
+  llvm::sort(Data, InstrProfValueDataGTComparer());
+  llvm::annotateValueSite(M, CB, Data, Sum,
+                          InstrProfValueKind::IPVK_IndirectCallTarget,
+                          Data.size());
+  LLVM_DEBUG(dbgs() << "[ctxprof] flat indirect call prof: " << CB
+                    << CB.getMetadata(LLVMContext::MD_prof) << "\n");
+}
+
+// We normally return a "Changed" bool, but the calling pass' run assumes
+// something will change - some profile will be added - so this won't add much
+// by returning false when applicable.
+void annotateIndCalls(Module &M, const CtxProfAnalysis::Result &CtxProf) {
+  const auto FlatIndCalls = CtxProf.flattenVirtCalls();
+  for (auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+    auto FlatProfIter = FlatIndCalls.find(AssignGUIDPass::getGUID(F));
+    if (FlatProfIter == FlatIndCalls.end())
+      continue;
+    const auto &FlatProf = FlatProfIter->second;
+    for (auto &BB : F) {
+      for (auto &I : BB) {
+        auto *CB = dyn_cast<CallBase>(&I);
+        if (!CB || !CB->isIndirectCall())
+          continue;
+        if (auto *Ins = CtxProfAnalysis::getCallsiteInstrumentation(*CB))
+          annotateIndirectCall(M, *CB, FlatProf, *Ins);
+      }
+    }
+  }
+}
+
 } // namespace
 
 PreservedAnalyses PGOCtxProfFlatteningPass::run(Module &M,
@@ -437,6 +492,8 @@ PreservedAnalyses PGOCtxProfFlatteningPass::run(Module &M,
   if (!IsPreThinlink && !CtxProf.isInSpecializedModule())
     return PreservedAnalyses::none();
 
+  if (IsPreThinlink)
+    annotateIndCalls(M, CtxProf);
   const auto FlattenedProfile = CtxProf.flatten();
 
   for (auto &F : M) {

llvm/test/Analysis/CtxProfAnalysis/flatten-insert-icp-mdprof.ll

@@ -0,0 +1,50 @@
+; REQUIRES:x86_64-linux
+
+; Test flattening indirect calls into "VP" MD_prof metadata, in prelink.
+
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromYAML --input %t/profile.yaml --output %t/profile.ctxprofdata
+; RUN: opt -passes=ctx-prof-flatten-prethinlink %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -S -o - | FileCheck %s --check-prefix=PRELINK
+
+; PRELINK:      call void @llvm.instrprof.callsite(ptr @foo, i64 1234, i32 2, i32 0, ptr %p)
+; PRELINK-NEXT: call void %p(), !prof ![[VPPROF:[0-9]+]]
+; PRELINK-NEXT: call void @llvm.instrprof.callsite(ptr @foo, i64 1234, i32 2, i32 1, ptr @bar)
+; PRELINK-NEXT: call void @bar(){{$}}
+; PRELINK:      ![[VPPROF]] = !{!"VP", i32 0, i64 5, i64 5678, i64 4, i64 5555, i64 1}
+
+; RUN: cp %t/example.ll %t/1234.ll
+; RUN: opt -passes=ctx-prof-flatten %t/1234.ll -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -S -o - | FileCheck %s --check-prefix=POSTLINK
+; RUN: opt -passes=ctx-prof-flatten %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -S -o - | FileCheck %s --check-prefix=POSTLINK
+
+; POSTLINK-NOT: call void %p(), !prof
+;--- example.ll
+
+declare !guid !0 void @bar()
+
+define void @foo(ptr %p) !guid !1 {
+  call void @llvm.instrprof.increment(ptr @foo, i64 1234, i32 1, i32 0)
+  call void @llvm.instrprof.callsite(ptr @foo, i64 1234, i32 2, i32 0, ptr %p)
+  call void %p()
+  call void @llvm.instrprof.callsite(ptr @foo, i64 1234, i32 2, i32 1, ptr @bar)
+  call void @bar()
+  ret void
+}
+
+!0 = !{i64 8888}
+!1 = !{i64 1234}
+
+;--- profile.yaml
+Contexts:
+  - Guid: 1234
+    TotalRootEntryCount: 5
+    Counters: [5]
+    Callsites:
+      - - Guid: 5555
+          Counters: [1]
+        - Guid: 5678
+          Counters: [4]
+      - - Guid: 8888
+          Counters: [5]