[BatchMode] Add -driver-batch-count to allow overriding batch count inferred by -j

Author: graydon

include/swift/Driver/Compilation.h

@@ -174,6 +174,9 @@ class Compilation {
   /// Provides a randomization seed to batch-mode partitioning, for debugging.
   const unsigned BatchSeed;
 
+  /// Overrides parallelism level as count of batches, if in batch-mode.
+  const Optional<unsigned> BatchCount;
+
   /// In order to test repartitioning, set to true if
   /// -driver-force-one-batch-repartition is present.
   const bool ForceOneBatchRepartition = false;
@@ -228,6 +231,7 @@ class Compilation {
               bool EnableIncrementalBuild = false,
               bool EnableBatchMode = false,
               unsigned BatchSeed = 0,
+              Optional<unsigned> BatchCount = None,
               bool ForceOneBatchRepartition = false,
               bool SaveTemps = false,
               bool ShowDriverTimeCompilation = false,

include/swift/Driver/Driver.h

@@ -176,12 +176,6 @@ 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;
-
-  /// Forces a repartition for testing.
-  bool DriverForceOneBatchRepartition = false;
-
 public:
   Driver(StringRef DriverExecutable, StringRef Name,
          ArrayRef<const char *> Args, DiagnosticEngine &Diags);

include/swift/Option/Options.td

@@ -110,6 +110,9 @@ def driver_filelist_threshold_EQ : Joined<["-"], "driver-filelist-threshold=">,
 def driver_batch_seed : Separate<["-"], "driver-batch-seed">,
   InternalDebugOpt,
   HelpText<"Use the given seed value to randomize batch-mode partitions">;
+def driver_batch_count : Separate<["-"], "driver-batch-count">,
+  InternalDebugOpt,
+  HelpText<"Use the given number of batch-mode partitions, rather than default parallelism level">;
 def driver_force_one_batch_repartition : Flag<["-"], "driver-force-one-batch-repartition">,
   InternalDebugOpt,
   HelpText<"Force one batch repartitioning for testing">;

lib/Driver/Compilation.cpp

@@ -111,6 +111,7 @@ Compilation::Compilation(DiagnosticEngine &Diags,
                          bool EnableIncrementalBuild,
                          bool EnableBatchMode,
                          unsigned BatchSeed,
+                         Optional<unsigned> BatchCount,
                          bool ForceOneBatchRepartition,
                          bool SaveTemps,
                          bool ShowDriverTimeCompilation,
@@ -130,6 +131,7 @@ Compilation::Compilation(DiagnosticEngine &Diags,
         OutputCompilationRecordForModuleOnlyBuild),
     EnableBatchMode(EnableBatchMode),
     BatchSeed(BatchSeed),
+    BatchCount(BatchCount),
     ForceOneBatchRepartition(ForceOneBatchRepartition),
     SaveTemps(SaveTemps),
     ShowDriverTimeCompilation(ShowDriverTimeCompilation),
@@ -880,7 +882,9 @@ namespace driver {
         return;
       }
 
-      size_t NumPartitions = TQ->getNumberOfParallelTasks();
+      size_t NumPartitions = (Comp.BatchCount.hasValue() ?
+                              Comp.BatchCount.getValue() :
+                              TQ->getNumberOfParallelTasks());
       CommandSetVector Batchable, NonBatchable;
       std::vector<const Job *> Batches;
       bool PretendTheCommandLineIsTooLongOnce =

lib/Driver/Driver.cpp

@@ -643,6 +643,35 @@ static bool getFilelistThreshold(DerivedArgList &Args, size_t &FilelistThreshold
   return false;
 }
 
+static unsigned
+getDriverBatchSeed(llvm::opt::InputArgList &ArgList,
+                   DiagnosticEngine &Diags) {
+  unsigned DriverBatchSeed = 0;
+  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());
+    }
+  }
+  return DriverBatchSeed;
+}
+
+static Optional<unsigned>
+getDriverBatchCount(llvm::opt::InputArgList &ArgList,
+                    DiagnosticEngine &Diags)
+{
+  if (const Arg *A = ArgList.getLastArg(options::OPT_driver_batch_count)) {
+    unsigned Count = 0;
+    if (StringRef(A->getValue()).getAsInteger(10, Count)) {
+      Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,
+                     A->getAsString(ArgList), A->getValue());
+    } else {
+      return Count;
+    }
+  }
+  return None;
+}
+
 std::unique_ptr<Compilation>
 Driver::buildCompilation(const ToolChain &TC,
                          std::unique_ptr<llvm::opt::InputArgList> ArgList) {
@@ -663,13 +692,9 @@ 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());
-    }
-  }
-  DriverForceOneBatchRepartition =
+  unsigned DriverBatchSeed = getDriverBatchSeed(*ArgList, Diags);
+  Optional<unsigned> DriverBatchCount = getDriverBatchCount(*ArgList, Diags);
+  bool DriverForceOneBatchRepartition =
       ArgList->hasArg(options::OPT_driver_force_one_batch_repartition);
 
   bool Incremental = ArgList->hasArg(options::OPT_incremental);
@@ -837,6 +862,7 @@ Driver::buildCompilation(const ToolChain &TC,
                       Incremental,
                       BatchMode,
                       DriverBatchSeed,
+                      DriverBatchCount,
                       DriverForceOneBatchRepartition,
                       SaveTemps,
                       ShowDriverTimeCompilation,