[llvm-profdata] Add option to cap profile output size

[huangjd]

Summary

D139603 (add option to llvm-profdata to reduce output profile size) contains test cases that are not cross-platform. Moving those tests to unit test and making sure the feature is callable from llvm library

[snehasish]

lgtm with some comments.

llvm/lib/ProfileData/SampleProfWriter.cpp

100

This will become unused in non-DEBUG build, maybe add the workaround here https://reviews.llvm.org/rG9f4a9d3f44501fa755eb71fe855e15cf0e59e8b8

Or wrap this in LLVM_DEBUG too.

Also IterationCount below as mentioned in this comment on the previous revision: https://reviews.llvm.org/D139603#inline-1365026

llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

42

I think we should replace report_fatal_error with

std::error_code EC = ReaderOrErr.getError();
ASSERT_FALSE(EC) << EC.message().c_str();

to let gtest handle the failure cleanly.

70

Perhaps TestWriteWithSizeLimit instead of TestOutputSizeLimit1 is a little more informative?

72

This should probably be EXPECT_LE.
The rationale is explained here https://stackoverflow.com/a/2565309

[huangjd]

Updating D141446: Fix to D139603(reverted) - moved size check to unit test so that it is cross-platform

[huangjd]

Updating D141446: Fix to D139603(reverted) - moved size check to unit test so that it is cross-platform

[vvereschaka]

Hi @huangjd ,

check-llvm gets failed with the following linker errors:

FAILED: unittests/tools/llvm-profdata/LLVMProfdataTests 
: && /usr/bin/c++ -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-maybe-uninitialized -Wno-class-memaccess -Wno-redundant-move -Wno-pessimizing-move -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wsuggest-override -Wno-comment -Wno-misleading-indentation -Wctad-maybe-unsupported -fdiagnostics-color -ffunction-sections -fdata-sections -O3 -DNDEBUG -fuse-ld=gold     -Wl,--gc-sections unittests/tools/llvm-profdata/CMakeFiles/LLVMProfdataTests.dir/OutputSizeLimitTest.cpp.o -o unittests/tools/llvm-profdata/LLVMProfdataTests  -Wl,-rpath,/home/buildbot/worker/as-builder-7/llvm-nvptx-nvidia-ubuntu/build/lib  lib/libLLVMProfileData.so.16git  lib/libllvm_gtest_main.so.16git  lib/libLLVMTestingSupport.so.16git  lib/libllvm_gtest.so.16git  lib/libLLVMSupport.so.16git  -Wl,-rpath-link,/home/buildbot/worker/as-builder-7/llvm-nvptx-nvidia-ubuntu/build/lib && :
unittests/tools/llvm-profdata/CMakeFiles/LLVMProfdataTests.dir/OutputSizeLimitTest.cpp.o:OutputSizeLimitTest.cpp:function TestOutputSizeLimit_TestOutputSizeLimit1_Test::TestBody() [clone .localalias]: error: undefined reference to 'llvm::LLVMContext::LLVMContext()'
unittests/tools/llvm-profdata/CMakeFiles/LLVMProfdataTests.dir/OutputSizeLimitTest.cpp.o:OutputSizeLimitTest.cpp:function TestOutputSizeLimit_TestOutputSizeLimit1_Test::TestBody() [clone .localalias]: error: undefined reference to 'llvm::LLVMContext::~LLVMContext()'
collect2: error: ld returned 1 exit status

https://lab.llvm.org/buildbot/#/builders/234/builds/2655

Would you take a look?

[thakis]

This unit-test based test also fails on non-linux, see e.g. http://45.33.8.238/macm1/52570/step_11.txt

I think that's because writeWithSizeLimitInternal calls Strategy->Erase(StringBuffer.size()); and that erases an element from an unordered_map, and it depends on the C++ standard library if an "important" function name gets removed. Depending on that, the error is either too_large (which the test expects) or truncated_name_table (which it doesn't).

Please take a look and revert for now if it takes a while to fix. If you do end up reverting, it'd be cool if you could revert 6be251352e6b4d9708a1b7b7b146ea199342de22 in the same commit (and reland it when you reland).

llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

38

https://github.com/google/googletest/blob/main/docs/advanced.md#propagating-fatal-failures FYI

[dyung]

We are also seeing an assertion failure in the added unit test on our internal Windows builder similar to what this bot is hitting:

https://lab.llvm.org/buildbot/#/builders/231/builds/7174

[ RUN      ] TestOutputSizeLimit.TestOutputSizeLimit1
LLVMProfdataTests: /home/buildbots/ppc64be-clang-test-suite/clang-ppc64be-test-suite/llvm-project/llvm/lib/ProfileData/SampleProfWriter.cpp:85: virtual void llvm::sampleprof::DefaultFunctionPruningStrategy::Erase(size_t): Assertion `NumToRemove <= SortedFunctions.size()' failed.

[vitalybuka]

This bot is still broken https://lab.llvm.org/buildbot/#/builders/168/builds/11266

[huangjd]

Updated unit test

[huangjd]

Updated unit test

[vvereschaka]

I have tested the latest changes on llvm-nvptx-nvidia-ubuntu (https://lab.llvm.org/buildbot/#/builders/234) builder locally and the unit tests get built and run successfully.
Thank you @huangjd.

[huangjd]

Update unit test, previously may crash on mac OS

[huangjd]

Refactored to fix bugs on non linux platforms

Moved all the tests into unit test since

1. Output size limit doesn't require the preservation of specific functions if a different pruning strategy is used in the future, so CHECK-NEXT is too restrictive. In unit test each rewritten sample is checked against the original sample to confirm all its fields are identical

2. Windows use '\r\n' line ending so the behavior of output size limit in text mode is different (more functions being pruned than linux for the same limit because extra char per line). This breaks lit test

3. Write once for many tests with repetative output

[huangjd]

cleanup

[huangjd]

use llvm::unittest::TempFile instead

[huangjd]

@snehasish Please re-review since I introduced a major change to the code.

[snehasish]

lgtm

llvm/lib/ProfileData/SampleProfWriter.cpp

111

Is this leaking the raw_svector_ostream objects?

Can we rewrite this as the following?

raw_svector_ostream OS(StringBuffer);
OutputStream.reset(&OS);

Also might be worth it to use a separate OutputStream object within the loop scope so that we don't have to worry about the swap before and after and the lifetimes of the raw_vector_ostream objects.

llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

71

I think these parameters name should be capitalized based on the guide.
https://llvm.org/docs/CodingStandards.html

Also consider moving this to the implementation of FunctionSamples since it seems generally useful to have operator== implemented?

173

VAR_RETURN_IF_ERROR can be used here?

197

EXPECT_THAT_EXPECTED is probably better to continue with other test cases rather than aborting the test on the first failure?

[huangjd]

llvm/lib/ProfileData/SampleProfWriter.cpp

111

unique_ptr.reset(new obj) is the correct usage.
https://en.cppreference.com/w/cpp/memory/unique_ptr/reset

[snehasish]

llvm/lib/ProfileData/SampleProfWriter.cpp

111

Thanks, knowing the type makes it clear! Perhaps that is additional motivation for a separate, local OutputStream var?

[huangjd]

llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

197

The callee returns an error only if there's something wrong with I/O or profile reading, which is not expected to happen at all, so ASSERT is used properly here. This test checks the correctness of the sample map after size reduction, which is checked with EXPECT_EQ

[huangjd]

Cleanup unit test

[huangjd]

update tests because API change from main

llvm/include/llvm/ProfileData/SampleProfWriter.h

enum SectionLayout {
   DefaultLayout,
   // The layout splits profile with context information from profile without
   // context information. When Thinlto is enabled, ThinLTO postlink phase only
   // has to load profile with context information and can skip the other part.
   CtxSplitLayout,
   NumOfLayout,
 };
 
+/// When writing a profile with size limit, user may want to use a different
+/// strategy to reduce function count other than dropping functions with fewest
+/// samples first. In this case a class implementing the same interfaces should
+/// be provided to SampleProfileWriter::writeWithSizeLimit().
+class FunctionPruningStrategy {
+protected:
+  SampleProfileMap &ProfileMap;
+  size_t OutputSizeLimit;
+
+public:
+  /// \p ProfileMap A reference to the original profile map. It will be modified
+  /// by Erase().
+  /// \p OutputSizeLimit Size limit in bytes of the output profile. This is
+  /// necessary to estimate how many functions to remove.
+  FunctionPruningStrategy(SampleProfileMap &ProfileMap, size_t OutputSizeLimit)
+      : ProfileMap(ProfileMap), OutputSizeLimit(OutputSizeLimit) {}
+
+  virtual ~FunctionPruningStrategy() = default;
+
+  /// SampleProfileWriter::writeWithSizeLimit() calls this after every write
+  /// iteration if the output size still exceeds the limit. This function
+  /// should erase some functions from the profile map so that the writer tries
+  /// to write the profile again with fewer functions. At least 1 entry from the
+  /// profile map must be erased.
+  ///
+  /// \p CurrentOutputSize Number of bytes in the output if current profile map
+  /// is written.
+  virtual void Erase(size_t CurrentOutputSize) = 0;
+};
+
+class DefaultFunctionPruningStrategy : public FunctionPruningStrategy {
+  std::vector<NameFunctionSamples> SortedFunctions;
+
+public:
+  DefaultFunctionPruningStrategy(SampleProfileMap &ProfileMap,
+                                 size_t OutputSizeLimit);
+
+  /// In this default implementation, functions with fewest samples are dropped
+  /// first. Since the exact size of the output cannot be easily calculated due
+  /// to compression, we use a heuristic to remove as many functions as
+  /// necessary but not too many, aiming to minimize the number of write
+  /// iterations.
+  /// Empirically, functions with larger total sample count contain linearly
+  /// more sample entries, meaning it takes linearly more space to write them.
+  /// The cumulative length is therefore quadratic if all functions are sorted
+  /// by total sample count.
+  /// TODO: Find better heuristic.
+  void Erase(size_t CurrentOutputSize) override;
+};
+
 /// Sample-based profile writer. Base class.
 class SampleProfileWriter {
 public:
   virtual ~SampleProfileWriter() = default;
 
   /// Write sample profiles in \p S.
   ///
   /// \returns status code of the file update operation.
   virtual std::error_code writeSample(const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
   /// \returns status code of the file update operation.
   virtual std::error_code write(const SampleProfileMap &ProfileMap);
 
+  /// Write sample profiles up to given size limit, using the pruning strategy
+  /// to drop some functions if necessary.
+  ///
+  /// \returns status code of the file update operation.
+  template <typename FunctionPruningStrategy = DefaultFunctionPruningStrategy>
+  std::error_code writeWithSizeLimit(SampleProfileMap &ProfileMap,
+                                     size_t OutputSizeLimit) {
+    FunctionPruningStrategy Strategy(ProfileMap, OutputSizeLimit);
+    return writeWithSizeLimitInternal(ProfileMap, OutputSizeLimit, &Strategy);
+  }
+
   raw_ostream &getOutputStream() { return *OutputStream; }
 
   /// Profile writer factory.
   ///
   /// Create a new file writer based on the value of \p Format.
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(StringRef Filename, SampleProfileFormat Format);
 
SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
       : OutputStream(std::move(OS)) {}
 
   /// Write a file header for the profile file.
   virtual std::error_code writeHeader(const SampleProfileMap &ProfileMap) = 0;
 
   // Write function profiles to the profile file.
   virtual std::error_code writeFuncProfiles(const SampleProfileMap &ProfileMap);
 
+  std::error_code writeWithSizeLimitInternal(SampleProfileMap &ProfileMap,
+                                             size_t OutputSizeLimit,
+                                             FunctionPruningStrategy *Strategy);
+
+  /// For writeWithSizeLimit in text mode, each newline takes 1 additional byte
+  /// on Windows when actually written to the file, but not written to a memory
+  /// buffer. This needs to be accounted for when rewriting the profile.
+  size_t LineCount;
+
   /// Output stream where to emit the profile to.
   std::unique_ptr<raw_ostream> OutputStream;
 
   /// Profile summary.
   std::unique_ptr<ProfileSummary> Summary;
 
   /// Compute summary for this profile.
   void computeSummary(const SampleProfileMap &ProfileMap);
 
   /// Profile format.
   SampleProfileFormat Format = SPF_None;
 };
 
 /// Sample-based profile writer (text format).
 class SampleProfileWriterText : public SampleProfileWriter {
 public:
   std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
   SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS), Indent(0) {}
 
   std::error_code writeHeader(const SampleProfileMap &ProfileMap) override {
+    LineCount = 0;
     return sampleprof_error::success;
   }
 
 private:
   /// Indent level to use when writing.
   ///
   /// This is used when printing inlined callees.
   unsigned Indent;

llvm/lib/ProfileData/SampleProfWriter.cpp

 #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>
+#include <cmath>
 #include <cstdint>
 #include <memory>
 #include <set>
 #include <system_error>
 #include <utility>
 #include <vector>
 
+#define DEBUG_TYPE "llvm-profdata"
+
 using namespace llvm;
 using namespace sampleprof;
 
+namespace llvm {
+namespace support {
+namespace endian {
+namespace {
+
+// Adapter class to llvm::support::endian::Writer for pwrite().
+struct SeekableWriter {
+  raw_pwrite_stream &OS;
+  endianness Endian;
+  SeekableWriter(raw_pwrite_stream &OS, endianness Endian)
+      : OS(OS), Endian(Endian) {}
+
+  template <typename ValueType>
+  void pwrite(ValueType Val, size_t Offset) {
+    std::string StringBuf;
+    raw_string_ostream SStream(StringBuf);
+    Writer(SStream, Endian).write(Val);
+    OS.pwrite(StringBuf.data(), StringBuf.size(), Offset);
+  }
+};
+
+} // namespace
+} // namespace endian
+} // namespace support
+} // namespace llvm
+
+DefaultFunctionPruningStrategy::DefaultFunctionPruningStrategy(
+    SampleProfileMap &ProfileMap, size_t OutputSizeLimit)
+    : FunctionPruningStrategy(ProfileMap, OutputSizeLimit) {
+  sortFuncProfiles(ProfileMap, SortedFunctions);
+}
+
+void DefaultFunctionPruningStrategy::Erase(size_t CurrentOutputSize) {
+  double D = (double)OutputSizeLimit / CurrentOutputSize;
+  size_t NewSize = (size_t)round(ProfileMap.size() * D * D);
+  size_t NumToRemove = ProfileMap.size() - NewSize;
+  if (NumToRemove < 1)
+    NumToRemove = 1;
+
+  assert(NumToRemove <= SortedFunctions.size());
+  llvm::for_each(
+      llvm::make_range(SortedFunctions.begin() + SortedFunctions.size() -
+                           NumToRemove,
+                       SortedFunctions.end()),
+      [&](const NameFunctionSamples &E) { ProfileMap.erase(E.first); });
+  SortedFunctions.resize(SortedFunctions.size() - NumToRemove);
+}
+
+std::error_code SampleProfileWriter::writeWithSizeLimitInternal(
+    SampleProfileMap &ProfileMap, size_t OutputSizeLimit,
+    FunctionPruningStrategy *Strategy) {
+  if (OutputSizeLimit == 0)
+    return write(ProfileMap);
+
+  size_t OriginalFunctionCount = ProfileMap.size();
+
+  std::unique_ptr<raw_ostream> OriginalOutputStream;
+  OutputStream.swap(OriginalOutputStream);
+
+  size_t IterationCount = 0;
+  size_t TotalSize;
+
+  SmallVector<char> StringBuffer;
+  do {
+    StringBuffer.clear();
+    OutputStream.reset(new raw_svector_ostream(StringBuffer));
+    if (std::error_code EC = write(ProfileMap))
+      return EC;
+
+    TotalSize = StringBuffer.size();
+    // On Windows every "\n" is actually written as "\r\n" to disk but not to
+    // memory buffer, this difference should be added when considering the total
+    // output size.
+#ifdef _WIN32
+    if (Format == SPF_Text)
+      TotalSize += LineCount;
+#endif
+    if (TotalSize <= OutputSizeLimit)
+      break;
+
+    Strategy->Erase(TotalSize);
+    IterationCount++;
+  } while (ProfileMap.size() != 0);
+
+  if (ProfileMap.size() == 0)
+    return sampleprof_error::too_large;
+
+  OutputStream.swap(OriginalOutputStream);
+  OutputStream->write(StringBuffer.data(), StringBuffer.size());
+  LLVM_DEBUG(dbgs() << "Profile originally has " << OriginalFunctionCount
+                    << " functions, reduced to " << ProfileMap.size() << " in "
+                    << IterationCount << " iterations\n");
+  // Silence warning on Release build.
+  (void)OriginalFunctionCount;
+  (void)IterationCount;
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileWriter::writeFuncProfiles(const SampleProfileMap &ProfileMap) {
   std::vector<NameFunctionSamples> V;
   sortFuncProfiles(ProfileMap, V);
   for (const auto &I : V) {
     if (std::error_code EC = writeSample(*I.second))
       return EC;
   }
std::error_code SampleProfileWriterExtBinaryBase::addNewSection(
   }
   SecHdrTable.push_back({Type, Entry.Flags, SectionStart - FileStart,
                          OutputStream->tell() - SectionStart, LayoutIdx});
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileWriterExtBinaryBase::write(const SampleProfileMap &ProfileMap) {
+  // When calling write on a different profile map, existing states should be
+  // cleared.
+  NameTable.clear();
+  CSNameTable.clear();
+  SecHdrTable.clear();
+
   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;
 
std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   if (FunctionSamples::ProfileIsCS)
     OS << "[" << S.getContext().toString() << "]:" << S.getTotalSamples();
   else
     OS << S.getName() << ":" << S.getTotalSamples();
 
   if (Indent == 0)
     OS << ":" << S.getHeadSamples();
   OS << "\n";
+  LineCount++;
 
   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
   for (const auto &I : SortedSamples.get()) {
     LineLocation Loc = I->first;
     const SampleRecord &Sample = I->second;
     OS.indent(Indent + 1);
     if (Loc.Discriminator == 0)
       OS << Loc.LineOffset << ": ";
     else
       OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
 
     OS << Sample.getSamples();
 
     for (const auto &J : Sample.getSortedCallTargets())
       OS << " " << J.first << ":" << J.second;
     OS << "\n";
+    LineCount++;
   }
 
   SampleSorter<LineLocation, FunctionSamplesMap> SortedCallsiteSamples(
       S.getCallsiteSamples());
   Indent += 1;
   for (const auto &I : SortedCallsiteSamples.get())
     for (const auto &FS : I->second) {
       LineLocation Loc = I->first;
       const FunctionSamples &CalleeSamples = FS.second;
       OS.indent(Indent);
       if (Loc.Discriminator == 0)
         OS << Loc.LineOffset << ": ";
       else
         OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
       if (std::error_code EC = writeSample(CalleeSamples))
         return EC;
     }
   Indent -= 1;
 
   if (FunctionSamples::ProfileIsProbeBased) {
     OS.indent(Indent + 1);
     OS << "!CFGChecksum: " << S.getFunctionHash() << "\n";
+    LineCount++;
   }
 
   if (S.getContext().getAllAttributes()) {
     OS.indent(Indent + 1);
     OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
+    LineCount++;
   }
 
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
   assert(!Context.hasContext() && "cs profile is not supported");
std::error_code SampleProfileWriterBinary::writeNameTable() {
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
   auto &OS = *OutputStream;
 
   // Fill the slot remembered by TableOffset with the offset of FuncOffsetTable.
-  auto &OFS = static_cast<raw_fd_ostream &>(OS);
   uint64_t FuncOffsetTableStart = OS.tell();
-  if (OFS.seek(TableOffset) == (uint64_t)-1)
-    return sampleprof_error::ostream_seek_unsupported;
-  support::endian::Writer Writer(*OutputStream, support::little);
-  Writer.write(FuncOffsetTableStart);
-  if (OFS.seek(FuncOffsetTableStart) == (uint64_t)-1)
-    return sampleprof_error::ostream_seek_unsupported;
+  support::endian::SeekableWriter Writer(static_cast<raw_pwrite_stream &>(OS),
+                                         support::little);
+  Writer.pwrite(FuncOffsetTableStart, TableOffset);
 
   // Write out the table size.
   encodeULEB128(FuncOffsetTable.size(), OS);
 
   // Write out FuncOffsetTable.
   for (auto Entry : FuncOffsetTable) {
     if (std::error_code EC = writeNameIdx(Entry.first))
       return EC;
     encodeULEB128(Entry.second, OS);
   }
+  FuncOffsetTable.clear();
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(NameTable, V);
 
SampleProfileWriterBinary::writeMagicIdent(SampleProfileFormat Format) {
   // Write file magic identifier.
   encodeULEB128(SPMagic(Format), OS);
   encodeULEB128(SPVersion(), OS);
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileWriterBinary::writeHeader(const SampleProfileMap &ProfileMap) {
+  // When calling write on a different profile map, existing names should be
+  // cleared.
+  NameTable.clear();
+
   writeMagicIdent(Format);
 
   computeSummary(ProfileMap);
   if (auto EC = writeSummary())
     return EC;
 
   // Generate the name table for all the functions referenced in the profile.
   for (const auto &I : ProfileMap) {
for (uint32_t i = 0; i < SectionHdrLayout.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);
-
   assert(SecHdrTable.size() == SectionHdrLayout.size() &&
          "SecHdrTable entries doesn't match SectionHdrLayout");
   SmallVector<uint32_t, 16> IndexMap(SecHdrTable.size(), -1);
   for (uint32_t TableIdx = 0; TableIdx < SecHdrTable.size(); TableIdx++) {
     IndexMap[SecHdrTable[TableIdx].LayoutIndex] = TableIdx;
   }
 
   // Write the section header table in the order specified in
   // SectionHdrLayout. SectionHdrLayout specifies the sections
   // order in which profile reader expect to read, so the section
   // header table should be written in the order in SectionHdrLayout.
   // Note that the section order in SecHdrTable may be different
   // from the order in SectionHdrLayout, for example, SecFuncOffsetTable
   // needs to be computed after SecLBRProfile (the order in SecHdrTable),
   // but it needs to be read before SecLBRProfile (the order in
   // SectionHdrLayout). So we use IndexMap above to switch the order.
+  support::endian::SeekableWriter Writer(
+      static_cast<raw_pwrite_stream &>(*OutputStream), support::little);
   for (uint32_t LayoutIdx = 0; LayoutIdx < SectionHdrLayout.size();
        LayoutIdx++) {
     assert(IndexMap[LayoutIdx] < SecHdrTable.size() &&
            "Incorrect LayoutIdx in SecHdrTable");
     auto Entry = SecHdrTable[IndexMap[LayoutIdx]];
-    Writer.write(static_cast<uint64_t>(Entry.Type));
-    Writer.write(static_cast<uint64_t>(Entry.Flags));
-    Writer.write(static_cast<uint64_t>(Entry.Offset));
-    Writer.write(static_cast<uint64_t>(Entry.Size));
+    Writer.pwrite(static_cast<uint64_t>(Entry.Type),
+                  SecHdrTableOffset + 4 * LayoutIdx * sizeof(uint64_t));
+    Writer.pwrite(static_cast<uint64_t>(Entry.Flags),
+                  SecHdrTableOffset + (4 * LayoutIdx + 1) * sizeof(uint64_t));
+    Writer.pwrite(static_cast<uint64_t>(Entry.Offset),
+                  SecHdrTableOffset + (4 * LayoutIdx + 2) * sizeof(uint64_t));
+    Writer.pwrite(static_cast<uint64_t>(Entry.Size),
+                  SecHdrTableOffset + (4 * LayoutIdx + 3) * sizeof(uint64_t));
   }
 
-  // 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 SampleProfileMap &ProfileMap) {
   auto &OS = *OutputStream;
   FileStart = OS.tell();
   writeMagicIdent(Format);

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

 
 static void
 mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
                    StringRef OutputFilename, ProfileFormat OutputFormat,
                    StringRef ProfileSymbolListFile, bool CompressAllSections,
                    bool UseMD5, bool GenPartialProfile, bool GenCSNestedProfile,
                    bool SampleMergeColdContext, bool SampleTrimColdContext,
                    bool SampleColdContextFrameDepth, FailureMode FailMode,
-                   bool DropProfileSymbolList) {
+                   bool DropProfileSymbolList, size_t OutputSizeLimit) {
   using namespace sampleprof;
   SampleProfileMap ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
   LLVMContext Context;
   sampleprof::ProfileSymbolList WriterList;
   std::optional<bool> ProfileIsProbeBased;
   std::optional<bool> ProfileIsCS;
   for (const auto &Input : 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);
   handleExtBinaryWriter(*Writer, OutputFormat, Buffer.get(), WriterList,
                         CompressAllSections, UseMD5, GenPartialProfile);
-  if (std::error_code EC = Writer->write(ProfileMap))
+
+  // If OutputSizeLimit is 0 (default), it is the same as write().
+  if (std::error_code EC =
+          Writer->writeWithSizeLimit(ProfileMap, OutputSizeLimit))
     exitWithErrorCode(std::move(EC));
 }
 
 static WeightedFile parseWeightedFile(const StringRef &WeightedFilename) {
   StringRef WeightStr, FileName;
   std::tie(WeightStr, FileName) = WeightedFilename.split(',');
 
   uint64_t Weight;
static int merge_main(int argc, const char *argv[]) {
   cl::opt<bool> SampleTrimColdContext(
       "sample-trim-cold-context", cl::init(false), cl::Hidden,
       cl::desc(
           "Trim context sample profiles whose count is below cold threshold"));
   cl::opt<uint32_t> SampleColdContextFrameDepth(
       "sample-frame-depth-for-cold-context", cl::init(1),
       cl::desc("Keep the last K frames while merging cold profile. 1 means the "
                "context-less base profile"));
+  cl::opt<size_t> OutputSizeLimit(
+      "output-size-limit", cl::init(0), cl::Hidden,
+      cl::desc("Trim cold functions until profile size is below specified "
+               "limit in bytes. This uses a heursitic and functions may be "
+               "excessively trimmed"));
   cl::opt<bool> GenPartialProfile(
       "gen-partial-profile", cl::init(false), cl::Hidden,
       cl::desc("Generate a partial profile (only meaningful for -extbinary)"));
   cl::opt<std::string> SupplInstrWithSample(
       "supplement-instr-with-sample", cl::init(""), cl::Hidden,
       cl::desc("Supplement an instr profile with sample profile, to correct "
                "the profile unrepresentativeness issue. The sample "
                "profile is the input of the flag. Output will be in instr "
if (ProfileKind == instr)
     mergeInstrProfile(WeightedInputs, DebugInfoFilename, Remapper.get(),
                       OutputFilename, OutputFormat, OutputSparse, NumThreads,
                       FailureMode, ProfiledBinary);
   else
     mergeSampleProfile(
         WeightedInputs, Remapper.get(), OutputFilename, OutputFormat,
         ProfileSymbolListFile, CompressAllSections, UseMD5, GenPartialProfile,
         GenCSNestedProfile, SampleMergeColdContext, SampleTrimColdContext,
-        SampleColdContextFrameDepth, FailureMode, DropProfileSymbolList);
+        SampleColdContextFrameDepth, FailureMode, DropProfileSymbolList,
+        OutputSizeLimit);
   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,
                                 raw_fd_ostream &OS, bool IsCS) {

llvm/unittests/tools/CMakeLists.txt

 if(LLVM_TARGETS_TO_BUILD MATCHES "X86")
   add_subdirectory(
     llvm-cfi-verify
   )
 endif()
 
 add_subdirectory(
   llvm-exegesis
 )
+add_subdirectory(llvm-profdata)
 add_subdirectory(llvm-profgen)
 add_subdirectory(llvm-mca)

llvm/unittests/tools/llvm-profdata/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  ProfileData
+  Support
+  )
+
+add_llvm_unittest(LLVMProfdataTests
+    OutputSizeLimitTest.cpp
+  )
+
+target_link_libraries(LLVMProfdataTests PRIVATE LLVMTestingSupport)
+
+set_property(TARGET LLVMProfdataTests PROPERTY FOLDER "Tests/UnitTests/ToolTests")

llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

+//===- llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ProfileData/SampleProfReader.h"
+#include "llvm/ProfileData/SampleProfWriter.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using llvm::unittest::TempFile;
+
+std::string Input1 = R"(main:184019:0
+ 4: 534
+ 4.2: 534
+ 5: 1075
+ 5.1: 1075
+ 6: 2080
+ 7: 534
+ 9: 2064 _Z3bari:1471 _Z3fooi:631
+ 10: inline1:1000
+  1: 1000
+ 10: inline2:2000
+  1: 2000
+_Z3bari:20301:1437
+ 1: 1437
+_Z3fooi:7711:610
+ 1: 610)";
+
+const char EmptyProfile[18] = "\xff\xe5\xd0\xb1\xf4\xc9\x94\xa8\x53\x67";
+
+/// sys::fs and SampleProf mix Error and error_code, making an adapter class
+/// to keep code elegant.
+template <typename T> class ExpectedErrorOr : public Expected<T> {
+public:
+  ExpectedErrorOr(T &&Obj) : Expected<T>(Obj) {}
+
+  ExpectedErrorOr(std::error_code EC) : Expected<T>(errorCodeToError(EC)) {}
+
+  ExpectedErrorOr(Error &&E) : Expected<T>(std::move(E)) {}
+
+  template <typename U>
+  ExpectedErrorOr(ErrorOr<U> &&E)
+      : Expected<T>(errorCodeToError(E.getError())) {}
+
+  template <typename U>
+  ExpectedErrorOr(Expected<U> &&E) : Expected<T>(E.takeError()) {}
+};
+
+#define DEF_VAR_RETURN_IF_ERROR(Var, Value)                                    \
+  auto Var##OrErr = Value;                                                     \
+  if (!Var##OrErr)                                                             \
+    return Var##OrErr;                                                         \
+  auto Var = std::move(Var##OrErr.get())
+
+#define VAR_RETURN_IF_ERROR(Var, Value)                                        \
+  Var##OrErr = Value;                                                          \
+  if (!Var##OrErr)                                                             \
+    return Var##OrErr;                                                         \
+  Var = std::move(Var##OrErr.get())
+
+#define RETURN_IF_ERROR(Value)                                                 \
+  if (auto E = Value)                                                          \
+  return std::move(E)
+
+static bool operator==(const FunctionSamples &a, const FunctionSamples &b) {
+  const BodySampleMap &BodySamplesA = a.getBodySamples();
+  const BodySampleMap &BodySamplesB = b.getBodySamples();
+
+  const CallsiteSampleMap &CallsiteSamplesA = a.getCallsiteSamples();
+  const CallsiteSampleMap &CallsiteSamplesB = b.getCallsiteSamples();
+
+  if (a.getTotalSamples() != b.getTotalSamples() ||
+      BodySamplesA.size() != BodySamplesB.size() ||
+      CallsiteSamplesA.size() != CallsiteSamplesB.size())
+    return false;
+
+  for (auto &BodySampleA : BodySamplesA) {
+    auto BodySampleB = BodySamplesB.find(BodySampleA.first);
+    if (BodySampleB == BodySamplesB.end() ||
+        BodySampleA.second.getSamples() != BodySampleB->second.getSamples() ||
+        BodySampleA.second.getCallTargets() !=
+            BodySampleB->second.getCallTargets())
+      return false;
+  }
+
+  for (auto &CallsiteA : CallsiteSamplesA) {
+    auto CallsiteB = CallsiteSamplesB.find(CallsiteA.first);
+    if (CallsiteB == CallsiteSamplesB.end())
+      return false;
+
+    auto FunctionsA = CallsiteA.second;
+    auto FunctionsB = CallsiteB->second;
+    if (FunctionsA.size() != FunctionsB.size())
+      return false;
+    for (auto &FunctionA : FunctionsA) {
+      auto FunctionB = FunctionsB.find(FunctionA.first);
+      if (FunctionB == FunctionsB.end() ||
+          !(FunctionA.second == FunctionB->second))
+        return false;
+    }
+  }
+  return true;
+}
+
+/// The main testing routine. After rewriting profiles with size limit, check
+/// the following:
+/// 1. The file size of the new profile is within the size limit.
+/// 2. The new profile is a subset of the old profile, and the content of every
+/// sample in the new profile is unchanged.
+/// Note that even though by default samples with fewest total count are dropped
+/// first, this is not a requirement. Samples can be dropped by any order.
+static ExpectedErrorOr<void *> RunTest(StringRef Input, size_t SizeLimit,
+                                       SampleProfileFormat Format) {
+  // Read Input profile.
+  LLVMContext Context;
+  auto InputBuffer = MemoryBuffer::getMemBuffer(Input);
+  DEF_VAR_RETURN_IF_ERROR(Reader,
+                          SampleProfileReader::create(InputBuffer, Context));
+  RETURN_IF_ERROR(Reader->read());
+  SampleProfileMap OldProfiles = Reader->getProfiles();
+
+  // Rewrite it to a temp file with size limit.
+  TempFile Temp("profile", "afdo");
+  bool isEmpty = false;
+  {
+    DEF_VAR_RETURN_IF_ERROR(Writer,
+                            SampleProfileWriter::create(Temp.path(), Format));
+    std::error_code EC = Writer->writeWithSizeLimit(OldProfiles, SizeLimit);
+    // too_large means no sample could be written because SizeLimit is too
+    // small. Otherwise any other error code indicates unexpected failure.
+    if (EC == sampleprof_error::too_large)
+      isEmpty = true;
+    else if (EC)
+      return EC;
+  }
+
+  // Read the temp file to get new profiles. Use the default empty profile if
+  // temp file was not written because size limit is too small.
+  SampleProfileMap NewProfiles;
+  InputBuffer = MemoryBuffer::getMemBuffer(StringRef(EmptyProfile, 17));
+  DEF_VAR_RETURN_IF_ERROR(NewReader,
+                          SampleProfileReader::create(InputBuffer, Context));
+  if (!isEmpty) {
+    VAR_RETURN_IF_ERROR(
+        NewReader, SampleProfileReader::create(Temp.path().str(), Context));
+    RETURN_IF_ERROR(NewReader->read());
+    NewProfiles = NewReader->getProfiles();
+  }
+
+  // Check temp file is actually within size limit.
+  uint64_t FileSize;
+  RETURN_IF_ERROR(sys::fs::file_size(Temp.path(), FileSize));
+  EXPECT_LE(FileSize, SizeLimit);
+
+  // For compact binary format, function names are stored as MD5, so we cannot
+  // directly match the samples of the new profile with the old profile. A
+  // simple way is to convert the old profile to compact binary format and read
+  // it back
+  if (Format == llvm::sampleprof::SPF_Compact_Binary) {
+    TempFile CompBinary("compbinary", "afdo");
+    {
+      DEF_VAR_RETURN_IF_ERROR(
+          Writer, SampleProfileWriter::create(
+                      CompBinary.path(), llvm::sampleprof::SPF_Compact_Binary));
+      RETURN_IF_ERROR(Writer->write(OldProfiles));
+    }
+    ReaderOrErr = SampleProfileReader::create(CompBinary.path().str(), Context);
+    if (!ReaderOrErr)
+      return ReaderOrErr;
+    Reader = std::move(ReaderOrErr.get());
+    RETURN_IF_ERROR(Reader->read());
+    OldProfiles = Reader->getProfiles();
+  }
+
+  // For every sample in the new profile, confirm it is in the old profile and
+  // unchanged.
+  for (auto Sample : NewProfiles) {
+    auto FindResult = OldProfiles.find(Sample.first);
+    EXPECT_NE(FindResult, OldProfiles.end());
+    if (FindResult != OldProfiles.end()) {
+      EXPECT_EQ(Sample.second.getHeadSamples(),
+                FindResult->second.getHeadSamples());
+      EXPECT_TRUE(Sample.second == FindResult->second);
+    }
+  }
+  return nullptr;
+}
+
+TEST(TestOutputSizeLimit, TestOutputSizeLimitExtBinary) {
+  for (size_t OutputSizeLimit : {490, 489, 488, 475, 474, 459, 400})
+    ASSERT_THAT_EXPECTED(
+        RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Ext_Binary),
+        Succeeded());
+}
+
+TEST(TestOutputSizeLimit, TestOutputSizeLimitBinary) {
+  for (size_t OutputSizeLimit : {250, 249, 248, 237, 236, 223, 200})
+    ASSERT_THAT_EXPECTED(
+        RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Binary),
+        Succeeded());
+}
+
+TEST(TestOutputSizeLimit, TestOutputSizeLimitCompBinary) {
+  for (size_t OutputSizeLimit : {277, 276, 275, 264, 263, 250, 200})
+    ASSERT_THAT_EXPECTED(
+        RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Compact_Binary),
+        Succeeded());
+}
+
+TEST(TestOutputSizeLimit, TestOutputSizeLimitText) {
+  for (size_t OutputSizeLimit :
+       {229, 228, 227, 213, 212, 211, 189, 188, 187, 186, 150})
+    ASSERT_THAT_EXPECTED(
+        RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Text),
+        Succeeded());
+}