[ctxprof] Profile format support for flat profiles

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -21,12 +21,6 @@ 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 =
-    std::map<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
-
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
   friend class CtxProfAnalysis;
@@ -38,7 +32,7 @@ class PGOContextualProfile {
     PGOCtxProfContext Index;
     FunctionInfo(StringRef Name) : Name(Name) {}
   };
-  std::optional<PGOCtxProfContext::CallTargetMapTy> Profiles;
+  std::optional<PGOCtxProfile> Profile;
   // For the GUIDs in this module, associate metadata about each function which
   // we'll need when we maintain the profiles during IPO transformations.
   std::map<GlobalValue::GUID, FunctionInfo> FuncInfo;
@@ -56,12 +50,12 @@ class PGOContextualProfile {
   PGOContextualProfile(const PGOContextualProfile &) = delete;
   PGOContextualProfile(PGOContextualProfile &&) = default;
 
-  operator bool() const { return Profiles.has_value(); }
-
-  const PGOCtxProfContext::CallTargetMapTy &profiles() const {
-    return *Profiles;
+  const PGOCtxProfContext::CallTargetMapTy &contexts() const {
+    return Profile->Contexts;
   }
 
+  const PGOCtxProfile &profiles() const { return *Profile; }
+
   bool isFunctionKnown(const Function &F) const {
     return getDefinedFunctionGUID(F) != 0;
   }

llvm/include/llvm/ProfileData/PGOCtxProfReader.h

@@ -164,6 +164,24 @@ class PGOCtxProfContext final : public internal::IndexNode {
   }
 };
 
+// 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 =
+    std::map<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
+
+using CtxProfContextualProfiles =
+    std::map<GlobalValue::GUID, PGOCtxProfContext>;
+struct PGOCtxProfile {
+  CtxProfContextualProfiles Contexts;
+  CtxProfFlatProfile FlatProfiles;
+
+  PGOCtxProfile() = default;
+  PGOCtxProfile(const PGOCtxProfile &) = delete;
+  PGOCtxProfile(PGOCtxProfile &&) = default;
+  PGOCtxProfile &operator=(PGOCtxProfile &&) = default;
+};
+
 class PGOCtxProfileReader final {
   StringRef Magic;
   BitstreamCursor Cursor;
@@ -173,18 +191,21 @@ class PGOCtxProfileReader final {
   Error unsupported(const Twine &);
 
   Expected<std::pair<std::optional<uint32_t>, PGOCtxProfContext>>
-  readContext(bool ExpectIndex);
-  bool canReadContext();
+  readProfile(PGOCtxProfileBlockIDs Kind);
+
+  bool canEnterBlockWithID(PGOCtxProfileBlockIDs ID);
+  Error enterBlockWithID(PGOCtxProfileBlockIDs ID);
+
+  Error loadContexts(CtxProfContextualProfiles &);
 
 public:
   PGOCtxProfileReader(StringRef Buffer)
       : Magic(Buffer.substr(0, PGOCtxProfileWriter::ContainerMagic.size())),
         Cursor(Buffer.substr(PGOCtxProfileWriter::ContainerMagic.size())) {}
 
-  Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>> loadContexts();
+  Expected<PGOCtxProfile> loadProfiles();
 };
 
-void convertCtxProfToYaml(raw_ostream &OS,
-                          const PGOCtxProfContext::CallTargetMapTy &);
+void convertCtxProfToYaml(raw_ostream &OS, const PGOCtxProfile &);
 } // namespace llvm
 #endif

llvm/include/llvm/ProfileData/PGOCtxProfWriter.h

@@ -23,7 +23,9 @@ enum PGOCtxProfileRecords { Invalid = 0, Version, Guid, CalleeIndex, Counters };
 
 enum PGOCtxProfileBlockIDs {
   ProfileMetadataBlockID = bitc::FIRST_APPLICATION_BLOCKID,
-  ContextNodeBlockID = ProfileMetadataBlockID + 1
+  ContextsSectionBlockID = ProfileMetadataBlockID + 1,
+  ContextRootBlockID = ContextsSectionBlockID + 1,
+  ContextNodeBlockID = ContextRootBlockID + 1,
 };
 
 /// Write one or more ContextNodes to the provided raw_fd_stream.
@@ -60,10 +62,13 @@ enum PGOCtxProfileBlockIDs {
 /// like value profiling - which would appear as additional records. For
 /// example, value profiling would produce a new record with a new record ID,
 /// containing the profiled values (much like the counters)
-class PGOCtxProfileWriter final {
+class PGOCtxProfileWriter : public ctx_profile::ProfileWriter {
+  enum class EmptyContextCriteria { None, EntryIsZero, AllAreZero };
+
   BitstreamWriter Writer;
 
-  void writeCounters(const ctx_profile::ContextNode &Node);
+  void writeGuid(ctx_profile::GUID Guid);
+  void writeCounters(ArrayRef<uint64_t> Counters);
   void writeImpl(std::optional<uint32_t> CallerIndex,
                  const ctx_profile::ContextNode &Node);
 
@@ -72,11 +77,13 @@ class PGOCtxProfileWriter final {
                       std::optional<unsigned> VersionOverride = std::nullopt);
   ~PGOCtxProfileWriter() { Writer.ExitBlock(); }
 
-  void write(const ctx_profile::ContextNode &);
+  void startContextSection() override;
+  void writeContextual(const ctx_profile::ContextNode &RootNode) override;
+  void endContextSection() override;
 
   // constants used in writing which a reader may find useful.
   static constexpr unsigned CodeLen = 2;
-  static constexpr uint32_t CurrentVersion = 1;
+  static constexpr uint32_t CurrentVersion = 2;
   static constexpr unsigned VBREncodingBits = 6;
   static constexpr StringRef ContainerMagic = "CTXP";
 };

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -88,7 +88,7 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
     return {};
   }
   PGOCtxProfileReader Reader(MB.get()->getBuffer());
-  auto MaybeCtx = Reader.loadContexts();
+  auto MaybeCtx = Reader.loadProfiles();
   if (!MaybeCtx) {
     M.getContext().emitError("contextual profile file is invalid: " +
                              toString(MaybeCtx.takeError()));
@@ -99,16 +99,16 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
   for (const auto &F : M)
     if (!F.isDeclaration())
       if (auto GUID = AssignGUIDPass::getGUID(F);
-          MaybeCtx->find(GUID) != MaybeCtx->end())
+          MaybeCtx->Contexts.find(GUID) != MaybeCtx->Contexts.end())
         ProfileRootsInModule.insert(GUID);
 
   // Trim first the roots that aren't in this module.
-  for (auto &[RootGuid, _] : llvm::make_early_inc_range(*MaybeCtx))
+  for (auto &[RootGuid, _] : llvm::make_early_inc_range(MaybeCtx->Contexts))
     if (!ProfileRootsInModule.contains(RootGuid))
-      MaybeCtx->erase(RootGuid);
+      MaybeCtx->Contexts.erase(RootGuid);
   // If none of the roots are in the module, we have no profile (for this
   // module)
-  if (MaybeCtx->empty())
+  if (MaybeCtx->Contexts.empty())
     return {};
 
   // OK, so we have a valid profile and it's applicable to roots in this module.
@@ -146,7 +146,7 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
   }
   // If we made it this far, the Result is valid - which we mark by setting
   // .Profiles.
-  Result.Profiles = std::move(*MaybeCtx);
+  Result.Profile = std::move(*MaybeCtx);
   Result.initIndex();
   return Result;
 }
@@ -164,7 +164,7 @@ CtxProfAnalysisPrinterPass::CtxProfAnalysisPrinterPass(raw_ostream &OS)
 PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
                                                   ModuleAnalysisManager &MAM) {
   CtxProfAnalysis::Result &C = MAM.getResult<CtxProfAnalysis>(M);
-  if (!C) {
+  if (C.contexts().empty()) {
     OS << "No contextual profile was provided.\n";
     return PreservedAnalyses::all();
   }
@@ -245,7 +245,7 @@ void PGOContextualProfile::initIndex() {
   for (auto &[Guid, FI] : FuncInfo)
     InsertionPoints[Guid] = &FI.Index;
   preorderVisit<PGOCtxProfContext::CallTargetMapTy, PGOCtxProfContext>(
-      *Profiles, [&](PGOCtxProfContext &Ctx) {
+      Profile->Contexts, [&](PGOCtxProfContext &Ctx) {
         auto InsertIt = InsertionPoints.find(Ctx.guid());
         if (InsertIt == InsertionPoints.end())
           return;
@@ -270,7 +270,7 @@ void PGOContextualProfile::update(Visitor V, const Function &F) {
 void PGOContextualProfile::visit(ConstVisitor V, const Function *F) const {
   if (!F)
     return preorderVisit<const PGOCtxProfContext::CallTargetMapTy,
-                         const PGOCtxProfContext>(*Profiles, V);
+                         const PGOCtxProfContext>(Profile->Contexts, V);
   assert(isFunctionKnown(*F));
   GlobalValue::GUID G = getDefinedFunctionGUID(*F);
   for (const auto *Node = FuncInfo.find(G)->second.Index.Next; Node;
@@ -279,11 +279,11 @@ void PGOContextualProfile::visit(ConstVisitor V, const Function *F) const {
 }
 
 const CtxProfFlatProfile PGOContextualProfile::flatten() const {
-  assert(Profiles.has_value());
+  assert(Profile.has_value());
   CtxProfFlatProfile Flat;
   preorderVisit<const PGOCtxProfContext::CallTargetMapTy,
                 const PGOCtxProfContext>(
-      *Profiles, [&](const PGOCtxProfContext &Ctx) {
+      Profile->Contexts, [&](const PGOCtxProfContext &Ctx) {
         auto [It, Ins] = Flat.insert({Ctx.guid(), {}});
         if (Ins) {
           llvm::append_range(It->second, Ctx.counters());

llvm/lib/ProfileData/PGOCtxProfReader.cpp

@@ -58,26 +58,33 @@ Error PGOCtxProfileReader::unsupported(const Twine &Msg) {
   return make_error<InstrProfError>(instrprof_error::unsupported_version, Msg);
 }
 
-bool PGOCtxProfileReader::canReadContext() {
+bool PGOCtxProfileReader::canEnterBlockWithID(PGOCtxProfileBlockIDs ID) {
   auto Blk = advance();
   if (!Blk) {
     consumeError(Blk.takeError());
     return false;
   }
-  return Blk->Kind == BitstreamEntry::SubBlock &&
-         Blk->ID == PGOCtxProfileBlockIDs::ContextNodeBlockID;
+  return Blk->Kind == BitstreamEntry::SubBlock && Blk->ID == ID;
+}
+Error PGOCtxProfileReader::enterBlockWithID(PGOCtxProfileBlockIDs ID) {
+  RET_ON_ERR(Cursor.EnterSubBlock(ID));
+  return Error::success();
 }
 
 Expected<std::pair<std::optional<uint32_t>, PGOCtxProfContext>>
-PGOCtxProfileReader::readContext(bool ExpectIndex) {
-  RET_ON_ERR(Cursor.EnterSubBlock(PGOCtxProfileBlockIDs::ContextNodeBlockID));
+PGOCtxProfileReader::readProfile(PGOCtxProfileBlockIDs Kind) {
+  assert((Kind == PGOCtxProfileBlockIDs::ContextRootBlockID ||
+          Kind == PGOCtxProfileBlockIDs::ContextNodeBlockID) &&
+         "Unexpected profile kind");
+  RET_ON_ERR(enterBlockWithID(Kind));
 
   std::optional<ctx_profile::GUID> Guid;
   std::optional<SmallVector<uint64_t, 16>> Counters;
   std::optional<uint32_t> CallsiteIndex;
 
   SmallVector<uint64_t, 1> RecordValues;
 
+  const bool ExpectIndex = Kind == PGOCtxProfileBlockIDs::ContextNodeBlockID;
   // We don't prescribe the order in which the records come in, and we are ok
   // if other unsupported records appear. We seek in the current subblock until
   // we get all we know.
@@ -121,8 +128,8 @@ PGOCtxProfileReader::readContext(bool ExpectIndex) {
 
   PGOCtxProfContext Ret(*Guid, std::move(*Counters));
 
-  while (canReadContext()) {
-    EXPECT_OR_RET(SC, readContext(true));
+  while (canEnterBlockWithID(PGOCtxProfileBlockIDs::ContextNodeBlockID)) {
+    EXPECT_OR_RET(SC, readProfile(PGOCtxProfileBlockIDs::ContextNodeBlockID));
     auto &Targets = Ret.callsites()[*SC->first];
     auto [_, Inserted] =
         Targets.insert({SC->second.guid(), std::move(SC->second)});
@@ -168,16 +175,23 @@ Error PGOCtxProfileReader::readMetadata() {
   return Error::success();
 }
 
-Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>>
-PGOCtxProfileReader::loadContexts() {
-  std::map<GlobalValue::GUID, PGOCtxProfContext> Ret;
-  RET_ON_ERR(readMetadata());
-  while (canReadContext()) {
-    EXPECT_OR_RET(E, readContext(false));
-    auto Key = E->second.guid();
-    if (!Ret.insert({Key, std::move(E->second)}).second)
-      return wrongValue("Duplicate roots");
+Error PGOCtxProfileReader::loadContexts(CtxProfContextualProfiles &P) {
+  if (canEnterBlockWithID(PGOCtxProfileBlockIDs::ContextsSectionBlockID)) {
+    RET_ON_ERR(enterBlockWithID(PGOCtxProfileBlockIDs::ContextsSectionBlockID));
+    while (canEnterBlockWithID(PGOCtxProfileBlockIDs::ContextRootBlockID)) {
+      EXPECT_OR_RET(E, readProfile(PGOCtxProfileBlockIDs::ContextRootBlockID));
+      auto Key = E->second.guid();
+      if (!P.insert({Key, std::move(E->second)}).second)
+        return wrongValue("Duplicate roots");
+    }
   }
+  return Error::success();
+}
+
+Expected<PGOCtxProfile> PGOCtxProfileReader::loadProfiles() {
+  RET_ON_ERR(readMetadata());
+  PGOCtxProfile Ret;
+  RET_ON_ERR(loadContexts(Ret.Contexts));
   return std::move(Ret);
 }
 
@@ -225,41 +239,54 @@ void toYaml(yaml::Output &Out,
   Out.endSequence();
 }
 
-void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx) {
+void toYaml(yaml::Output &Out, GlobalValue::GUID Guid,
+            const SmallVectorImpl<uint64_t> &Counters,
+            const PGOCtxProfContext::CallsiteMapTy &Callsites) {
   yaml::EmptyContext Empty;
   Out.beginMapping();
   void *SaveInfo = nullptr;
   bool UseDefault = false;
   {
     Out.preflightKey("Guid", /*Required=*/true, /*SameAsDefault=*/false,
                      UseDefault, SaveInfo);
-    auto Guid = Ctx.guid();
     yaml::yamlize(Out, Guid, true, Empty);
     Out.postflightKey(nullptr);
   }
   {
     Out.preflightKey("Counters", true, false, UseDefault, SaveInfo);
     Out.beginFlowSequence();
-    for (size_t I = 0U, E = Ctx.counters().size(); I < E; ++I) {
+    for (size_t I = 0U, E = Counters.size(); I < E; ++I) {
       Out.preflightFlowElement(I, SaveInfo);
-      uint64_t V = Ctx.counters()[I];
+      uint64_t V = Counters[I];
       yaml::yamlize(Out, V, true, Empty);
       Out.postflightFlowElement(SaveInfo);
     }
     Out.endFlowSequence();
     Out.postflightKey(nullptr);
   }
-  if (!Ctx.callsites().empty()) {
+  if (!Callsites.empty()) {
     Out.preflightKey("Callsites", true, false, UseDefault, SaveInfo);
-    toYaml(Out, Ctx.callsites());
+    toYaml(Out, Callsites);
     Out.postflightKey(nullptr);
   }
   Out.endMapping();
 }
+void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx) {
+  toYaml(Out, Ctx.guid(), Ctx.counters(), Ctx.callsites());
+}
+
 } // namespace
 
-void llvm::convertCtxProfToYaml(
-    raw_ostream &OS, const PGOCtxProfContext::CallTargetMapTy &Profiles) {
+void llvm::convertCtxProfToYaml(raw_ostream &OS,
+                                const PGOCtxProfile &Profiles) {
   yaml::Output Out(OS);
-  toYaml(Out, Profiles);
+  void *SaveInfo = nullptr;
+  bool UseDefault = false;
+  Out.beginMapping();
+  if (!Profiles.Contexts.empty()) {
+    Out.preflightKey("Contexts", false, false, UseDefault, SaveInfo);
+    toYaml(Out, Profiles.Contexts);
+    Out.postflightKey(nullptr);
+  }
+  Out.endMapping();
 }