[ctxprof] Profile section for flat profiles (llvm#129932)

A section for flat profiles (i.e. non-contextual). This is useful for debugging or for intentional cases where a root isn't identified.

This patch adds the reader/writer support. `compiler-rt` changes follow in a subsequent change.

Author: mtrofin

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;
@@ -192,10 +193,12 @@ class PGOCtxProfileReader final {
   Expected<std::pair<std::optional<uint32_t>, PGOCtxProfContext>>
   readProfile(PGOCtxProfileBlockIDs Kind);
 
+  bool tryGetNextKnownBlockID(PGOCtxProfileBlockIDs &ID);
   bool canEnterBlockWithID(PGOCtxProfileBlockIDs ID);
   Error enterBlockWithID(PGOCtxProfileBlockIDs ID);
 
   Error loadContexts(CtxProfContextualProfiles &P);
+  Error loadFlatProfiles(CtxProfFlatProfile &P);
 
 public:
   PGOCtxProfileReader(StringRef Buffer)

llvm/include/llvm/ProfileData/PGOCtxProfWriter.h

@@ -22,10 +22,14 @@ namespace llvm {
 enum PGOCtxProfileRecords { Invalid = 0, Version, Guid, CalleeIndex, Counters };
 
 enum PGOCtxProfileBlockIDs {
-  ProfileMetadataBlockID = bitc::FIRST_APPLICATION_BLOCKID,
+  FIRST_VALID = bitc::FIRST_APPLICATION_BLOCKID,
+  ProfileMetadataBlockID = FIRST_VALID,
   ContextsSectionBlockID = ProfileMetadataBlockID + 1,
   ContextRootBlockID = ContextsSectionBlockID + 1,
   ContextNodeBlockID = ContextRootBlockID + 1,
+  FlatProfilesSectionBlockID = ContextNodeBlockID + 1,
+  FlatProfileBlockID = FlatProfilesSectionBlockID + 1,
+  LAST_VALID = FlatProfileBlockID
 };
 
 /// Write one or more ContextNodes to the provided raw_fd_stream.
@@ -83,6 +87,11 @@ class PGOCtxProfileWriter final : public ctx_profile::ProfileWriter {
   void writeContextual(const ctx_profile::ContextNode &RootNode) override;
   void endContextSection() override;
 
+  void startFlatSection();
+  void writeFlatSection(ctx_profile::GUID Guid, const uint64_t *Buffer,
+                        size_t BufferSize);
+  void endFlatSection();
+
   // constants used in writing which a reader may find useful.
   static constexpr unsigned CodeLen = 2;
   static constexpr uint32_t CurrentVersion = 2;

llvm/lib/ProfileData/PGOCtxProfReader.cpp

@@ -57,24 +57,39 @@ Error PGOCtxProfileReader::unsupported(const Twine &Msg) {
   return make_error<InstrProfError>(instrprof_error::unsupported_version, Msg);
 }
 
-bool PGOCtxProfileReader::canEnterBlockWithID(PGOCtxProfileBlockIDs ID) {
+bool PGOCtxProfileReader::tryGetNextKnownBlockID(PGOCtxProfileBlockIDs &ID) {
   auto Blk = advance();
   if (!Blk) {
     consumeError(Blk.takeError());
     return false;
   }
-  return Blk->Kind == BitstreamEntry::SubBlock && Blk->ID == ID;
+  if (Blk->Kind != BitstreamEntry::SubBlock)
+    return false;
+  if (PGOCtxProfileBlockIDs::FIRST_VALID > Blk->ID ||
+      PGOCtxProfileBlockIDs::LAST_VALID < Blk->ID)
+    return false;
+  ID = static_cast<PGOCtxProfileBlockIDs>(Blk->ID);
+  return true;
+}
+
+bool PGOCtxProfileReader::canEnterBlockWithID(PGOCtxProfileBlockIDs ID) {
+  PGOCtxProfileBlockIDs Test = {};
+  return tryGetNextKnownBlockID(Test) && Test == ID;
 }
 
 Error PGOCtxProfileReader::enterBlockWithID(PGOCtxProfileBlockIDs ID) {
   RET_ON_ERR(Cursor.EnterSubBlock(ID));
   return Error::success();
 }
 
+// Note: we use PGOCtxProfContext for flat profiles also, as the latter are
+// structurally similar. Alternative modeling here seems a bit overkill at the
+// moment.
 Expected<std::pair<std::optional<uint32_t>, PGOCtxProfContext>>
 PGOCtxProfileReader::readProfile(PGOCtxProfileBlockIDs Kind) {
   assert((Kind == PGOCtxProfileBlockIDs::ContextRootBlockID ||
-          Kind == PGOCtxProfileBlockIDs::ContextNodeBlockID) &&
+          Kind == PGOCtxProfileBlockIDs::ContextNodeBlockID ||
+          Kind == PGOCtxProfileBlockIDs::FlatProfileBlockID) &&
          "Unexpected profile kind");
   RET_ON_ERR(enterBlockWithID(Kind));
 
@@ -176,22 +191,44 @@ Error PGOCtxProfileReader::readMetadata() {
 }
 
 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");
-    }
+  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();
+}
+
+Error PGOCtxProfileReader::loadFlatProfiles(CtxProfFlatProfile &P) {
+  RET_ON_ERR(
+      enterBlockWithID(PGOCtxProfileBlockIDs::FlatProfilesSectionBlockID));
+  while (canEnterBlockWithID(PGOCtxProfileBlockIDs::FlatProfileBlockID)) {
+    EXPECT_OR_RET(E, readProfile(PGOCtxProfileBlockIDs::FlatProfileBlockID));
+    auto Guid = E->second.guid();
+    if (!P.insert({Guid, std::move(E->second.counters())}).second)
+      return wrongValue("Duplicate flat profile entries");
   }
   return Error::success();
 }
 
 Expected<PGOCtxProfile> PGOCtxProfileReader::loadProfiles() {
   RET_ON_ERR(readMetadata());
   PGOCtxProfile Ret;
-  RET_ON_ERR(loadContexts(Ret.Contexts));
+  PGOCtxProfileBlockIDs Test = {};
+  for (auto I = 0; I < 2; ++I) {
+    if (!tryGetNextKnownBlockID(Test))
+      break;
+    if (Test == PGOCtxProfileBlockIDs::ContextsSectionBlockID) {
+      RET_ON_ERR(loadContexts(Ret.Contexts));
+    } else if (Test == PGOCtxProfileBlockIDs::FlatProfilesSectionBlockID) {
+      RET_ON_ERR(loadFlatProfiles(Ret.FlatProfiles));
+    } else {
+      return wrongValue("Unexpected section");
+    }
+  }
+
   return std::move(Ret);
 }
 
@@ -287,5 +324,17 @@ void llvm::convertCtxProfToYaml(raw_ostream &OS, const PGOCtxProfile &Profile) {
     toYaml(Out, Profile.Contexts);
     Out.postflightKey(nullptr);
   }
+  if (!Profile.FlatProfiles.empty()) {
+    Out.preflightKey("FlatProfiles", false, false, UseDefault, SaveInfo);
+    Out.beginSequence();
+    size_t ElemID = 0;
+    for (const auto &[Guid, Counters] : Profile.FlatProfiles) {
+      Out.preflightElement(ElemID++, SaveInfo);
+      toYaml(Out, Guid, Counters, {});
+      Out.postflightElement(nullptr);
+    }
+    Out.endSequence();
+    Out.postflightKey(nullptr);
+  }
   Out.endMapping();
 }

llvm/lib/ProfileData/PGOCtxProfWriter.cpp

@@ -13,6 +13,7 @@
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
 #include "llvm/Bitstream/BitCodeEnums.h"
 #include "llvm/ProfileData/CtxInstrContextNode.h"
+#include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/YAMLTraits.h"
@@ -59,6 +60,11 @@ PGOCtxProfileWriter::PGOCtxProfileWriter(
     DescribeRecord(PGOCtxProfileRecords::Guid, "GUID");
     DescribeRecord(PGOCtxProfileRecords::CalleeIndex, "CalleeIndex");
     DescribeRecord(PGOCtxProfileRecords::Counters, "Counters");
+    DescribeBlock(PGOCtxProfileBlockIDs::FlatProfilesSectionBlockID,
+                  "FlatProfiles");
+    DescribeBlock(PGOCtxProfileBlockIDs::FlatProfileBlockID, "Flat");
+    DescribeRecord(PGOCtxProfileRecords::Guid, "GUID");
+    DescribeRecord(PGOCtxProfileRecords::Counters, "Counters");
   }
   Writer.ExitBlock();
   Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, CodeLen);
@@ -108,12 +114,27 @@ void PGOCtxProfileWriter::startContextSection() {
   Writer.EnterSubblock(PGOCtxProfileBlockIDs::ContextsSectionBlockID, CodeLen);
 }
 
+void PGOCtxProfileWriter::startFlatSection() {
+  Writer.EnterSubblock(PGOCtxProfileBlockIDs::FlatProfilesSectionBlockID,
+                       CodeLen);
+}
+
 void PGOCtxProfileWriter::endContextSection() { Writer.ExitBlock(); }
+void PGOCtxProfileWriter::endFlatSection() { Writer.ExitBlock(); }
 
 void PGOCtxProfileWriter::writeContextual(const ContextNode &RootNode) {
   writeImpl(std::nullopt, RootNode);
 }
 
+void PGOCtxProfileWriter::writeFlatSection(ctx_profile::GUID Guid,
+                                           const uint64_t *Buffer,
+                                           size_t Size) {
+  Writer.EnterSubblock(PGOCtxProfileBlockIDs::FlatProfileBlockID, CodeLen);
+  writeGuid(Guid);
+  writeCounters({Buffer, Size});
+  Writer.ExitBlock();
+}
+
 namespace {
 
 /// Representation of the context node suitable for yaml serialization /
@@ -123,8 +144,13 @@ struct SerializableCtxRepresentation {
   std::vector<uint64_t> Counters;
   std::vector<std::vector<SerializableCtxRepresentation>> Callsites;
 };
+
+using SerializableFlatProfileRepresentation =
+    std::pair<ctx_profile::GUID, std::vector<uint64_t>>;
+
 struct SerializableProfileRepresentation {
   std::vector<SerializableCtxRepresentation> Contexts;
+  std::vector<SerializableFlatProfileRepresentation> FlatProfiles;
 };
 
 ctx_profile::ContextNode *
@@ -164,6 +190,7 @@ createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(SerializableCtxRepresentation)
 LLVM_YAML_IS_SEQUENCE_VECTOR(std::vector<SerializableCtxRepresentation>)
+LLVM_YAML_IS_SEQUENCE_VECTOR(SerializableFlatProfileRepresentation)
 template <> struct yaml::MappingTraits<SerializableCtxRepresentation> {
   static void mapping(yaml::IO &IO, SerializableCtxRepresentation &SCR) {
     IO.mapRequired("Guid", SCR.Guid);
@@ -175,6 +202,15 @@ template <> struct yaml::MappingTraits<SerializableCtxRepresentation> {
 template <> struct yaml::MappingTraits<SerializableProfileRepresentation> {
   static void mapping(yaml::IO &IO, SerializableProfileRepresentation &SPR) {
     IO.mapOptional("Contexts", SPR.Contexts);
+    IO.mapOptional("FlatProfiles", SPR.FlatProfiles);
+  }
+};
+
+template <> struct yaml::MappingTraits<SerializableFlatProfileRepresentation> {
+  static void mapping(yaml::IO &IO,
+                      SerializableFlatProfileRepresentation &SFPR) {
+    IO.mapRequired("Guid", SFPR.first);
+    IO.mapRequired("Counters", SFPR.second);
   }
 };
 
@@ -201,6 +237,12 @@ Error llvm::createCtxProfFromYAML(StringRef Profile, raw_ostream &Out) {
     }
     Writer.endContextSection();
   }
+  if (!SPR.FlatProfiles.empty()) {
+    Writer.startFlatSection();
+    for (const auto &[Guid, Counters] : SPR.FlatProfiles)
+      Writer.writeFlatSection(Guid, Counters.data(), Counters.size());
+    Writer.endFlatSection();
+  }
   if (EC)
     return createStringError(EC, "failed to write output");
   return Error::success();

llvm/test/tools/llvm-ctxprof-util/Inputs/invalid-flat.yaml

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

llvm/test/tools/llvm-ctxprof-util/Inputs/valid-ctx-only.yaml

@@ -0,0 +1,14 @@
+
+Contexts:
+  - Guid:            1000
+    Counters:        [ 1, 2, 3 ]
+    Callsites:
+      - [  ]
+      - - Guid:            2000
+          Counters:        [ 4, 5 ]
+        - Guid:            18446744073709551613
+          Counters:        [ 6, 7, 8 ]
+      - - Guid:            3000
+          Counters:        [ 40, 50 ]
+  - Guid:            18446744073709551612
+    Counters:        [ 5, 9, 10 ]

llvm/test/tools/llvm-ctxprof-util/Inputs/valid-flat-first.yaml

@@ -0,0 +1,19 @@
+
+FlatProfiles:
+  - Guid:            1234
+    Counters:        [ 5, 6, 7 ]
+  - Guid:            5555
+    Counters:        [ 1 ]
+Contexts:
+  - Guid:            1000
+    Counters:        [ 1, 2, 3 ]
+    Callsites:
+      - [  ]
+      - - Guid:            2000
+          Counters:        [ 4, 5 ]
+        - Guid:            18446744073709551613
+          Counters:        [ 6, 7, 8 ]
+      - - Guid:            3000
+          Counters:        [ 40, 50 ]
+  - Guid:            18446744073709551612
+    Counters:        [ 5, 9, 10 ]

llvm/test/tools/llvm-ctxprof-util/Inputs/valid-flat-only.yaml

@@ -0,0 +1,6 @@
+
+FlatProfiles:
+  - Guid:            1234
+    Counters:        [ 5, 6, 7 ]
+  - Guid:            5555
+    Counters:        [ 1 ]

llvm/test/tools/llvm-ctxprof-util/Inputs/valid.yaml

@@ -12,3 +12,8 @@ Contexts:
           Counters:        [ 40, 50 ]
   - Guid:            18446744073709551612
     Counters:        [ 5, 9, 10 ]
+FlatProfiles:
+  - Guid:            1234
+    Counters:        [ 5, 6, 7 ]
+  - Guid:            5555
+    Counters:        [ 1 ]

llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util-negative.test

@@ -9,6 +9,7 @@
 ; RUN: not llvm-ctxprof-util fromYAML --input %S/Inputs/invalid-no-ctx.yaml 2>&1 | FileCheck %s --check-prefix=NO_CTX
 ; RUN: not llvm-ctxprof-util fromYAML --input %S/Inputs/invalid-no-counters.yaml 2>&1 | FileCheck %s --check-prefix=NO_COUNTERS
 ; RUN: not llvm-ctxprof-util fromYAML --input %S/Inputs/invalid-bad-subctx.yaml 2>&1 | FileCheck %s --check-prefix=BAD_SUBCTX
+; RUN: not llvm-ctxprof-util fromYAML --input %S/Inputs/invalid-flat.yaml 2>&1 | FileCheck %s --check-prefix=BAD_FLAT
 ; RUN: rm -rf %t
 ; RUN: not llvm-ctxprof-util fromYAML --input %S/Inputs/valid.yaml --output %t/output.bitstream 2>&1 | FileCheck %s --check-prefix=NO_DIR
 
@@ -21,4 +22,5 @@
 ; NO_CTX: YAML:1:1: error: not a mapping
 ; NO_COUNTERS: YAML:2:5: error: missing required key 'Counters'
 ; BAD_SUBCTX: YAML:4:18: error: not a sequence
+; BAD_FLAT: YAML:2:5: error: missing required key 'Counters'
 ; NO_DIR: failed to open output

llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util.test

@@ -8,6 +8,21 @@
 ; RUN: llvm-ctxprof-util toYAML -input %t/valid.bitstream -output %t/valid2.yaml
 ; RUN: diff %t/valid2.yaml %S/Inputs/valid.yaml
 
+
+; RUN: llvm-ctxprof-util fromYAML -input %S/Inputs/valid-ctx-only.yaml -output %t/valid-ctx-only.bitstream
+; RUN: llvm-ctxprof-util toYAML -input %t/valid-ctx-only.bitstream -output %t/valid-ctx-only.yaml
+; RUN: diff %t/valid-ctx-only.yaml %S/Inputs/valid-ctx-only.yaml
+
+; RUN: llvm-ctxprof-util fromYAML -input %S/Inputs/valid-flat-only.yaml -output %t/valid-flat-only.bitstream
+; RUN: llvm-ctxprof-util toYAML -input %t/valid-flat-only.bitstream -output %t/valid-flat-only.yaml
+; RUN: diff %t/valid-flat-only.yaml %S/Inputs/valid-flat-only.yaml
+
+; This case is the "valid.yaml" case but with the flat profile first.
+; The output, though, should match valid.yaml
+; RUN: llvm-ctxprof-util fromYAML -input %S/Inputs/valid-flat-first.yaml -output %t/valid-flat-first.bitstream
+; RUN: llvm-ctxprof-util toYAML -input %t/valid-flat-first.bitstream -output %t/valid-flat-first.yaml
+; RUN: diff %t/valid-flat-first.yaml %S/Inputs/valid.yaml
+
 ; For the valid case, check against a reference output.
 ; Note that uint64_t are printed as signed values by llvm-bcanalyzer:
 ;  * 18446744073709551613 in yaml is -3 in the output
@@ -22,7 +37,7 @@
 ; EMPTY-NEXT: </Metadata>
 
 ; VALID:      <BLOCKINFO_BLOCK/>
-; VALID-NEXT: <Metadata NumWords=33 BlockCodeSize=2>
+; VALID-NEXT: <Metadata NumWords=45 BlockCodeSize=2>
 ; VALID-NEXT:   <Version op0=2/>
 ; VALID-NEXT:   <Contexts NumWords=29 BlockCodeSize=2>
 ; VALID-NEXT:     <Root NumWords=20 BlockCodeSize=2>
@@ -49,4 +64,14 @@
 ; VALID-NEXT:       <Counters op0=5 op1=9 op2=10/>
 ; VALID-NEXT:     </Root>
 ; VALID-NEXT:   </Contexts>
+; VALID-NEXT:   <FlatProfiles NumWords=10 BlockCodeSize=2>
+; VALID-NEXT:       <Flat NumWords=3 BlockCodeSize=2>
+; VALID-NEXT:         <GUID op0=1234/>
+; VALID-NEXT:         <Counters op0=5 op1=6 op2=7/>
+; VALID-NEXT:       </Flat>
+; VALID-NEXT:       <Flat NumWords=2 BlockCodeSize=2>
+; VALID-NEXT:         <GUID op0=5555/>
+; VALID-NEXT:         <Counters op0=1/>
+; VALID-NEXT:       </Flat>
+; VALID-NEXT:   </FlatProfiles>
 ; VALID-NEXT: </Metadata>