Merge pull request #85440 from rjmansfield/fix-wmo-opt-record-paths

rjmansfield · web-flow · commit e60ca698c62c · 2025-12-08T16:14:53.000-05:00
Fix optimization record paths in multi-threaded WMO mode
diff --git a/lib/Frontend/ArgsToFrontendOutputsConverter.cpp b/lib/Frontend/ArgsToFrontendOutputsConverter.cpp
@@ -286,23 +286,41 @@ SupplementaryOutputPathsComputer::computeOutputPaths() const {
   // For other cases, process the paths normally
   std::vector<SupplementaryOutputPaths> outputPaths;
   unsigned i = 0;
-  bool hadError = InputsAndOutputs.forEachInputProducingSupplementaryOutput(
-      [&](const InputFile &input) -> bool {
-        if (auto suppPaths = computeOutputPathsForOneInput(
-                OutputFiles[i], (*pathsFromUser)[i], input)) {
-          ++i;
-          outputPaths.push_back(*suppPaths);
-          return false;
-        }
-        return true;
-      });
+  bool hadError = false;
+
+  // In multi-threaded WMO with supplementary output file map, we have paths
+  // for all inputs, so process them all through computeOutputPathsForOneInput
+  if (!InputsAndOutputs.hasPrimaryInputs() && OutputFiles.size() > 1 &&
+      pathsFromUser->size() == InputsAndOutputs.inputCount()) {
+    hadError = InputsAndOutputs.forEachInput([&](const InputFile &input) -> bool {
+      if (auto suppPaths = computeOutputPathsForOneInput(
+              OutputFiles[i], (*pathsFromUser)[i], input)) {
+        ++i;
+        outputPaths.push_back(*suppPaths);
+        return false;
+      }
+      return true;
+    });
+  } else {
+    // Standard path: process inputs that produce supplementary output
+    hadError = InputsAndOutputs.forEachInputProducingSupplementaryOutput(
+        [&](const InputFile &input) -> bool {
+          if (auto suppPaths = computeOutputPathsForOneInput(
+                  OutputFiles[i], (*pathsFromUser)[i], input)) {
+            ++i;
+            outputPaths.push_back(*suppPaths);
+            return false;
+          }
+          return true;
+        });
+  }
   if (hadError)
     return std::nullopt;
 
-  // In WMO mode compute supplementary output paths for optimization record
-  // data (opt-remarks). We need one path per LLVMModule that will be created as
-  // part of wmo.
-  if (!InputsAndOutputs.hasPrimaryInputs()) {
+  // In WMO mode without supplementary output file map, compute supplementary
+  // output paths for optimization records for inputs beyond the first one.
+  if (!InputsAndOutputs.hasPrimaryInputs() && OutputFiles.size() > 1 &&
+      pathsFromUser->size() != InputsAndOutputs.inputCount()) {
     unsigned i = 0;
     InputsAndOutputs.forEachInput([&](const InputFile &input) -> bool {
       // First input is already computed.
@@ -448,8 +466,9 @@ SupplementaryOutputPathsComputer::getSupplementaryOutputPathsFromArguments()
     sop.APIDescriptorOutputPath = (*apiDescriptorOutput)[moduleIndex];
     sop.ConstValuesOutputPath = (*constValuesOutput)[moduleIndex];
     sop.ModuleSemanticInfoOutputPath = (*moduleSemanticInfoOutput)[moduleIndex];
-    sop.YAMLOptRecordPath = (*optRecordOutput)[moduleIndex];
-    sop.BitstreamOptRecordPath = (*optRecordOutput)[moduleIndex];
+    // Optimization record paths are per-file in multi-threaded WMO, like IR
+    sop.YAMLOptRecordPath = (*optRecordOutput)[i];
+    sop.BitstreamOptRecordPath = (*optRecordOutput)[i];
     sop.SILOutputPath = (*silOutput)[silOutputIndex];
     sop.LLVMIROutputPath = (*irOutput)[i];
     result.push_back(sop);
@@ -478,18 +497,28 @@ SupplementaryOutputPathsComputer::getSupplementaryFilenamesFromArguments(
       paths.emplace_back();
     return paths;
   }
-  // Special handling for SIL and IR output paths: allow multiple paths per file
-  // type
-  else if ((pathID == options::OPT_sil_output_path ||
-            pathID == options::OPT_ir_output_path) &&
+  // Special handling for IR and optimization record output paths: allow multiple paths per file
+  // type. Note: SIL is NOT included here because in WMO mode, SIL is generated once
+  // per module, not per source file.
+  else if ((pathID == options::OPT_ir_output_path ||
+            pathID == options::OPT_save_optimization_record_path) &&
            paths.size() > N) {
-    // For parallel compilation, we can have multiple SIL/IR output paths
+    // For parallel compilation, we can have multiple IR/opt-record output paths
     // so return all the paths.
     return paths;
   }
 
-  if (paths.empty())
+  if (paths.empty()) {
+    // For IR and optimization records in multi-threaded WMO, we need one entry per input file.
+    // Check if WMO is enabled and we have multiple output files (multi-threaded WMO).
+    if ((pathID == options::OPT_ir_output_path ||
+         pathID == options::OPT_save_optimization_record_path) &&
+        Args.hasArg(options::OPT_whole_module_optimization) &&
+        OutputFiles.size() > 1) {
+      return std::vector<std::string>(OutputFiles.size(), std::string());
+    }
     return std::vector<std::string>(N, std::string());
+  }
 
   Diags.diagnose(SourceLoc(), diag::error_wrong_number_of_arguments,
                  Args.getLastArg(pathID)->getOption().getPrefixedName(), N,
@@ -864,5 +893,28 @@ SupplementaryOutputPathsComputer::readSupplementaryOutputFileMap() const {
   if (hadError)
     return std::nullopt;
 
+  // In multi-threaded WMO mode, we need to read supplementary output paths
+  // for all inputs beyond the first one (which was already processed above).
+  // Entries in the map are optional, so if an input is missing, the regular
+  // WMO path generation logic will handle it.
+  if (!InputsAndOutputs.hasPrimaryInputs() && OutputFiles.size() > 1) {
+    InputsAndOutputs.forEachInput([&](const InputFile &input) -> bool {
+      // Skip the first input, which was already processed above
+      if (InputsAndOutputs.firstInput().getFileName() == input.getFileName()) {
+        return false;
+      }
+
+      // Check if this input has an entry in the supplementary output file map
+      const TypeToPathMap *mapForInput =
+          OFM->getOutputMapForInput(input.getFileName());
+      if (mapForInput) {
+        // Entry exists, use it
+        outputPaths.push_back(createFromTypeToPathMap(mapForInput));
+      }
+      // If no entry exists, skip - the regular WMO logic will generate paths
+      return false;
+    });
+  }
+
   return outputPaths;
 }
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -3231,8 +3231,12 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
   if (const Arg *A = Args.getLastArg(OPT_save_optimization_record_passes))
     Opts.OptRecordPasses = A->getValue();
 
-  if (const Arg *A = Args.getLastArg(OPT_save_optimization_record_path))
-    Opts.OptRecordFile = A->getValue();
+  // Only use getLastArg for single -save-optimization-record-path.
+  // With multiple paths (multi-threaded WMO), FrontendTool will populate
+  // OptRecordFile and AuxOptRecordFiles from command-line arguments.
+  auto allOptRecordPaths = Args.getAllArgValues(OPT_save_optimization_record_path);
+  if (allOptRecordPaths.size() == 1)
+    Opts.OptRecordFile = allOptRecordPaths[0];
 
   // If any of the '-g<kind>', except '-gnone', is given,
   // tell the SILPrinter to print debug info as well
diff --git a/lib/FrontendTool/FrontendTool.cpp b/lib/FrontendTool/FrontendTool.cpp
@@ -731,7 +731,8 @@ static bool writeAPIDescriptorIfNeeded(CompilerInstance &Instance) {
 static bool performCompileStepsPostSILGen(
     CompilerInstance &Instance, std::unique_ptr<SILModule> SM,
     ModuleOrSourceFile MSF, const PrimarySpecificPaths &PSPs, int &ReturnValue,
-    FrontendObserver *observer, ArrayRef<const char *> CommandLineArgs);
+    FrontendObserver *observer, ArrayRef<const char *> CommandLineArgs,
+    ArrayRef<PrimarySpecificPaths> auxPSPs = {});
 
 bool swift::performCompileStepsPostSema(
     CompilerInstance &Instance, int &ReturnValue, FrontendObserver *observer,
@@ -749,14 +750,48 @@ bool swift::performCompileStepsPostSema(
           PSPs.SupplementaryOutputs.YAMLOptRecordPath :
           PSPs.SupplementaryOutputs.BitstreamOptRecordPath;
     }
+
+    auto populateOptRecordPathsFromCmdLine = [&]() {
+      auto optRecordPaths = collectSupplementaryOutputPaths(
+          CommandLineArgs, options::OPT_save_optimization_record_path);
+      if (!optRecordPaths.empty()) {
+        // Set the first path. With multiple paths, CompilerInvocation leaves
+        // OptRecordFile empty, so we populate it here along with aux paths.
+        SILOpts.OptRecordFile = optRecordPaths[0];
+        if (optRecordPaths.size() > 1) {
+          SILOpts.AuxOptRecordFiles.assign(optRecordPaths.begin() + 1,
+                                           optRecordPaths.end());
+        }
+      }
+    };
+
     if (!auxPSPs.empty()) {
       assert(SILOpts.AuxOptRecordFiles.empty());
+      // Check if ALL auxPSPs have optimization record paths populated
+      bool allHaveOptRecordPaths = true;
       for (const auto &auxFile: auxPSPs) {
-        SILOpts.AuxOptRecordFiles.push_back(
-          SILOpts.OptRecordFormat == llvm::remarks::Format::YAML ?
-          auxFile.SupplementaryOutputs.YAMLOptRecordPath :
-          auxFile.SupplementaryOutputs.BitstreamOptRecordPath);
+        bool hasPath = SILOpts.OptRecordFormat == llvm::remarks::Format::YAML ?
+          !auxFile.SupplementaryOutputs.YAMLOptRecordPath.empty() :
+          !auxFile.SupplementaryOutputs.BitstreamOptRecordPath.empty();
+        if (!hasPath) {
+          allHaveOptRecordPaths = false;
+          break;
+        }
       }
+
+      if (allHaveOptRecordPaths) {
+        for (const auto &auxFile: auxPSPs) {
+          SILOpts.AuxOptRecordFiles.push_back(
+            SILOpts.OptRecordFormat == llvm::remarks::Format::YAML ?
+            auxFile.SupplementaryOutputs.YAMLOptRecordPath :
+            auxFile.SupplementaryOutputs.BitstreamOptRecordPath);
+        }
+      } else {
+        populateOptRecordPathsFromCmdLine();
+      }
+    } else {
+      assert(SILOpts.AuxOptRecordFiles.empty());
+      populateOptRecordPathsFromCmdLine();
     }
     return SILOpts;
   };
@@ -781,7 +816,7 @@ bool swift::performCompileStepsPostSema(
                                  &irgenOpts);
     return performCompileStepsPostSILGen(Instance, std::move(SM), mod, PSPs,
                                          ReturnValue, observer,
-                                         CommandLineArgs);
+                                         CommandLineArgs, auxPSPs);
   }
 
 
@@ -1983,7 +2018,8 @@ static bool generateCode(CompilerInstance &Instance, StringRef OutputFilename,
 static bool performCompileStepsPostSILGen(
     CompilerInstance &Instance, std::unique_ptr<SILModule> SM,
     ModuleOrSourceFile MSF, const PrimarySpecificPaths &PSPs, int &ReturnValue,
-    FrontendObserver *observer, ArrayRef<const char *> CommandLineArgs) {
+    FrontendObserver *observer, ArrayRef<const char *> CommandLineArgs,
+    ArrayRef<PrimarySpecificPaths> auxPSPs) {
   const auto &Invocation = Instance.getInvocation();
   const auto &opts = Invocation.getFrontendOptions();
   FrontendOptions::ActionType Action = opts.RequestedAction;
@@ -2156,10 +2192,38 @@ static bool performCompileStepsPostSILGen(
   std::vector<std::string> ParallelOutputFilenames =
       opts.InputsAndOutputs.copyOutputFilenames();
 
-  // Collect IR output paths from command line arguments
-  std::vector<std::string> ParallelIROutputFilenames =
-      collectSupplementaryOutputPaths(CommandLineArgs,
-                                      options::OPT_ir_output_path);
+  // Collect IR output paths - check if supplementary output file map has paths,
+  // otherwise fall back to command line arguments
+  std::vector<std::string> ParallelIROutputFilenames;
+  if (!auxPSPs.empty()) {
+    // Check if the first file (PSPs) and ALL auxPSPs have IR output paths populated
+    bool allHaveIRPaths = !PSPs.SupplementaryOutputs.LLVMIROutputPath.empty();
+    if (allHaveIRPaths) {
+      for (const auto &auxFile : auxPSPs) {
+        if (auxFile.SupplementaryOutputs.LLVMIROutputPath.empty()) {
+          allHaveIRPaths = false;
+          break;
+        }
+      }
+    }
+
+    if (allHaveIRPaths) {
+      // Paths are in supplementary output file map - include first file + aux files
+      ParallelIROutputFilenames.push_back(PSPs.SupplementaryOutputs.LLVMIROutputPath);
+      for (const auto &auxFile : auxPSPs) {
+        ParallelIROutputFilenames.push_back(
+            auxFile.SupplementaryOutputs.LLVMIROutputPath);
+      }
+    } else {
+      // Fall back to command line arguments
+      ParallelIROutputFilenames = collectSupplementaryOutputPaths(
+          CommandLineArgs, options::OPT_ir_output_path);
+    }
+  } else {
+    // No auxPSPs, use command line arguments
+    ParallelIROutputFilenames = collectSupplementaryOutputPaths(
+        CommandLineArgs, options::OPT_ir_output_path);
+  }
 
   llvm::GlobalVariable *HashGlobal;
   cas::SwiftCASOutputBackend *casBackend =
@@ -2310,10 +2374,13 @@ collectSupplementaryOutputPaths(ArrayRef<const char *> Args,
     StringRef arg = Args[i];
     StringRef optionName;
 
-    if (OptionID == options::OPT_sil_output_path) {
-      optionName = "-sil-output-path";
-    } else if (OptionID == options::OPT_ir_output_path) {
+    // Note: SIL is NOT included here because in WMO mode, SIL is generated once
+    // per module, not per source file. Only IR and optimization records can have
+    // per-file outputs in multi-threaded WMO.
+    if (OptionID == options::OPT_ir_output_path) {
       optionName = "-ir-output-path";
+    } else if (OptionID == options::OPT_save_optimization_record_path) {
+      optionName = "-save-optimization-record-path";
     } else {
       continue;
     }
diff --git a/test/Frontend/multi-threaded-wmo-opt-record-supplementary-map.swift b/test/Frontend/multi-threaded-wmo-opt-record-supplementary-map.swift
@@ -0,0 +1,37 @@
+// Test that frontend properly handles supplementary output file maps with
+// optimization records in multi-threaded WMO mode
+
+// RUN: %empty-directory(%t)
+
+// RUN: echo 'public func funcA() -> Int { return 42 }' > %t/file_a.swift
+// RUN: echo 'public func funcB() -> String { return "hello" }' > %t/file_b.swift
+
+// RUN: echo '{' > %t/output-file-map.json
+// RUN: echo '  "%/t/file_a.swift": {' >> %t/output-file-map.json
+// RUN: echo '    "object": "%/t/file_a.o",' >> %t/output-file-map.json
+// RUN: echo '    "yaml-opt-record": "%/t/file_a.opt.yaml",' >> %t/output-file-map.json
+// RUN: echo '    "llvm-ir": "%/t/file_a.ll"' >> %t/output-file-map.json
+// RUN: echo '  },' >> %t/output-file-map.json
+// RUN: echo '  "%/t/file_b.swift": {' >> %t/output-file-map.json
+// RUN: echo '    "object": "%/t/file_b.o",' >> %t/output-file-map.json
+// RUN: echo '    "yaml-opt-record": "%/t/file_b.opt.yaml",' >> %t/output-file-map.json
+// RUN: echo '    "llvm-ir": "%/t/file_b.ll"' >> %t/output-file-map.json
+// RUN: echo '  }' >> %t/output-file-map.json
+// RUN: echo '}' >> %t/output-file-map.json
+
+// RUN: %target-swift-frontend -c %/t/file_a.swift %/t/file_b.swift \
+// RUN:   -wmo -num-threads 2 -O -module-name TestModule \
+// RUN:   -supplementary-output-file-map %t/output-file-map.json \
+// RUN:   -o %t/file_a.o -o %t/file_b.o
+
+// RUN: ls %t/file_a.o
+// RUN: ls %t/file_b.o
+// RUN: ls %t/file_a.opt.yaml
+// RUN: ls %t/file_b.opt.yaml
+// RUN: ls %t/file_a.ll
+// RUN: ls %t/file_b.ll
+
+// RUN: grep -q "funcA" %t/file_a.opt.yaml
+// RUN: grep -q "funcB" %t/file_b.opt.yaml
+// RUN: grep -q "funcA" %t/file_a.ll
+// RUN: grep -q "funcB" %t/file_b.ll
diff --git a/test/Frontend/multi-threaded-wmo-supplementary-outputs.swift b/test/Frontend/multi-threaded-wmo-supplementary-outputs.swift
@@ -0,0 +1,38 @@
+// Test that frontend properly handles multiple supplementary output paths
+// using command line options in multi-threaded WMO mode.
+
+// RUN: %empty-directory(%t)
+
+// RUN: echo 'public func functionA() -> Int { return 42 }' > %t/FileA.swift
+// RUN: echo 'public func functionB() -> String { return "hello" }' > %t/FileB.swift
+
+// RUN: %target-swift-frontend -c %t/FileA.swift %t/FileB.swift \
+// RUN:   -wmo -num-threads 2 -O -module-name TestModule \
+// RUN:   -save-optimization-record-path %t/FileA.opt.yaml \
+// RUN:   -save-optimization-record-path %t/FileB.opt.yaml \
+// RUN:   -ir-output-path %t/FileA.ll \
+// RUN:   -ir-output-path %t/FileB.ll \
+// RUN:   -sil-output-path %t/TestModule.sil \
+// RUN:   -o %t/FileA.o -o %t/FileB.o
+
+// RUN: ls %t/FileA.opt.yaml
+// RUN: ls %t/FileB.opt.yaml
+
+// RUN: ls %t/FileA.ll
+// RUN: ls %t/FileB.ll
+
+// RUN: ls %t/TestModule.sil
+
+// RUN: ls %t/FileA.o
+// RUN: ls %t/FileB.o
+
+// RUN: grep -q "functionA" %t/FileA.ll
+// RUN: grep -q "functionB" %t/FileB.ll
+
+// In multi-threaded WMO, each source file should generate its own optimization record file
+// RUN: grep -q "functionA" %t/FileA.opt.yaml
+// RUN: grep -q "functionB" %t/FileB.opt.yaml
+
+// Verify the SIL output contains both functions (whole module)
+// RUN: grep -q "functionA" %t/TestModule.sil
+// RUN: grep -q "functionB" %t/TestModule.sil
