[ctx_prof] CtxProfAnalysis: populate module data

mtrofin · mtrofin · commit 6f23165ee03a · 2024-08-12T09:32:06.000-07:00
Continuing from llvm#102084, which introduced the analysis, we now populate it with info about functions contained in the module. When we will update the profile due to e.g. inlined callsites, we'll ingest the callee's counters and callsites to the caller. We'll move those to the caller's respective index space (counter and callers), so we need to know and maintain where those currently end. We also don't need to keep profiles not pertinent to this module. This patch also introduces an arguably much simpler way to track the GUID of a function from the frontend compilation, through ThinLTO, and into the post-thinlink compilation step, which doesn't rely on keeping names around. A separate RFC and patches will discuss extending this to the current PGO (instrumented and sampled) and other consumers as an infrastructural component.
diff --git a/llvm/include/llvm/Analysis/CtxProfAnalysis.h b/llvm/include/llvm/Analysis/CtxProfAnalysis.h
@@ -9,23 +9,38 @@
 #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)) {}
+  GlobalValue::GUID getKnownGUID(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 +50,18 @@ class PGOContextualProfile {
     return *Profiles;
   }
 
+  bool isFunctionKnown(const Function &F) const { return getKnownGUID(F) != 0; }
+
+  uint32_t allocateNextCounterIndex(const Function &F) {
+    assert(isFunctionKnown(F));
+    return FuncInfo.find(getKnownGUID(F))->second.NextCounterIndex++;
+  }
+
+  uint32_t allocateNextCallsiteIndex(const Function &F) {
+    assert(isFunctionKnown(F));
+    return FuncInfo.find(getKnownGUID(F))->second.NextCallsiteIndex++;
+  }
+
   bool invalidate(Module &, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &) {
     // Check whether the analysis has been explicitly invalidated. Otherwise,
diff --git a/llvm/include/llvm/Transforms/Instrumentation/PGOCtxProfLowering.h b/llvm/include/llvm/Transforms/Instrumentation/PGOCtxProfLowering.h
@@ -24,5 +24,26 @@ 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). It's
+/// arguably a much simpler mechanism than PGO's current GV-based one, and can
+/// be made available more broadly.
+
+// 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
diff --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/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,55 @@ 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);
+    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::getKnownGUID(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 +141,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();
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/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);
diff --git a/llvm/lib/Passes/PassRegistry.def b/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())
diff --git a/llvm/lib/Transforms/Instrumentation/PGOCtxProfLowering.cpp b/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);
+  return cast<ConstantInt>(cast<ConstantAsMetadata>(MD->getOperand(0))
+                               ->getValue()
+                               ->stripPointerCasts())
+      ->getZExtValue();
+}
diff --git a/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll b/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
+  }
+]
diff --git a/llvm/test/Analysis/CtxProfAnalysis/load.ll b/llvm/test/Analysis/CtxProfAnalysis/load.ll
diff --git a/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll b/llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll
diff --git a/llvm/test/Transforms/PGOProfile/ctx-prof-use-prelink.ll b/llvm/test/Transforms/PGOProfile/ctx-prof-use-prelink.ll
