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

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.

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

llvm-svn: 373914

include/llvm/ProfileData/SampleProf.h

@@ -145,11 +145,25 @@ static inline std::string getSecName(SecType Type) {
 // and SampleProfileExtBinaryBaseWriter.
 struct SecHdrTableEntry {
   SecType Type;
-  uint64_t Flag;
+  uint64_t Flags;
   uint64_t Offset;
   uint64_t Size;
 };
 
+enum SecFlags { SecFlagInValid = 0, SecFlagCompress = (1 << 0) };
+
+static inline void addSecFlags(SecHdrTableEntry &Entry, uint64_t Flags) {
+  Entry.Flags |= Flags;
+}
+
+static inline void removeSecFlags(SecHdrTableEntry &Entry, uint64_t Flags) {
+  Entry.Flags &= ~Flags;
+}
+
+static inline bool hasSecFlag(SecHdrTableEntry &Entry, SecFlags Flag) {
+  return Entry.Flags & Flag;
+}
+
 /// Represents the relative location of an instruction.
 ///
 /// Instruction locations are specified by the line offset from the
@@ -643,9 +657,9 @@ public:
   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;
 

include/llvm/ProfileData/SampleProfReader.h

@@ -488,6 +488,14 @@ public:
 /// 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;
@@ -518,7 +526,7 @@ 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,

include/llvm/ProfileData/SampleProfWriter.h

@@ -143,14 +143,16 @@ class SampleProfileWriterRawBinary : public 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;
@@ -158,34 +160,52 @@ protected:
   // 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 addSectionFlags(SecType Type, SecFlags Flags);
+  SecHdrTableEntry &getEntryInLayout(SecType Type);
+  std::error_code compressAndOutput();
 
+  // We will swap the raw_ostream held by LocalBufStream and that
+  // held 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 held 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;

lib/ProfileData/SampleProf.cpp

@@ -15,7 +15,6 @@
 #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"
@@ -198,66 +197,34 @@ FunctionSamples::findFunctionSamples(const DILocation *DIL) const {
 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;
 }
 

lib/ProfileData/SampleProfReader.cpp

@@ -26,6 +26,7 @@
 #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"
@@ -471,6 +472,7 @@ 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())
@@ -487,7 +489,7 @@ SampleProfileReaderExtBinary::readOneSection(const uint8_t *Start,
     }
     break;
   case SecProfileSymbolList:
-    if (std::error_code EC = readProfileSymbolList())
+    if (std::error_code EC = readProfileSymbolList(Size))
       return EC;
     break;
   default:
@@ -496,27 +498,43 @@ SampleProfileReaderExtBinary::readOneSection(const uint8_t *Start,
   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 (!llvm::zlib::isAvailable())
+    return sampleprof_error::zlib_unavailable;
 
-  if (std::error_code EC =
-          ProfSymList->read(*CompressSize, *UncompressSize, Data))
-    return EC;
-
-  // 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;
 }
 
@@ -528,11 +546,35 @@ std::error_code SampleProfileReaderExtBinaryBase::read() {
     // 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;
+      SecStart = DecompressBuf;
+      SecSize = DecompressBufSize;
+    }
+
+    if (std::error_code EC = readOneSection(SecStart, SecSize, Entry.Type))
       return EC;
-    if (Data != SecStart + Entry.Size)
+    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;
@@ -621,10 +663,10 @@ std::error_code SampleProfileReaderExtBinaryBase::readSecHdrTableEntry() {
     return EC;
   Entry.Type = static_cast<SecType>(*Type);
 
-  auto Flag = readUnencodedNumber<uint64_t>();
-  if (std::error_code EC = Flag.getError())
+  auto Flags = readUnencodedNumber<uint64_t>();
+  if (std::error_code EC = Flags.getError())
     return EC;
-  Entry.Flag = *Flag;
+  Entry.Flags = *Flags;
 
   auto Offset = readUnencodedNumber<uint64_t>();
   if (std::error_code EC = Offset.getError())

lib/ProfileData/SampleProfWriter.cpp

@@ -21,6 +21,7 @@
 #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"
@@ -72,21 +73,58 @@ SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
   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();
+SecHdrTableEntry &
+SampleProfileWriterExtBinaryBase::getEntryInLayout(SecType Type) {
+  auto SecIt = std::find_if(
+      SectionLayout.begin(), SectionLayout.end(),
+      [=](const auto &Entry) -> bool { return Entry.Type == Type; });
+  return *SecIt;
 }
 
-/// Add a new section into section header table. Return the position
-/// of SectionEnd.
-uint64_t
-SampleProfileWriterExtBinaryBase::addNewSection(SecType Sec,
+/// Return the current position and prepare to use it as the start
+/// position of a section.
+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.
+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.Flags, SectionStart - FileStart,
+                         OutputStream->tell() - SectionStart});
+  return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::write(
@@ -94,6 +132,8 @@ std::error_code SampleProfileWriterExtBinaryBase::write(
   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;
 
@@ -105,28 +145,38 @@ std::error_code SampleProfileWriterExtBinaryBase::write(
 
 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;
 }
@@ -308,6 +358,23 @@ std::error_code SampleProfileWriterBinary::writeHeader(
   return sampleprof_error::success;
 }
 
+void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
+  for (auto &Entry : SectionLayout)
+    addSecFlags(Entry, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::setToCompressSection(SecType Type) {
+  addSectionFlags(Type, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::addSectionFlags(SecType Type,
+                                                       SecFlags Flags) {
+  for (auto &Entry : SectionLayout) {
+    if (Entry.Type == Type)
+      addSecFlags(Entry, Flags);
+  }
+}
+
 void SampleProfileWriterExtBinaryBase::allocSecHdrTable() {
   support::endian::Writer Writer(*OutputStream, support::little);
 
@@ -342,9 +409,9 @@ std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
   // 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].Flags));
     Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Offset));
     Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Size));
   }
@@ -362,7 +429,6 @@ std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
   FileStart = OS.tell();
   writeMagicIdent(Format);
 
-  initSectionLayout();
   allocSecHdrTable();
   return sampleprof_error::success;
 }

test/Transforms/SampleProfile/compressed-profile-symbol-list.ll

@@ -1,5 +1,5 @@
 ; REQUIRES: zlib
 ; Append inline.prof with profile symbol list and save it after compression.
-; RUN: llvm-profdata merge --sample --prof-sym-list=%S/Inputs/profile-symbol-list.text --compress-prof-sym-list=true --extbinary %S/Inputs/inline.prof --output=%t.profdata
+; RUN: llvm-profdata merge --sample --prof-sym-list=%S/Inputs/profile-symbol-list.text --compress-all-sections=true --extbinary %S/Inputs/inline.prof --output=%t.profdata
 ; RUN: opt < %S/Inputs/profile-symbol-list.ll -sample-profile -profile-accurate-for-symsinlist -sample-profile-file=%t.profdata -S | FileCheck %S/Inputs/profile-symbol-list.ll
 ; RUN: opt < %S/Inputs/profile-symbol-list.ll -passes=sample-profile -profile-accurate-for-symsinlist -sample-profile-file=%t.profdata -S | FileCheck %S/Inputs/profile-symbol-list.ll

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

@@ -0,0 +1,123 @@
+; 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)

test/Transforms/SampleProfile/uncompressed-profile-symbol-list.ll

@@ -1,4 +1,4 @@
 ; Append inline.prof with profile symbol list and save it without compression.
-; RUN: llvm-profdata merge --sample --prof-sym-list=%S/Inputs/profile-symbol-list.text --compress-prof-sym-list=false --extbinary %S/Inputs/inline.prof --output=%t.profdata
+; RUN: llvm-profdata merge --sample --prof-sym-list=%S/Inputs/profile-symbol-list.text --compress-all-sections=false --extbinary %S/Inputs/inline.prof --output=%t.profdata
 ; RUN: opt < %S/Inputs/profile-symbol-list.ll -sample-profile -profile-accurate-for-symsinlist -sample-profile-file=%t.profdata -S | FileCheck %S/Inputs/profile-symbol-list.ll
 ; RUN: opt < %S/Inputs/profile-symbol-list.ll -passes=sample-profile -profile-accurate-for-symsinlist -sample-profile-file=%t.profdata -S | FileCheck %S/Inputs/profile-symbol-list.ll

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

@@ -0,0 +1,6 @@
+REQUIRES: zlib
+; RUN: llvm-profdata merge -sample -extbinary -compress-all-sections -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-all-sections -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-all-sections %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

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

@@ -0,0 +1,10 @@
+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

tools/llvm-profdata/llvm-profdata.cpp

@@ -439,12 +439,35 @@ static void populateProfileSymbolList(MemoryBuffer *Buffer,
     PSL.add(symbol);
 }
 
+static void handleExtBinaryWriter(sampleprof::SampleProfileWriter &Writer,
+                                  ProfileFormat OutputFormat,
+                                  MemoryBuffer *Buffer,
+                                  sampleprof::ProfileSymbolList &WriterList,
+                                  bool CompressAllSections) {
+  populateProfileSymbolList(Buffer, WriterList);
+  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-section is ignored. Specify -extbinary to enable it");
+    } 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, FailureMode FailMode) {
+                               bool CompressAllSections, FailureMode FailMode) {
   using namespace sampleprof;
   StringMap<FunctionSamples> ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
@@ -496,17 +519,12 @@ static void mergeSampleProfile(const WeightedFileVector &Inputs,
   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,
+                        CompressAllSections);
   Writer->write(ProfileMap);
 }
 
@@ -630,9 +648,10 @@ static int merge_main(int argc, const char *argv[]) {
       "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");
 
@@ -666,8 +685,8 @@ static int merge_main(int argc, const char *argv[]) {
                       OutputFormat, OutputSparse, NumThreads, FailureMode);
   else
     mergeSampleProfile(WeightedInputs, Remapper.get(), OutputFilename,
-                       OutputFormat, ProfileSymbolListFile,
-                       CompressProfSymList, FailureMode);
+                       OutputFormat, ProfileSymbolListFile, CompressAllSections,
+                       FailureMode);
 
   return 0;
 }