[WIP] Whole module optimization mode

Author: DougGregor

Sources/SwiftDriver/Jobs/CompileJob.swift

@@ -54,8 +54,10 @@ extension Driver {
       }
       commandLine.append(.path(input.file))
 
-      // If there is a primary output, add it.
-      if isPrimary, let compilerOutputType = compilerOutputType {
+      // If there is a primary output or we are doing multithreaded compiles,
+      // add an output for the input.
+      if isPrimary || numThreads > 0,
+          let compilerOutputType = compilerOutputType {
         let output = (outputFileMap ?? OutputFileMap()).getOutput(
           inputFile: input.file,
           outputType: compilerOutputType
@@ -64,6 +66,15 @@ extension Driver {
       }
     }
 
+    // When not using primary file inputs or multithreading, add a single output.
+    if !usesPrimaryFileInputs && numThreads == 0,
+        let outputType = compilerOutputType {
+      let existingOutputPath = outputFileMap?.existingOutputForSingleInput(
+          outputType: outputType)
+      let output = existingOutputPath ?? VirtualPath.temporary(.init(moduleName.appendingFileTypeExtension(outputType)))
+      primaryOutputs.append(TypedVirtualPath(file: output, type: outputType))
+    }
+
     return primaryOutputs
   }
 

Sources/SwiftDriver/Jobs/EmitModuleJob.swift

@@ -32,7 +32,7 @@ extension Driver {
     try addCommonFrontendOptions(commandLine: &commandLine)
     // FIXME: Add MSVC runtime library flags
 
-    // Add suppplementable outputs.
+    // Add supplemental outputs.
     func addSupplementalOutput(path: VirtualPath?, flag: String, type: FileType) {
       guard let path = path else { return }
 

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -164,40 +164,75 @@ extension Driver {
   mutating func addFrontendSupplementaryOutputArguments(commandLine: inout [Job.ArgTemplate], primaryInputs: [TypedVirtualPath]) throws -> [TypedVirtualPath] {
     var outputs: [TypedVirtualPath] = []
 
-    func addOutputsOfType(outputType: FileType, input: VirtualPath, flag: String) {
+    /// Add output of a particular type, if needed.
+    func addOutputOfType(
+        outputType: FileType, finalOutputPath: VirtualPath?,
+        input: VirtualPath?, flag: String
+    ) {
+      // If there is no final output, there's nothing to do.
+      guard let finalOutputPath = finalOutputPath else { return }
+
+      // If the whole of the compiler output is this type, there's nothing to
+      // do.
       if outputType == compilerOutputType { return }
 
+      // Add the appropriate flag.
       commandLine.appendFlag(flag)
 
-      let path = (outputFileMap ?? OutputFileMap())
-        .getOutput(inputFile: input, outputType: outputType)
-      commandLine.append(.path(path))
-      outputs.append(TypedVirtualPath(file: path, type: outputType))
-    }
-
-    for input in primaryInputs {
-      if moduleOutput != nil && !forceEmitModuleInSingleInvocation {
-        addOutputsOfType(outputType: .swiftModule, input: input.file, flag: "-emit-module-path")
+      // Compute the output path based on the input path (if there is one), or
+      // use the final output.
+      let outputPath: VirtualPath
+      if let input = input {
+        outputPath = (outputFileMap ?? OutputFileMap())
+          .getOutput(inputFile: input, outputType: outputType)
+      } else {
+        outputPath = finalOutputPath
       }
 
-      if moduleDocOutputPath != nil && !forceEmitModuleInSingleInvocation {
-        addOutputsOfType(outputType: .swiftDocumentation, input: input.file, flag: "-emit-module-doc-path")
-      }
+      commandLine.append(.path(outputPath))
+      outputs.append(TypedVirtualPath(file: outputPath, type: outputType))
+    }
 
-      if dependenciesFilePath != nil && !forceEmitModuleInSingleInvocation {
-        addOutputsOfType(outputType: .dependencies, input: input.file, flag: "-emit-dependencies-path")
+    /// Add all of the outputs needed for a given input.
+    func addAllOutputsFor(input: VirtualPath?) {
+      if !forceEmitModuleInSingleInvocation {
+        addOutputOfType(
+            outputType: .swiftModule,
+            finalOutputPath: moduleOutput?.outputPath,
+            input: input,
+            flag: "-emit-module-path")
+        addOutputOfType(
+            outputType: .swiftDocumentation,
+            finalOutputPath: moduleDocOutputPath,
+            input: input,
+            flag: "-emit-module-doc-path")
+        addOutputOfType(
+            outputType: .dependencies,
+            finalOutputPath: dependenciesFilePath,
+            input: input,
+            flag: "-emit-dependencies-path")
       }
 
-      if optimizationRecordPath != nil {
-        addOutputsOfType(outputType: .optimizationRecord, input: input.file, flag: "-save-optimization-record-path")
-      }
+      addOutputOfType(
+          outputType: .optimizationRecord,
+          finalOutputPath: optimizationRecordPath,
+          input: input,
+          flag: "-save-optimization-record-path")
 
       #if false
       // FIXME: handle -update-code
-      addOutputsOfType(outputType: .remap, input: input.file, flag: "-emit-remap-file-path")
+      addOutputOfType(outputType: .remap, input: input.file, flag: "-emit-remap-file-path")
       #endif
     }
 
+    if compilerMode.usesPrimaryFileInputs {
+      for input in primaryInputs {
+        addAllOutputsFor(input: input.file)
+      }
+    } else {
+      addAllOutputsFor(input: nil)
+    }
+
     return outputs
   }
 }

Sources/SwiftDriver/Jobs/Planning.swift

@@ -40,15 +40,33 @@ extension Driver {
     }
 
     let partitions: BatchPartitions?
-    if case let .batchCompile(batchInfo) = compilerMode {
+    switch compilerMode {
+    case .batchCompile(let batchInfo):
       partitions = batchPartitions(batchInfo)
-    } else {
+
+    case .immediate, .repl:
+      fatalError("immediate and REPL modes are currently unsupported")
+
+    case .singleCompile:
+      // Create a single compile job for all of the files, none of which
+      // are primary.
+      var jobOutputs: [TypedVirtualPath] = []
+      let job = try compileJob(primaryInputs: [], outputType: compilerOutputType, allOutputs: &jobOutputs)
+      jobs.append(job)
+      addJobOutputs(jobOutputs)
+
+      partitions = nil
+
+    case .standardCompile:
       partitions = nil
     }
 
     for input in inputFiles {
       switch input.type {
       case .swift, .sil, .sib:
+        // Generate a compile job for primary inputs here.
+        guard compilerMode.usesPrimaryFileInputs else { break }
+
         var primaryInputs: [TypedVirtualPath]
         if let partitions = partitions, let partitionIdx = partitions.assignment[input] {
           // We have a partitioning for batch mode. If this input file isn't the first
@@ -94,7 +112,7 @@ extension Driver {
     }
 
     // Plan the merge-module job, if there are module inputs.
-    if moduleOutput != nil && !moduleInputs.isEmpty {
+    if moduleOutput != nil && !moduleInputs.isEmpty && compilerMode.usesPrimaryFileInputs {
       jobs.append(try mergeModuleJob(inputs: moduleInputs))
     }
 
@@ -126,10 +144,10 @@ extension Driver {
   public mutating func planBuild() throws -> [Job] {
     // Plan the build.
     switch compilerMode {
-    case .immediate, .repl, .singleCompile:
+    case .immediate, .repl:
       fatalError("Not yet supported")
 
-    case .standardCompile, .batchCompile:
+    case .standardCompile, .batchCompile, .singleCompile:
       return try planStandardCompile()
     }
   }

Tests/SwiftDriverTests/SwiftDriverTests.swift

@@ -721,6 +721,33 @@ final class SwiftDriverTests: XCTestCase {
     XCTAssertEqual(plannedJobs[3].outputs.first!.file, VirtualPath.relative(RelativePath("Test")))
   }
 
+  func testSingleThreadedWholeModuleOptimizationCompiles() throws {
+    var driver1 = try Driver(args: ["swiftc", "-whole-module-optimization", "foo.swift", "bar.swift", "-module-name", "Test", "-target", "x86_64-apple-macosx10.15"])
+    let plannedJobs = try driver1.planBuild()
+    XCTAssertEqual(plannedJobs.count, 2)
+    XCTAssertEqual(plannedJobs[0].kind, .compile)
+    XCTAssertEqual(plannedJobs[0].outputs.count, 1)
+    XCTAssertEqual(plannedJobs[0].outputs.first!.file, VirtualPath.temporary(RelativePath("Test.o")))
+    XCTAssert(!plannedJobs[0].commandLine.contains(.flag("-primary-file")))
+
+    XCTAssertEqual(plannedJobs[1].kind, .link)
+  }
+
+  func testMultiThreadedWholeModuleOptimizationCompiles() throws {
+    var driver1 = try Driver(args: ["swiftc", "-whole-module-optimization", "foo.swift", "bar.swift", "wibble.swift", "-module-name", "Test",
+        "-num-threads", "4", "-target", "x86_64-apple-macosx10.15"])
+    let plannedJobs = try driver1.planBuild()
+    XCTAssertEqual(plannedJobs.count, 2)
+    XCTAssertEqual(plannedJobs[0].kind, .compile)
+    XCTAssertEqual(plannedJobs[0].outputs.count, 3)
+    XCTAssertEqual(plannedJobs[0].outputs[0].file, VirtualPath.temporary(RelativePath("foo.o")))
+    XCTAssertEqual(plannedJobs[0].outputs[1].file, VirtualPath.temporary(RelativePath("bar.o")))
+    XCTAssertEqual(plannedJobs[0].outputs[2].file, VirtualPath.temporary(RelativePath("wibble.o")))
+    XCTAssert(!plannedJobs[0].commandLine.contains(.flag("-primary-file")))
+
+    XCTAssertEqual(plannedJobs[1].kind, .link)
+  }
+
   func testMergeModulesOnly() throws {
     do {
       var driver = try Driver(args: ["swiftc", "foo.swift", "bar.swift", "-module-name", "Test", "-emit-module", "-import-objc-header", "TestInputHeader.h", "-emit-dependencies", "-emit-module-doc-path", "/foo/bar/Test.swiftdoc"])