Merge pull request #14732 from graydon/batch-mode-driver-work

swift-ci · web-flow · commit 2e2d16150e9f · 2018-02-20T22:33:43.000-08:00
diff --git a/include/swift/Driver/Compilation.h b/include/swift/Driver/Compilation.h
@@ -161,6 +161,9 @@ class Compilation {
   /// redundant work.
   bool EnableBatchMode;
 
+  /// Provides a randomization seed to batch-mode partitioning, for debugging.
+  unsigned BatchSeed;
+
   /// True if temporary files should not be deleted.
   bool SaveTemps;
 
@@ -203,6 +206,7 @@ class Compilation {
               unsigned NumberOfParallelCommands = 1,
               bool EnableIncrementalBuild = false,
               bool EnableBatchMode = false,
+              unsigned BatchSeed = 0,
               bool SkipTaskExecution = false,
               bool SaveTemps = false,
               bool ShowDriverTimeCompilation = false,
diff --git a/include/swift/Driver/Driver.h b/include/swift/Driver/Driver.h
@@ -156,6 +156,9 @@ class Driver {
   /// Indicates whether the driver should check that the input files exist.
   bool CheckInputFilesExist = true;
 
+  /// Provides a randomization seed to batch-mode partitioning, for debugging.
+  unsigned DriverBatchSeed = 0;
+
 public:
   Driver(StringRef DriverExecutable, StringRef Name,
          ArrayRef<const char *> Args, DiagnosticEngine &Diags);
diff --git a/include/swift/Option/Options.td b/include/swift/Option/Options.td
@@ -98,6 +98,9 @@ def driver_show_job_lifecycle : Flag<["-"], "driver-show-job-lifecycle">,
   HelpText<"Show every step in the lifecycle of driver jobs">;
 def driver_use_filelists : Flag<["-"], "driver-use-filelists">,
   InternalDebugOpt, HelpText<"Pass input files as filelists whenever possible">;
+def driver_batch_seed : Separate<["-"], "driver-batch-seed">,
+  InternalDebugOpt,
+  HelpText<"Use the given seed value to randomize batch-mode partitions">;
 
 def driver_always_rebuild_dependents :
   Flag<["-"], "driver-always-rebuild-dependents">, InternalDebugOpt,
diff --git a/lib/Driver/Compilation.cpp b/lib/Driver/Compilation.cpp
@@ -42,6 +42,10 @@
 
 #include "CompilationRecord.h"
 
+// Batch-mode has a sub-mode for testing that randomizes batch partitions,
+// by user-provided seed. That is the only thing randomized here.
+#include <random>
+
 using namespace swift;
 using namespace swift::sys;
 using namespace swift::driver;
@@ -97,6 +101,7 @@ Compilation::Compilation(DiagnosticEngine &Diags,
                          unsigned NumberOfParallelCommands,
                          bool EnableIncrementalBuild,
                          bool EnableBatchMode,
+                         unsigned BatchSeed,
                          bool SkipTaskExecution,
                          bool SaveTemps,
                          bool ShowDriverTimeCompilation,
@@ -112,6 +117,7 @@ Compilation::Compilation(DiagnosticEngine &Diags,
     SkipTaskExecution(SkipTaskExecution),
     EnableIncrementalBuild(EnableIncrementalBuild),
     EnableBatchMode(EnableBatchMode),
+    BatchSeed(BatchSeed),
     SaveTemps(SaveTemps),
     ShowDriverTimeCompilation(ShowDriverTimeCompilation),
     Stats(std::move(StatsReporter)) {
@@ -120,7 +126,8 @@ Compilation::Compilation(DiagnosticEngine &Diags,
 static bool writeFilelistIfNecessary(const Job *job, DiagnosticEngine &diags);
 
 using CommandSet = llvm::SmallPtrSet<const Job *, 16>;
-
+using CommandSetVector = llvm::SetVector<const Job*>;
+using BatchPartition = std::vector<std::vector<const Job*>>;
 
 using InputInfoMap = llvm::SmallMapVector<const llvm::opt::Arg *,
                                           CompileJobAction::InputInfo, 16>;
@@ -140,7 +147,7 @@ namespace driver {
     /// A temporary buffer to hold commands that were scheduled but haven't been
     /// added to the Task Queue yet, because we might try batching them together
     /// first.
-    CommandSet PendingExecution;
+    CommandSetVector PendingExecution;
 
     /// Set of synthetic BatchJobs that serve to cluster subsets of jobs waiting
     /// in PendingExecution. Also used to identify (then unpack) BatchJobs back
@@ -680,7 +687,8 @@ namespace driver {
 
     /// Insert all jobs in \p Cmds (of descriptive name \p Kind) to the \c
     /// TaskQueue, and clear \p Cmds.
-    void transferJobsToTaskQueue(CommandSet &Cmds, StringRef Kind) {
+    template <typename Container>
+    void transferJobsToTaskQueue(Container &Cmds, StringRef Kind) {
       for (const Job *Cmd : Cmds) {
         if (Comp.ShowJobLifecycle)
           llvm::outs() << "Adding " << Kind
@@ -694,8 +702,8 @@ namespace driver {
     /// Partition the jobs in \c PendingExecution into those that are \p
     /// Batchable and those that are \p NonBatchable, clearing \p
     /// PendingExecution.
-    void getPendingBatchableJobs(CommandSet &Batchable,
-                                 CommandSet &NonBatchable) {
+    void getPendingBatchableJobs(CommandSetVector &Batchable,
+                                 CommandSetVector &NonBatchable) {
       for (const Job *Cmd : PendingExecution) {
         if (Comp.getToolChain().jobIsBatchable(Comp, Cmd)) {
           if (Comp.ShowJobLifecycle)
@@ -710,49 +718,83 @@ namespace driver {
       PendingExecution.clear();
     }
 
-    /// If \p CurrentBatch is nonempty, construct a new \c BatchJob from its
+    /// If \p Batch is nonempty, construct a new \c BatchJob from its
     /// contents by calling \p ToolChain::constructBatchJob, then insert the
-    /// new \c BatchJob into \p Batches and clear \p CurrentBatch.
+    /// new \c BatchJob into \p Batches.
     void
-    formBatchJobFromCurrentBatch(CommandSet &Batches,
-                                 llvm::SetVector<const Job *> &CurrentBatch) {
-      if (CurrentBatch.empty())
+    formBatchJobFromPartitionBatch(std::vector<const Job *> &Batches,
+                                   std::vector<const Job *> const &Batch) {
+      if (Batch.empty())
         return;
       if (Comp.ShowJobLifecycle)
         llvm::outs() << "Forming batch job from "
-                     << CurrentBatch.size() << " constituents\n";
+                     << Batch.size() << " constituents\n";
       auto const &TC = Comp.getToolChain();
-      auto J = TC.constructBatchJob(CurrentBatch.getArrayRef(), Comp);
+      auto J = TC.constructBatchJob(Batch, Comp);
       if (J)
-        Batches.insert(Comp.addJob(std::move(J)));
-      CurrentBatch.clear();
+        Batches.push_back(Comp.addJob(std::move(J)));
     }
 
-    /// Return true iff \p Cmd can be expanded by \p CurrentBatch, meaning
-    /// that \p CurrentBatch is smaller than \p TargetBatchSize and \p Cmd
-    /// is batch-combinable with the equivalence class of \p CurrentBatch
-    /// (as represented by element 0 of \p CurrentBatch).
-    bool canExpandBatch(const Job *Cmd,
-                        llvm::SetVector<const Job *> &CurrentBatch,
-                        size_t TargetBatchSize) {
-      auto const &TC = Comp.getToolChain();
-      return (CurrentBatch.empty() ||
-              (TC.jobsAreBatchCombinable(Comp, Cmd, CurrentBatch[0]) &&
-               CurrentBatch.size() < TargetBatchSize));
+    /// Inspect current batch \p i of the \p Partition currently being built
+    /// and, if that batch is "full" (in the sense of holding an evenly-divided
+    /// portion of NumJobs) then advance \p i to the next batch index in the
+    /// partition.
+    void maybeAdvanceToNextPartition(size_t &i,
+                                     BatchPartition const &Partition,
+                                     size_t NumJobs) {
+      assert(i < Partition.size());
+      size_t Remainder = NumJobs % Partition.size();
+      size_t TargetSize = NumJobs / Partition.size();
+      // Spread remainder evenly across partitions by adding 1 to the target
+      // size of the first Remainder of them.
+      if (i < Remainder)
+        TargetSize++;
+      if (Partition[i].size() >= TargetSize)
+        ++i;
+      assert(i < Partition.size());
     }
 
-    /// If \p CurrentBatch can't be expanded with \p Cmd, form a new \c BatchJob
-    /// from \p CurrentBatch, add it to \p Batches, and reset\p CurrentBatch;
-    /// then in either case, insert \p Cmd into \p CurrentBatch.
-    void expandBatch(const Job *Cmd,
-                     CommandSet &Batches,
-                     llvm::SetVector<const Job *> &CurrentBatch,
-                     size_t TargetBatchSize) {
-      if (!canExpandBatch(Cmd, CurrentBatch, TargetBatchSize)) {
-        formBatchJobFromCurrentBatch(Batches, CurrentBatch);
+    /// Shuffle \p Batchable if -driver-batch-seed is nonzero.
+    void maybeShuffleBatchable(std::vector<const Job *> &Batchable) {
+      if (Comp.BatchSeed != 0) {
+        std::minstd_rand gen(Comp.BatchSeed);
+        std::shuffle(Batchable.begin(), Batchable.end(), gen);
+      }
+    }
+
+    /// Create \c NumberOfParallelCommands batches and assign each job to a
+    /// batch either filling each partition in order or, if seeded with a
+    /// nonzero value, pseudo-randomly (but determinstically and nearly-evenly).
+    void partitionIntoBatches(std::vector<const Job *> Batchable,
+                              BatchPartition &Partition) {
+      if (Comp.ShowJobLifecycle) {
+        llvm::outs() << "Found " << Batchable.size() << " batchable jobs\n";
+        llvm::outs() << "Forming into " << Partition.size() << " batches\n";
+      }
+
+      assert(Partition.size() > 0);
+      maybeShuffleBatchable(Batchable);
+
+      size_t i = 0;
+      auto const &TC = Comp.getToolChain();
+      for (const Job *Cmd : Batchable) {
+        maybeAdvanceToNextPartition(i, Partition, Batchable.size());
+        std::vector<const Job*> &P = Partition[i];
+        if (P.empty() || TC.jobsAreBatchCombinable(Comp, P[0], Cmd)) {
+          if (Comp.ShowJobLifecycle)
+            llvm::outs() << "Adding " << LogJob(Cmd)
+                         << " to batch " << i << '\n';
+          P.push_back(Cmd);
+        } else {
+          // Strange but theoretically possible that we have a batchable job
+          // that's not combinable with others; tack a new batch on for it.
+          if (Comp.ShowJobLifecycle)
+            llvm::outs() << "Adding " << LogJob(Cmd)
+                         << " to new batch " << Partition.size() << '\n';
+          Partition.push_back(std::vector<const Job*>());
+          Partition.back().push_back(Cmd);
+        }
       }
-      llvm::outs() << "Adding to batch: " << LogJob(Cmd) << "\n";
-      CurrentBatch.insert(Cmd);
     }
 
     /// Select jobs that are batch-combinable from \c PendingExecution, combine
@@ -768,25 +810,18 @@ namespace driver {
         return;
       }
 
-      // Partition the pending jobs.
-      CommandSet Batchable, NonBatchable, Batches;
+      // Split the batchable from non-batchable pending jobs.
+      CommandSetVector Batchable, NonBatchable;
       getPendingBatchableJobs(Batchable, NonBatchable);
-      size_t TargetBatchSize = Batchable.size() / Comp.NumberOfParallelCommands;
 
-      if (Comp.ShowJobLifecycle) {
-        llvm::outs() << "Found " << Batchable.size() << " batchable jobs\n";
-        llvm::outs() << "Aiming for batch size " << TargetBatchSize << '\n';
-      }
-
-      // Batch the batchable jobs.
-      llvm::SetVector<const Job *> CurrentBatch;
-      for (const Job *Cmd : Batchable) {
-        expandBatch(Cmd, Batches, CurrentBatch, TargetBatchSize);
-      }
+      // Partition the batchable jobs into sets.
+      BatchPartition Partition(Comp.NumberOfParallelCommands);
+      partitionIntoBatches(Batchable.takeVector(), Partition);
 
-      // Form a residual incomplete batch if any jobs remain.
-      if (!CurrentBatch.empty()) {
-        formBatchJobFromCurrentBatch(Batches, CurrentBatch);
+      // Construct a BatchJob from each batch in the partition.
+      std::vector<const Job *> Batches;
+      for (auto const &Batch : Partition) {
+        formBatchJobFromPartitionBatch(Batches, Batch);
       }
 
       // Save batches so we can locate and decompose them on task-exit.
diff --git a/lib/Driver/Driver.cpp b/lib/Driver/Driver.cpp
@@ -522,6 +522,12 @@ Driver::buildCompilation(const ToolChain &TC,
     ArgList->hasArg(options::OPT_driver_show_incremental);
   bool ShowJobLifecycle =
     ArgList->hasArg(options::OPT_driver_show_job_lifecycle);
+  if (const Arg *A = ArgList->getLastArg(options::OPT_driver_batch_seed)) {
+    if (StringRef(A->getValue()).getAsInteger(10, DriverBatchSeed)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(*ArgList), A->getValue());
+    }
+  }
 
   bool Incremental = ArgList->hasArg(options::OPT_incremental);
   if (ArgList->hasArg(options::OPT_whole_module_optimization)) {
@@ -684,6 +690,7 @@ Driver::buildCompilation(const ToolChain &TC,
                                                  NumberOfParallelCommands,
                                                  Incremental,
                                                  BatchMode,
+                                                 DriverBatchSeed,
                                                  DriverSkipExecution,
                                                  SaveTemps,
                                                  ShowDriverTimeCompilation,
diff --git a/lib/Driver/Job.cpp b/lib/Driver/Job.cpp
@@ -244,36 +244,44 @@ void Job::printSummary(raw_ostream &os) const {
   // from a JobAction that itself has InputActions sources, then we collect
   // those up. Otherwise it's more correct to talk about our inputs as the
   // outputs of our input-jobs.
-  SmallVector<std::string, 4> Inputs;
+  SmallVector<StringRef, 4> Inputs;
+  SmallVector<StringRef, 4> Outputs = getOutput().getPrimaryOutputFilenames();
 
   for (const Action *A : getSource().getInputs())
     if (const auto *IA = dyn_cast<InputAction>(A))
       Inputs.push_back(IA->getInputArg().getValue());
 
   for (const Job *J : getInputs())
-    for (const std::string &f : J->getOutput().getPrimaryOutputFilenames())
+    for (StringRef f : J->getOutput().getPrimaryOutputFilenames())
       Inputs.push_back(f);
 
   size_t limit = 3;
-  size_t actual = Inputs.size();
-  if (actual > limit) {
+  size_t actual_in = Inputs.size();
+  size_t actual_out = Outputs.size();
+  if (actual_in > limit) {
     Inputs.erase(Inputs.begin() + limit, Inputs.end());
   }
+  if (actual_out > limit) {
+    Outputs.erase(Outputs.begin() + limit, Outputs.end());
+  }
 
   os << "{" << getSource().getClassName() << ": ";
-  interleave(getOutput().getPrimaryOutputFilenames(),
+  interleave(Outputs,
              [&](const std::string &Arg) {
                os << llvm::sys::path::filename(Arg);
              },
              [&] { os << ' '; });
+  if (actual_out > limit) {
+    os << " ... " << (actual_out-limit) << " more";
+  }
   os << " <= ";
   interleave(Inputs,
              [&](const std::string &Arg) {
                os << llvm::sys::path::filename(Arg);
              },
              [&] { os << ' '; });
-  if (actual > limit) {
-    os << " ... " << (actual-limit) << " more";
+  if (actual_in > limit) {
+    os << " ... " << (actual_in-limit) << " more";
   }
   os << "}";
 }
diff --git a/test/Driver/batch_mode.swift b/test/Driver/batch_mode.swift
