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

[huangjd]

Summary

Allow user to specify --output-size-limit=n to cap the size of generated profile to be strictly under n. Functions with the lowest total sample count are dropped first if necessary. Due to using a heuristic, excessive functions may be dropped to satisfy the size requirement

[huangjd]

https://discourse.llvm.org/t/rfc-add-option-to-limit-llvm-profdata-profile-output-size/67036 for RFC

[davidxl]

Add a test case for completeness.

llvm/lib/ProfileData/SampleProfWriter.cpp

631

unneeded change

639

unneeded change.

665

The clear call changes the behavior. Why is it needed?

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

1068

Extract the following to a helper function.

[wenlei]

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

1080–1082

This heuristic can be quite inaccurate. The number of body sample + call site sample entries can be much better proxy for actual profile size and yet still easily accessible.

functionSamples.getBodySamples().size() + functionSamples.getCallsiteSamples().size()

[wenlei]

cc @hoy @wlei

[huangjd]

llvm/lib/ProfileData/SampleProfWriter.cpp

665

Name table should only contains names exist in the current ProfileMap. The original implementation adds to the name table when writing a new profile, and the old names are never cleared, which is actually a bug. (However SampleProfileWriter is single use, an instance never calls write twice, so this bug was not showing up until this new feature)

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

1080–1082

The new revision performs the iterations on a string buffer so the performance of the heuristic is not a big concern (and now it should not overshoot by reducing too many functions). Although using a simple proportional heuristic is still too slow because it ends up reducing one function at the tail at a time

[huangjd]

Refactored code structure
Use a string buffer to rewrite files
Added API for potential new strategy to reduce profile size

[huangjd]

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

1080–1082

Better heuristics can be added in a later patch. Need more real world profile data (industrial use) to confirm which model is the best

[huangjd]

Added output size check

[xur]

I have some high level questions:
(1) have you considered removing samples with smaller values across the program-- it's like downsampling. Comparing to removing functions with smaller total counts, I think that results in a more consistent profile.
(2) if we choose to remove function, and you sort the function with total count, should we find the exactly place to cut to satisfy the size limit? In theory it should as the profile organized in unit of function. You can keep writing to the buffer until it reaches the limits. Of cause there are some section data for extbinary and summary, but they should be able to compute. Using heuristic to guess the function to remove and doing it iteratively does not seem to be appealing here.

[huangjd]

The biggest challenge to compute the number of functions accurately is the compression in extbinary, because the compressed size is non-linear to the original size. Since profile samples and function names are written to different sections (and in CS profile the names are split into two sections and samples can also be split into two sections), there is no way to predict ahead the offset between them. Based on use cases, the current heuristic is under estimating how many functions to prune (and the last iteration typically converges to pruning 1 function) so it's unlikely to remove too many functions. (Note: Also tried using cubic equation for heuristic but that will remove too many functions, so the optimal heuristic is between O(n^2) and O(n^3))

[huangjd]

As for down sampling, having a sample count of 0 vs not having a sample means differently to the compiler, so that may change the branch basic block placement on hot functions, not sure if good idea.

[snehasish]

In D139603#4000123, @huangjd wrote:

The biggest challenge to compute the number of functions accurately is the compression in extbinary, because the compressed size is non-linear to the original size. Since profile samples and function names are written to different sections (and in CS profile the names are split into two sections and samples can also be split into two sections), there is no way to predict ahead the offset between them. Based on use cases, the current heuristic is under estimating how many functions to prune (and the last iteration typically converges to pruning 1 function) so it's unlikely to remove too many functions. (Note: Also tried using cubic equation for heuristic but that will remove too many functions, so the optimal heuristic is between O(n^2) and O(n^3))

I think it makes sense to start with straight-forward implementation and look for opportunities to refine it further in the future.

llvm/include/llvm/ProfileData/SampleProfWriter.h

133

It looks like derived classes must always call this function. Can you add a comment here for the future? Maybe something like // This function must always be called by the overridden implementation?

llvm/lib/ProfileData/SampleProfWriter.cpp

58

I think this should be ValueType based on the guidance here:
https://llvm.org/docs/CodingStandards.html#name-types-functions-variables-and-enumerators-properly

llvm/test/tools/llvm-profdata/output-size-limit.test

6

I don't think we should use *-DAG here since that means the test will pass if the lines are reordered. However, if we reorder the symbol line with its contents the text format will be incorrect.

I think the -DAG directive is useful if there is non-determinism in the text format where profile information for a whole symbol may appear before another symbol. In this case we should separate out the CHECKs like the example in [1]. Though we should fix non-deterministic output if this is the case.

[1] https://llvm.org/docs/CommandGuide/FileCheck.html#the-check-dag-directive

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

981

How about moving this to SampleProfileWriter so that we can use this API in other tooling where we don't invoke llvm-profdata?

1028

Perhaps wrap this in DEBUG(), I'm not sure whether it's useful to have this message all the time.

1077

An if statement with initialization would make the intent clearer here:

if (EC = File.error(); EC) {
}

1211

nit: Just heuristic instead of "heuristic algorithm"?

[huangjd]

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

[huangjd]

llvm/include/llvm/ProfileData/SampleProfWriter.h

133

It is not called by derived classes. It is called by llvm-profdata if the writer needs to be reused

[snehasish]

llvm/include/llvm/ProfileData/SampleProfWriter.h

133

My comment was to add some documentation for future implementations which derive from SampleProfileWriter. For example, in this patch

SampleProfileWriterText::reset calls SampleProfileWriter::reset(OS) on L110.
SampleProfileWriterBinary::reset calls SampleProfileWriter::reset(OS) on L142.

I hope that clarifies the suggestion.

[huangjd]

Add comment clarifying reset() usage

[snehasish]

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

981

Re-opening since I think we could move the RewriteProfileSizeLimit and CalculateNumFunctionsToRemove method to SampleProfileWriter as well so that this heuristic (and subsequent updates to it) can be reused in internal tooling directly.

[huangjd]

Refactor: moved implementation to llvm lib so that it can be used by other tools

[snehasish]

lgtm

Please wait a bit to see if others have additional comments. Thanks!

llvm/include/llvm/ProfileData/SampleProfWriter.h

45

Perhaps specify that this is in bytes either in the variable name or a comment?

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

1209

I think this is generally useful and we should make it visible.

[huangjd]

Clarified comments

[thakis]

This breaks tests on Mac: http://45.33.8.238/macm1/52334/step_11.txt

Please take a look and revert for now if it takes a while to fix.

[thakis]

Windows too: http://45.33.8.238/win/72992/step_11.txt

[dyung]

@huangjd the test you added seems to be failing on Windows. Can you take a look and revert if you need time to investigate?

https://lab.llvm.org/buildbot/#/builders/216/builds/15534

[chapuni]

llvm/lib/ProfileData/SampleProfWriter.cpp

125

FYI, OriginalFunctionCount was unused but fixed in rG9f4a9d3f4450

127

IterationCount is used only here.

[dyung]

llvm/test/tools/llvm-profdata/output-size-limit.test

62

On Windows this seems to be expanded in a way you probably did not expect:
(https://lab.llvm.org/buildbot/#/builders/216/builds/15553/steps/7/logs/FAIL__LLVM__output-size-limit_test)
"$(stat" "-c" "%s" "Z:\test\build\test\tools\llvm-profdata\Output\output-size-limit.test.tmp.output)"

[huangjd]

In D139603#4038245, @dyung wrote:

@huangjd the test you added seems to be failing on Windows. Can you take a look and revert if you need time to investigate?

https://lab.llvm.org/buildbot/#/builders/216/builds/15534

How to do size check in a cross platform way?

llvm/include/llvm/ProfileData/SampleProfWriter.h

45

Clarified comment in constructor

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

981

CalculateNumFunctionToRemove is moved inside FunctionPruningStrategy since it won't be used anywhere else. It can be overriden if necessary

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.
+  /// \p OutputSizeLimit Size limit 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
+  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);
 
   /// Create a new stream writer based on the value of \p Format.
   /// For testing.
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(std::unique_ptr<raw_ostream> &OS, SampleProfileFormat Format);
 
   virtual void setProfileSymbolList(ProfileSymbolList *PSL) {}
   virtual void setToCompressAllSections() {}
   virtual void setUseMD5() {}
   virtual void setPartialProfile() {}
   virtual void resetSecLayout(SectionLayout SL) {}
 
 protected:
   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;
 

 SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
   assert(!Context.hasContext() && "cs profile is not supported");
   return writeNameIdx(Context.getName());
 }
 
 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
   auto &NTable = getNameTable();
   const auto &Ret = NTable.find(FName);
   if (Ret == NTable.end())
     return sampleprof_error::truncated_name_table;
   encodeULEB128(Ret->second, *OutputStream);
   return sampleprof_error::success;
 }
 
 void SampleProfileWriterBinary::addName(StringRef FName) {
   auto &NTable = getNameTable();
   NTable.insert(std::make_pair(FName, 0));
 }
 
 void SampleProfileWriterBinary::addContext(const SampleContext &Context) {
   addName(Context.getName());
 }
 
 void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
   // Add all the names in indirect call targets.
for (const auto &FS : J.second) {
       const FunctionSamples &CalleeSamples = FS.second;
       addName(CalleeSamples.getName());
       addNames(CalleeSamples);
     }
 }
 
 void SampleProfileWriterExtBinaryBase::addContext(
     const SampleContext &Context) {
   if (Context.hasContext()) {
     for (auto &Callsite : Context.getContextFrames())
       SampleProfileWriterBinary::addName(Callsite.FuncName);
     CSNameTable.insert(std::make_pair(Context, 0));
   } else {
     SampleProfileWriterBinary::addName(Context.getName());
   }
 }
 
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.
+
+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: test $(stat -c %%s %t.output) -le 489
+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: test $(stat -c %%s %t.output) -le 488
+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: test $(stat -c %%s %t.output) -le 474
+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) {
     auto ReaderOrErr = SampleProfileReader::create(Input.Filename, Context,
                                                    FSDiscriminatorPassOption);
     if (std::error_code EC = ReaderOrErr.getError()) {
       warnOrExitGivenError(FailMode, EC, Input.Filename);
       continue;
     }
 
     // We need to keep the readers around until after all the files are
     // read so that we do not lose the function names stored in each
     // reader's memory. The function names are needed to write out the
     // merged profile map.
     Readers.push_back(std::move(ReaderOrErr.get()));
     const auto Reader = Readers.back().get();
     if (std::error_code EC = Reader->read()) {
       warnOrExitGivenError(FailMode, EC, Input.Filename);
       Readers.pop_back();
       continue;
     }
if (!DropProfileSymbolList) {
         WriterList.merge(*ReaderList);
     }
   }
 
   if (ProfileIsCS && (SampleMergeColdContext || SampleTrimColdContext)) {
     // Use threshold calculated from profile summary unless specified.
     SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
     auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
     uint64_t SampleProfColdThreshold =
         ProfileSummaryBuilder::getColdCountThreshold(
             (Summary->getDetailedSummary()));
 
     // Trim and merge cold context profile using cold threshold above;
     SampleContextTrimmer(ProfileMap)
         .trimAndMergeColdContextProfiles(
             SampleProfColdThreshold, SampleTrimColdContext,
             SampleMergeColdContext, SampleColdContextFrameDepth, false);
   }
 
   if (ProfileIsCS && GenCSNestedProfile) {
     CSProfileConverter CSConverter(ProfileMap);
     CSConverter.convertProfiles();
     ProfileIsCS = FunctionSamples::ProfileIsCS = false;
   }
 
+  // If limiting the output size, write to a string buffer first, and drop
+  // functions if the output size exceeds limit. This iterates multiple times
+  // until the limit is satisfied.
+  SmallVector<char> StringBuffer;
+  std::unique_ptr<raw_ostream> BufferStream(
+      new raw_svector_ostream(StringBuffer));
+
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
   auto Writer = std::move(WriterOrErr.get());
   // WriterList will have StringRef refering to string in Buffer.
   // Make sure Buffer lives as long as WriterList.
   auto Buffer = getInputFileBuf(ProfileSymbolListFile);
   handleExtBinaryWriter(*Writer, OutputFormat, Buffer.get(), WriterList,
                         CompressAllSections, UseMD5, GenPartialProfile);
-  if (std::error_code EC = Writer->write(ProfileMap))
+
+  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;
   if (WeightStr.getAsInteger(10, Weight) || Weight < 1)
     exitWithError("input weight must be a positive integer");
 
   return {std::string(FileName), Weight};
 }
 
 static void addWeightedInput(WeightedFileVector &WNI, const WeightedFile &WF) {
   StringRef Filename = WF.Filename;
   uint64_t Weight = WF.Weight;
 
   // If it's STDIN just pass it on.
   if (Filename == "-") {
     WNI.push_back({std::string(Filename), Weight});
     return;
   }
 
   llvm::sys::fs::file_status Status;
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) {