Fix optimization record paths in multi-threaded WMO mode

lib/Frontend/ArgsToFrontendOutputsConverter.cpp

@@ -478,12 +478,13 @@ 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;
   }

lib/FrontendTool/FrontendTool.cpp

@@ -749,14 +749,41 @@ 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 (overriding what might have been set by getLastArg
+        SILOpts.OptRecordFile = optRecordPaths[0];
+        if (optRecordPaths.size() > 1) {
+          SILOpts.AuxOptRecordFiles.assign(optRecordPaths.begin() + 1,
+                                           optRecordPaths.end());
+        }
+      }
+    };
+
     if (!auxPSPs.empty()) {
       assert(SILOpts.AuxOptRecordFiles.empty());
-      for (const auto &auxFile: auxPSPs) {
-        SILOpts.AuxOptRecordFiles.push_back(
-          SILOpts.OptRecordFormat == llvm::remarks::Format::YAML ?
-          auxFile.SupplementaryOutputs.YAMLOptRecordPath :
-          auxFile.SupplementaryOutputs.BitstreamOptRecordPath);
+      // Check if the first auxPSP has optimization record paths populated
+      bool hasOptRecordPaths = SILOpts.OptRecordFormat == llvm::remarks::Format::YAML ?
+        !auxPSPs[0].SupplementaryOutputs.YAMLOptRecordPath.empty() :
+        !auxPSPs[0].SupplementaryOutputs.BitstreamOptRecordPath.empty();
+
+      if (hasOptRecordPaths) {
+        // Paths are in supplementary output file map
+        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;
   };
@@ -2310,6 +2337,8 @@ collectSupplementaryOutputPaths(ArrayRef<const char *> Args,
       optionName = "-sil-output-path";
     } else 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;
     }

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 multithreaded 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