Add code to create llvm::RemarkStreamer objects for all the LLVMModules in WMO mode

aschwaighofer · aschwaighofer · commit b1edbac45bad · 2025-06-30T11:06:04.000-07:00
rdar://154403078
diff --git a/include/swift/AST/SILOptions.h b/include/swift/AST/SILOptions.h
@@ -316,6 +316,10 @@ class SILOptions {
   /// records.
   std::string OptRecordFile;
 
+  /// The names of the auxiliar files to which the backend should save optimization
+  /// records for the remaining (other than the main one) LLVMModules.
+  std::vector<std::string> AuxOptRecordFiles;
+
   /// The regex that filters the passes that should be saved to the optimization
   /// records.
   std::string OptRecordPasses;
diff --git a/include/swift/Frontend/FrontendInputsAndOutputs.h b/include/swift/Frontend/FrontendInputsAndOutputs.h
@@ -253,6 +253,9 @@ class FrontendInputsAndOutputs {
   const PrimarySpecificPaths &
   getPrimarySpecificPathsForAtMostOnePrimary() const;
 
+  const PrimarySpecificPaths &
+  getPrimarySpecificPathsForRemaining(unsigned i) const;
+
   const PrimarySpecificPaths &
       getPrimarySpecificPathsForPrimary(StringRef) const;
 
diff --git a/lib/Frontend/ArgsToFrontendOutputsConverter.cpp b/lib/Frontend/ArgsToFrontendOutputsConverter.cpp
@@ -298,6 +298,44 @@ SupplementaryOutputPathsComputer::computeOutputPaths() const {
       });
   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() && OutputFiles.size() > 1) {
+    unsigned i = 0;
+    InputsAndOutputs.forEachInput([&](const InputFile &input) -> bool {
+      // First input is already computed.
+      if (InputsAndOutputs.firstInput().getFileName() == input.getFileName()) {
+        ++i;
+        return false;
+      }
+      SupplementaryOutputPaths outputs;
+
+      // Compute auxiliar opt record paths.
+      StringRef defaultSupplementaryOutputPathExcludingExtension =
+          deriveDefaultSupplementaryOutputPathExcludingExtension(OutputFiles[i], input);
+
+      auto YAMLOptRecordPath = determineSupplementaryOutputFilename(
+        options::OPT_save_optimization_record_path,
+        "",
+        file_types::TY_YAMLOptRecord, "",
+        defaultSupplementaryOutputPathExcludingExtension, true);
+      outputs.YAMLOptRecordPath = YAMLOptRecordPath;
+
+      auto bitstreamOptRecordPath = determineSupplementaryOutputFilename(
+        options::OPT_save_optimization_record_path,
+        "",
+        file_types::TY_BitstreamOptRecord, "",
+        defaultSupplementaryOutputPathExcludingExtension, true);
+
+      outputs.BitstreamOptRecordPath = bitstreamOptRecordPath;
+
+      outputPaths.emplace_back(std::move(outputs));
+      ++i;
+      return false;
+    });
+  }
   return outputPaths;
 }
 
@@ -615,7 +653,21 @@ std::string
 SupplementaryOutputPathsComputer::determineSupplementaryOutputFilename(
     options::ID emitOpt, std::string pathFromArguments, file_types::ID type,
     StringRef mainOutputIfUsable,
-    StringRef defaultSupplementaryOutputPathExcludingExtension) const {
+    StringRef defaultSupplementaryOutputPathExcludingExtension,
+    bool forceDefaultSupplementaryOutputPathExcludingExtension) const {
+
+  auto computeDefaultSupplementaryOutputPathExcludingExtension =
+    [&] () -> std::string {
+      llvm::SmallString<128> path(
+        defaultSupplementaryOutputPathExcludingExtension);
+
+      llvm::sys::path::replace_extension(path, file_types::getExtension(type));
+      return path.str().str();
+    };
+
+  if (forceDefaultSupplementaryOutputPathExcludingExtension) {
+    return computeDefaultSupplementaryOutputPathExcludingExtension();
+  }
 
   if (!pathFromArguments.empty())
     return pathFromArguments;
@@ -627,9 +679,7 @@ SupplementaryOutputPathsComputer::determineSupplementaryOutputFilename(
     return mainOutputIfUsable.str();
   }
 
-  llvm::SmallString<128> path(defaultSupplementaryOutputPathExcludingExtension);
-  llvm::sys::path::replace_extension(path, file_types::getExtension(type));
-  return path.str().str();
+  return computeDefaultSupplementaryOutputPathExcludingExtension();
 }
 
 void SupplementaryOutputPathsComputer::deriveModulePathParameters(
diff --git a/lib/Frontend/ArgsToFrontendOutputsConverter.h b/lib/Frontend/ArgsToFrontendOutputsConverter.h
@@ -180,7 +180,8 @@ class SupplementaryOutputPathsComputer {
   std::string determineSupplementaryOutputFilename(
       options::ID emitOpt, std::string pathFromArgumentsOrFilelists,
       file_types::ID type, StringRef mainOutputIfUsable,
-      StringRef defaultSupplementaryOutputPathExcludingExtension) const;
+      StringRef defaultSupplementaryOutputPathExcludingExtension,
+      bool forceDefaultSupplementaryOutputPathExcludingExtension = false) const;
 
   void deriveModulePathParameters(StringRef mainOutputFile,
                                   options::ID &emitOption,
diff --git a/lib/Frontend/FrontendInputsAndOutputs.cpp b/lib/Frontend/FrontendInputsAndOutputs.cpp
@@ -380,21 +380,32 @@ void FrontendInputsAndOutputs::setMainAndSupplementaryOutputs(
     }
     return;
   }
-  assert(supplementaryOutputs.size() == 1 &&
-         "WMO only ever produces one set of supplementary outputs");
+
+  assert(supplementaryOutputs.size() == 1 ||
+         supplementaryOutputs.size() == AllInputs.size() &&
+             "WMO produces wrong number of sets of supplementary outputs");
   if (outputFiles.size() == 1) {
-    AllInputs.front().setPrimarySpecificPaths(PrimarySpecificPaths(
-        outputFiles.front(), outputFilesForIndexUnits.front(),
-        firstInputProducingOutput().getFileName(),
-        supplementaryOutputs.front()));
+    for (auto i : indices(AllInputs)) {
+      if (i == 0)
+        AllInputs[i].setPrimarySpecificPaths(PrimarySpecificPaths(
+            outputFiles.front(), outputFilesForIndexUnits.front(),
+            firstInputProducingOutput().getFileName(),
+            supplementaryOutputs.front()));
+      else
+        AllInputs[i].setPrimarySpecificPaths(PrimarySpecificPaths(
+            "", "", "",
+            supplementaryOutputs.size() == 1 ? SupplementaryOutputPaths()
+                                             : supplementaryOutputs[i]));
+    }
     return;
   }
   assert(outputFiles.size() == AllInputs.size() &&
          "Multi-threaded WMO requires one main output per input");
   for (auto i : indices(AllInputs))
     AllInputs[i].setPrimarySpecificPaths(PrimarySpecificPaths(
         outputFiles[i], outputFilesForIndexUnits[i], outputFiles[i],
-        i == 0 ? supplementaryOutputs.front() : SupplementaryOutputPaths()));
+        supplementaryOutputs.size() == 1 && i != 0 ? SupplementaryOutputPaths()
+                                                   : supplementaryOutputs[i]));
 }
 
 std::vector<std::string> FrontendInputsAndOutputs::copyOutputFilenames() const {
@@ -488,6 +499,14 @@ FrontendInputsAndOutputs::getPrimarySpecificPathsForAtMostOnePrimary() const {
                      : emptyPaths;
 }
 
+const PrimarySpecificPaths &
+FrontendInputsAndOutputs::getPrimarySpecificPathsForRemaining(unsigned i) const {
+  static auto emptyPaths = PrimarySpecificPaths();
+  unsigned firstProducingIdx = getIndexOfFirstOutputProducingInput();
+  return (hasInputs()  && i > 0) ?
+    AllInputs[firstProducingIdx+i].getPrimarySpecificPaths() : emptyPaths;
+}
+
 const PrimarySpecificPaths &
 FrontendInputsAndOutputs::getPrimarySpecificPathsForPrimary(
     StringRef filename) const {
diff --git a/lib/FrontendTool/FrontendTool.cpp b/lib/FrontendTool/FrontendTool.cpp
@@ -740,7 +740,8 @@ bool swift::performCompileStepsPostSema(CompilerInstance &Instance,
   const auto &Invocation = Instance.getInvocation();
   const FrontendOptions &opts = Invocation.getFrontendOptions();
 
-  auto getSILOptions = [&](const PrimarySpecificPaths &PSPs) -> SILOptions {
+  auto getSILOptions = [&](const PrimarySpecificPaths &PSPs,
+                           const std::vector<PrimarySpecificPaths> &auxPSPs) -> SILOptions {
     SILOptions SILOpts = Invocation.getSILOptions();
     if (SILOpts.OptRecordFile.empty()) {
       // Check if the record file path was passed via supplemental outputs.
@@ -749,6 +750,15 @@ bool swift::performCompileStepsPostSema(CompilerInstance &Instance,
           PSPs.SupplementaryOutputs.YAMLOptRecordPath :
           PSPs.SupplementaryOutputs.BitstreamOptRecordPath;
     }
+    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);
+      }
+    }
     return SILOpts;
   };
 
@@ -758,13 +768,24 @@ bool swift::performCompileStepsPostSema(CompilerInstance &Instance,
     // SILModule for the entire module.
     const PrimarySpecificPaths PSPs =
         Instance.getPrimarySpecificPathsForWholeModuleOptimizationMode();
-    SILOptions SILOpts = getSILOptions(PSPs);
+
+    std::vector<PrimarySpecificPaths> auxPSPs;
+    for (unsigned i = 1; i < opts.InputsAndOutputs.inputCount(); ++i) {
+      auto &auxPSP =
+        opts.InputsAndOutputs.getPrimarySpecificPathsForRemaining(i);
+      auxPSPs.push_back(auxPSP);
+    }
+
+    SILOptions SILOpts = getSILOptions(PSPs, auxPSPs);
     IRGenOptions irgenOpts = Invocation.getIRGenOptions();
     auto SM = performASTLowering(mod, Instance.getSILTypes(), SILOpts,
                                  &irgenOpts);
     return performCompileStepsPostSILGen(Instance, std::move(SM), mod, PSPs,
                                          ReturnValue, observer);
   }
+
+
+  std::vector<PrimarySpecificPaths> emptyAuxPSPs;
   // If there are primary source files, build a separate SILModule for
   // each source file, and run the remaining SILOpt-Serialize-IRGen-LLVM
   // once for each such input.
@@ -773,7 +794,7 @@ bool swift::performCompileStepsPostSema(CompilerInstance &Instance,
     for (auto *PrimaryFile : Instance.getPrimarySourceFiles()) {
       const PrimarySpecificPaths PSPs =
           Instance.getPrimarySpecificPathsForSourceFile(*PrimaryFile);
-      SILOptions SILOpts = getSILOptions(PSPs);
+      SILOptions SILOpts = getSILOptions(PSPs, emptyAuxPSPs);
     IRGenOptions irgenOpts = Invocation.getIRGenOptions();
       auto SM = performASTLowering(*PrimaryFile, Instance.getSILTypes(),
                                    SILOpts, &irgenOpts);
@@ -793,7 +814,7 @@ bool swift::performCompileStepsPostSema(CompilerInstance &Instance,
       if (opts.InputsAndOutputs.isInputPrimary(SASTF->getFilename())) {
         const PrimarySpecificPaths &PSPs =
             Instance.getPrimarySpecificPathsForPrimary(SASTF->getFilename());
-        SILOptions SILOpts = getSILOptions(PSPs);
+        SILOptions SILOpts = getSILOptions(PSPs, emptyAuxPSPs);
         auto SM = performASTLowering(*SASTF, Instance.getSILTypes(), SILOpts);
         result |= performCompileStepsPostSILGen(Instance, std::move(SM), mod,
                                                 PSPs, ReturnValue, observer);
diff --git a/lib/IRGen/IRGen.cpp b/lib/IRGen/IRGen.cpp
@@ -20,6 +20,7 @@
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/ObjectFile.h"
 #include "swift/AST/DiagnosticsIRGen.h"
+#include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/IRGenOptions.h"
 #include "swift/AST/IRGenRequests.h"
 #include "swift/AST/LinkLibrary.h"
@@ -71,6 +72,8 @@
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Remarks/Remark.h"
+#include "llvm/Remarks/RemarkStreamer.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -1143,7 +1146,7 @@ static void embedBitcode(llvm::Module *M, const IRGenOptions &Opts)
   NewUsed->setSection("llvm.metadata");
 }
 
-static void initLLVMModule(IRGenModule &IGM, SILModule &SIL) {
+static void initLLVMModule(IRGenModule &IGM, SILModule &SIL, std::optional<unsigned> idx = {}) {
   auto *Module = IGM.getModule();
   assert(Module && "Expected llvm:Module for IR generation!");
   
@@ -1186,17 +1189,61 @@ static void initLLVMModule(IRGenModule &IGM, SILModule &SIL) {
                                          llvm::ConstantAsMetadata::get(Value)}));
 
   if (auto *SILstreamer = SIL.getSILRemarkStreamer()) {
-    // Install RemarkStreamer into LLVM and keep the remarks file alive. This is
-    // required even if no LLVM remarks are enabled, because the AsmPrinter
-    // serializes meta information about the remarks into the object file.
-    IGM.RemarkStream = SILstreamer->releaseStream();
-    SILstreamer->intoLLVMContext(Context);
-    auto &RS = *IGM.getLLVMContext().getMainRemarkStreamer();
-    if (IGM.getOptions().AnnotateCondFailMessage) {
+    auto remarkStream = SILstreamer->releaseStream();
+    if (remarkStream) {
+      // Install RemarkStreamer into LLVM and keep the remarks file alive. This is
+      // required even if no LLVM remarks are enabled, because the AsmPrinter
+      // serializes meta information about the remarks into the object file.
+      IGM.RemarkStream = std::move(remarkStream);
+      SILstreamer->intoLLVMContext(Context);
+      auto &RS = *IGM.getLLVMContext().getMainRemarkStreamer();
+      if (IGM.getOptions().AnnotateCondFailMessage) {
+        Context.setLLVMRemarkStreamer(
+           std::make_unique<llvm::LLVMRemarkStreamer>(RS));
+        // FIXME: add a frontend flag to enable all LLVM remarks
+        cantFail(RS.setFilter("annotation-remarks"));
+      } else {
+        // Don't filter for now.
+        Context.setLLVMRemarkStreamer(
+            std::make_unique<llvm::LLVMRemarkStreamer>(RS));
+      }
+    } else {
+      assert(idx && "Not generating multiple output files?");
+
+      // Construct llvmremarkstreamer objects for LLVM remarks originating in
+      // the LLVM backend and install it in the remaining LLVMModule(s).
+      auto &SILOpts = SIL.getOptions();
+      assert(SILOpts.AuxOptRecordFiles.size() > (*idx - 1));
+
+      const auto &filename = SILOpts.AuxOptRecordFiles[*idx - 1];
+      auto &diagEngine = SIL.getASTContext().Diags;
+      std::error_code errorCode;
+      auto file = std::make_unique<llvm::raw_fd_ostream>(filename, errorCode,
+                                                         llvm::sys::fs::OF_None);
+      if (errorCode) {
+        diagEngine.diagnose(SourceLoc(), diag::cannot_open_file, filename,
+                            errorCode.message());
+        return;
+      }
+
+      const auto format = SILOpts.OptRecordFormat;
+      llvm::Expected<std::unique_ptr<llvm::remarks::RemarkSerializer>>
+        remarkSerializerOrErr = llvm::remarks::createRemarkSerializer(
+          format, llvm::remarks::SerializerMode::Separate, *file);
+      if (llvm::Error err = remarkSerializerOrErr.takeError()) {
+        diagEngine.diagnose(SourceLoc(), diag::error_creating_remark_serializer,
+                            toString(std::move(err)));
+        return;
+      }
+
+      auto auxRS = std::make_unique<llvm::remarks::RemarkStreamer>(
+        std::move(*remarkSerializerOrErr), filename);
+
+      Context.setMainRemarkStreamer(std::move(auxRS));
       Context.setLLVMRemarkStreamer(
-          std::make_unique<llvm::LLVMRemarkStreamer>(RS));
-      // FIXME: add a frontend flag to enable all LLVM remarks
-      cantFail(RS.setFilter("annotation-remarks"));
+        std::make_unique<llvm::LLVMRemarkStreamer>(
+          *Context.getMainRemarkStreamer()));
+      IGM.RemarkStream = std::move(file);
     }
   }
 }
@@ -1545,6 +1592,7 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
   auto &Ctx = M->getASTContext();
   // Create an IRGenModule for each source file.
   bool DidRunSILCodeGenPreparePasses = false;
+  unsigned idx = 0;
   for (auto *File : M->getFiles()) {
     auto nextSF = dyn_cast<SourceFile>(File);
     if (!nextSF)
@@ -1561,11 +1609,12 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
     if (!targetMachine) continue;
 
     // Create the IR emitter.
+    auto outputName = *OutputIter++;
     IRGenModule *IGM = new IRGenModule(
-        irgen, std::move(targetMachine), nextSF, desc.ModuleName, *OutputIter++,
+        irgen, std::move(targetMachine), nextSF, desc.ModuleName, outputName,
         nextSF->getFilename(), nextSF->getPrivateDiscriminator().str());
 
-    initLLVMModule(*IGM, *SILMod);
+    initLLVMModule(*IGM, *SILMod, idx++);
     if (!DidRunSILCodeGenPreparePasses) {
       // Run SIL level IRGen preparation passes on the module the first time
       // around.
diff --git a/test/Frontend/Inputs/opt-record-2.swift b/test/Frontend/Inputs/opt-record-2.swift
@@ -0,0 +1,24 @@
+open class C {
+    var x = 1
+    var y = 2
+
+    public init() {
+    }
+
+    public func method() {
+        print("x: \(x)")
+    }
+
+    public func method2() {
+        print("x2: \(y)")
+    }
+}
+
+@_assemblyVision
+@inline(never)
+public func runSomeTest(_ c: C) {
+  for i in 0..<100 {
+    c.method()
+    c.method2()
+  }
+}
diff --git a/test/Frontend/opt-record-supplemental.swift b/test/Frontend/opt-record-supplemental.swift
