[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

@@ -39,6 +39,8 @@
 
 using namespace llvm;
 
+#define DEBUG_TYPE "ctx_prof_flatten"
+
 namespace {
 
 class ProfileAnnotator final {
@@ -414,6 +416,52 @@ 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;
+  }
+  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 +485,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) {