[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

41

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.

69

Perhaps TestWriteWithSizeLimit instead of TestOutputSizeLimit1 is a little more informative?

71

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

72

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?

174

VAR_RETURN_IF_ERROR can be used here?

198

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

198

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);
+
   /// 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);

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();
+
+  SmallVector<char> StringBuffer;
+  std::unique_ptr<raw_ostream> BufferStream(
+      new raw_svector_ostream(StringBuffer));
+  OutputStream.swap(BufferStream);
+
+  if (std::error_code EC = write(ProfileMap))
+    return EC;
+  size_t IterationCount = 0;
+  while (StringBuffer.size() > OutputSizeLimit) {
+    Strategy->Erase(StringBuffer.size());
+
+    if (ProfileMap.size() == 0)
+      return sampleprof_error::too_large;
+
+    StringBuffer.clear();
+    OutputStream.reset(new raw_svector_ostream(StringBuffer));
+    if (std::error_code EC = write(ProfileMap))
+      return EC;
+    IterationCount++;
+  }
+
+  OutputStream.swap(BufferStream);
+  OutputStream->write(StringBuffer.data(), StringBuffer.size());
+  LLVM_DEBUG(dbgs() << "Profile originally has " << OriginalFunctionCount
+                    << " functions, reduced to " << ProfileMap.size() << " in "
+                    << IterationCount << " iterations\n");
+  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 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;
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/test/tools/llvm-profdata/output-size-limit.test

+Tests for output-size-limit option. Functions with least sample count are
+dropped. This test only checks the correctness of the output. The output size
+checks are in llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp because
+there isn't a cross platform way to get file size.
+
+1- No effect if output size limit >= original size
+RUN: llvm-profdata merge --sample --text --output-size-limit=212 %p/Inputs/sample-profile.proftext | FileCheck %s --check-prefix=TEST_TEXT1
+TEST_TEXT1: main:184019:0
+TEST_TEXT1-NEXT:  4: 534
+TEST_TEXT1-NEXT:  4.2: 534
+TEST_TEXT1-NEXT:  5: 1075
+TEST_TEXT1-NEXT:  5.1: 1075
+TEST_TEXT1-NEXT:  6: 2080
+TEST_TEXT1-NEXT:  7: 534
+TEST_TEXT1-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_TEXT1-NEXT:  10: inline1:1000
+TEST_TEXT1-NEXT:   1: 1000
+TEST_TEXT1-NEXT:  10: inline2:2000
+TEST_TEXT1-NEXT:   1: 2000
+TEST_TEXT1-NEXT: _Z3bari:20301:1437
+TEST_TEXT1-NEXT:  1: 1437
+TEST_TEXT1-NEXT: _Z3fooi:7711:610
+TEST_TEXT1-NEXT:  1: 610
+
+2- 1 function dropped
+RUN: llvm-profdata merge --sample --text --output-size-limit=211 %p/Inputs/sample-profile.proftext | FileCheck %s --check-prefix=TEST_TEXT2
+RUN: llvm-profdata merge --sample --text --output-size-limit=187 %p/Inputs/sample-profile.proftext | FileCheck %s --check-prefix=TEST_TEXT2
+TEST_TEXT2: main:184019:0
+TEST_TEXT2-NEXT:  4: 534
+TEST_TEXT2-NEXT:  4.2: 534
+TEST_TEXT2-NEXT:  5: 1075
+TEST_TEXT2-NEXT:  5.1: 1075
+TEST_TEXT2-NEXT:  6: 2080
+TEST_TEXT2-NEXT:  7: 534
+TEST_TEXT2-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_TEXT2-NEXT:  10: inline1:1000
+TEST_TEXT2-NEXT:   1: 1000
+TEST_TEXT2-NEXT:  10: inline2:2000
+TEST_TEXT2-NEXT:   1: 2000
+TEST_TEXT2-NEXT: _Z3bari:20301:1437
+TEST_TEXT2-NEXT:  1: 1437
+
+3- 2 functions dropped
+RUN: llvm-profdata merge --sample --text --output-size-limit=170 %p/Inputs/sample-profile.proftext | FileCheck %s --check-prefix=TEST_TEXT3
+TEST_TEXT3: main:184019:0
+TEST_TEXT3-NEXT:  4: 534
+TEST_TEXT3-NEXT:  4.2: 534
+TEST_TEXT3-NEXT:  5: 1075
+TEST_TEXT3-NEXT:  5.1: 1075
+TEST_TEXT3-NEXT:  6: 2080
+TEST_TEXT3-NEXT:  7: 534
+TEST_TEXT3-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_TEXT3-NEXT:  10: inline1:1000
+TEST_TEXT3-NEXT:   1: 1000
+TEST_TEXT3-NEXT:  10: inline2:2000
+TEST_TEXT3-NEXT:   1: 2000
+
+4- All functions dropped, should report an error
+RUN: not llvm-profdata merge --sample --text --output-size-limit=158 %p/Inputs/sample-profile.proftext 2>&1 | FileCheck %s --check-prefix=INVALID1
+INVALID1: error: Too much profile data
+
+5- ExtBinary form, no function dropped. Check output size and file content converted back to text
+RUN: llvm-profdata merge --sample --extbinary --output-size-limit=489  %p/Inputs/sample-profile.proftext -o %t.output
+RUN: llvm-profdata merge --sample --text %t.output | FileCheck %s --check-prefix=TEST_EXTBINARY1
+TEST_EXTBINARY1: main:184019:0
+TEST_EXTBINARY1-NEXT:  4: 534
+TEST_EXTBINARY1-NEXT:  4.2: 534
+TEST_EXTBINARY1-NEXT:  5: 1075
+TEST_EXTBINARY1-NEXT:  5.1: 1075
+TEST_EXTBINARY1-NEXT:  6: 2080
+TEST_EXTBINARY1-NEXT:  7: 534
+TEST_EXTBINARY1-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_EXTBINARY1-NEXT:  10: inline1:1000
+TEST_EXTBINARY1-NEXT:   1: 1000
+TEST_EXTBINARY1-NEXT:  10: inline2:2000
+TEST_EXTBINARY1-NEXT:   1: 2000
+TEST_EXTBINARY1-NEXT: _Z3bari:20301:1437
+TEST_EXTBINARY1-NEXT:  1: 1437
+TEST_EXTBINARY1-NEXT: _Z3fooi:7711:610
+TEST_EXTBINARY1-NEXT:  1: 610
+
+6- ExtBinary form, 1 function dropped
+RUN: llvm-profdata merge --sample --extbinary --output-size-limit=488  %p/Inputs/sample-profile.proftext -o %t.output
+RUN: llvm-profdata merge --sample --text %t.output | FileCheck %s --check-prefix=TEST_EXTBINARY2
+TEST_EXTBINARY2: main:184019:0
+TEST_EXTBINARY2-NEXT:  4: 534
+TEST_EXTBINARY2-NEXT:  4.2: 534
+TEST_EXTBINARY2-NEXT:  5: 1075
+TEST_EXTBINARY2-NEXT:  5.1: 1075
+TEST_EXTBINARY2-NEXT:  6: 2080
+TEST_EXTBINARY2-NEXT:  7: 534
+TEST_EXTBINARY2-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_EXTBINARY2-NEXT:  10: inline1:1000
+TEST_EXTBINARY2-NEXT:   1: 1000
+TEST_EXTBINARY2-NEXT:  10: inline2:2000
+TEST_EXTBINARY2-NEXT:   1: 2000
+TEST_EXTBINARY2-NEXT: _Z3bari:20301:1437
+TEST_EXTBINARY2-NEXT:  1: 1437
+
+7- ExtBinary form, 2 functions dropped
+RUN: llvm-profdata merge --sample --extbinary --output-size-limit=474  %p/Inputs/sample-profile.proftext -o %t.output
+RUN: llvm-profdata merge --sample --text %t.output | FileCheck %s --check-prefix=TEST_EXTBINARY3
+TEST_EXTBINARY3: main:184019:0
+TEST_EXTBINARY3-NEXT:  4: 534
+TEST_EXTBINARY3-NEXT:  4.2: 534
+TEST_EXTBINARY3-NEXT:  5: 1075
+TEST_EXTBINARY3-NEXT:  5.1: 1075
+TEST_EXTBINARY3-NEXT:  6: 2080
+TEST_EXTBINARY3-NEXT:  7: 534
+TEST_EXTBINARY3-NEXT:  9: 2064 _Z3bari:1471 _Z3fooi:631
+TEST_EXTBINARY3-NEXT:  10: inline1:1000
+TEST_EXTBINARY3-NEXT:   1: 1000
+TEST_EXTBINARY3-NEXT:  10: inline2:2000
+TEST_EXTBINARY3-NEXT:   1: 2000
+
+8- ExtBinary form, all functions dropped
+RUN: not llvm-profdata merge --sample --extbinary --output-size-limit=400  %p/Inputs/sample-profile.proftext 2>&1 | FileCheck %s --check-prefix=INVALID2
+INVALID2: error: Too much profile data

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
+//
+//===----------------------------------------------------------------------===//
+
+// Correctness check of SampleProfileWriter::writeWithSizeLimit() is done in
+// llvm lit tests (llvm/test/tools/llvm-profdata/output-size-limit.test). This
+// test checks the output size is in fact under the specified size limit.
+
+#include "llvm/ProfileData/SampleProfReader.h"
+#include "llvm/ProfileData/SampleProfWriter.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+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)";
+
+static SampleProfileMap ReadInput(StringRef Input) {
+  LLVMContext Context;
+  auto InputBuffer = MemoryBuffer::getMemBufferCopy(Input);
+  auto ReaderOrErr = SampleProfileReader::create(InputBuffer, Context);
+  if (std::error_code EC = ReaderOrErr.getError())
+    report_fatal_error(EC.message().c_str());
+  auto Reader = std::move(ReaderOrErr.get());
+  if (std::error_code EC = Reader->read())
+    report_fatal_error(EC.message().c_str());
+  return Reader->getProfiles();
+}
+
+static std::unique_ptr<SampleProfileWriter>
+CreateWriter(SmallVector<char> &OutputBuffer) {
+  std::unique_ptr<raw_ostream> BufferStream(
+      new raw_svector_ostream(OutputBuffer));
+  auto WriterOrErr = SampleProfileWriter::create(
+      BufferStream, llvm::sampleprof::SPF_Ext_Binary);
+  if (std::error_code EC = WriterOrErr.getError())
+    report_fatal_error(EC.message().c_str());
+  return std::move(WriterOrErr.get());
+}
+
+// Only returns the actual size of the written profile.
+static size_t WriteProfile(StringRef Input, size_t SizeLimit) {
+  SampleProfileMap Profiles = ReadInput(Input);
+  SmallVector<char> OutputBuffer;
+  auto Writer = CreateWriter(OutputBuffer);
+  if (std::error_code EC = Writer->writeWithSizeLimit(Profiles, SizeLimit))
+    report_fatal_error(EC.message().c_str());
+  return OutputBuffer.size();
+}
+
+TEST(TestOutputSizeLimit, TestOutputSizeLimit1) {
+  for (size_t OutputSizeLimit : {489, 488, 474})
+    ASSERT_LE(WriteProfile(Input1, OutputSizeLimit), OutputSizeLimit);
+}