[SampleFDO] Add compression support for any section in ExtBinary profile format

[wmi]

Summary

Previously ExtBinary profile format only supports compression using zlib for profile symbol list. In this patch, we extend the compression support to any section. User can select some or all of the sections to compress. In an experiment, for a 45M profile in ExtBinary format, compressing name table reduced its size to 24M, and compressing all the sections reduced its size to 11M.

[davidxl]

llvm/include/llvm/ProfileData/SampleProf.h

148

Flag->Flags?

155

may be name it 'addSecFlags' ? 'setSecFlags' implies Entry.Flag = Flags;

159

clearSecFlags or removeSecFlags?

llvm/include/llvm/ProfileData/SampleProfWriter.h

174

hold->held

178

held

llvm/tools/llvm-profdata/llvm-profdata.cpp

656

Should this imply compressProfSymlist when true? Or under what condition do we want to compressAllSections but not profsymList? If such case does not exist, we should remove compress symlist option.

[wmi]

llvm/include/llvm/ProfileData/SampleProf.h

148

Fixed.

155

Fixed.

159

Used removeSecFlags.

llvm/include/llvm/ProfileData/SampleProfWriter.h

174

Fixed.

178

Fixed.

llvm/tools/llvm-profdata/llvm-profdata.cpp

656

I removed compress symlist option.

[wmi]

Address David's comment.

[davidxl]

llvmprof command documentation also needs to be updated.

llvm/tools/llvm-profdata/llvm-profdata.cpp

456

'-compress-all-section is ignored. Specify -extbinary to enable it'

[wmi]

llvm/tools/llvm-profdata/llvm-profdata.cpp

456

Fixed.

[wmi]

Address David's comment

[davidxl]

lgtm

llvm/include/llvm/ProfileData/SampleProf.h

static inline std::string getSecName(SecType Type) {
   }
   llvm_unreachable("A SecType has no name for output");
 }
 
 // Entry type of section header table used by SampleProfileExtBinaryBaseReader
 // and SampleProfileExtBinaryBaseWriter.
 struct SecHdrTableEntry {
   SecType Type;
   uint64_t Flag;
   uint64_t Offset;
   uint64_t Size;
 };
 
+enum SecFlags { SecFlagInValid = 0, SecFlagCompress = (1 << 0) };
+
+static inline void setSecFlags(SecHdrTableEntry &Entry, uint64_t Flags) {
+  Entry.Flag |= Flags;
+}
+
+static inline void clearSecFlag(SecHdrTableEntry &Entry, uint64_t Flags) {
+  Entry.Flag &= ~Flags;
+}
+
+static inline bool hasSecFlag(SecHdrTableEntry &Entry, SecFlags Flag) {
+  return Entry.Flag & Flag;
+}
+
 /// Represents the relative location of an instruction.
 ///
 /// Instruction locations are specified by the line offset from the
 /// beginning of the function (marked by the line where the function
 /// header is) and the discriminator value within that line.
 ///
 /// The discriminator value is useful to distinguish instructions
 /// that are on the same line but belong to different basic blocks
public:
   void merge(const ProfileSymbolList &List) {
     for (auto Sym : List.Syms)
       add(Sym, true);
   }
 
   unsigned size() { return Syms.size(); }
 
   void setToCompress(bool TC) { ToCompress = TC; }
+  bool toCompress() { return ToCompress; }
 
-  std::error_code read(uint64_t CompressSize, uint64_t UncompressSize,
-                       const uint8_t *Data);
+  std::error_code read(const uint8_t *Data, uint64_t ListSize);
   std::error_code write(raw_ostream &OS);
   void dump(raw_ostream &OS = dbgs()) const;
 
 private:
   // Determine whether or not to compress the symbol list when
   // writing it into profile. The variable is unused when the symbol
   // list is read from an existing profile.
   bool ToCompress = false;
   DenseSet<StringRef> Syms;
   BumpPtrAllocator Allocator;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_SAMPLEPROF_H

llvm/include/llvm/ProfileData/SampleProfReader.h

 /// 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 {
+private:
+  std::error_code decompressSection(const uint8_t *SecStart,
+                                    const uint64_t SecSize,
+                                    const uint8_t *&DecompressBuf,
+                                    uint64_t &DecompressBufSize);
+
+  BumpPtrAllocator Allocator;
+
 protected:
   std::vector<SecHdrTableEntry> SecHdrTable;
   std::unique_ptr<ProfileSymbolList> ProfSymList;
   std::error_code readSecHdrTableEntry();
   std::error_code readSecHdrTable();
   virtual std::error_code readHeader() override;
   virtual std::error_code verifySPMagic(uint64_t Magic) override = 0;
   virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
public:
   virtual bool dumpSectionInfo(raw_ostream &OS = dbgs()) override;
 };
 
 class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
 private:
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
   virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
                                          SecType Type) override;
-  std::error_code readProfileSymbolList();
+  std::error_code readProfileSymbolList(uint64_t Size);
 
 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);

llvm/include/llvm/ProfileData/SampleProfWriter.h

 };
 
 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;
 
+  void setToCompressAllSections();
+  void setToCompressSection(SecType Type);
+
 protected:
-  uint64_t markSectionStart();
-  uint64_t addNewSection(SecType Sec, uint64_t SectionStart);
+  uint64_t markSectionStart(SecType Type);
+  std::error_code 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;
+  SmallVector<SecHdrTableEntry, 8> SectionLayout;
 
 private:
   void allocSecHdrTable();
   std::error_code writeSecHdrTable();
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
-
+  void setSectionFlag(SecType Type, SecFlags Flag);
+  SecHdrTableEntry &getEntryInLayout(SecType Type);
+  std::error_code compressAndOutput();
+
+  // We will swap the raw_ostream hold by LocalBufStream and that
+  // hold by OutputStream if we try to add a section which needs
+  // compression. After the swap, all the data written to output
+  // will be temporarily buffered into the underlying raw_string_ostream
+  // originally hold by LocalBufStream. After the data writing for the
+  // section is completed, compress the data in the local buffer,
+  // swap the raw_ostream back and write the compressed data to the
+  // real output.
+  std::unique_ptr<raw_ostream> LocalBufStream;
   // 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;
+  // Initial Section Flags setting.
   std::vector<SecHdrTableEntry> SecHdrTable;
 };
 
 class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
-  using SampleProfileWriterExtBinaryBase::SampleProfileWriterExtBinaryBase;
-
 public:
+  SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriterExtBinaryBase(OS) {
+    initSectionLayout();
+  }
+
   virtual void setProfileSymbolList(ProfileSymbolList *PSL) override {
     ProfSymList = PSL;
   };
 
 private:
   virtual void initSectionLayout() override {
-    SectionLayout = {SecProfSummary, SecNameTable, SecLBRProfile,
-                     SecProfileSymbolList};
+    SectionLayout = {{SecProfSummary},
+                     {SecNameTable},
+                     {SecLBRProfile},
+                     {SecProfileSymbolList}};
   };
   virtual std::error_code
   writeSections(const StringMap<FunctionSamples> &ProfileMap) override;
   ProfileSymbolList *ProfSymList = nullptr;
 };
 
 // CompactBinary is a compact format of binary profile which both reduces
 // the profile size and the load time needed when compiling. It has two

llvm/lib/ProfileData/SampleProf.cpp

 // sample profile data.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/Support/Compiler.h"
-#include "llvm/Support/Compression.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/raw_ostream.h"
 #include <string>
 #include <system_error>
FunctionSamples::findFunctionSamples(const DILocation *DIL) const {
   }
   return FS;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void FunctionSamples::dump() const { print(dbgs(), 0); }
 #endif
 
-std::error_code ProfileSymbolList::read(uint64_t CompressSize,
-                                        uint64_t UncompressSize,
-                                        const uint8_t *Data) {
+std::error_code ProfileSymbolList::read(const uint8_t *Data,
+                                        uint64_t ListSize) {
   const char *ListStart = reinterpret_cast<const char *>(Data);
-  // CompressSize being non-zero means the profile is compressed and
-  // needs to be uncompressed first.
-  if (CompressSize) {
-    if (!llvm::zlib::isAvailable())
-      return sampleprof_error::zlib_unavailable;
-
-    StringRef CompressedStrings(reinterpret_cast<const char *>(Data),
-                                CompressSize);
-    char *Buffer = Allocator.Allocate<char>(UncompressSize);
-    size_t UCSize = UncompressSize;
-    llvm::Error E = zlib::uncompress(CompressedStrings, Buffer, UCSize);
-    if (E)
-      return sampleprof_error::uncompress_failed;
-    ListStart = Buffer;
-  }
-
   uint64_t Size = 0;
-  while (Size < UncompressSize) {
+  while (Size < ListSize) {
     StringRef Str(ListStart + Size);
     add(Str);
     Size += Str.size() + 1;
   }
+  if (Size != ListSize)
+    return sampleprof_error::malformed;
   return sampleprof_error::success;
 }
 
 std::error_code ProfileSymbolList::write(raw_ostream &OS) {
-  // Sort the symbols before doing compression. It will make the
-  // compression much more effective.
+  // Sort the symbols before output. If doing compression.
+  // It will make the compression much more effective.
   std::vector<StringRef> SortedList;
   SortedList.insert(SortedList.begin(), Syms.begin(), Syms.end());
   llvm::sort(SortedList);
 
-  std::string UncompressedStrings;
+  std::string OutputString;
   for (auto &Sym : SortedList) {
-    UncompressedStrings.append(Sym.str());
-    UncompressedStrings.append(1, '\0');
+    OutputString.append(Sym.str());
+    OutputString.append(1, '\0');
   }
 
-  if (ToCompress) {
-    if (!llvm::zlib::isAvailable())
-      return sampleprof_error::zlib_unavailable;
-    SmallString<128> CompressedStrings;
-    llvm::Error E = zlib::compress(UncompressedStrings, CompressedStrings,
-                                   zlib::BestSizeCompression);
-    if (E)
-      return sampleprof_error::compress_failed;
-    encodeULEB128(UncompressedStrings.size(), OS);
-    encodeULEB128(CompressedStrings.size(), OS);
-    OS << CompressedStrings.str();
-  } else {
-    encodeULEB128(UncompressedStrings.size(), OS);
-    // If profile symbol list is not compressed, we will still save
-    // a compressed size value, but the value of the size is 0.
-    encodeULEB128(0, OS);
-    OS << UncompressedStrings;
-  }
+  OS << OutputString;
   return sampleprof_error::success;
 }
 
 void ProfileSymbolList::dump(raw_ostream &OS) const {
   OS << "======== Dump profile symbol list ========\n";
   std::vector<StringRef> SortedList;
   SortedList.insert(SortedList.begin(), Syms.begin(), Syms.end());
   llvm::sort(SortedList);
 
   for (auto &Sym : SortedList)
     OS << Sym << "\n";
 }

llvm/lib/ProfileData/SampleProfReader.cpp

 
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstddef>
std::error_code SampleProfileReaderBinary::read() {
 
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileReaderExtBinary::readOneSection(const uint8_t *Start,
                                              uint64_t Size, SecType Type) {
   Data = Start;
+  End = Start + Size;
   switch (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 < Start + Size) {
       if (std::error_code EC = readFuncProfile())
         return EC;
     }
     break;
   case SecProfileSymbolList:
-    if (std::error_code EC = readProfileSymbolList())
+    if (std::error_code EC = readProfileSymbolList(Size))
       return EC;
     break;
   default:
     break;
   }
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileReaderExtBinary::readProfileSymbolList() {
-  auto UncompressSize = readNumber<uint64_t>();
-  if (std::error_code EC = UncompressSize.getError())
+std::error_code
+SampleProfileReaderExtBinary::readProfileSymbolList(uint64_t Size) {
+  if (!ProfSymList)
+    ProfSymList = std::make_unique<ProfileSymbolList>();
+
+  if (std::error_code EC = ProfSymList->read(Data, Size))
     return EC;
 
+  Data = Data + Size;
+  return sampleprof_error::success;
+}
+
+std::error_code SampleProfileReaderExtBinaryBase::decompressSection(
+    const uint8_t *SecStart, const uint64_t SecSize,
+    const uint8_t *&DecompressBuf, uint64_t &DecompressBufSize) {
+  Data = SecStart;
+  End = SecStart + SecSize;
+  auto DecompressSize = readNumber<uint64_t>();
+  if (std::error_code EC = DecompressSize.getError())
+    return EC;
+  DecompressBufSize = *DecompressSize;
+
   auto CompressSize = readNumber<uint64_t>();
   if (std::error_code EC = CompressSize.getError())
     return EC;
 
-  if (!ProfSymList)
-    ProfSymList = std::make_unique<ProfileSymbolList>();
-
-  if (std::error_code EC =
-          ProfSymList->read(*CompressSize, *UncompressSize, Data))
-    return EC;
+  if (!llvm::zlib::isAvailable())
+    return sampleprof_error::zlib_unavailable;
 
-  // CompressSize is zero only when ProfileSymbolList is not compressed.
-  if (*CompressSize == 0)
-    Data = Data + *UncompressSize;
-  else
-    Data = Data + *CompressSize;
+  StringRef CompressedStrings(reinterpret_cast<const char *>(Data),
+                              *CompressSize);
+  char *Buffer = Allocator.Allocate<char>(DecompressBufSize);
+  llvm::Error E =
+      zlib::uncompress(CompressedStrings, Buffer, DecompressBufSize);
+  if (E)
+    return sampleprof_error::uncompress_failed;
+  DecompressBuf = reinterpret_cast<const uint8_t *>(Buffer);
   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;
+
     const uint8_t *SecStart = BufStart + Entry.Offset;
-    if (std::error_code EC = readOneSection(SecStart, Entry.Size, Entry.Type))
+    uint64_t SecSize = Entry.Size;
+
+    // If the section is compressed, decompress it into a buffer
+    // DecompressBuf before reading the actual data. The pointee of
+    // 'Data' will be changed to buffer hold by DecompressBuf
+    // temporarily when reading the actual data.
+    bool isCompressed = hasSecFlag(Entry, SecFlagCompress);
+    if (isCompressed) {
+      const uint8_t *DecompressBuf;
+      uint64_t DecompressBufSize;
+      if (std::error_code EC = decompressSection(
+              SecStart, SecSize, DecompressBuf, DecompressBufSize))
         return EC;
-    if (Data != SecStart + Entry.Size)
+      SecStart = DecompressBuf;
+      SecSize = DecompressBufSize;
+    }
+
+    if (std::error_code EC = readOneSection(SecStart, SecSize, Entry.Type))
+      return EC;
+    if (Data != SecStart + SecSize)
       return sampleprof_error::malformed;
+
+    // Change the pointee of 'Data' from DecompressBuf to original Buffer.
+    if (isCompressed) {
+      Data = BufStart + Entry.Offset;
+      End = BufStart + Buffer->getBufferSize();
+    }
   }
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderCompactBinary::read() {
   std::vector<uint64_t> OffsetsToUse;
   if (UseAllFuncs) {

llvm/lib/ProfileData/SampleProfWriter.cpp

 // supported formats.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
 
   return sampleprof_error::success;
 }
 
+SecHdrTableEntry &
+SampleProfileWriterExtBinaryBase::getEntryInLayout(SecType Type) {
+  auto SecIt = std::find_if(
+      SectionLayout.begin(), SectionLayout.end(),
+      [=](const auto &Entry) -> bool { return Entry.Type == Type; });
+  return *SecIt;
+}
+
 /// Return the current position and prepare to use it as the start
 /// position of a section.
-uint64_t SampleProfileWriterExtBinaryBase::markSectionStart() {
-  return OutputStream->tell();
+uint64_t SampleProfileWriterExtBinaryBase::markSectionStart(SecType Type) {
+  uint64_t SectionStart = OutputStream->tell();
+  auto &Entry = getEntryInLayout(Type);
+  // Use LocalBuf as a temporary output for writting data.
+  if (hasSecFlag(Entry, SecFlagCompress))
+    LocalBufStream.swap(OutputStream);
+  return SectionStart;
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::compressAndOutput() {
+  if (!llvm::zlib::isAvailable())
+    return sampleprof_error::zlib_unavailable;
+  std::string &UncompressedStrings =
+      static_cast<raw_string_ostream *>(LocalBufStream.get())->str();
+  if (UncompressedStrings.size() == 0)
+    return sampleprof_error::success;
+  auto &OS = *OutputStream;
+  SmallString<128> CompressedStrings;
+  llvm::Error E = zlib::compress(UncompressedStrings, CompressedStrings,
+                                 zlib::BestSizeCompression);
+  if (E)
+    return sampleprof_error::compress_failed;
+  encodeULEB128(UncompressedStrings.size(), OS);
+  encodeULEB128(CompressedStrings.size(), OS);
+  OS << CompressedStrings.str();
+  UncompressedStrings.clear();
+  return sampleprof_error::success;
 }
 
-/// Add a new section into section header table. Return the position
-/// of SectionEnd.
-uint64_t
-SampleProfileWriterExtBinaryBase::addNewSection(SecType Sec,
+/// Add a new section into section header table.
+std::error_code
+SampleProfileWriterExtBinaryBase::addNewSection(SecType Type,
                                                 uint64_t SectionStart) {
-  uint64_t SectionEnd = OutputStream->tell();
-  SecHdrTable.push_back(
-      {Sec, 0, SectionStart - FileStart, SectionEnd - SectionStart});
-  return SectionEnd;
+  auto Entry = getEntryInLayout(Type);
+  if (hasSecFlag(Entry, SecFlagCompress)) {
+    LocalBufStream.swap(OutputStream);
+    if (std::error_code EC = compressAndOutput())
+      return EC;
+  }
+  SecHdrTable.push_back({Type, Entry.Flag, SectionStart - FileStart,
+                         OutputStream->tell() - SectionStart});
+  return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::write(
     const StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
+  std::string LocalBuf;
+  LocalBufStream = std::make_unique<raw_string_ostream>(LocalBuf);
   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();
+  uint64_t SectionStart = markSectionStart(SecProfSummary);
   computeSummary(ProfileMap);
   if (auto EC = writeSummary())
     return EC;
-  SectionStart = addNewSection(SecProfSummary, SectionStart);
+  if (std::error_code EC = addNewSection(SecProfSummary, SectionStart))
+    return EC;
 
   // Generate the name table for all the functions referenced in the profile.
+  SectionStart = markSectionStart(SecNameTable);
   for (const auto &I : ProfileMap) {
     addName(I.first());
     addNames(I.second);
   }
   writeNameTable();
-  SectionStart = addNewSection(SecNameTable, SectionStart);
+  if (std::error_code EC = addNewSection(SecNameTable, SectionStart))
+    return EC;
 
+  SectionStart = markSectionStart(SecLBRProfile);
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
-  SectionStart = addNewSection(SecLBRProfile, SectionStart);
+  if (std::error_code EC = addNewSection(SecLBRProfile, SectionStart))
+    return EC;
+
+  if (ProfSymList && ProfSymList->toCompress())
+    setToCompressSection(SecProfileSymbolList);
 
+  SectionStart = markSectionStart(SecProfileSymbolList);
   if (ProfSymList && ProfSymList->size() > 0)
     if (std::error_code EC = ProfSymList->write(*OutputStream))
       return EC;
-  addNewSection(SecProfileSymbolList, SectionStart);
+  if (std::error_code EC = addNewSection(SecProfileSymbolList, SectionStart))
+    return EC;
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::write(
     const StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
     return EC;
for (const auto &I : ProfileMap) {
     addName(I.first());
     addNames(I.second);
   }
 
   writeNameTable();
   return sampleprof_error::success;
 }
 
+void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
+  for (auto &Entry : SectionLayout)
+    setSecFlags(Entry, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::setToCompressSection(SecType Type) {
+  setSectionFlag(Type, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::setSectionFlag(SecType Type,
+                                                      SecFlags Flag) {
+  for (auto &Entry : SectionLayout) {
+    if (Entry.Type == Type)
+      setSecFlags(Entry, Flag);
+  }
+}
+
 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));
std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
 
   // 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])];
+    uint32_t idx = IndexMap[static_cast<uint32_t>(SectionLayout[i].Type)];
     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);
   if (auto EC = SampleProfileWriterBinary::writeHeader(ProfileMap))

llvm/test/Transforms/SampleProfile/profile-format-compress.ll

+; REQUIRES: zlib
+; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/inline.prof -S | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline.prof -S | FileCheck %s
+; RUN: llvm-profdata merge -sample -extbinary -compress-all-sections %S/Inputs/inline.prof -o %t.compress.extbinary.afdo
+; RUN: opt < %s -sample-profile -sample-profile-file=%t.compress.extbinary.afdo -S | FileCheck %s
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t.compress.extbinary.afdo -S | FileCheck %s
+
+; Original C++ test case
+;
+; #include <stdio.h>
+;
+; int sum(int x, int y) {
+;   return x + y;
+; }
+;
+; int main() {
+;   int s, i = 0;
+;   while (i++ < 20000 * 20000)
+;     if (i != 100) s = sum(i, s); else s = 30;
+;   printf("sum is %d\n", s);
+;   return 0;
+; }
+;
+@.str = private unnamed_addr constant [11 x i8] c"sum is %d\0A\00", align 1
+
+; Check sample-profile phase using compressed 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]*]]
+; CHECK: = !{!"TotalCount", i64 26781}
+; CHECK: = !{!"MaxCount", i64 5553}
+; CHECK: ![[IDX1]] = !{!"branch_weights", i32 5392, i32 163}
+; CHECK: ![[IDX2]] = !{!"branch_weights", i32 5280, i32 113}
+; CHECK: ![[IDX3]] = !{!"branch_weights", i32 1}
+
+; Function Attrs: nounwind uwtable
+define i32 @_Z3sumii(i32 %x, i32 %y) !dbg !4 {
+entry:
+  %x.addr = alloca i32, align 4
+  %y.addr = alloca i32, align 4
+  store i32 %x, i32* %x.addr, align 4
+  store i32 %y, i32* %y.addr, align 4
+  %0 = load i32, i32* %x.addr, align 4, !dbg !11
+  %1 = load i32, i32* %y.addr, align 4, !dbg !11
+  %add = add nsw i32 %0, %1, !dbg !11
+  ret i32 %add, !dbg !11
+}
+
+; Function Attrs: uwtable
+define i32 @main() !dbg !7 {
+entry:
+  %retval = alloca i32, align 4
+  %s = alloca i32, align 4
+  %i = alloca i32, align 4
+  store i32 0, i32* %retval
+  store i32 0, i32* %i, align 4, !dbg !12
+  br label %while.cond, !dbg !13
+
+while.cond:                                       ; preds = %if.end, %entry
+  %0 = load i32, i32* %i, align 4, !dbg !14
+  %inc = add nsw i32 %0, 1, !dbg !14
+  store i32 %inc, i32* %i, align 4, !dbg !14
+  %cmp = icmp slt i32 %0, 400000000, !dbg !14
+  br i1 %cmp, label %while.body, label %while.end, !dbg !14
+
+while.body:                                       ; preds = %while.cond
+  %1 = load i32, i32* %i, align 4, !dbg !16
+  %cmp1 = icmp ne i32 %1, 100, !dbg !16
+  br i1 %cmp1, label %if.then, label %if.else, !dbg !16
+
+
+if.then:                                          ; preds = %while.body
+  %2 = load i32, i32* %i, align 4, !dbg !18
+  %3 = load i32, i32* %s, align 4, !dbg !18
+  %call = call i32 @_Z3sumii(i32 %2, i32 %3), !dbg !18
+  store i32 %call, i32* %s, align 4, !dbg !18
+  br label %if.end, !dbg !18
+
+if.else:                                          ; preds = %while.body
+  store i32 30, i32* %s, align 4, !dbg !20
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  br label %while.cond, !dbg !22
+
+while.end:                                        ; preds = %while.cond
+  %4 = load i32, i32* %s, align 4, !dbg !24
+  %call2 = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([11 x i8], [11 x i8]* @.str, i32 0, i32 0), i32 %4), !dbg !24
+  ret i32 0, !dbg !25
+}
+
+declare i32 @printf(i8*, ...) #2
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!8, !9}
+!llvm.ident = !{!10}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, producer: "clang version 3.5 ", isOptimized: false, emissionKind: NoDebug, file: !1, enums: !2, retainedTypes: !2, globals: !2, imports: !2)
+!1 = !DIFile(filename: "calls.cc", directory: ".")
+!2 = !{}
+!4 = distinct !DISubprogram(name: "sum", line: 3, isLocal: false, isDefinition: true, virtualIndex: 6, flags: DIFlagPrototyped, isOptimized: false, unit: !0, scopeLine: 3, file: !1, scope: !5, type: !6, retainedNodes: !2)
+!5 = !DIFile(filename: "calls.cc", directory: ".")
+!6 = !DISubroutineType(types: !2)
+!7 = distinct !DISubprogram(name: "main", line: 7, isLocal: false, isDefinition: true, virtualIndex: 6, flags: DIFlagPrototyped, isOptimized: false, unit: !0, scopeLine: 7, file: !1, scope: !5, type: !6, retainedNodes: !2)
+!8 = !{i32 2, !"Dwarf Version", i32 4}
+!9 = !{i32 1, !"Debug Info Version", i32 3}
+!10 = !{!"clang version 3.5 "}
+!11 = !DILocation(line: 4, scope: !4)
+!12 = !DILocation(line: 8, scope: !7)
+!13 = !DILocation(line: 9, scope: !7)
+!14 = !DILocation(line: 9, scope: !15)
+!15 = !DILexicalBlockFile(discriminator: 2, file: !1, scope: !7)
+!16 = !DILocation(line: 10, scope: !17)
+!17 = distinct !DILexicalBlock(line: 10, column: 0, file: !1, scope: !7)
+!18 = !DILocation(line: 10, scope: !19)
+!19 = !DILexicalBlockFile(discriminator: 2, file: !1, scope: !17)
+!20 = !DILocation(line: 10, scope: !21)
+!21 = !DILexicalBlockFile(discriminator: 4, file: !1, scope: !17)
+!22 = !DILocation(line: 10, scope: !23)
+!23 = !DILexicalBlockFile(discriminator: 6, file: !1, scope: !17)
+!24 = !DILocation(line: 11, scope: !7)
+!25 = !DILocation(line: 12, scope: !7)

llvm/test/tools/llvm-profdata/profile-symbol-list-compress.test

+REQUIRES: zlib
+; RUN: llvm-profdata merge -sample -extbinary -compress-prof-sym-list -prof-sym-list=%S/Inputs/profile-symbol-list-1.text %S/Inputs/sample-profile.proftext -o %t.1.output
+; RUN: llvm-profdata merge -sample -extbinary -compress-prof-sym-list -prof-sym-list=%S/Inputs/profile-symbol-list-2.text %S/Inputs/sample-profile.proftext -o %t.2.output
+; RUN: llvm-profdata merge -sample -extbinary -compress-prof-sym-list %t.1.output %t.2.output -o %t.3.output
+; RUN: llvm-profdata show -sample -show-prof-sym-list %t.3.output > %t.4.output
+; RUN: diff %S/Inputs/profile-symbol-list.expected %t.4.output

llvm/test/tools/llvm-profdata/roundtrip-compress.test

+REQUIRES: zlib
+# Round trip from text --> compressed extbinary --> text
+RUN: llvm-profdata merge --sample --extbinary -compress-all-sections -output=%t.1.profdata %S/Inputs/sample-profile.proftext
+RUN: llvm-profdata merge --sample --text -output=%t.1.proftext %t.1.profdata
+RUN: diff %t.1.proftext %S/Inputs/sample-profile.proftext
+# Round trip from text --> binary --> compressed extbinary --> text
+RUN: llvm-profdata merge --sample --binary -output=%t.2.profdata %S/Inputs/sample-profile.proftext
+RUN: llvm-profdata merge --sample --extbinary -compress-all-sections -output=%t.3.profdata %t.2.profdata
+RUN: llvm-profdata merge --sample --text -output=%t.2.proftext %t.3.profdata
+RUN: diff %t.2.proftext %S/Inputs/sample-profile.proftext

llvm/tools/llvm-profdata/llvm-profdata.cpp

static void populateProfileSymbolList(MemoryBuffer *Buffer,
   SmallVector<StringRef, 32> SymbolVec;
   StringRef Data = Buffer->getBuffer();
   Data.split(SymbolVec, '\n', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
 
   for (StringRef symbol : SymbolVec)
     PSL.add(symbol);
 }
 
-static void mergeSampleProfile(const WeightedFileVector &Inputs,
-                               SymbolRemapper *Remapper,
-                               StringRef OutputFilename,
+static void handleExtBinaryWriter(sampleprof::SampleProfileWriter &Writer,
                                   ProfileFormat OutputFormat,
-                               StringRef ProfileSymbolListFile,
-                               bool CompressProfSymList, FailureMode FailMode) {
+                                  MemoryBuffer *Buffer,
+                                  sampleprof::ProfileSymbolList &WriterList,
+                                  bool CompressProfSymList,
+                                  bool CompressAllSections) {
+  populateProfileSymbolList(Buffer, WriterList);
+  WriterList.setToCompress(CompressProfSymList);
+  if (WriterList.size() > 0 && OutputFormat != PF_Ext_Binary)
+    warn("Profile Symbol list is not empty but the output format is not "
+         "ExtBinary format. The list will be lost in the output. ");
+
+  Writer.setProfileSymbolList(&WriterList);
+
+  if (CompressAllSections) {
+    if (OutputFormat != PF_Ext_Binary) {
+      warn("-compress-all-sections is only meaningful when -extbinary is "
+           "used. ");
+    } else {
+      auto ExtBinaryWriter =
+          static_cast<sampleprof::SampleProfileWriterExtBinary *>(&Writer);
+      ExtBinaryWriter->setToCompressAllSections();
+    }
+  }
+}
+
+static void
+mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
+                   StringRef OutputFilename, ProfileFormat OutputFormat,
+                   StringRef ProfileSymbolListFile, bool CompressProfSymList,
+                   bool CompressAllSections, FailureMode FailMode) {
   using namespace sampleprof;
   StringMap<FunctionSamples> ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
   LLVMContext Context;
   sampleprof::ProfileSymbolList WriterList;
   for (const auto &Input : Inputs) {
     auto ReaderOrErr = SampleProfileReader::create(Input.Filename, Context);
     if (std::error_code EC = ReaderOrErr.getError()) {
for (const auto &Input : Inputs) {
     if (ReaderList)
       WriterList.merge(*ReaderList);
   }
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
+  auto Writer = std::move(WriterOrErr.get());
   // WriterList will have StringRef refering to string in Buffer.
   // Make sure Buffer lives as long as WriterList.
   auto Buffer = getInputFileBuf(ProfileSymbolListFile);
-  populateProfileSymbolList(Buffer.get(), WriterList);
-  WriterList.setToCompress(CompressProfSymList);
-  if (WriterList.size() > 0 && OutputFormat != PF_Ext_Binary)
-    warn("Profile Symbol list is not empty but the output format is not "
-         "ExtBinary format. The list will be lost in the output. ");
-
-  auto Writer = std::move(WriterOrErr.get());
-  Writer->setProfileSymbolList(&WriterList);
+  handleExtBinaryWriter(*Writer, OutputFormat, Buffer.get(), WriterList,
+                        CompressProfSymList, CompressAllSections);
   Writer->write(ProfileMap);
 }
 
 static WeightedFile parseWeightedFile(const StringRef &WeightedFilename) {
   StringRef WeightStr, FileName;
   std::tie(WeightStr, FileName) = WeightedFilename.split(',');
 
   uint64_t Weight;
cl::alias NumThreadsA("j", cl::desc("Alias for --num-threads"),
                         cl::aliasopt(NumThreads));
   cl::opt<std::string> ProfileSymbolListFile(
       "prof-sym-list", cl::init(""),
       cl::desc("Path to file containing the list of function symbols "
                "used to populate profile symbol list"));
   cl::opt<bool> CompressProfSymList(
       "compress-prof-sym-list", cl::init(false), cl::Hidden,
       cl::desc("Compress profile symbol list before write it into profile. "));
+  cl::opt<bool> CompressAllSections(
+      "compress-all-sections", cl::init(false), cl::Hidden,
+      cl::desc("Compress all sections when writing the profile (only "
+               "meaningful for -extbinary)"));
 
   cl::ParseCommandLineOptions(argc, argv, "LLVM profile data merger\n");
 
   WeightedFileVector WeightedInputs;
   for (StringRef Filename : InputFilenames)
     addWeightedInput(WeightedInputs, {Filename, 1});
   for (StringRef WeightedFilename : WeightedInputFilenames)
     addWeightedInput(WeightedInputs, parseWeightedFile(WeightedFilename));
static int merge_main(int argc, const char *argv[]) {
   if (!RemappingFile.empty())
     Remapper = SymbolRemapper::create(RemappingFile);
 
   if (ProfileKind == instr)
     mergeInstrProfile(WeightedInputs, Remapper.get(), OutputFilename,
                       OutputFormat, OutputSparse, NumThreads, FailureMode);
   else
     mergeSampleProfile(WeightedInputs, Remapper.get(), OutputFilename,
-                       OutputFormat, ProfileSymbolListFile,
-                       CompressProfSymList, FailureMode);
+                       OutputFormat, ProfileSymbolListFile, CompressProfSymList,
+                       CompressAllSections, FailureMode);
 
   return 0;
 }
 
 /// Computer the overlap b/w profile BaseFilename and profile TestFilename.
 static void overlapInstrProfile(const std::string &BaseFilename,
                                 const std::string &TestFilename,
                                 const OverlapFuncFilters &FuncFilter,