[BatchMode] Preserve invariant that shuffled batches are subsequences of inputs.

Author: graydon

lib/Driver/Compilation.cpp

@@ -15,6 +15,7 @@
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsDriver.h"
 #include "swift/Basic/Program.h"
+#include "swift/Basic/STLExtras.h"
 #include "swift/Basic/Statistic.h"
 #include "swift/Basic/TaskQueue.h"
 #include "swift/Basic/Version.h"
@@ -724,31 +725,31 @@ namespace driver {
         Batches.push_back(Comp.addJob(std::move(J)));
     }
 
-    /// 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());
-    }
-
-    /// Shuffle \p Batchable if -driver-batch-seed is nonzero.
-    void maybeShuffleBatchable(std::vector<const Job *> &Batchable) {
+    /// Build a vector of partition indices, one per Job: the i'th index says
+    /// which batch of the partition the i'th Job will be assigned to. If we are
+    /// shuffling due to -driver-batch-seed, the returned indices will not be
+    /// arranged in contiguous runs. We shuffle partition-indices here, not
+    /// elements themselves, to preserve the invariant that each batch is a
+    /// subsequence of the full set of inputs, not just a subset.
+    std::vector<size_t>
+    assignJobsToPartitions(size_t PartitionSize,
+                           size_t NumJobs) {
+      size_t Remainder = NumJobs % PartitionSize;
+      size_t TargetSize = NumJobs / PartitionSize;
+      std::vector<size_t> PartitionIndex;
+      PartitionIndex.reserve(NumJobs);
+      for (size_t P = 0; P < PartitionSize; ++P) {
+        // Spread remainder evenly across partitions by adding 1 to the target
+        // size of the first Remainder of them.
+        size_t FillCount = TargetSize + ((P < Remainder) ? 1 : 0);
+        std::fill_n(std::back_inserter(PartitionIndex), FillCount, P);
+      }
       if (Comp.BatchSeed != 0) {
         std::minstd_rand gen(Comp.BatchSeed);
-        std::shuffle(Batchable.begin(), Batchable.end(), gen);
+        std::shuffle(PartitionIndex.begin(), PartitionIndex.end(), gen);
       }
+      assert(PartitionIndex.size() == NumJobs);
+      return PartitionIndex;
     }
 
     /// Create \c NumberOfParallelCommands batches and assign each job to a
@@ -762,28 +763,28 @@ namespace driver {
       }
 
       assert(!Partition.empty());
-      maybeShuffleBatchable(Batchable);
-
-      size_t i = 0;
+      auto PartitionIndex = assignJobsToPartitions(Partition.size(),
+                                                   Batchable.size());
+      assert(PartitionIndex.size() == Batchable.size());
       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);
-        }
-      }
+      for_each(Batchable, PartitionIndex, [&](const Job *Cmd, size_t Idx) {
+          assert(Idx < Partition.size());
+          std::vector<const Job*> &P = Partition[Idx];
+          if (P.empty() || TC.jobsAreBatchCombinable(Comp, P[0], Cmd)) {
+            if (Comp.ShowJobLifecycle)
+              llvm::outs() << "Adding " << LogJob(Cmd)
+                           << " to batch " << Idx << '\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);
+          }
+        });
     }
 
     // FIXME: at the moment we're not passing OutputFileMaps to frontends, so

lib/Driver/ToolChain.cpp

@@ -280,6 +280,35 @@ mergeBatchInputs(ArrayRef<const Job *> jobs,
   return false;
 }
 
+/// Debugging only: return whether the set of \p jobs is an ordered subsequence
+/// of the sequence of top-level input files in the \c Compilation \p C.
+static bool
+jobsAreSubsequenceOfCompilationInputs(ArrayRef<const Job *> jobs,
+                                      Compilation &C) {
+  llvm::SmallVector<const Job *, 16> sortedJobs;
+  llvm::StringMap<const Job *> jobsByInput;
+  for (const Job *J : jobs) {
+    const CompileJobAction *CJA = cast<CompileJobAction>(&J->getSource());
+    const InputAction* IA = findSingleSwiftInput(CJA);
+    auto R = jobsByInput.insert(std::make_pair(IA->getInputArg().getValue(),
+                                               J));
+    assert(R.second);
+  }
+  for (const InputPair &P : C.getInputFiles()) {
+    auto I = jobsByInput.find(P.second->getValue());
+    if (I != jobsByInput.end()) {
+      sortedJobs.push_back(I->second);
+    }
+  }
+  if (sortedJobs.size() != jobs.size())
+    return false;
+  for (size_t i = 0; i < sortedJobs.size(); ++i) {
+    if (sortedJobs[i] != jobs[i])
+      return false;
+  }
+  return true;
+}
+
 /// Construct a \c BatchJob by merging the constituent \p jobs' CommandOutput,
 /// input \c Job and \c Action members. Call through to \c constructInvocation
 /// on \p BatchJob, to build the \c InvocationInfo.
@@ -290,6 +319,8 @@ ToolChain::constructBatchJob(ArrayRef<const Job *> jobs,
   if (jobs.empty())
     return nullptr;
 
+  assert(jobsAreSubsequenceOfCompilationInputs(jobs, C));
+
   // Synthetic OutputInfo is a slightly-modified version of the initial
   // compilation's OI.
   auto OI = C.getOutputInfo();

test/Driver/batch_mode_aux_file_order.swift

@@ -0,0 +1,20 @@
+// When multiple additional-outputs on the same command-line are no longer
+// supported (i.e. when we've moved to mandatory use of output file maps for
+// communicating multiple additional-outputs to frontends) this test will no
+// longer make sense, and should be removed.
+//
+// RUN: %empty-directory(%t)
+// RUN: touch %t/file-01.swift %t/file-02.swift %t/file-03.swift %t/file-04.swift %t/file-05.swift
+// RUN: touch %t/file-06.swift %t/file-07.swift %t/file-08.swift %t/file-09.swift
+//
+// RUN: %swiftc_driver -enable-batch-mode -driver-batch-seed 1 -driver-print-jobs -driver-skip-execution -j 3 -emit-module -module-name foo %t/file-01.swift %t/file-02.swift %t/file-03.swift %t/file-04.swift %t/file-05.swift %t/file-06.swift %t/file-07.swift %t/file-08.swift %t/file-09.swift >%t/out.txt
+// RUN: %FileCheck %s <%t/out.txt
+//
+// Each batch should get 3 primaries; check that each has 3 modules _in the same numeric order_.
+//
+// CHECK: {{.*}}/swift {{.*}}-primary-file {{[^ ]*}}/file-[[A1:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[A2:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[A3:[0-9]+]].swift
+// CHECK-SAME: -o {{.*}}/file-[[A1]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[A2]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[A3]]-{{[a-z0-9]+}}.swiftmodule
+// CHECK: {{.*}}/swift {{.*}}-primary-file {{[^ ]*}}/file-[[B1:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[B2:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[B3:[0-9]+]].swift
+// CHECK-SAME: -o {{.*}}/file-[[B1]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[B2]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[B3]]-{{[a-z0-9]+}}.swiftmodule
+// CHECK: {{.*}}/swift {{.*}}-primary-file {{[^ ]*}}/file-[[C1:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[C2:[0-9]+]].swift {{.*}}-primary-file {{[^ ]*}}/file-[[C3:[0-9]+]].swift
+// CHECK-SAME: -o {{.*}}/file-[[C1]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[C2]]-{{[a-z0-9]+}}.swiftmodule -o {{.*}}/file-[[C3]]-{{[a-z0-9]+}}.swiftmodule