[llvm-profdata] Rewrite algorithm to reduce profile size limit

[huangjd]

Summary

Using the redesigned interface to take a range of NameFunctionSamples in
SampleProfileWriter, the performance of the default algorithm to reduce
the profile is greatly improved. Now use binary search to converge faster
since the container is not modified each iteration.

[snehasish]

Can you add some numbers on the speedup to the patch description? Also how much is profile size reduction wall time as a fraction of the entire time?

[huangjd]

In D155770#4519278, @snehasish wrote:

Can you add some numbers on the speedup to the patch description? Also how much is profile size reduction wall time as a fraction of the entire time?

Reduce 50 MB compressed extbinary profile to 12 MB, previously takes 1 minute, after refactoring takes 34s.

What do you mean by profile size reduction as a fraction of the entire time?
Profile size reduction is usually done manually to satisfy certain objectives, it is not used automatically by the compiler. However we still want the time spent to be reasonable to the user.

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 {
+class DefaultFunctionPruningStrategy {
   std::vector<NameFunctionSamples> SortedFunctions;
+  size_t OutputSizeLimit;
+  size_t Index;
+  size_t Step;
 
 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;
+  operator const llvm::ArrayRef<NameFunctionSamples>() const {
+    return ArrayRef<NameFunctionSamples>(SortedFunctions.data(), Index);
+  }
+
+  bool operator()(size_t CurrentOutputSize);
 };
 
 /// Sample-based profile writer. Base class.
 class SampleProfileWriter {
 public:
   using FunctionSamplesVector = const llvm::ArrayRef<NameFunctionSamples>;
 
   virtual ~SampleProfileWriter() = default;
public:
   std::error_code write(const SampleProfileMap &ProfileMap) {
     std::vector<NameFunctionSamples> SortedProfiles;
     sortFuncProfiles(ProfileMap, SortedProfiles);
     return write(SortedProfiles);
   }
 
   virtual std::error_code write(FunctionSamplesVector &ProfileMap);
 
-  /// Write sample profiles up to given size limit, using the pruning strategy
-  /// to drop some functions if necessary.
+  /// Write sample profiles up to a given size limit, using a strategy to
+  /// iteratively reduce the amount of content being written.
+  /// The default strategy uses binary search on the vector of function samples
+  /// sorted by their total count, to find a threshold that only function
+  /// samples having a total count greater than which are written.
+  /// The strategy should be a class providing the following interfaces:
+  /// (1) A constructor taking a SampleProfileMap and a size_t for output size
+  /// limit. The profile map may be modified by this class.
+  /// (2) A type conversion function returning a representation of the current
+  /// sample profiles to be written.
+  /// (3) bool operator()(size_t), called after each iteration to write the
+  /// current sample profiles with the argument being the size written. This
+  /// function should either modify the profile map or the class' internal state
+  /// so that in the next iteration (2) returns a profile map (or a slice of it)
+  /// with fewer contents and return true, or return false to indicate it cannot
+  /// modify the profile map further.
+  /// (2) and (3) are called iteratively until either the strategy terminates,
+  /// or all function samples are removed. The strategy should ensure itself
+  /// will terminate.
+  /// Generally the strategy should sort the profile map and return an ArrayRef
+  /// of the sorted NameFunctionSamples for (2), as sample profiles are always
+  /// writen in sorted order, and by doing this it greatly improves performance
+  /// by not sorting the sample profiles in every iteration. However this
+  /// is not a mandatory requirement.
   ///
   /// \returns status code of the file update operation.
-  template <typename FunctionPruningStrategy = DefaultFunctionPruningStrategy>
+  template <typename Strategy = DefaultFunctionPruningStrategy>
   std::error_code writeWithSizeLimit(SampleProfileMap &ProfileMap,
                                      size_t OutputSizeLimit) {
-    FunctionPruningStrategy Strategy(ProfileMap, OutputSizeLimit);
-    return writeWithSizeLimitInternal(ProfileMap, OutputSizeLimit, &Strategy);
+    if (OutputSizeLimit == 0)
+      return write(ProfileMap);
+
+    Strategy S(ProfileMap, OutputSizeLimit);
+
+    size_t TotalSize;
+    std::unique_ptr<raw_ostream> OriginalOutputStream;
+    OutputStream.swap(OriginalOutputStream);
+    SmallVector<char> StringBuffer;
+    do {
+      // Write the current slice of the profile map to a memory buffer to check
+      // the output size.
+      StringBuffer.clear();
+      OutputStream.reset(new raw_svector_ostream(StringBuffer));
+      if (std::error_code EC = write(S))
+        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 calculating the
+      // total output size.
+#ifdef _WIN32
+      if (Format == SPF_Text)
+        TotalSize += LineCount;
+#endif
+
+      // If no function can be written, exit with error to indicate the output
+      // file is not generated, unless we started with an empty profile map.
+      if (ProfileMap.empty())
+        break;
+      else if ((Format == SPF_Text && LineCount == 0) ||
+               ((Format == SPF_Binary || Format == SPF_Ext_Binary)
+                && Summary->getNumFunctions() == 0))
+        return sampleprof_error::too_large;
+
+      // Update the strategy with the output size from this iteration.
+    } while (S(TotalSize));
+
+    // If the strategy cannot reduce the sample profiles within the size limit,
+    // signal an error.
+    if (TotalSize > OutputSizeLimit)
+      return sampleprof_error::too_large;
+
+    // Write the reduced sample profile back to the actual output.
+    OutputStream.swap(OriginalOutputStream);
+    OutputStream->write(StringBuffer.data(), StringBuffer.size());
+    return sampleprof_error::success;
   }
 
   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>>
SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
       : OutputStream(std::move(OS)) {}
 
   /// Write a file header for the profile file.
   virtual std::error_code writeHeader(FunctionSamplesVector &ProfileMap) = 0;
 
   // Write function profiles to the profile file.
   virtual std::error_code writeFuncProfiles(FunctionSamplesVector &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;
 

llvm/lib/ProfileData/SampleProfWriter.cpp

 
 } // namespace
 } // namespace endian
 } // namespace support
 } // namespace llvm
 
 DefaultFunctionPruningStrategy::DefaultFunctionPruningStrategy(
     SampleProfileMap &ProfileMap, size_t OutputSizeLimit)
-    : FunctionPruningStrategy(ProfileMap, OutputSizeLimit) {
+  : OutputSizeLimit(OutputSizeLimit),
+    Index(ProfileMap.size()), Step(std::numeric_limits<size_t>::max()) {
   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();
+bool DefaultFunctionPruningStrategy::operator()(size_t CurrentOutputSize) {
+  assert(SortedFunctions.size() != 0);
 
-  std::unique_ptr<raw_ostream> OriginalOutputStream;
-  OutputStream.swap(OriginalOutputStream);
-
-  size_t IterationCount = 0;
-  size_t TotalSize;
+  // Hitting the exact limit, and we can immediately exit.
+  if (CurrentOutputSize == OutputSizeLimit)
+    return false;
+
+  // For the first iteration we want to descend as much as possible  before
+  // switching to binary search, because writing a large sample profile is
+  // costly. We use an empirical model where the number of sample entries in a
+  // function is linear to its sample count, so the total length of the entire
+  // profile map is quadratic to the number of functions.
+  if (Index == SortedFunctions.size()) {
+    if (CurrentOutputSize < OutputSizeLimit)
+      return false;
+
+    // Usually specified size limit is much smaller than the output size of the
+    // full profile, so we estimate a cutoff point near the head of the sorted
+    // functions where the length of them is still above limit. In this case
+    // binary search only needs to be appled on a much smaller range.
+    size_t Estimate = (size_t)round(
+        SortedFunctions.size() *
+        (1.0 - sqrt(1.0 - (double)OutputSizeLimit / CurrentOutputSize)));
+    Estimate = std::min(Estimate, SortedFunctions.size() - 1);
+    Index = llvm::bit_floor(Estimate);
+    Step = std::numeric_limits<size_t>::max();
+    return true;
+  }
+
+  // This only happens after the first iteration.
+  if (Step == std::numeric_limits<size_t>::max()) {
+    // We only need to search in the subrange of [0..Index).
+    if (CurrentOutputSize > OutputSizeLimit)
+      Step = Index;
+    else {
+      // We overestimated CurrentOutputSize, so we need to reperform binary
+      // search on the entire sample map.
+      Index = 0;
+      Step = llvm::bit_floor(SortedFunctions.size()) * 2;
+    }
+  }
 
-  SmallVector<char> StringBuffer;
+  if (Step == 0)
+    return false;
+
+  // Binary search on each iteration.
+  if (CurrentOutputSize > OutputSizeLimit)
+    Index -= Step;
   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;
+    Step /= 2;
+  } while (Index + Step >= SortedFunctions.size());
+  Index += Step;
+  return true;
 }
 
 std::error_code
 SampleProfileWriter::writeFuncProfiles(FunctionSamplesVector &ProfileMap) {
   for (const auto &I : ProfileMap) {
     if (std::error_code EC = writeSample(*I.second))
       return EC;
   }