[ctxprof] Prepare profile format for flat profiles

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -54,6 +54,8 @@ class PGOContextualProfile {
     return Profiles.Contexts;
   }
 
+  const PGOCtxProfile &profiles() const { return Profiles; }
+
   bool isFunctionKnown(const Function &F) const {
     return getDefinedFunctionGUID(F) != 0;
   }

llvm/include/llvm/ProfileData/PGOCtxProfReader.h

@@ -174,6 +174,7 @@ using CtxProfContextualProfiles =
     std::map<GlobalValue::GUID, PGOCtxProfContext>;
 struct PGOCtxProfile {
   CtxProfContextualProfiles Contexts;
+  CtxProfFlatProfile FlatProfiles;
 
   PGOCtxProfile() = default;
   PGOCtxProfile(const PGOCtxProfile &) = delete;
@@ -190,8 +191,12 @@ 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)
@@ -201,7 +206,6 @@ class PGOCtxProfileReader final {
   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

@@ -180,7 +180,7 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
 
   if (Mode == PrintMode::Everything)
     OS << "\nCurrent Profile:\n";
-  convertCtxProfToYaml(OS, C.contexts());
+  convertCtxProfToYaml(OS, C.profiles());
   OS << "\n";
   if (Mode == PrintMode::YAML)
     return PreservedAnalyses::all();

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,15 +175,23 @@ Error PGOCtxProfileReader::readMetadata() {
   return Error::success();
 }
 
+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() {
-  PGOCtxProfile Ret;
   RET_ON_ERR(readMetadata());
-  while (canReadContext()) {
-    EXPECT_OR_RET(E, readContext(false));
-    auto Key = E->second.guid();
-    if (!Ret.Contexts.insert({Key, std::move(E->second)}).second)
-      return wrongValue("Duplicate roots");
-  }
+  PGOCtxProfile Ret;
+  RET_ON_ERR(loadContexts(Ret.Contexts));
   return std::move(Ret);
 }
 
@@ -224,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();
 }

llvm/lib/ProfileData/PGOCtxProfWriter.cpp

@@ -13,14 +13,19 @@
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
 #include "llvm/Bitstream/BitCodeEnums.h"
 #include "llvm/ProfileData/CtxInstrContextNode.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 using namespace llvm::ctx_profile;
 
+static cl::opt<bool> IncludeEmpty(
+    "ctx-prof-include-empty", cl::init(false),
+    cl::desc(
+        "Also write profiles with all-zero counters. Mostly for debugging."));
+
 PGOCtxProfileWriter::PGOCtxProfileWriter(
     raw_ostream &Out, std::optional<unsigned> VersionOverride)
     : Writer(Out, 0) {
@@ -44,9 +49,11 @@ PGOCtxProfileWriter::PGOCtxProfileWriter(
     DescribeBlock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, "Metadata");
     DescribeRecord(PGOCtxProfileRecords::Version, "Version");
     DescribeBlock(PGOCtxProfileBlockIDs::ContextNodeBlockID, "Context");
+    DescribeBlock(PGOCtxProfileBlockIDs::ContextRootBlockID, "Root");
     DescribeRecord(PGOCtxProfileRecords::Guid, "GUID");
     DescribeRecord(PGOCtxProfileRecords::CalleeIndex, "CalleeIndex");
     DescribeRecord(PGOCtxProfileRecords::Counters, "Counters");
+    DescribeBlock(PGOCtxProfileBlockIDs::ContextsSectionBlockID, "Contexts");
   }
   Writer.ExitBlock();
   Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, CodeLen);
@@ -55,12 +62,16 @@ PGOCtxProfileWriter::PGOCtxProfileWriter(
                     SmallVector<unsigned, 1>({Version}));
 }
 
-void PGOCtxProfileWriter::writeCounters(const ContextNode &Node) {
+void PGOCtxProfileWriter::writeCounters(ArrayRef<uint64_t> Counters) {
   Writer.EmitCode(bitc::UNABBREV_RECORD);
   Writer.EmitVBR(PGOCtxProfileRecords::Counters, VBREncodingBits);
-  Writer.EmitVBR(Node.counters_size(), VBREncodingBits);
-  for (uint32_t I = 0U; I < Node.counters_size(); ++I)
-    Writer.EmitVBR64(Node.counters()[I], VBREncodingBits);
+  Writer.EmitVBR(Counters.size(), VBREncodingBits);
+  for (uint32_t I = 0U; I < Counters.size(); ++I)
+    Writer.EmitVBR64(Counters[I], VBREncodingBits);
+}
+
+void PGOCtxProfileWriter::writeGuid(ctx_profile::GUID Guid) {
+  Writer.EmitRecord(PGOCtxProfileRecords::Guid, SmallVector<uint64_t, 1>{Guid});
 }
 
 // recursively write all the subcontexts. We do need to traverse depth first to
@@ -69,21 +80,30 @@ void PGOCtxProfileWriter::writeCounters(const ContextNode &Node) {
 // keep the implementation simple.
 void PGOCtxProfileWriter::writeImpl(std::optional<uint32_t> CallerIndex,
                                     const ContextNode &Node) {
-  Writer.EnterSubblock(PGOCtxProfileBlockIDs::ContextNodeBlockID, CodeLen);
-  Writer.EmitRecord(PGOCtxProfileRecords::Guid,
-                    SmallVector<uint64_t, 1>{Node.guid()});
+  if (!IncludeEmpty && Node.entrycount() == 0)
+    return;
+  Writer.EnterSubblock(CallerIndex ? PGOCtxProfileBlockIDs::ContextNodeBlockID
+                                   : PGOCtxProfileBlockIDs::ContextRootBlockID,
+                       CodeLen);
+  writeGuid(Node.guid());
   if (CallerIndex)
     Writer.EmitRecord(PGOCtxProfileRecords::CalleeIndex,
                       SmallVector<uint64_t, 1>{*CallerIndex});
-  writeCounters(Node);
+  writeCounters({Node.counters(), Node.counters_size()});
   for (uint32_t I = 0U; I < Node.callsites_size(); ++I)
     for (const auto *Subcontext = Node.subContexts()[I]; Subcontext;
          Subcontext = Subcontext->next())
       writeImpl(I, *Subcontext);
   Writer.ExitBlock();
 }
 
-void PGOCtxProfileWriter::write(const ContextNode &RootNode) {
+void PGOCtxProfileWriter::startContextSection() {
+  Writer.EnterSubblock(PGOCtxProfileBlockIDs::ContextsSectionBlockID, CodeLen);
+}
+
+void PGOCtxProfileWriter::endContextSection() { Writer.ExitBlock(); }
+
+void PGOCtxProfileWriter::writeContextual(const ContextNode &RootNode) {
   writeImpl(std::nullopt, RootNode);
 }
 
@@ -96,6 +116,9 @@ struct SerializableCtxRepresentation {
   std::vector<uint64_t> Counters;
   std::vector<std::vector<SerializableCtxRepresentation>> Callsites;
 };
+struct SerializableProfileRepresentation {
+  std::vector<SerializableCtxRepresentation> Contexts;
+};
 
 ctx_profile::ContextNode *
 createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
@@ -142,23 +165,34 @@ template <> struct yaml::MappingTraits<SerializableCtxRepresentation> {
   }
 };
 
+template <> struct yaml::MappingTraits<SerializableProfileRepresentation> {
+  static void mapping(yaml::IO &IO, SerializableProfileRepresentation &SPR) {
+    IO.mapOptional("Contexts", SPR.Contexts);
+  }
+};
+
 Error llvm::createCtxProfFromYAML(StringRef Profile, raw_ostream &Out) {
   yaml::Input In(Profile);
-  std::vector<SerializableCtxRepresentation> DCList;
-  In >> DCList;
+  SerializableProfileRepresentation SPR;
+  In >> SPR;
   if (In.error())
     return createStringError(In.error(), "incorrect yaml content");
   std::vector<std::unique_ptr<char[]>> Nodes;
   std::error_code EC;
   if (EC)
     return createStringError(EC, "failed to open output");
   PGOCtxProfileWriter Writer(Out);
-  for (const auto &DC : DCList) {
-    auto *TopList = createNode(Nodes, DC);
-    if (!TopList)
-      return createStringError(
-          "Unexpected error converting internal structure to ctx profile");
-    Writer.write(*TopList);
+
+  if (!SPR.Contexts.empty()) {
+    Writer.startContextSection();
+    for (const auto &DC : SPR.Contexts) {
+      auto *TopList = createNode(Nodes, DC);
+      if (!TopList)
+        return createStringError(
+            "Unexpected error converting internal structure to ctx profile");
+      Writer.writeContextual(*TopList);
+    }
+    Writer.endContextSection();
   }
   if (EC)
     return createStringError(EC, "failed to write output");

llvm/test/tools/llvm-ctxprof-util/Inputs/invalid-bad-subctx.yaml

@@ -1,4 +1,5 @@
-- Guid: 123
-  Counters: [1, 2]
-  Callsites: - Guid: 1
+Contexts:
+  - Guid: 123
+    Counters: [1, 2]
+    Callsites: - Guid: 1
 

llvm/test/tools/llvm-ctxprof-util/Inputs/invalid-no-counters.yaml

@@ -1 +1,2 @@
-- Guid: 1231
+Contexts:
+  - Guid: 1231

llvm/test/tools/llvm-ctxprof-util/Inputs/invalid-no-section.yaml

@@ -0,0 +1 @@
+- Guid: 1