[llvm-profdata] Add support for weighted merge of profile data

This change adds support for an optional weight when merging profile data with the llvm-profdata tool.
Weights are specified by adding an option ':<weight>' suffix to the input file names.

Adding support for arbitrary weighting of input profile data allows for relative importance to be placed on the
input data from multiple training runs.

Both sampled and instrumented profiles are supported.

Reviewers: dnovillo, bogner, davidxl

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D14547

llvm-svn: 254669

1 files changed, 54 insertions, 16 deletions

diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 10b6855233d..56c80f518ec 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/ProfileData/InstrProfReader.h"
@@ -27,6 +28,7 @@
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/raw_ostream.h"
+#include <tuple>
 
 using namespace llvm;
 
@@ -93,7 +95,17 @@ static void handleMergeWriterError(std::error_code &Error,
   }
 }
 
-static void mergeInstrProfile(const cl::list<std::string> &Inputs,
+struct WeightedFile {
+  StringRef Filename;
+  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) {
   if (OutputFilename.compare("-") == 0)
@@ -109,21 +121,21 @@ static void mergeInstrProfile(const cl::list<std::string> &Inputs,
 
   InstrProfWriter Writer;
   SmallSet<std::error_code, 4> WriterErrorCodes;
-  for (const auto &Filename : Inputs) {
-    auto ReaderOrErr = InstrProfReader::create(Filename);
+  for (const auto &Input : Inputs) {
+    auto ReaderOrErr = InstrProfReader::create(Input.Filename);
     if (std::error_code ec = ReaderOrErr.getError())
-      exitWithErrorCode(ec, Filename);
+      exitWithErrorCode(ec, Input.Filename);
 
     auto Reader = std::move(ReaderOrErr.get());
     for (auto &I : *Reader) {
-      if (std::error_code EC = Writer.addRecord(std::move(I))) {
+      if (std::error_code EC = Writer.addRecord(std::move(I), Input.Weight)) {
         // Only show hint the first time an error occurs.
         bool firstTime = WriterErrorCodes.insert(EC).second;
-        handleMergeWriterError(EC, Filename, I.Name, firstTime);
+        handleMergeWriterError(EC, Input.Filename, I.Name, firstTime);
       }
     }
     if (Reader->hasError())
-      exitWithErrorCode(Reader->getError(), Filename);
+      exitWithErrorCode(Reader->getError(), Input.Filename);
   }
   if (OutputFormat == PF_Text)
     Writer.writeText(Output);
@@ -135,7 +147,7 @@ static sampleprof::SampleProfileFormat FormatMap[] = {
     sampleprof::SPF_None, sampleprof::SPF_Text, sampleprof::SPF_Binary,
     sampleprof::SPF_GCC};
 
-static void mergeSampleProfile(const cl::list<std::string> &Inputs,
+static void mergeSampleProfile(const WeightedFileVector &Inputs,
                                StringRef OutputFilename,
                                ProfileFormat OutputFormat) {
   using namespace sampleprof;
@@ -147,11 +159,11 @@ static void mergeSampleProfile(const cl::list<std::string> &Inputs,
   auto Writer = std::move(WriterOrErr.get());
   StringMap<FunctionSamples> ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
-  for (const auto &Filename : Inputs) {
+  for (const auto &Input : Inputs) {
     auto ReaderOrErr =
-        SampleProfileReader::create(Filename, getGlobalContext());
+        SampleProfileReader::create(Input.Filename, getGlobalContext());
     if (std::error_code EC = ReaderOrErr.getError())
-      exitWithErrorCode(EC, Filename);
+      exitWithErrorCode(EC, Input.Filename);
 
     // 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
@@ -160,7 +172,7 @@ static void mergeSampleProfile(const cl::list<std::string> &Inputs,
     Readers.push_back(std::move(ReaderOrErr.get()));
     const auto Reader = Readers.back().get();
     if (std::error_code EC = Reader->read())
-      exitWithErrorCode(EC, Filename);
+      exitWithErrorCode(EC, Input.Filename);
 
     StringMap<FunctionSamples> &Profiles = Reader->getProfiles();
     for (StringMap<FunctionSamples>::iterator I = Profiles.begin(),
@@ -168,15 +180,38 @@ static void mergeSampleProfile(const cl::list<std::string> &Inputs,
          I != E; ++I) {
       StringRef FName = I->first();
       FunctionSamples &Samples = I->second;
-      ProfileMap[FName].merge(Samples);
+      ProfileMap[FName].merge(Samples, Input.Weight);
     }
   }
   Writer->write(ProfileMap);
 }
 
+static void parseInputFiles(const cl::list<std::string> &Inputs,
+                            WeightedFileVector &WeightedInputs) {
+  WeightedInputs.reserve(Inputs.size());
+
+  for (StringRef Input : Inputs) {
+    StringRef FileName;
+    StringRef WeightStr;
+    std::tie(FileName, WeightStr) = Input.rsplit(':');
+    if (WeightStr.empty() || sys::fs::exists(Input)) {
+      // No weight specified or valid path containing delimiter.
+      WeightedInputs.push_back(WeightedFile(Input, 1));
+    } else {
+      // Input weight specified.
+      uint64_t Weight;
+      if (WeightStr.getAsInteger(10, Weight) || Weight < 1) {
+        // Invalid input weight.
+        exitWithError("Input weight must be a positive integer.");
+      }
+      WeightedInputs.push_back(WeightedFile(FileName, Weight));
+    }
+  }
+}
+
 static int merge_main(int argc, const char *argv[]) {
   cl::list<std::string> Inputs(cl::Positional, cl::Required, cl::OneOrMore,
-                               cl::desc("<filenames...>"));
+                               cl::desc("<filename[:weight]...>"));
 
   cl::opt<std::string> OutputFilename("output", cl::value_desc("output"),
                                       cl::init("-"), cl::Required,
@@ -198,10 +233,13 @@ static int merge_main(int argc, const char *argv[]) {
 
   cl::ParseCommandLineOptions(argc, argv, "LLVM profile data merger\n");
 
+  WeightedFileVector WeightedInputs;
+  parseInputFiles(Inputs, WeightedInputs);
+
   if (ProfileKind == instr)
-    mergeInstrProfile(Inputs, OutputFilename, OutputFormat);
+    mergeInstrProfile(WeightedInputs, OutputFilename, OutputFormat);
   else
-    mergeSampleProfile(Inputs, OutputFilename, OutputFormat);
+    mergeSampleProfile(WeightedInputs, OutputFilename, OutputFormat);
 
   return 0;
 }