[SampleFDO] Add ExtBinary format to support extension of binary profile.

This is a patch split from https://reviews.llvm.org/D66374. It tries to add
a new format of profile called ExtBinary. The format adds a section header
table to the profile and organize the profile in sections, so the future
extension like adding a new section or extending an existing section will be
easier while keeping backward compatiblity feasible.

Differential Revision: https://reviews.llvm.org/D66513

llvm-svn: 369798

include/llvm/ProfileData/SampleProf.h

@@ -84,6 +84,7 @@ enum SampleProfileFormat {
   SPF_Text = 0x1,
   SPF_Compact_Binary = 0x2,
   SPF_GCC = 0x3,
+  SPF_Ext_Binary = 0x4,
   SPF_Binary = 0xff
 };
 
@@ -106,6 +107,27 @@ static inline StringRef getRepInFormat(StringRef Name,
 
 static inline uint64_t SPVersion() { return 103; }
 
+// Section Type used by SampleProfileExtBinaryBaseReader and
+// SampleProfileExtBinaryBaseWriter. Never change the existing
+// value of enum. Only append new ones.
+enum SecType {
+  SecInValid = 0,
+  SecProfSummary = 1,
+  SecNameTable = 2,
+  // marker for the first type of profile.
+  SecFuncProfileFirst = 32,
+  SecLBRProfile = SecFuncProfileFirst
+};
+
+// Entry type of section header table used by SampleProfileExtBinaryBaseReader
+// and SampleProfileExtBinaryBaseWriter.
+struct SecHdrTableEntry {
+  SecType Type;
+  uint64_t Flag;
+  uint64_t Offset;
+  uint64_t Size;
+};
+
 /// Represents the relative location of an instruction.
 ///
 /// Instruction locations are specified by the line offset from the

include/llvm/ProfileData/SampleProfReader.h

@@ -416,38 +416,89 @@ protected:
   /// Read the contents of the given profile instance.
   std::error_code readProfile(FunctionSamples &FProfile);
 
+  /// Read the contents of Magic number and Version number.
+  std::error_code readMagicIdent();
+
+  /// Read profile summary.
+  std::error_code readSummary();
+
+  /// Read the whole name table.
+  virtual std::error_code readNameTable();
+
   /// Points to the current location in the buffer.
   const uint8_t *Data = nullptr;
 
   /// Points to the end of the buffer.
   const uint8_t *End = nullptr;
 
+  /// Function name table.
+  std::vector<StringRef> NameTable;
+
+  /// Read a string indirectly via the name table.
+  virtual ErrorOr<StringRef> readStringFromTable();
+
 private:
   std::error_code readSummaryEntry(std::vector<ProfileSummaryEntry> &Entries);
   virtual std::error_code verifySPMagic(uint64_t Magic) = 0;
-
-  /// Read profile summary.
-  std::error_code readSummary();
-
-  /// Read the whole name table.
-  virtual std::error_code readNameTable() = 0;
-
-  /// Read a string indirectly via the name table.
-  virtual ErrorOr<StringRef> readStringFromTable() = 0;
 };
 
 class SampleProfileReaderRawBinary : public SampleProfileReaderBinary {
 private:
-  /// Function name table.
-  std::vector<StringRef> NameTable;
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
-  virtual std::error_code readNameTable() override;
-  /// Read a string indirectly via the name table.
-  virtual ErrorOr<StringRef> readStringFromTable() override;
 
 public:
-  SampleProfileReaderRawBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C)
-      : SampleProfileReaderBinary(std::move(B), C, SPF_Binary) {}
+  SampleProfileReaderRawBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
+                               SampleProfileFormat Format = SPF_Binary)
+      : SampleProfileReaderBinary(std::move(B), C, Format) {}
+
+  /// \brief Return true if \p Buffer is in the format supported by this class.
+  static bool hasFormat(const MemoryBuffer &Buffer);
+};
+
+/// SampleProfileReaderExtBinaryBase/SampleProfileWriterExtBinaryBase defines
+/// the basic structure of the extensible binary format.
+/// The format is organized in sections except the magic and version number
+/// at the beginning. There is a section table before all the sections, and
+/// each entry in the table describes the entry type, start, size and
+/// attributes. The format in each section is defined by the section itself.
+///
+/// It is easy to add a new section while maintaining the backward
+/// compatibility of the profile. Nothing extra needs to be done. If we want
+/// to extend an existing section, like add cache misses information in
+/// addition to the sample count in the profile body, we can add a new section
+/// with the extension and retire the existing section, and we could choose
+/// to keep the parser of the old section if we want the reader to be able
+/// to read both new and old format profile.
+///
+/// SampleProfileReaderExtBinary/SampleProfileWriterExtBinary define the
+/// commonly used sections of a profile in extensible binary format. It is
+/// possible to define other types of profile inherited from
+/// SampleProfileReaderExtBinaryBase/SampleProfileWriterExtBinaryBase.
+class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
+protected:
+  std::vector<SecHdrTableEntry> SecHdrTable;
+  std::error_code readSecHdrTableEntry();
+  std::error_code readSecHdrTable();
+  virtual std::error_code readHeader() override;
+  virtual std::error_code verifySPMagic(uint64_t Magic) = 0;
+
+public:
+  SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
+                                   LLVMContext &C, SampleProfileFormat Format)
+      : SampleProfileReaderBinary(std::move(B), C, Format) {}
+
+  /// Read sample profiles in extensible format from the associated file.
+  std::error_code read() override;
+};
+
+class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
+private:
+  virtual std::error_code verifySPMagic(uint64_t Magic) override;
+
+public:
+  SampleProfileReaderExtBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
+                               SampleProfileFormat Format = SPF_Ext_Binary)
+      : SampleProfileReaderExtBinaryBase(std::move(B), C, Format) {}
 
   /// \brief Return true if \p Buffer is in the format supported by this class.
   static bool hasFormat(const MemoryBuffer &Buffer);

include/llvm/ProfileData/SampleProfWriter.h

@@ -36,7 +36,7 @@ public:
   /// Write sample profiles in \p S.
   ///
   /// \returns status code of the file update operation.
-  virtual std::error_code write(const FunctionSamples &S) = 0;
+  virtual std::error_code writeSample(const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
@@ -64,6 +64,10 @@ protected:
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) = 0;
 
+  // Write function profiles to the profile file.
+  virtual std::error_code
+  writeFuncProfiles(const StringMap<FunctionSamples> &ProfileMap);
+
   /// Output stream where to emit the profile to.
   std::unique_ptr<raw_ostream> OutputStream;
 
@@ -72,12 +76,15 @@ protected:
 
   /// Compute summary for this profile.
   void computeSummary(const StringMap<FunctionSamples> &ProfileMap);
+
+  /// Profile format.
+  SampleProfileFormat Format;
 };
 
 /// Sample-based profile writer (text format).
 class SampleProfileWriterText : public SampleProfileWriter {
 public:
-  std::error_code write(const FunctionSamples &S) override;
+  std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
   SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
@@ -102,13 +109,14 @@ private:
 /// Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
-  virtual std::error_code write(const FunctionSamples &S) override;
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
+
+protected:
   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS) {}
 
-protected:
-  virtual std::error_code writeNameTable() = 0;
-  virtual std::error_code writeMagicIdent() = 0;
+  virtual std::error_code writeMagicIdent(SampleProfileFormat Format);
+  virtual std::error_code writeNameTable();
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeSummary();
@@ -118,10 +126,10 @@ protected:
 
   MapVector<StringRef, uint32_t> NameTable;
 
-private:
   void addName(StringRef FName);
   void addNames(const FunctionSamples &S);
 
+private:
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
@@ -129,10 +137,50 @@ private:
 
 class SampleProfileWriterRawBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
+};
+
+class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
+  using SampleProfileWriterBinary::SampleProfileWriterBinary;
+
+public:
+  virtual std::error_code
+  write(const StringMap<FunctionSamples> &ProfileMap) override;
 
 protected:
-  virtual std::error_code writeNameTable() override;
-  virtual std::error_code writeMagicIdent() override;
+  uint64_t markSectionStart();
+  uint64_t addNewSection(SecType Sec, uint64_t SectionStart);
+  virtual void initSectionLayout() = 0;
+  virtual std::error_code
+  writeSections(const StringMap<FunctionSamples> &ProfileMap) = 0;
+
+  // Specifiy the section layout in the profile. Note that the order in
+  // SecHdrTable (order to collect sections) may be different from the
+  // order in SectionLayout (order to write out sections into profile).
+  SmallVector<SecType, 8> SectionLayout;
+
+private:
+  void allocSecHdrTable();
+  std::error_code writeSecHdrTable();
+  virtual std::error_code
+  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+
+  // The location where the output stream starts.
+  uint64_t FileStart;
+  // The location in the output stream where the SecHdrTable should be
+  // written to.
+  uint64_t SecHdrTableOffset;
+  std::vector<SecHdrTableEntry> SecHdrTable;
+};
+
+class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
+  using SampleProfileWriterExtBinaryBase::SampleProfileWriterExtBinaryBase;
+
+private:
+  virtual void initSectionLayout() {
+    SectionLayout = {SecProfSummary, SecNameTable, SecLBRProfile};
+  };
+  virtual std::error_code
+  writeSections(const StringMap<FunctionSamples> &ProfileMap) override;
 };
 
 // CompactBinary is a compact format of binary profile which both reduces
@@ -169,7 +217,7 @@ class SampleProfileWriterCompactBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 
 public:
-  virtual std::error_code write(const FunctionSamples &S) override;
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
   virtual std::error_code
   write(const StringMap<FunctionSamples> &ProfileMap) override;
 
@@ -181,7 +229,6 @@ protected:
   /// towards profile start.
   uint64_t TableOffset;
   virtual std::error_code writeNameTable() override;
-  virtual std::error_code writeMagicIdent() override;
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeFuncOffsetTable();

lib/ProfileData/SampleProfReader.cpp

@@ -345,7 +345,7 @@ inline ErrorOr<uint32_t> SampleProfileReaderBinary::readStringIndex(T &Table) {
   return *Idx;
 }
 
-ErrorOr<StringRef> SampleProfileReaderRawBinary::readStringFromTable() {
+ErrorOr<StringRef> SampleProfileReaderBinary::readStringFromTable() {
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
@@ -467,6 +467,40 @@ std::error_code SampleProfileReaderBinary::read() {
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileReaderExtBinaryBase::read() {
+  const uint8_t *BufStart =
+      reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
+
+  for (auto &Entry : SecHdrTable) {
+    // Skip empty section.
+    if (!Entry.Size)
+      continue;
+    Data = BufStart + Entry.Offset;
+    switch (Entry.Type) {
+    case SecProfSummary:
+      if (std::error_code EC = readSummary())
+        return EC;
+      break;
+    case SecNameTable:
+      if (std::error_code EC = readNameTable())
+        return EC;
+      break;
+    case SecLBRProfile:
+      while (Data < BufStart + Entry.Offset + Entry.Size) {
+        if (std::error_code EC = readFuncProfile())
+          return EC;
+      }
+      break;
+    default:
+      continue;
+    }
+    if (Data != BufStart + Entry.Offset + Entry.Size)
+      return sampleprof_error::malformed;
+  }
+
+  return sampleprof_error::success;
+}
+
 std::error_code SampleProfileReaderCompactBinary::read() {
   std::vector<uint64_t> OffsetsToUse;
   if (UseAllFuncs) {
@@ -501,6 +535,12 @@ std::error_code SampleProfileReaderRawBinary::verifySPMagic(uint64_t Magic) {
   return sampleprof_error::bad_magic;
 }
 
+std::error_code SampleProfileReaderExtBinary::verifySPMagic(uint64_t Magic) {
+  if (Magic == SPMagic(SPF_Ext_Binary))
+    return sampleprof_error::success;
+  return sampleprof_error::bad_magic;
+}
+
 std::error_code
 SampleProfileReaderCompactBinary::verifySPMagic(uint64_t Magic) {
   if (Magic == SPMagic(SPF_Compact_Binary))
@@ -508,7 +548,7 @@ SampleProfileReaderCompactBinary::verifySPMagic(uint64_t Magic) {
   return sampleprof_error::bad_magic;
 }
 
-std::error_code SampleProfileReaderRawBinary::readNameTable() {
+std::error_code SampleProfileReaderBinary::readNameTable() {
   auto Size = readNumber<uint32_t>();
   if (std::error_code EC = Size.getError())
     return EC;
@@ -537,10 +577,60 @@ std::error_code SampleProfileReaderCompactBinary::readNameTable() {
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderBinary::readHeader() {
-  Data = reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
-  End = Data + Buffer->getBufferSize();
+std::error_code SampleProfileReaderExtBinaryBase::readSecHdrTableEntry() {
+  SecHdrTableEntry Entry;
+  auto Type = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = Type.getError())
+    return EC;
+  Entry.Type = static_cast<SecType>(*Type);
 
+  auto Flag = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = Flag.getError())
+    return EC;
+  Entry.Flag = *Flag;
+
+  auto Offset = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = Offset.getError())
+    return EC;
+  Entry.Offset = *Offset;
+
+  auto Size = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = Size.getError())
+    return EC;
+  Entry.Size = *Size;
+
+  SecHdrTable.push_back(std::move(Entry));
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderExtBinaryBase::readSecHdrTable() {
+  auto EntryNum = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = EntryNum.getError())
+    return EC;
+
+  for (uint32_t i = 0; i < (*EntryNum); i++)
+    if (std::error_code EC = readSecHdrTableEntry())
+      return EC;
+
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderExtBinaryBase::readHeader() {
+  const uint8_t *BufStart =
+      reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
+  Data = BufStart;
+  End = BufStart + Buffer->getBufferSize();
+
+  if (std::error_code EC = readMagicIdent())
+    return EC;
+
+  if (std::error_code EC = readSecHdrTable())
+    return EC;
+
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderBinary::readMagicIdent() {
   // Read and check the magic identifier.
   auto Magic = readNumber<uint64_t>();
   if (std::error_code EC = Magic.getError())
@@ -555,6 +645,16 @@ std::error_code SampleProfileReaderBinary::readHeader() {
   else if (*Version != SPVersion())
     return sampleprof_error::unsupported_version;
 
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderBinary::readHeader() {
+  Data = reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
+  End = Data + Buffer->getBufferSize();
+
+  if (std::error_code EC = readMagicIdent())
+    return EC;
+
   if (std::error_code EC = readSummary())
     return EC;
 
@@ -674,6 +774,13 @@ bool SampleProfileReaderRawBinary::hasFormat(const MemoryBuffer &Buffer) {
   return Magic == SPMagic();
 }
 
+bool SampleProfileReaderExtBinary::hasFormat(const MemoryBuffer &Buffer) {
+  const uint8_t *Data =
+      reinterpret_cast<const uint8_t *>(Buffer.getBufferStart());
+  uint64_t Magic = decodeULEB128(Data);
+  return Magic == SPMagic(SPF_Ext_Binary);
+}
+
 bool SampleProfileReaderCompactBinary::hasFormat(const MemoryBuffer &Buffer) {
   const uint8_t *Data =
       reinterpret_cast<const uint8_t *>(Buffer.getBufferStart());
@@ -1023,6 +1130,8 @@ SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
   std::unique_ptr<SampleProfileReader> Reader;
   if (SampleProfileReaderRawBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderRawBinary(std::move(B), C));
+  else if (SampleProfileReaderExtBinary::hasFormat(*B))
+    Reader.reset(new SampleProfileReaderExtBinary(std::move(B), C));
   else if (SampleProfileReaderCompactBinary::hasFormat(*B))
     Reader.reset(new SampleProfileReaderCompactBinary(std::move(B), C));
   else if (SampleProfileReaderGCC::hasFormat(*B))
@@ -1033,8 +1142,9 @@ SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
     return sampleprof_error::unrecognized_format;
 
   FunctionSamples::Format = Reader->getFormat();
-  if (std::error_code EC = Reader->readHeader())
+  if (std::error_code EC = Reader->readHeader()) {
     return EC;
+  }
 
   return std::move(Reader);
 }

lib/ProfileData/SampleProfWriter.cpp

@@ -39,11 +39,8 @@
 using namespace llvm;
 using namespace sampleprof;
 
-std::error_code
-SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
-  if (std::error_code EC = writeHeader(ProfileMap))
-    return EC;
-
+std::error_code SampleProfileWriter::writeFuncProfiles(
+    const StringMap<FunctionSamples> &ProfileMap) {
   // Sort the ProfileMap by total samples.
   typedef std::pair<StringRef, const FunctionSamples *> NameFunctionSamples;
   std::vector<NameFunctionSamples> V;
@@ -58,12 +55,77 @@ SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
       });
 
   for (const auto &I : V) {
-    if (std::error_code EC = write(*I.second))
+    if (std::error_code EC = writeSample(*I.second))
       return EC;
   }
   return sampleprof_error::success;
 }
 
+std::error_code
+SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
+  if (std::error_code EC = writeHeader(ProfileMap))
+    return EC;
+
+  if (std::error_code EC = writeFuncProfiles(ProfileMap))
+    return EC;
+
+  return sampleprof_error::success;
+}
+
+/// Return the current position and prepare to use it as the start
+/// position of a section.
+uint64_t SampleProfileWriterExtBinaryBase::markSectionStart() {
+  return OutputStream->tell();
+}
+
+/// Add a new section into section header table. Return the position
+/// of SectionEnd.
+uint64_t
+SampleProfileWriterExtBinaryBase::addNewSection(SecType Sec,
+                                                uint64_t SectionStart) {
+  uint64_t SectionEnd = OutputStream->tell();
+  SecHdrTable.push_back(
+      {Sec, 0, SectionStart - FileStart, SectionEnd - SectionStart});
+  return SectionEnd;
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::write(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  if (std::error_code EC = writeHeader(ProfileMap))
+    return EC;
+
+  if (std::error_code EC = writeSections(ProfileMap))
+    return EC;
+
+  if (std::error_code EC = writeSecHdrTable())
+    return EC;
+
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileWriterExtBinary::writeSections(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  uint64_t SectionStart = markSectionStart();
+  computeSummary(ProfileMap);
+  if (auto EC = writeSummary())
+    return EC;
+  SectionStart = addNewSection(SecProfSummary, SectionStart);
+
+  // Generate the name table for all the functions referenced in the profile.
+  for (const auto &I : ProfileMap) {
+    addName(I.first());
+    addNames(I.second);
+  }
+  writeNameTable();
+  SectionStart = addNewSection(SecNameTable, SectionStart);
+
+  if (std::error_code EC = writeFuncProfiles(ProfileMap))
+    return EC;
+  addNewSection(SecLBRProfile, SectionStart);
+
+  return sampleprof_error::success;
+}
+
 std::error_code SampleProfileWriterCompactBinary::write(
     const StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
@@ -81,7 +143,7 @@ std::error_code SampleProfileWriterCompactBinary::write(
 ///
 /// The format used here is more structured and deliberate because
 /// it needs to be parsed by the SampleProfileReaderText class.
-std::error_code SampleProfileWriterText::write(const FunctionSamples &S) {
+std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   auto &OS = *OutputStream;
   OS << S.getName() << ":" << S.getTotalSamples();
   if (Indent == 0)
@@ -117,7 +179,7 @@ std::error_code SampleProfileWriterText::write(const FunctionSamples &S) {
         OS << Loc.LineOffset << ": ";
       else
         OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
-      if (std::error_code EC = write(CalleeSamples))
+      if (std::error_code EC = writeSample(CalleeSamples))
         return EC;
     }
   Indent -= 1;
@@ -163,7 +225,7 @@ void SampleProfileWriterBinary::stablizeNameTable(std::set<StringRef> &V) {
     NameTable[N] = i++;
 }
 
-std::error_code SampleProfileWriterRawBinary::writeNameTable() {
+std::error_code SampleProfileWriterBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(V);
@@ -214,25 +276,18 @@ std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterRawBinary::writeMagicIdent() {
+std::error_code
+SampleProfileWriterBinary::writeMagicIdent(SampleProfileFormat Format) {
   auto &OS = *OutputStream;
   // Write file magic identifier.
-  encodeULEB128(SPMagic(), OS);
-  encodeULEB128(SPVersion(), OS);
-  return sampleprof_error::success;
-}
-
-std::error_code SampleProfileWriterCompactBinary::writeMagicIdent() {
-  auto &OS = *OutputStream;
-  // Write file magic identifier.
-  encodeULEB128(SPMagic(SPF_Compact_Binary), OS);
+  encodeULEB128(SPMagic(Format), OS);
   encodeULEB128(SPVersion(), OS);
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterBinary::writeHeader(
     const StringMap<FunctionSamples> &ProfileMap) {
-  writeMagicIdent();
+  writeMagicIdent(Format);
 
   computeSummary(ProfileMap);
   if (auto EC = writeSummary())
@@ -248,6 +303,65 @@ std::error_code SampleProfileWriterBinary::writeHeader(
   return sampleprof_error::success;
 }
 
+void SampleProfileWriterExtBinaryBase::allocSecHdrTable() {
+  support::endian::Writer Writer(*OutputStream, support::little);
+
+  Writer.write(static_cast<uint64_t>(SectionLayout.size()));
+  SecHdrTableOffset = OutputStream->tell();
+  for (uint32_t i = 0; i < SectionLayout.size(); i++) {
+    Writer.write(static_cast<uint64_t>(-1));
+    Writer.write(static_cast<uint64_t>(-1));
+    Writer.write(static_cast<uint64_t>(-1));
+    Writer.write(static_cast<uint64_t>(-1));
+  }
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
+  auto &OFS = static_cast<raw_fd_ostream &>(*OutputStream);
+  uint64_t Saved = OutputStream->tell();
+
+  // Set OutputStream to the location saved in SecHdrTableOffset.
+  if (OFS.seek(SecHdrTableOffset) == (uint64_t)-1)
+    return sampleprof_error::ostream_seek_unsupported;
+  support::endian::Writer Writer(*OutputStream, support::little);
+
+  DenseMap<uint32_t, uint32_t> IndexMap;
+  for (uint32_t i = 0; i < SecHdrTable.size(); i++) {
+    IndexMap.insert({static_cast<uint32_t>(SecHdrTable[i].Type), i});
+  }
+
+  // Write the sections in the order specified in SectionLayout.
+  // That is the sections order Reader will see. Note that the
+  // sections order in which Reader expects to read may be different
+  // from the order in which Writer is able to write, so we need
+  // to adjust the order in SecHdrTable to be consistent with
+  // SectionLayout when we write SecHdrTable to the memory.
+  for (uint32_t i = 0; i < SectionLayout.size(); i++) {
+    uint32_t idx = IndexMap[static_cast<uint32_t>(SectionLayout[i])];
+    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Type));
+    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Flag));
+    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Offset));
+    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Size));
+  }
+
+  // Reset OutputStream.
+  if (OFS.seek(Saved) == (uint64_t)-1)
+    return sampleprof_error::ostream_seek_unsupported;
+
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  auto &OS = *OutputStream;
+  FileStart = OS.tell();
+  writeMagicIdent(Format);
+
+  initSectionLayout();
+  allocSecHdrTable();
+  return sampleprof_error::success;
+}
+
 std::error_code SampleProfileWriterCompactBinary::writeHeader(
     const StringMap<FunctionSamples> &ProfileMap) {
   support::endian::Writer Writer(*OutputStream, support::little);
@@ -324,13 +438,14 @@ std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
 /// Write samples of a top-level function to a binary file.
 ///
 /// \returns true if the samples were written successfully, false otherwise.
-std::error_code SampleProfileWriterBinary::write(const FunctionSamples &S) {
+std::error_code
+SampleProfileWriterBinary::writeSample(const FunctionSamples &S) {
   encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(S);
 }
 
 std::error_code
-SampleProfileWriterCompactBinary::write(const FunctionSamples &S) {
+SampleProfileWriterCompactBinary::writeSample(const FunctionSamples &S) {
   uint64_t Offset = OutputStream->tell();
   StringRef Name = S.getName();
   FuncOffsetTable[Name] = Offset;
@@ -349,7 +464,8 @@ ErrorOr<std::unique_ptr<SampleProfileWriter>>
 SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) {
   std::error_code EC;
   std::unique_ptr<raw_ostream> OS;
-  if (Format == SPF_Binary || Format == SPF_Compact_Binary)
+  if (Format == SPF_Binary || Format == SPF_Ext_Binary ||
+      Format == SPF_Compact_Binary)
     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_None));
   else
     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_Text));
@@ -374,6 +490,8 @@ SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
 
   if (Format == SPF_Binary)
     Writer.reset(new SampleProfileWriterRawBinary(OS));
+  else if (Format == SPF_Ext_Binary)
+    Writer.reset(new SampleProfileWriterExtBinary(OS));
   else if (Format == SPF_Compact_Binary)
     Writer.reset(new SampleProfileWriterCompactBinary(OS));
   else if (Format == SPF_Text)
@@ -386,6 +504,7 @@ SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
   if (EC)
     return EC;
 
+  Writer->Format = Format;
   return std::move(Writer);
 }
 

test/Transforms/SampleProfile/profile-format.ll

@@ -2,6 +2,8 @@
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline.prof -S | FileCheck %s
 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/inline.compactbinary.afdo -S | FileCheck %s
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline.compactbinary.afdo -S | FileCheck %s
+; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/inline.extbinary.afdo -S | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline.extbinary.afdo -S | FileCheck %s
 
 ; Original C++ test case
 ;
@@ -21,8 +23,8 @@
 ;
 @.str = private unnamed_addr constant [11 x i8] c"sum is %d\0A\00", align 1
 
-; Check sample-profile phase using compactbinary format profile will annotate
-; the IR with exactly the same result as using text format.
+; Check sample-profile phase using compactbinary or extbinary format profile
+; will annotate the IR with exactly the same result as using text format.
 ; CHECK: br i1 %cmp, label %while.body, label %while.end{{.*}} !prof ![[IDX1:[0-9]*]]
 ; CHECK: br i1 %cmp1, label %if.then, label %if.else{{.*}} !prof ![[IDX2:[0-9]*]]
 ; CHECK: call i32 (i8*, ...) @printf{{.*}} !prof ![[IDX3:[0-9]*]]

test/tools/llvm-profdata/roundtrip.test

@@ -6,4 +6,13 @@ RUN: llvm-profdata show -o %t.1.proftext -all-functions -text %t.1.profdata
 RUN: diff %t.1.proftext %S/Inputs/IR_profile.proftext
 RUN: llvm-profdata merge --sample --binary -output=%t.2.profdata %S/Inputs/sample-profile.proftext
 RUN: llvm-profdata merge --sample --text -output=%t.2.proftext %t.2.profdata
-RUN: diff %t.2.proftext %S/Inputs/sample-profile.proftext
+RUN: diff %t.2.proftext %S/Inputs/sample-profile.proftext
+# Round trip from text --> extbinary --> text
+RUN: llvm-profdata merge --sample --extbinary -output=%t.3.profdata %S/Inputs/sample-profile.proftext
+RUN: llvm-profdata merge --sample --text -output=%t.3.proftext %t.3.profdata
+RUN: diff %t.3.proftext %S/Inputs/sample-profile.proftext
+# Round trip from text --> binary --> extbinary --> text
+RUN: llvm-profdata merge --sample --binary -output=%t.4.profdata %S/Inputs/sample-profile.proftext
+RUN: llvm-profdata merge --sample --extbinary -output=%t.5.profdata %t.4.profdata
+RUN: llvm-profdata merge --sample --text -output=%t.4.proftext %t.5.profdata
+RUN: diff %t.4.proftext %S/Inputs/sample-profile.proftext

tools/llvm-profdata/llvm-profdata.cpp

@@ -37,6 +37,7 @@ enum ProfileFormat {
   PF_None = 0,
   PF_Text,
   PF_Compact_Binary,
+  PF_Ext_Binary,
   PF_GCC,
   PF_Binary
 };
@@ -314,7 +315,7 @@ static void mergeInstrProfile(const WeightedFileVector &Inputs,
     exitWithError("Cannot write indexed profdata format to stdout.");
 
   if (OutputFormat != PF_Binary && OutputFormat != PF_Compact_Binary &&
-      OutputFormat != PF_Text)
+      OutputFormat != PF_Ext_Binary && OutputFormat != PF_Text)
     exitWithError("Unknown format is specified.");
 
   std::mutex ErrorLock;
@@ -425,8 +426,12 @@ remapSamples(const sampleprof::FunctionSamples &Samples,
 }
 
 static sampleprof::SampleProfileFormat FormatMap[] = {
-    sampleprof::SPF_None, sampleprof::SPF_Text, sampleprof::SPF_Compact_Binary,
-    sampleprof::SPF_GCC, sampleprof::SPF_Binary};
+    sampleprof::SPF_None,
+    sampleprof::SPF_Text,
+    sampleprof::SPF_Compact_Binary,
+    sampleprof::SPF_Ext_Binary,
+    sampleprof::SPF_GCC,
+    sampleprof::SPF_Binary};
 
 static void mergeSampleProfile(const WeightedFileVector &Inputs,
                                SymbolRemapper *Remapper,
@@ -583,12 +588,14 @@ static int merge_main(int argc, const char *argv[]) {
                  clEnumVal(sample, "Sample profile")));
   cl::opt<ProfileFormat> OutputFormat(
       cl::desc("Format of output profile"), cl::init(PF_Binary),
-      cl::values(clEnumValN(PF_Binary, "binary", "Binary encoding (default)"),
-                 clEnumValN(PF_Compact_Binary, "compbinary",
-                            "Compact binary encoding"),
-                 clEnumValN(PF_Text, "text", "Text encoding"),
-                 clEnumValN(PF_GCC, "gcc",
-                            "GCC encoding (only meaningful for -sample)")));
+      cl::values(
+          clEnumValN(PF_Binary, "binary", "Binary encoding (default)"),
+          clEnumValN(PF_Compact_Binary, "compbinary",
+                     "Compact binary encoding"),
+          clEnumValN(PF_Ext_Binary, "extbinary", "Extensible binary encoding"),
+          clEnumValN(PF_Text, "text", "Text encoding"),
+          clEnumValN(PF_GCC, "gcc",
+                     "GCC encoding (only meaningful for -sample)")));
   cl::opt<bool> OutputSparse("sparse", cl::init(false),
       cl::desc("Generate a sparse profile (only meaningful for -instr)"));
   cl::opt<unsigned> NumThreads(

unittests/ProfileData/SampleProfTest.cpp

@@ -285,6 +285,10 @@ TEST_F(SampleProfTest, roundtrip_compact_binary_profile) {
   testRoundTrip(SampleProfileFormat::SPF_Compact_Binary, false);
 }
 
+TEST_F(SampleProfTest, roundtrip_ext_binary_profile) {
+  testRoundTrip(SampleProfileFormat::SPF_Ext_Binary, false);
+}
+
 TEST_F(SampleProfTest, remap_text_profile) {
   testRoundTrip(SampleProfileFormat::SPF_Text, true);
 }
@@ -293,6 +297,10 @@ TEST_F(SampleProfTest, remap_raw_binary_profile) {
   testRoundTrip(SampleProfileFormat::SPF_Binary, true);
 }
 
+TEST_F(SampleProfTest, remap_ext_binary_profile) {
+  testRoundTrip(SampleProfileFormat::SPF_Ext_Binary, true);
+}
+
 TEST_F(SampleProfTest, sample_overflow_saturation) {
   const uint64_t Max = std::numeric_limits<uint64_t>::max();
   sampleprof_error Result;