[ctx_prof] CtxProfAnalysis (llvm#102084)

This is an immutable analysis that loads and makes the contextual profile available to other passes. This patch introduces the analysis and an analysis printer pass. Subsequent patches will introduce the APIs that IPO passes will call to modify the profile as result of their changes.

Change-Id: I1a1047a73074f143a2df79afdbc06b075473feb8

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -9,47 +9,23 @@
 #ifndef LLVM_ANALYSIS_CTXPROFANALYSIS_H
 #define LLVM_ANALYSIS_CTXPROFANALYSIS_H
 
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/GlobalValue.h"
-#include "llvm/IR/InstrTypes.h"
-#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
+#include <map>
 
 namespace llvm {
 
 class CtxProfAnalysis;
 
-// Setting initial capacity to 1 because all contexts must have at least 1
-// counter, and then, because all contexts belonging to a function have the same
-// size, there'll be at most one other heap allocation.
-using CtxProfFlatProfile =
-    DenseMap<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
-
 /// 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;
-
-  /// 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:
+  explicit PGOContextualProfile(PGOCtxProfContext::CallTargetMapTy &&Profiles)
+      : Profiles(std::move(Profiles)) {}
+  PGOContextualProfile() = default;
   PGOContextualProfile(const PGOContextualProfile &) = delete;
   PGOContextualProfile(PGOContextualProfile &&) = default;
 
@@ -59,22 +35,6 @@ 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++;
-  }
-
-  const CtxProfFlatProfile flatten() const;
-
   bool invalidate(Module &, const PreservedAnalyses &PA,
                   ModuleAnalysisManager::Invalidator &) {
     // Check whether the analysis has been explicitly invalidated. Otherwise,
@@ -94,8 +54,6 @@ class CtxProfAnalysis : public AnalysisInfoMixin<CtxProfAnalysis> {
   using Result = PGOContextualProfile;
 
   PGOContextualProfile run(Module &M, ModuleAnalysisManager &MAM);
-
-  static InstrProfCallsite *getCallsiteInstrumentation(CallBase &CB);
 };
 
 class CtxProfAnalysisPrinterPass
@@ -108,27 +66,5 @@ class CtxProfAnalysisPrinterPass
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
   static bool isRequired() { return true; }
 };
-
-/// 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 AssignGUIDPass : public PassInfoMixin<AssignGUIDPass> {
-public:
-  explicit AssignGUIDPass() = default;
-
-  /// Assign a GUID *if* one is not already assign, as a function metadata named
-  /// `GUIDMetadataName`.
-  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_ANALYSIS_CTXPROFANALYSIS_H

llvm/lib/Analysis/CMakeLists.txt

@@ -46,6 +46,7 @@ add_llvm_component_library(LLVMAnalysis
   CostModel.cpp
   CodeMetrics.cpp
   ConstantFolding.cpp
+  CtxProfAnalysis.cpp
   CycleAnalysis.cpp
   DDG.cpp
   DDGPrinter.cpp

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -0,0 +1,95 @@
+//===- CtxProfAnalysis.cpp - contextual profile analysis ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of the contextual profile analysis, which maintains contextual
+// profiling info through IPO passes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/CtxProfAnalysis.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/ProfileData/PGOCtxProfReader.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
+
+#define DEBUG_TYPE "ctx_prof"
+
+namespace llvm {
+namespace json {
+Value toJSON(const PGOCtxProfContext &P) {
+  Object Ret;
+  Ret["Guid"] = P.guid();
+  Ret["Counters"] = Array(P.counters());
+  if (P.callsites().empty())
+    return Ret;
+  auto AllCS =
+      ::llvm::map_range(P.callsites(), [](const auto &P) { return P.first; });
+  auto MaxIt = ::llvm::max_element(AllCS);
+  assert(MaxIt != AllCS.end() && "We should have a max value because the "
+                                 "callsites collection is not empty.");
+  Array CSites;
+  // Iterate to, and including, the maximum index.
+  for (auto I = 0U, Max = *MaxIt; I <= Max; ++I) {
+    CSites.push_back(Array());
+    Array &Targets = *CSites.back().getAsArray();
+    if (P.hasCallsite(I))
+      for (const auto &[_, Ctx] : P.callsite(I))
+        Targets.push_back(toJSON(Ctx));
+  }
+  Ret["Callsites"] = std::move(CSites);
+
+  return Ret;
+}
+
+Value toJSON(const PGOCtxProfContext::CallTargetMapTy &P) {
+  Array Ret;
+  for (const auto &[_, Ctx] : P)
+    Ret.push_back(toJSON(Ctx));
+  return Ret;
+}
+} // namespace json
+} // namespace llvm
+
+using namespace llvm;
+
+AnalysisKey CtxProfAnalysis::Key;
+
+CtxProfAnalysis::Result 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: " +
+                             EC.message());
+    return {};
+  }
+  PGOCtxProfileReader Reader(MB.get()->getBuffer());
+  auto MaybeCtx = Reader.loadContexts();
+  if (!MaybeCtx) {
+    M.getContext().emitError("contextual profile file is invalid: " +
+                             toString(MaybeCtx.takeError()));
+    return {};
+  }
+  return Result(std::move(*MaybeCtx));
+}
+
+PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
+                                                  ModuleAnalysisManager &MAM) {
+  CtxProfAnalysis::Result &C = MAM.getResult<CtxProfAnalysis>(M);
+  if (!C) {
+    M.getContext().emitError("Invalid CtxProfAnalysis");
+    return PreservedAnalyses::all();
+  }
+  const auto JSONed = ::llvm::json::toJSON(C.profiles());
+
+  OS << formatv("{0:2}", JSONed);
+  OS << "\n";
+  return PreservedAnalyses::all();
+}

llvm/lib/Passes/PassBuilder.cpp

@@ -28,6 +28,7 @@
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallPrinter.h"
 #include "llvm/Analysis/CostModel.h"
+#include "llvm/Analysis/CtxProfAnalysis.h"
 #include "llvm/Analysis/CycleAnalysis.h"
 #include "llvm/Analysis/DDG.h"
 #include "llvm/Analysis/DDGPrinter.h"

llvm/lib/Passes/PassRegistry.def

@@ -22,6 +22,7 @@ MODULE_ANALYSIS("callgraph", CallGraphAnalysis())
 MODULE_ANALYSIS("collector-metadata", CollectorMetadataAnalysis())
 MODULE_ANALYSIS("dxil-metadata", DXILMetadataAnalysis())
 MODULE_ANALYSIS("dxil-resource", DXILResourceAnalysis())
+MODULE_ANALYSIS("ctx-prof-analysis", CtxProfAnalysis(UseCtxProfile))
 MODULE_ANALYSIS("inline-advisor", InlineAdvisorAnalysis())
 MODULE_ANALYSIS("ir-similarity", IRSimilarityAnalysis())
 MODULE_ANALYSIS("lcg", LazyCallGraphAnalysis())
@@ -84,6 +85,7 @@ MODULE_PASS("insert-gcov-profiling", GCOVProfilerPass())
 MODULE_PASS("instrorderfile", InstrOrderFilePass())
 MODULE_PASS("instrprof", InstrProfilingLoweringPass())
 MODULE_PASS("ctx-instr-lower", PGOCtxProfLoweringPass())
+MODULE_PASS("print<ctx-prof-analysis>", CtxProfAnalysisPrinterPass(dbgs()))
 MODULE_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 MODULE_PASS("iroutliner", IROutlinerPass())
 MODULE_PASS("jmc-instrumenter", JMCInstrumenterPass())

llvm/test/Analysis/CtxProfAnalysis/load.ll

@@ -0,0 +1,56 @@
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: not opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   %t/empty.ll -S 2>&1 | FileCheck %s --check-prefix=NO-FILE
+
+; RUN: not opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   -use-ctx-profile=does_not_exist.ctxprofdata %t/empty.ll -S 2>&1 | FileCheck %s --check-prefix=NO-FILE
+
+; RUN: opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   -use-ctx-profile=%t/profile.ctxprofdata %t/empty.ll -S 2> %t/output.json
+; RUN: diff %t/profile.json %t/output.json
+
+; NO-FILE: error: could not open contextual profile file
+;
+; This is the reference profile, laid out in the format the json formatter will
+; output it from opt.
+;--- profile.json
+[
+  {
+    "Callsites": [
+      [],
+      [
+        {
+          "Counters": [
+            4,
+            5
+          ],
+          "Guid": 2000
+        },
+        {
+          "Counters": [
+            6,
+            7,
+            8
+          ],
+          "Guid": 18446744073709551613
+        }
+      ]
+    ],
+    "Counters": [
+      1,
+      2,
+      3
+    ],
+    "Guid": 1000
+  },
+  {
+    "Counters": [
+      5,
+      9,
+      10
+    ],
+    "Guid": 18446744073709551612
+  }
+]
+;--- empty.ll

revert_patches.txt

@@ -86,7 +86,6 @@ contact: RonL
 ---
 Reverts:
 1488fb415336 [PAC][AArch64] Lower ptrauth constants in code (#96879)  
-[ctx_prof] CtxProfAnalysis (#102084)
 aca01bff0 [ctx_prof] CtxProfAnalysis: populate module data (#102930) 
 0c876a486eb [nfc][ctx_prof] Remove the need for `PassBuilder` to know about `UseCtxProfile` (#104492)
 1e70122cbc18 [ctx_prof] API to get the instrumentation of a BB (#105468)