[ProfileData] Fix data racing in merging indexed profiles

[xur]

Summary

There are some data races when merging multiple indexed profiles with multiple threads.
(1) the static variable RecordIndex in index profile reader needs to be thread_local.
Otherwise, multiple function entries with different hashes in the profile lead to segfault.
(2) thread partition in loading the profile does not seem to be right: we are submitting a separated thread for each input and multiple threads can access the same WriteContext.
This patch fixes the above issues.

[vsk]

lib/ProfileData/InstrProfReader.cpp

736

Why isn't this just a member variable?

If we used "thread_local", wouldn't there still be a bug if you instantiated and used two IndexedInstrProfReaders?

tools/llvm-profdata/llvm-profdata.cpp

230

I don't understand the problem here. Each time a WriterContext is accessed, the accessing thread immediately acquires a unique lock on the context. Where does the racing access occur?

[xur]

Integrated the review comments from vsk.

[vsk]

Thanks! Lgtm.

lib/ProfileData/InstrProfReader.cpp

Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
   if (Error E = Record.takeError())
     return error(std::move(E));
 
   Counts = Record.get().Counts;
   return success();
 }
 
 Error IndexedInstrProfReader::readNextRecord(NamedInstrProfRecord &Record) {
-  static unsigned RecordIndex = 0;
+  thread_local static unsigned RecordIndex = 0;
 
   ArrayRef<NamedInstrProfRecord> Data;
 
   Error E = Index->getRecords(Data);
   if (E)
     return error(std::move(E));
 
   Record = Data[RecordIndex++];
   if (RecordIndex >= Data.size()) {
     Index->advanceToNextKey();
     RecordIndex = 0;
   }
   return success();
 }

test/tools/llvm-profdata/Inputs/multiple-profdata-merge.proftext

+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
[phone]
+# Num Counters:
+4
+# Counter Values:
+700000
+700000
+0
+0
+
+foo
+# Func Hash:
[phone]
+# Num Counters:
+6
+# Counter Values:
+400000
+400000
+0
+0
+0
+0
+
+foo
+# Func Hash:
[phone]
+# Num Counters:
+3
+# Counter Values:
+200000
+200000
+0
+
+foo
+# Func Hash:
[phone]
+# Num Counters:
+6
+# Counter Values:
+0
+100000
+0
+0
+0
+0
+
+bar
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0
+
+foo2
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0
+
+foo3
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0
+
+foo4
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0
+
+foo5
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0
+
+foo1
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+100000
+
+main
+# Func Hash:
[phone]
+# Num Counters:
+2
+# Counter Values:
+1400000
+14
+

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

+Test multi-thread merge of multiple profdata files.
+
+RUN: llvm-profdata merge %p/Inputs/multiple-profdata-merge.proftext -o %t
+RUN: llvm-profdata merge -j 4 %t %t %t %t -o %t_2
+RUN: llvm-profdata show %t_2 | FileCheck %s
+
+; CHECK:Total functions: 11
+; CHECK:Maximum function count: 5600000
+; CHECK:Maximum internal block count: 2800000
+
+

tools/llvm-profdata/llvm-profdata.cpp

static void mergeInstrProfile(const WeightedFileVector &Inputs,
 
   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;
+    unsigned ChunkSize = Inputs.size() / NumThreads;
+    unsigned ChunkRemainder = Inputs.size() % NumThreads;
+    for (unsigned Ctx = 0; Ctx < NumThreads; Ctx++) {
+      Pool.async(
+          [&](unsigned Ctx) {
+            for (unsigned I = 0; I < ChunkSize; I++)
+              loadInput(Inputs[Ctx * ChunkSize + I], Contexts[Ctx].get());
+            if (Ctx < ChunkRemainder)
+              loadInput(Inputs[NumThreads * ChunkSize + Ctx],
+                        Contexts[Ctx].get());
+          },
+          Ctx);
     }
     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 {