[ctx_prof] Make the profile output analyzable by llvm-bcanalyzer (#99563)

This requires output-ing a "Magic" 4-byte header. We also emit a block info block, to describe our blocks and records. The output of `llvm-bcanalyzer` would look like:

```
<BLOCKINFO_BLOCK/>
<Metadata NumWords=17 BlockCodeSize=2>
  <Version op0=1/>
  <Context NumWords=13 BlockCodeSize=2>
    <GUID op0=2/>
    <Counters op0=1 op1=2 op2=3/>
```

Instead of having `Unknown` for block and record IDs.

Author: mtrofin

llvm/include/llvm/ProfileData/PGOCtxProfReader.h

@@ -73,7 +73,8 @@ class PGOContextualProfile final {
 };
 
 class PGOCtxProfileReader final {
-  BitstreamCursor &Cursor;
+  StringRef Magic;
+  BitstreamCursor Cursor;
   Expected<BitstreamEntry> advance();
   Error readMetadata();
   Error wrongValue(const Twine &);
@@ -84,7 +85,9 @@ class PGOCtxProfileReader final {
   bool canReadContext();
 
 public:
-  PGOCtxProfileReader(BitstreamCursor &Cursor) : Cursor(Cursor) {}
+  PGOCtxProfileReader(StringRef Buffer)
+      : Magic(Buffer.substr(0, PGOCtxProfileWriter::ContainerMagic.size())),
+        Cursor(Buffer.substr(PGOCtxProfileWriter::ContainerMagic.size())) {}
 
   Expected<std::map<GlobalValue::GUID, PGOContextualProfile>> loadContexts();
 };

llvm/include/llvm/ProfileData/PGOCtxProfWriter.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_PROFILEDATA_PGOCTXPROFWRITER_H_
 #define LLVM_PROFILEDATA_PGOCTXPROFWRITER_H_
 
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Bitstream/BitCodeEnums.h"
 #include "llvm/Bitstream/BitstreamWriter.h"
 #include "llvm/ProfileData/CtxInstrContextNode.h"
@@ -68,15 +69,7 @@ class PGOCtxProfileWriter final {
 
 public:
   PGOCtxProfileWriter(raw_ostream &Out,
-                      std::optional<unsigned> VersionOverride = std::nullopt)
-      : Writer(Out, 0) {
-    Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID,
-                         CodeLen);
-    const auto Version = VersionOverride ? *VersionOverride : CurrentVersion;
-    Writer.EmitRecord(PGOCtxProfileRecords::Version,
-                      SmallVector<unsigned, 1>({Version}));
-  }
-
+                      std::optional<unsigned> VersionOverride = std::nullopt);
   ~PGOCtxProfileWriter() { Writer.ExitBlock(); }
 
   void write(const ctx_profile::ContextNode &);
@@ -85,6 +78,7 @@ class PGOCtxProfileWriter final {
   static constexpr unsigned CodeLen = 2;
   static constexpr uint32_t CurrentVersion = 1;
   static constexpr unsigned VBREncodingBits = 6;
+  static constexpr StringRef ContainerMagic = "CTXP";
 };
 
 } // namespace llvm

llvm/lib/ProfileData/PGOCtxProfReader.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ProfileData/PGOCtxProfReader.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Bitstream/BitCodeEnums.h"
 #include "llvm/Bitstream/BitstreamReader.h"
 #include "llvm/ProfileData/InstrProf.h"
@@ -139,6 +140,20 @@ PGOCtxProfileReader::readContext(bool ExpectIndex) {
 }
 
 Error PGOCtxProfileReader::readMetadata() {
+  if (Magic.size() < PGOCtxProfileWriter::ContainerMagic.size() ||
+      Magic != PGOCtxProfileWriter::ContainerMagic)
+    return make_error<InstrProfError>(instrprof_error::invalid_prof,
+                                      "Invalid magic");
+
+  BitstreamEntry Entry;
+  RET_ON_ERR(Cursor.advance().moveInto(Entry));
+  if (Entry.Kind != BitstreamEntry::SubBlock ||
+      Entry.ID != bitc::BLOCKINFO_BLOCK_ID)
+    return unsupported("Expected Block ID");
+  // We don't need the blockinfo to read the rest, it's metadata usable for e.g.
+  // llvm-bcanalyzer.
+  RET_ON_ERR(Cursor.SkipBlock());
+
   EXPECT_OR_RET(Blk, advance());
   if (Blk->Kind != BitstreamEntry::SubBlock)
     return unsupported("Expected Version record");

llvm/lib/ProfileData/PGOCtxProfWriter.cpp

@@ -16,6 +16,40 @@
 using namespace llvm;
 using namespace llvm::ctx_profile;
 
+PGOCtxProfileWriter::PGOCtxProfileWriter(
+    raw_ostream &Out, std::optional<unsigned> VersionOverride)
+    : Writer(Out, 0) {
+  static_assert(ContainerMagic.size() == 4);
+  Out.write(ContainerMagic.data(), ContainerMagic.size());
+  Writer.EnterBlockInfoBlock();
+  {
+    auto DescribeBlock = [&](unsigned ID, StringRef Name) {
+      Writer.EmitRecord(bitc::BLOCKINFO_CODE_SETBID,
+                        SmallVector<unsigned, 1>{ID});
+      Writer.EmitRecord(bitc::BLOCKINFO_CODE_BLOCKNAME,
+                        llvm::arrayRefFromStringRef(Name));
+    };
+    SmallVector<uint64_t, 16> Data;
+    auto DescribeRecord = [&](unsigned RecordID, StringRef Name) {
+      Data.clear();
+      Data.push_back(RecordID);
+      llvm::append_range(Data, Name);
+      Writer.EmitRecord(bitc::BLOCKINFO_CODE_SETRECORDNAME, Data);
+    };
+    DescribeBlock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, "Metadata");
+    DescribeRecord(PGOCtxProfileRecords::Version, "Version");
+    DescribeBlock(PGOCtxProfileBlockIDs::ContextNodeBlockID, "Context");
+    DescribeRecord(PGOCtxProfileRecords::Guid, "GUID");
+    DescribeRecord(PGOCtxProfileRecords::CalleeIndex, "CalleeIndex");
+    DescribeRecord(PGOCtxProfileRecords::Counters, "Counters");
+  }
+  Writer.ExitBlock();
+  Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID, CodeLen);
+  const auto Version = VersionOverride ? *VersionOverride : CurrentVersion;
+  Writer.EmitRecord(PGOCtxProfileRecords::Version,
+                    SmallVector<unsigned, 1>({Version}));
+}
+
 void PGOCtxProfileWriter::writeCounters(const ContextNode &Node) {
   Writer.EmitCode(bitc::UNABBREV_RECORD);
   Writer.EmitVBR(PGOCtxProfileRecords::Counters, VBREncodingBits);

llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Bitstream/BitstreamReader.h"
+#include "llvm/Bitcode/BitcodeAnalyzer.h"
 #include "llvm/ProfileData/CtxInstrContextNode.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
@@ -106,8 +106,20 @@ TEST_F(PGOCtxProfRWTest, RoundTrip) {
         MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+
+    // Check it's analyzable by the BCAnalyzer
+    BitcodeAnalyzer BA((*MB)->getBuffer());
+    std::string AnalyzerDump;
+    raw_string_ostream OS(AnalyzerDump);
+    BCDumpOptions Opts(OS);
+
+    // As in, expect no error.
+    EXPECT_FALSE(BA.analyze(Opts));
+    EXPECT_TRUE(AnalyzerDump.find("<Metadata BlockID") != std::string::npos);
+    EXPECT_TRUE(AnalyzerDump.find("<Context BlockID") != std::string::npos);
+    EXPECT_TRUE(AnalyzerDump.find("<CalleeIndex codeid") != std::string::npos);
+
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     ASSERT_TRUE(!!Expected);
     auto &Ctxes = *Expected;
@@ -143,25 +155,22 @@ TEST_F(PGOCtxProfRWTest, InvalidCounters) {
     auto MB = MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
   }
 }
 
 TEST_F(PGOCtxProfRWTest, Empty) {
-  BitstreamCursor Cursor("");
-  PGOCtxProfileReader Reader(Cursor);
+  PGOCtxProfileReader Reader("");
   auto Expected = Reader.loadContexts();
   EXPECT_FALSE(Expected);
   consumeError(Expected.takeError());
 }
 
 TEST_F(PGOCtxProfRWTest, Invalid) {
-  BitstreamCursor Cursor("Surely this is not valid");
-  PGOCtxProfileReader Reader(Cursor);
+  PGOCtxProfileReader Reader("Surely this is not valid");
   auto Expected = Reader.loadContexts();
   EXPECT_FALSE(Expected);
   consumeError(Expected.takeError());
@@ -182,8 +191,8 @@ TEST_F(PGOCtxProfRWTest, ValidButEmpty) {
     auto MB = MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     EXPECT_TRUE(!!Expected);
     EXPECT_TRUE(Expected->empty());
@@ -204,8 +213,8 @@ TEST_F(PGOCtxProfRWTest, WrongVersion) {
     auto MB = MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
@@ -228,8 +237,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateRoots) {
     auto MB = MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
@@ -255,8 +263,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateTargets) {
     auto MB = MemoryBuffer::getFile(ProfileFile.path());
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
-    BitstreamCursor Cursor((*MB)->getBuffer());
-    PGOCtxProfileReader Reader(Cursor);
+    PGOCtxProfileReader Reader((*MB)->getBuffer());
     auto Expected = Reader.loadContexts();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());