[llvm-profdata] Speed up merging by using a thread pool

[vsk]

Summary

Add a "-j" option to llvm-profdata to control the number of threads
used. Auto-detect NumThreads when it isn't specified, and avoid spawning
threads when they wouldn't be beneficial.

I tested this patch using a raw profile produced by clang (147MB). Here is the
time taken to merge 4 copies together on my laptop:

build	user	system	cpu	total
No thread pool	112.87s	5.92s	97%	2:01.08
With 2 threads	134.99s	26.54s	164%	1:33.31

[silvas]

The way you arrived at sqrt(N) seems sketchy to me.
Assuming infinite parallelism, the first step performs N/2 merges and then you recurse on the N/2 merged profiles:

numMerges(N) = N/2 + numMerges(N/2)
numSerialMerges(N) = 1 + numSerialMerges(N/2)

Which must work out to something like log(N) serial merges. I.e. a binary reduction.

[davidxl]

There isn't any reduction of the number of merges compared with serial case. Say there are T threads, the total number of merges is

T * f(N/T) + T - 1, so f(N) is still N-1.

I suspect a much larger speed up can be obtained by unifying symbol table creation across profiles (that are generated from the same binary -- as their symtab is identical). I was planning doing that but did not find time. If you can help with that, that will be a great win.

David

tools/llvm-profdata/llvm-profdata.cpp

196

Make this a regular function and move it out of line.

197

This does not look right. Just make it a controllable option (default with 2 or 4 etc).

202

Move struct definition out of the function.

[vsk]

In D22438#486320, @silvas wrote:

The way you arrived at sqrt(N) seems sketchy to me.

I didn't consider the option of splitting the input by twos. You're right that it's possible to do better than T - 1 + (N/T) serial merges.

Assuming infinite parallelism, the first step performs N/2 merges and then you recurse on the N/2 merged profiles [...]. Which must work out to something like log(N) serial merges.

In practice (dropping the assumption of infinite parallelism), I don't think we can get to log(N) serial merges. You have to at least visit each input once (using any one of T threads), so N/T serial operations looks like the lower bound.

That said, your suggested approach seems better and I'll take a shot at it.

In D22438#486320, @davidxl wrote:

There isn't any reduction of the number of merges compared with serial case. Say there are T threads, the total number of merges is: T * f(N/T) + T - 1, so f(N) is still N-1.

No disagreement here. We will count the same (or higher) number of merges in total when using T threads. It's strictly meant as a speed improvement.

I suspect a much larger speed up can be obtained by unifying symbol table creation across profiles (that are generated from the same binary -- as their symtab is identical). I was planning doing that but did not find time. If you can help with that, that will be a great win.

That sounds like a great idea. Did you have in mind some sort of cache of InstrProfSymtabs which can optionally be passed to the raw profile reader? I can investigate that after this patch is done.

[vsk]

• Address Sean and David's review comments.
• Add tests for {even, odd} number of threads x {even, odd} number of inputs.

[vsk]

Mark some inline comments as done.

[vsk]

• Remove stray semi-colon.

[davidxl]

lgtm

Probably also add a test as a follow up to test option -num-threads=1 and some other non-default value.

tools/llvm-profdata/llvm-profdata.cpp

173

--> Merge Src context into Dst.

[silvas]

LGTM.

[vsk]

In D22438#487497, @davidxl wrote:

lgtm

Probably also add a test as a follow up to test option -num-threads=1 and some other non-default value.

Thanks. I'll add the test.

docs/CommandGuide/llvm-profdata.rst

  Emit the profile using GCC's gcov format (Not yet supported).
 
 .. option:: -sparse[=true|false]
 
  Do not emit function records with 0 execution count. Can only be used in
  conjunction with -instr. Defaults to false, since it can inhibit compiler
  optimization during PGO.
 
+.. option:: -num-threads=N, -j=N
+
+ Use N threads to perform profile merging. When N=0, llvm-profdata auto-detects
+ an appropriate number of threads to use. This is the default.
+
 EXAMPLES
 ^^^^^^^^
 Basic Usage
 +++++++++++
 Merge three profiles:
 
 ::
 

include/llvm/ProfileData/InstrProfWriter.h

 public:
   InstrProfWriter(bool Sparse = false);
   ~InstrProfWriter();
 
   /// Add function counts for the given function. If there are already counts
   /// for this function and the hash and number of counts match, each counter is
   /// summed. Optionally scale counts by \p Weight.
   Error addRecord(InstrProfRecord &&I, uint64_t Weight = 1);
+  /// Merge existing function counts from the given writer.
+  Error mergeRecordsFromWriter(InstrProfWriter &&IPW);
   /// Write the profile to \c OS
   void write(raw_fd_ostream &OS);
   /// Write the profile in text format to \c OS
   void writeText(raw_fd_ostream &OS);
   /// Write \c Record in text format to \c OS
   static void writeRecordInText(const InstrProfRecord &Record,
                                 InstrProfSymtab &Symtab, raw_fd_ostream &OS);
   /// Write the profile, returning the raw data. For testing.

lib/ProfileData/InstrProfWriter.cpp

     Dest.merge(I, Weight);
   }
 
   Dest.sortValueData();
 
   return Dest.takeError();
 }
 
+Error InstrProfWriter::mergeRecordsFromWriter(InstrProfWriter &&IPW) {
+  for (auto &I : IPW.FunctionData)
+    for (auto &Func : I.getValue())
+      if (Error E = addRecord(std::move(Func.second), 1))
+        return E;
+  return Error::success();
+}
+
 bool InstrProfWriter::shouldEncodeData(const ProfilingData &PD) {
   if (!Sparse)
     return true;
   for (const auto &Func : PD) {
     const InstrProfRecord &IPR = Func.second;
     if (std::any_of(IPR.Counts.begin(), IPR.Counts.end(),
                     [](uint64_t Count) { return Count > 0; }))
       return true;

test/tools/llvm-profdata/multiple-inputs.test

 DISJOINT-1: Block counts: [2, 3]
 DISJOINT-2: bar:
 DISJOINT-2: Counters: 3
 DISJOINT-2: Function count: 1
 DISJOINT-2: Block counts: [2, 3]
 DISJOINT: Total functions: 2
 DISJOINT: Maximum function count: 1
 DISJOINT: Maximum internal block count: 3
+
+RUN: llvm-profdata merge %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     -num-threads 2 -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FOO4
+RUN: llvm-profdata merge %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     -j 3 -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FOO4
+FOO4: foo:
+FOO4: Counters: 3
+FOO4: Function count: 4
+FOO4: Block counts: [8, 12]
+FOO4: Total functions: 1
+FOO4: Maximum function count: 4
+FOO4: Maximum internal block count: 12
+
+RUN: llvm-profdata merge %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext -j 2 -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FOO5
+RUN: llvm-profdata merge %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext -j 3 -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FOO5
+RUN: llvm-profdata merge %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext %p/Inputs/foo3-1.proftext \
+RUN:                     %p/Inputs/foo3-1.proftext -o %t
+RUN: llvm-profdata show %t -all-functions -counts | FileCheck %s --check-prefix=FOO5
+FOO5: foo:
+FOO5: Counters: 3
+FOO5: Function count: 5
+FOO5: Block counts: [10, 15]
+FOO5: Total functions: 1
+FOO5: Maximum function count: 5
+FOO5: Maximum internal block count: 15

tools/llvm-profdata/llvm-profdata.cpp

 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 
 using namespace llvm;
 
 enum ProfileFormat { PF_None = 0, PF_Text, PF_Binary, PF_GCC };
 
 static void exitWithError(const Twine &Message, StringRef Whence = "",
struct WeightedFile {
   uint64_t Weight;
 
   WeightedFile() {}
 
   WeightedFile(StringRef F, uint64_t W) : Filename{F}, Weight{W} {}
 };
 typedef SmallVector<WeightedFile, 5> WeightedFileVector;
 
-static void mergeInstrProfile(const WeightedFileVector &Inputs,
-                              StringRef OutputFilename,
-                              ProfileFormat OutputFormat, bool OutputSparse) {
-  if (OutputFilename.compare("-") == 0)
-    exitWithError("Cannot write indexed profdata format to stdout.");
+/// Keep track of merged data and reported errors.
+struct WriterContext {
+  std::mutex Lock;
+  InstrProfWriter Writer;
+  Error Err;
+  StringRef ErrWhence;
+  std::mutex &ErrLock;
+  SmallSet<instrprof_error, 4> &WriterErrorCodes;
+
+  WriterContext(bool IsSparse, std::mutex &ErrLock,
+                SmallSet<instrprof_error, 4> &WriterErrorCodes)
+      : Lock(), Writer(IsSparse), Err(Error::success()), ErrWhence(""),
+        ErrLock(ErrLock), WriterErrorCodes(WriterErrorCodes) {}
+};
 
-  if (OutputFormat != PF_Binary && OutputFormat != PF_Text)
-    exitWithError("Unknown format is specified.");
+/// Load an input into a writer context.
+static void loadInput(const WeightedFile &Input, WriterContext *WC) {
+  std::unique_lock<std::mutex> CtxLock{WC->Lock};
 
-  std::error_code EC;
-  raw_fd_ostream Output(OutputFilename.data(), EC, sys::fs::F_None);
-  if (EC)
-    exitWithErrorCode(EC, OutputFilename);
+  // If there's a pending hard error, don't do more work.
+  if (WC->Err)
+    return;
+
+  WC->ErrWhence = Input.Filename;
 
-  InstrProfWriter Writer(OutputSparse);
-  SmallSet<instrprof_error, 4> WriterErrorCodes;
-  for (const auto &Input : Inputs) {
-    auto ReaderOrErr = InstrProfReader::create(Input.Filename);
-    if (Error E = ReaderOrErr.takeError())
-      exitWithError(std::move(E), Input.Filename);
+  if ((WC->Err = ReaderOrErr.takeError()))
+    return;
 
-    auto Reader = std::move(ReaderOrErr.get());
-    bool IsIRProfile = Reader->isIRLevelProfile();
-    if (Writer.setIsIRLevelProfile(IsIRProfile))
-      exitWithError("Merge IR generated profile with Clang generated profile.");
+  if (WC->Writer.setIsIRLevelProfile(IsIRProfile)) {
+    WC->Err = make_error<StringError>(
+        "Merge IR generated profile with Clang generated profile.",
+        std::error_code());
+    return;
+  }
 
-    for (auto &I : *Reader) {
-      if (Error E = Writer.addRecord(std::move(I), Input.Weight)) {
-        // Only show hint the first time an error occurs.
-        instrprof_error IPE = InstrProfError::take(std::move(E));
-        bool firstTime = WriterErrorCodes.insert(IPE).second;
-        handleMergeWriterError(make_error<InstrProfError>(IPE), Input.Filename,
-                               I.Name, firstTime);
-      }
-    }
-    if (Reader->hasError())
-      exitWithError(Reader->getError(), Input.Filename);
+    if (Error E = WC->Writer.addRecord(std::move(I), Input.Weight)) {
+      std::unique_lock<std::mutex> ErrorGuard{WC->ErrLock};
+      bool firstTime = WC->WriterErrorCodes.insert(IPE).second;
+    WC->Err = Reader->getError();
+}
+
+/// Merge two writer contexts together.
+static void mergeWriterContexts(WriterContext *Dst, WriterContext *Src) {
+  if (Error E = Dst->Writer.mergeRecordsFromWriter(std::move(Src->Writer)))
+    Dst->Err = std::move(E);
+};
+
+static void mergeInstrProfile(const WeightedFileVector &Inputs,
+                              StringRef OutputFilename,
+                              ProfileFormat OutputFormat, bool OutputSparse,
+                              unsigned NumThreads) {
+  if (OutputFilename.compare("-") == 0)
+    exitWithError("Cannot write indexed profdata format to stdout.");
+
+  if (OutputFormat != PF_Binary && OutputFormat != PF_Text)
+    exitWithError("Unknown format is specified.");
+
+  std::error_code EC;
+  raw_fd_ostream Output(OutputFilename.data(), EC, sys::fs::F_None);
+  if (EC)
+    exitWithErrorCode(EC, OutputFilename);
+
+  std::mutex ErrorLock;
+  SmallSet<instrprof_error, 4> WriterErrorCodes;
+
+  // If NumThreads is not specified, auto-detect a good default.
+  if (NumThreads == 0)
+    NumThreads = std::max(1U, std::min(std::thread::hardware_concurrency(),
+                                       unsigned(Inputs.size() / 2)));
+
+  // Initialize the writer contexts.
+  SmallVector<std::unique_ptr<WriterContext>, 4> Contexts;
+  for (unsigned I = 0; I < NumThreads; ++I)
+    Contexts.emplace_back(llvm::make_unique<WriterContext>(
+        OutputSparse, ErrorLock, WriterErrorCodes));
+
+  if (NumThreads == 1) {
+    for (const auto &Input : Inputs)
+      loadInput(Input, Contexts[0].get());
+  } else {
+    ThreadPool Pool(NumThreads);
+
+    // Load the inputs in parallel (N/NumThreads serial steps).
+    unsigned Ctx = 0;
+    for (const auto &Input : Inputs) {
+      Pool.async(loadInput, Input, Contexts[Ctx].get());
+      Ctx = (Ctx + 1) % NumThreads;
     }
+    Pool.wait();
+
+    // Merge the writer contexts together (lg(NumThreads) serial steps).
+    unsigned Mid = Contexts.size() / 2;
+    unsigned End = Contexts.size();
+    assert(Mid > 0 && "Expected more than one context");
+    do {
+      for (unsigned I = 0; I < Mid; ++I)
+        Pool.async(mergeWriterContexts, Contexts[I].get(),
+                   Contexts[I + Mid].get());
+      if (End & 1)
+        Pool.async(mergeWriterContexts, Contexts[0].get(),
+                   Contexts[End - 1].get());
+      Pool.wait();
+      End = Mid;
+      Mid /= 2;
+    } while (Mid > 0);
+  }
+
+  // Handle deferred hard errors encountered during merging.
+  for (std::unique_ptr<WriterContext> &WC : Contexts)
+    if (WC->Err)
+      exitWithError(std::move(WC->Err), WC->ErrWhence);
+
+  InstrProfWriter &Writer = Contexts[0]->Writer;
   if (OutputFormat == PF_Text)
     Writer.writeText(Output);
   else
     Writer.write(Output);
 }
 
 static sampleprof::SampleProfileFormat FormatMap[] = {
     sampleprof::SPF_None, sampleprof::SPF_Text, sampleprof::SPF_Binary,
cl::opt<ProfileFormat> OutputFormat(
       cl::desc("Format of output profile"), cl::init(PF_Binary),
       cl::values(clEnumValN(PF_Binary, "binary", "Binary encoding (default)"),
                  clEnumValN(PF_Text, "text", "Text encoding"),
                  clEnumValN(PF_GCC, "gcc",
                             "GCC encoding (only meaningful for -sample)"),
                  clEnumValEnd));
   cl::opt<bool> OutputSparse("sparse", cl::init(false),
       cl::desc("Generate a sparse profile (only meaningful for -instr)"));
+  cl::opt<unsigned> NumThreads(
+      "num-threads", cl::init(0),
+      cl::desc("Number of merge threads to use (default: autodetect)"));
+  cl::alias NumThreadsA("j", cl::desc("Alias for --num-threads"),
+                        cl::aliasopt(NumThreads));
 
   cl::ParseCommandLineOptions(argc, argv, "LLVM profile data merger\n");
 
   WeightedFileVector WeightedInputs;
   for (StringRef Filename : InputFilenames)
     WeightedInputs.emplace_back(Filename, 1);
   for (StringRef WeightedFilename : WeightedInputFilenames)
     WeightedInputs.emplace_back(parseWeightedFile(WeightedFilename));
static int merge_main(int argc, const char *argv[]) {
   if (DumpInputFileList) {
     for (auto &WF : WeightedInputs)
       outs() << WF.Weight << "," << WF.Filename << "\n";
     return 0;
   }
 
   if (ProfileKind == instr)
     mergeInstrProfile(WeightedInputs, OutputFilename, OutputFormat,
-                      OutputSparse);
+                      OutputSparse, NumThreads);
   else
     mergeSampleProfile(WeightedInputs, OutputFilename, OutputFormat);
 
   return 0;
 }
 
 static int showInstrProfile(const std::string &Filename, bool ShowCounts,
                             bool ShowIndirectCallTargets,

unittests/ProfileData/InstrProfTest.cpp

   MD = M.getProfileSummary();
   ASSERT_TRUE(MD);
   PSFromMD = ProfileSummary::getFromMD(MD);
   ASSERT_TRUE(PSFromMD);
   VerifySummary(*PSFromMD);
   delete PSFromMD;
 }
 
+TEST_F(InstrProfTest, test_writer_merge) {
+  InstrProfRecord Record1("func1", 0x1234, {42});
+  NoError(Writer.addRecord(std::move(Record1)));
+
+  InstrProfWriter Writer2;
+  InstrProfRecord Record2("func2", 0x1234, {0, 0});
+  NoError(Writer2.addRecord(std::move(Record2)));
+
+  NoError(Writer.mergeRecordsFromWriter(std::move(Writer2)));
+
+  auto Profile = Writer.writeBuffer();
+  readProfile(std::move(Profile));
+
+  Expected<InstrProfRecord> R = Reader->getInstrProfRecord("func1", 0x1234);
+  ASSERT_TRUE(NoError(R.takeError()));
+  ASSERT_EQ(1U, R->Counts.size());
+  ASSERT_EQ(42U, R->Counts[0]);
+
+  R = Reader->getInstrProfRecord("func2", 0x1234);
+  ASSERT_TRUE(NoError(R.takeError()));
+  ASSERT_EQ(2U, R->Counts.size());
+  ASSERT_EQ(0U, R->Counts[0]);
+  ASSERT_EQ(0U, R->Counts[1]);
+}
+
 static const char callee1[] = "callee1";
 static const char callee2[] = "callee2";
 static const char callee3[] = "callee3";
 static const char callee4[] = "callee4";
 static const char callee5[] = "callee5";
 static const char callee6[] = "callee6";
 
 TEST_P(MaybeSparseInstrProfTest, get_icall_data_read_write) {