[memprof] Add MemProf version (#86414)

This patch adds a version field to the MemProf section of the indexed
profile format, calling the new version "version 1".  The existing
version is called "version 0".

The writer supports both versions via a command-line option:

  llvm-profdata merge --memprof-version=1 ...

The reader supports both versions by automatically detecting the
version from the header.

5 files changed, 91 insertions, 18 deletions

diff --git a/llvm/include/llvm/ProfileData/InstrProfWriter.h b/llvm/include/llvm/ProfileData/InstrProfWriter.h
index f70574d1f756..d2156c867872 100644
--- a/llvm/include/llvm/ProfileData/InstrProfWriter.h
+++ b/llvm/include/llvm/ProfileData/InstrProfWriter.h
@@ -75,11 +75,14 @@ private:
   // deployment of newer versions of llvm-profdata.
   bool WritePrevVersion = false;
 
+  // The MemProf version we should write.
+  memprof::IndexedVersion MemProfVersionRequested;
+
 public:
-  InstrProfWriter(bool Sparse = false,
-                  uint64_t TemporalProfTraceReservoirSize = 0,
-                  uint64_t MaxTemporalProfTraceLength = 0,
-                  bool WritePrevVersion = false);
+  InstrProfWriter(
+      bool Sparse = false, uint64_t TemporalProfTraceReservoirSize = 0,
+      uint64_t MaxTemporalProfTraceLength = 0, bool WritePrevVersion = false,
+      memprof::IndexedVersion MemProfVersionRequested = memprof::Version0);
   ~InstrProfWriter();
 
   StringMap<ProfilingData> &getProfileData() { return FunctionData; }
diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index cd0c148b4f2f..ff00900a1466 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -16,6 +16,20 @@
 namespace llvm {
 namespace memprof {
 
+// The versions of the indexed MemProf format
+enum IndexedVersion : uint64_t {
+  // Version 0: This version didn't have a version field.
+  Version0 = 0,
+  // Version 1: Added a version field to the header.
+  Version1 = 1,
+};
+
+constexpr uint64_t MinimumSupportedVersion = Version0;
+constexpr uint64_t MaximumSupportedVersion = Version1;
+
+// Verify that the minimum and maximum satisfy the obvious constraint.
+static_assert(MinimumSupportedVersion <= MaximumSupportedVersion);
+
 enum class Meta : uint64_t {
   Start = 0,
 #define MIBEntryDef(NameTag, Name, Type) NameTag,
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index 4f786a23f354..a275d4852c15 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -24,6 +24,7 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/Support/VirtualFileSystem.h"
@@ -1230,10 +1231,39 @@ Error IndexedInstrProfReader::readHeader() {
             Header->MemProfOffset);
 
     const unsigned char *Ptr = Start + MemProfOffset;
-    // The value returned from RecordTableGenerator.Emit.
-    const uint64_t RecordTableOffset =
+
+    // Read the first 64-bit word, which may be RecordTableOffset in
+    // memprof::MemProfVersion0 or the MemProf version number in
+    // memprof::MemProfVersion1.
+    const uint64_t FirstWord =
         support::endian::readNext<uint64_t, llvm::endianness::little,
                                   unaligned>(Ptr);
+
+    memprof::IndexedVersion Version = memprof::Version0;
+    if (FirstWord == memprof::Version1) {
+      // Everything is good.  We can proceed to deserialize the rest.
+      Version = memprof::Version1;
+    } else if (FirstWord >= 24) {
+      // This is a heuristic/hack to detect memprof::MemProfVersion0,
+      // which does not have a version field in the header.
+      // In memprof::MemProfVersion0, FirstWord will be RecordTableOffset,
+      // which should be at least 24 because of the MemProf header size.
+      Version = memprof::Version0;
+    } else {
+      return make_error<InstrProfError>(
+          instrprof_error::unsupported_version,
+          formatv("MemProf version {} not supported; "
+                  "requires version between {} and {}, inclusive",
+                  FirstWord, memprof::MinimumSupportedVersion,
+                  memprof::MaximumSupportedVersion));
+    }
+
+    // The value returned from RecordTableGenerator.Emit.
+    const uint64_t RecordTableOffset =
+        Version == memprof::Version0
+            ? FirstWord
+            : support::endian::readNext<uint64_t, llvm::endianness::little,
+                                        unaligned>(Ptr);
     // The offset in the stream right before invoking
     // FrameTableGenerator.Emit.
     const uint64_t FramePayloadOffset =
diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index d9fe88a00bdf..8f067f8d05e2 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -22,6 +22,7 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/OnDiskHashTable.h"
 #include "llvm/Support/raw_ostream.h"
@@ -179,14 +180,15 @@ public:
 
 } // end namespace llvm
 
-InstrProfWriter::InstrProfWriter(bool Sparse,
-                                 uint64_t TemporalProfTraceReservoirSize,
-                                 uint64_t MaxTemporalProfTraceLength,
-                                 bool WritePrevVersion)
+InstrProfWriter::InstrProfWriter(
+    bool Sparse, uint64_t TemporalProfTraceReservoirSize,
+    uint64_t MaxTemporalProfTraceLength, bool WritePrevVersion,
+    memprof::IndexedVersion MemProfVersionRequested)
     : Sparse(Sparse), MaxTemporalProfTraceLength(MaxTemporalProfTraceLength),
       TemporalProfTraceReservoirSize(TemporalProfTraceReservoirSize),
       InfoObj(new InstrProfRecordWriterTrait()),
-      WritePrevVersion(WritePrevVersion) {}
+      WritePrevVersion(WritePrevVersion),
+      MemProfVersionRequested(MemProfVersionRequested) {}
 
 InstrProfWriter::~InstrProfWriter() { delete InfoObj; }
 
@@ -516,6 +518,7 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
 
   // Write the MemProf profile data if we have it. This includes a simple schema
   // with the format described below followed by the hashtable:
+  // uint64_t Version
   // uint64_t RecordTableOffset = RecordTableGenerator.Emit
   // uint64_t FramePayloadOffset = Stream offset before emitting the frame table
   // uint64_t FrameTableOffset = FrameTableGenerator.Emit
@@ -528,7 +531,21 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   // OnDiskChainedHashTable MemProfFrameData
   uint64_t MemProfSectionStart = 0;
   if (static_cast<bool>(ProfileKind & InstrProfKind::MemProf)) {
+    if (MemProfVersionRequested < memprof::MinimumSupportedVersion ||
+        MemProfVersionRequested > memprof::MaximumSupportedVersion) {
+      return make_error<InstrProfError>(
+          instrprof_error::unsupported_version,
+          formatv("MemProf version {} not supported; "
+                  "requires version between {} and {}, inclusive",
+                  MemProfVersionRequested, memprof::MinimumSupportedVersion,
+                  memprof::MaximumSupportedVersion));
+    }
+
     MemProfSectionStart = OS.tell();
+
+    if (MemProfVersionRequested >= memprof::Version1)
+      OS.write(MemProfVersionRequested);
+
     OS.write(0ULL); // Reserve space for the memprof record table offset.
     OS.write(0ULL); // Reserve space for the memprof frame payload offset.
     OS.write(0ULL); // Reserve space for the memprof frame table offset.
@@ -570,12 +587,13 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
 
     uint64_t FrameTableOffset = FrameTableGenerator.Emit(OS.OS, *FrameWriter);
 
-    PatchItem PatchItems[] = {
-        {MemProfSectionStart, &RecordTableOffset, 1},
-        {MemProfSectionStart + sizeof(uint64_t), &FramePayloadOffset, 1},
-        {MemProfSectionStart + 2 * sizeof(uint64_t), &FrameTableOffset, 1},
-    };
-    OS.patch(PatchItems);
+    uint64_t Header[] = {RecordTableOffset, FramePayloadOffset,
+                         FrameTableOffset};
+    uint64_t HeaderUpdatePos = MemProfSectionStart;
+    if (MemProfVersionRequested >= memprof::Version1)
+      // The updates go just after the version field.
+      HeaderUpdatePos += sizeof(uint64_t);
+    OS.patch({{HeaderUpdatePos, Header, std::size(Header)}});
   }
 
   // BinaryIdSection has two parts:
diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 3a7bd061d3d2..e8ee3c238194 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -300,6 +300,13 @@ cl::opt<bool> DoWritePrevVersion(
     cl::desc("Write the previous version of indexed format, to enable "
              "some forward compatibility."));
 
+cl::opt<memprof::IndexedVersion> MemProfVersionRequested(
+    "memprof-version", cl::Hidden, cl::sub(MergeSubcommand),
+    cl::desc("Specify the version of the memprof format to use"),
+    cl::init(memprof::Version0),
+    cl::values(clEnumValN(memprof::Version0, "0", "version 0"),
+               clEnumValN(memprof::Version1, "1", "version 1")));
+
 // Options specific to overlap subcommand.
 cl::opt<std::string> BaseFilename(cl::Positional, cl::Required,
                                   cl::desc("<base profile file>"),
@@ -588,7 +595,8 @@ struct WriterContext {
   WriterContext(bool IsSparse, std::mutex &ErrLock,
                 SmallSet<instrprof_error, 4> &WriterErrorCodes,
                 uint64_t ReservoirSize = 0, uint64_t MaxTraceLength = 0)
-      : Writer(IsSparse, ReservoirSize, MaxTraceLength, DoWritePrevVersion),
+      : Writer(IsSparse, ReservoirSize, MaxTraceLength, DoWritePrevVersion,
+               MemProfVersionRequested),
         ErrLock(ErrLock), WriterErrorCodes(WriterErrorCodes) {}
 };