[ctx_prof] CtxProfAnalysis: populate module data

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -9,23 +9,39 @@
 #ifndef LLVM_ANALYSIS_CTXPROFANALYSIS_H
 #define LLVM_ANALYSIS_CTXPROFANALYSIS_H
 
+#include "llvm/ADT/StringMap.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
-#include <map>
 
 namespace llvm {
 
 class CtxProfAnalysis;
 
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
+  friend class CtxProfAnalysis;
+  friend class CtxProfAnalysisPrinterPass;
+  struct FunctionInfo {
+    uint32_t NextCounterIndex = 0;
+    uint32_t NextCallsiteIndex = 0;
+    const std::string Name;
+
+    FunctionInfo(StringRef Name) : Name(Name) {}
+  };
   std::optional<PGOCtxProfContext::CallTargetMapTy> Profiles;
+  // For the GUIDs in this module, associate metadata about each function which
+  // we'll need when we maintain the profiles during IPO transformations.
+  DenseMap<GlobalValue::GUID, FunctionInfo> FuncInfo;
 
-public:
-  explicit PGOContextualProfile(PGOCtxProfContext::CallTargetMapTy &&Profiles)
-      : Profiles(std::move(Profiles)) {}
+  /// Get the GUID of this Function if it's defined in this module.
+  GlobalValue::GUID getDefinedFunctionGUID(const Function &F) const;
+
+  // This is meant to be constructed from CtxProfAnalysis, which will also set
+  // its state piecemeal.
   PGOContextualProfile() = default;
+
+public:
   PGOContextualProfile(const PGOContextualProfile &) = delete;
   PGOContextualProfile(PGOContextualProfile &&) = default;
 
@@ -35,6 +51,20 @@ class PGOContextualProfile {
     return *Profiles;
   }
 
+  bool isFunctionKnown(const Function &F) const {
+    return getDefinedFunctionGUID(F) != 0;
+  }
+
+  uint32_t allocateNextCounterIndex(const Function &F) {
+    assert(isFunctionKnown(F));
+    return FuncInfo.find(getDefinedFunctionGUID(F))->second.NextCounterIndex++;
+  }
+
+  uint32_t allocateNextCallsiteIndex(const Function &F) {
+    assert(isFunctionKnown(F));
+    return FuncInfo.find(getDefinedFunctionGUID(F))->second.NextCallsiteIndex++;
+  }
+
   bool invalidate(Module &, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &) {
     // Check whether the analysis has been explicitly invalidated. Otherwise,

llvm/include/llvm/Transforms/Instrumentation/PGOCtxProfLowering.h

@@ -24,5 +24,24 @@ class PGOCtxProfLoweringPass : public PassInfoMixin<PGOCtxProfLoweringPass> {
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 };
+
+/// Assign a GUID to functions as metadata. GUID calculation takes linkage into
+/// account, which may change especially through and after thinlto. By
+/// pre-computing and assigning as metadata, this mechanism is resilient to such
+/// changes (as well as name changes e.g. suffix ".llvm." additions).
+
+// FIXME(mtrofin): we can generalize this mechanism to calculate a GUID early in
+// the pass pipeline, associate it with any Global Value, and then use it for
+// PGO and ThinLTO.
+// At that point, this should be moved elsewhere.
+class AssignUniqueIDPass : public PassInfoMixin<AssignUniqueIDPass> {
+public:
+  explicit AssignUniqueIDPass() = default;
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+  static const char *GUIDMetadataName;
+  // This should become GlobalValue::getGUID
+  static uint64_t getGUID(const Function &F);
+};
+
 } // namespace llvm
 #endif

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -14,12 +14,14 @@
 #include "llvm/Analysis/CtxProfAnalysis.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Analysis.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Transforms/Instrumentation/PGOCtxProfLowering.h"
 
 #define DEBUG_TYPE "ctx_prof"
 
@@ -66,8 +68,8 @@ Value toJSON(const PGOCtxProfContext::CallTargetMapTy &P) {
 
 AnalysisKey CtxProfAnalysis::Key;
 
-CtxProfAnalysis::Result CtxProfAnalysis::run(Module &M,
-                                             ModuleAnalysisManager &MAM) {
+PGOContextualProfile CtxProfAnalysis::run(Module &M,
+                                          ModuleAnalysisManager &MAM) {
   ErrorOr<std::unique_ptr<MemoryBuffer>> MB = MemoryBuffer::getFile(Profile);
   if (auto EC = MB.getError()) {
     M.getContext().emitError("could not open contextual profile file: " +
@@ -81,7 +83,56 @@ CtxProfAnalysis::Result CtxProfAnalysis::run(Module &M,
                              toString(MaybeCtx.takeError()));
     return {};
   }
-  return Result(std::move(*MaybeCtx));
+
+  PGOContextualProfile Result;
+
+  for (const auto &F : M) {
+    if (F.isDeclaration())
+      continue;
+    auto GUID = AssignUniqueIDPass::getGUID(F);
+    assert(GUID && "guid not found for defined function");
+    const auto &Entry = F.begin();
+    uint32_t MaxCounters = 0; // we expect at least a counter.
+    for (const auto &I : *Entry)
+      if (auto *C = dyn_cast<InstrProfIncrementInst>(&I)) {
+        MaxCounters =
+            static_cast<uint32_t>(C->getNumCounters()->getZExtValue());
+        break;
+      }
+    if (!MaxCounters)
+      continue;
+    uint32_t MaxCallsites = 0;
+    for (const auto &BB : F)
+      for (const auto &I : BB)
+        if (auto *C = dyn_cast<InstrProfCallsite>(&I)) {
+          MaxCallsites =
+              static_cast<uint32_t>(C->getNumCounters()->getZExtValue());
+          break;
+        }
+    auto [It, Ins] = Result.FuncInfo.insert(
+        {GUID, PGOContextualProfile::FunctionInfo(F.getName())});
+    (void)Ins;
+    assert(Ins);
+    It->second.NextCallsiteIndex = MaxCallsites;
+    It->second.NextCounterIndex = MaxCounters;
+  }
+  // If we made it this far, the Result is valid - which we mark by setting
+  // .Profiles.
+  // Trim first the roots that aren't in this module.
+  DenseSet<GlobalValue::GUID> ProfiledGUIDs;
+  for (auto &[RootGuid, Tree] : llvm::make_early_inc_range(*MaybeCtx))
+    if (!Result.FuncInfo.contains(RootGuid))
+      MaybeCtx->erase(RootGuid);
+  Result.Profiles = std::move(*MaybeCtx);
+  return Result;
+}
+
+GlobalValue::GUID
+PGOContextualProfile::getDefinedFunctionGUID(const Function &F) const {
+  if (auto It = FuncInfo.find(AssignUniqueIDPass::getGUID(F));
+      It != FuncInfo.end())
+    return It->first;
+  return 0;
 }
 
 PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
@@ -91,8 +142,16 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
     M.getContext().emitError("Invalid CtxProfAnalysis");
     return PreservedAnalyses::all();
   }
+
+  OS << "Function Info:\n";
+  for (const auto &[Guid, FuncInfo] : C.FuncInfo)
+    OS << Guid << " : " << FuncInfo.Name
+       << ". MaxCounterID: " << FuncInfo.NextCounterIndex
+       << ". MaxCallsiteID: " << FuncInfo.NextCallsiteIndex << "\n";
+
   const auto JSONed = ::llvm::json::toJSON(C.profiles());
 
+  OS << "\nCurrent Profile:\n";
   OS << formatv("{0:2}", JSONed);
   OS << "\n";
   return PreservedAnalyses::all();

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -1196,6 +1196,7 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
     // In pre-link, we just want the instrumented IR. We use the contextual
     // profile in the post-thinlink phase.
     // The instrumentation will be removed in post-thinlink after IPO.
+    MPM.addPass(AssignUniqueIDPass());
     if (IsCtxProfUse)
       return MPM;
     addPostPGOLoopRotation(MPM, Level);

llvm/lib/Passes/PassRegistry.def

@@ -46,6 +46,7 @@ MODULE_PASS("always-inline", AlwaysInlinerPass())
 MODULE_PASS("annotation2metadata", Annotation2MetadataPass())
 MODULE_PASS("attributor", AttributorPass())
 MODULE_PASS("attributor-light", AttributorLightPass())
+MODULE_PASS("assign-guid", AssignUniqueIDPass())
 MODULE_PASS("called-value-propagation", CalledValuePropagationPass())
 MODULE_PASS("canonicalize-aliases", CanonicalizeAliasesPass())
 MODULE_PASS("check-debugify", NewPMCheckDebugifyPass())

llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp

@@ -16,6 +16,7 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/CommandLine.h"
 #include <utility>
 
@@ -223,8 +224,8 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
       assert(Mark->getIndex()->isZero());
 
       IRBuilder<> Builder(Mark);
-      // FIXME(mtrofin): use InstrProfSymtab::getCanonicalName
-      Guid = Builder.getInt64(F.getGUID());
+
+      Guid = Builder.getInt64(AssignUniqueIDPass::getGUID(F));
       // The type of the context of this function is now knowable since we have
       // NrCallsites and NrCounters. We delcare it here because it's more
       // convenient - we have the Builder.
@@ -349,3 +350,34 @@ bool CtxInstrumentationLowerer::lowerFunction(Function &F) {
         F.getName());
   return true;
 }
+
+const char *AssignUniqueIDPass::GUIDMetadataName = "unique_id";
+
+PreservedAnalyses AssignUniqueIDPass::run(Module &M,
+                                          ModuleAnalysisManager &MAM) {
+  for (auto &F : M.functions()) {
+    if (F.isDeclaration())
+      continue;
+    const GlobalValue::GUID GUID = F.getGUID();
+    assert(!F.getMetadata(GUIDMetadataName) ||
+           GUID == AssignUniqueIDPass::getGUID(F));
+    F.setMetadata(GUIDMetadataName,
+                  MDNode::get(M.getContext(),
+                              {ConstantAsMetadata::get(ConstantInt::get(
+                                  Type::getInt64Ty(M.getContext()), GUID))}));
+  }
+  return PreservedAnalyses::none();
+}
+
+GlobalValue::GUID AssignUniqueIDPass::getGUID(const Function &F) {
+  if (F.isDeclaration()) {
+    assert(GlobalValue::isExternalLinkage(F.getLinkage()));
+    return GlobalValue::getGUID(F.getGlobalIdentifier());
+  }
+  auto *MD = F.getMetadata(GUIDMetadataName);
+  assert(MD && "unique_id not found for defined function");
+  return cast<ConstantInt>(cast<ConstantAsMetadata>(MD->getOperand(0))
+                               ->getValue()
+                               ->stripPointerCasts())
+      ->getZExtValue();
+}

llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll

@@ -0,0 +1,110 @@
+; REQUIRES: x86_64-linux
+
+; RUN: split-file %s %t
+;
+; Test that the GUID metadata survives through thinlink.
+;
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+;
+; Disable pre-inline to avoid losing the trivial functions to the preinliner.
+;
+; RUN: opt -module-summary -disable-preinline -passes='thinlto-pre-link<O2>' -use-ctx-profile=%t/profile.ctxprofdata -o %t/m1.bc %t/m1.ll
+; RUN: opt -module-summary -disable-preinline -passes='thinlto-pre-link<O2>' -use-ctx-profile=%t/profile.ctxprofdata -o %t/m2.bc %t/m2.ll
+;
+; RUN: rm -rf %t/postlink
+; RUN: mkdir %t/postlink
+; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc -o %t/ -thinlto-distributed-indexes \
+; RUN:  -r %t/m1.bc,f1,plx \
+; RUN:  -r %t/m2.bc,f1 \
+; RUN:  -r %t/m2.bc,entrypoint,plx
+; RUN: opt --passes='function-import,require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:  -summary-file=%t/m2.bc.thinlto.bc -use-ctx-profile=%t/profile.ctxprofdata %t/m2.bc \
+; RUN:  -S -o %t/m2.post.ll 2> %t/profile.txt
+; RUN: diff %t/expected.txt %t/profile.txt
+;--- m1.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+define private void @f2() {
+  ret void
+}
+
+define void @f1() {
+  call void @f2()
+  ret void
+}
+
+;--- m2.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+declare void @f1()
+
+define void @entrypoint() {
+  call void @f1()
+  ret void
+}
+;--- profile.json
+[
+  {
+    "Callsites": [
+      [
+        {
+          "Callsites": [
+            [
+              {
+                "Counters": [
+                  10
+                ],
+                "Guid": 5967942613276634709
+              }
+            ]
+          ],
+          "Counters": [
+            7
+          ],
+          "Guid": 2072045998141807037
+        }
+      ]
+    ],
+    "Counters": [
+      1
+    ],
+    "Guid": 10507721908651011566
+  }
+]
+;--- expected.txt
+Function Info:
+5967942613276634709 : f2.llvm.0. MaxCounterID: 1. MaxCallsiteID: 0
+10507721908651011566 : entrypoint. MaxCounterID: 1. MaxCallsiteID: 1
+2072045998141807037 : f1. MaxCounterID: 1. MaxCallsiteID: 1
+
+Current Profile:
+[
+  {
+    "Callsites": [
+      [
+        {
+          "Callsites": [
+            [
+              {
+                "Counters": [
+                  10
+                ],
+                "Guid": 5967942613276634709
+              }
+            ]
+          ],
+          "Counters": [
+            7
+          ],
+          "Guid": 2072045998141807037
+        }
+      ]
+    ],
+    "Counters": [
+      1
+    ],
+    "Guid": 10507721908651011566
+  }
+]