Lazily Associated a SILRemarkStreamer with an LLVMContext at IRGen time

Corrects a mistake introduced in #31106

I was under the impression that the LLVMContext for an instance of
llvm::remarks::RemarkStreamer was somehow just scratch-space. It turns
out the ASMPrinters don't gather remarks data from modules, they gather
it from the remark streamer associated with the module's context.

Thus, we cannot have the module's context be distinct from whatever
context the streamer is eventually associated with.

In order to bring these two systems back into harmony, introduce a simple
ownership contract to SILRemarkStreamer. That is, it starts owned by
a SILModule, in which case the SILRemarkStreamer holds onto the
underlying LLVM object as the optimizer emits remarks. When it comes
time to IRGen the module, then and only then do we install the streamer
on the module's context. This tranfers ownership of the underlying LLVM
streamer to LLVM itself, so it acts as a consuming operation. When we
are about to perform IR Generation, the SILModule will be expiring
anyways, so the streamer was already about to be destroyed.

We're just putting it to better use doing its job.

Author: CodaFi

include/swift/SIL/SILModule.h

@@ -523,6 +523,9 @@ class SILModule {
   swift::SILRemarkStreamer *getSILRemarkStreamer() {
     return silRemarkStreamer.get();
   }
+  std::unique_ptr<swift::SILRemarkStreamer> takeSILRemarkStreamer() {
+    return std::move(silRemarkStreamer);
+  }
   void installSILRemarkStreamer();
 
   // This is currently limited to VarDecl because the visibility of global

include/swift/SIL/SILRemarkStreamer.h

@@ -26,8 +26,11 @@ namespace swift {
 
 class SILRemarkStreamer {
 private:
-  /// The \c LLVMContext the underlying streamer uses for scratch space.
-  std::unique_ptr<llvm::LLVMContext> streamerContext;
+  /// The underlying LLVM streamer.
+  ///
+  /// If owned by a SILModule, this will be non-null.
+  std::unique_ptr<llvm::remarks::RemarkStreamer> streamer;
+
   /// The remark output stream used to record SIL remarks to a file.
   std::unique_ptr<llvm::raw_fd_ostream> remarkStream;
 
@@ -53,6 +56,12 @@ class SILRemarkStreamer {
 
   const ASTContext &getASTContext() const { return ctx; }
 
+public:
+  /// Perform a one-time ownership transfer to associate the underlying
+  /// \c llvm::remarks::RemarkStreamer with the given \c LLVMContext.
+  void intoLLVMContext(llvm::LLVMContext &Ctx) &;
+
+public:
   /// Emit a remark through the streamer.
   template <typename RemarkT>
   void emit(const OptRemark::Remark<RemarkT> &remark);

lib/IRGen/IRGen.cpp

@@ -34,6 +34,7 @@
 #include "swift/LLVMPasses/Passes.h"
 #include "swift/LLVMPasses/PassesFwd.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SIL/SILRemarkStreamer.h"
 #include "swift/SILOptimizer/PassManager/PassManager.h"
 #include "swift/SILOptimizer/PassManager/PassPipeline.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
@@ -884,7 +885,7 @@ static void embedBitcode(llvm::Module *M, const IRGenOptions &Opts)
   NewUsed->setSection("llvm.metadata");
 }
 
-static void initLLVMModule(const IRGenModule &IGM, ModuleDecl &M) {
+static void initLLVMModule(const IRGenModule &IGM, SILModule &SIL) {
   auto *Module = IGM.getModule();
   assert(Module && "Expected llvm:Module for IR generation!");
 
@@ -902,12 +903,17 @@ static void initLLVMModule(const IRGenModule &IGM, ModuleDecl &M) {
 
   auto *MDNode = IGM.getModule()->getOrInsertNamedMetadata("swift.module.flags");
   auto &Context = IGM.getModule()->getContext();
-  auto *Value = M.isStdlibModule() ? llvm::ConstantInt::getTrue(Context)
-                                   : llvm::ConstantInt::getFalse(Context);
+  auto *Value = SIL.getSwiftModule()->isStdlibModule()
+              ? llvm::ConstantInt::getTrue(Context)
+              : llvm::ConstantInt::getFalse(Context);
   MDNode->addOperand(llvm::MDTuple::get(Context,
                                         {llvm::MDString::get(Context,
                                                              "standard-library"),
                                          llvm::ConstantAsMetadata::get(Value)}));
+
+  if (auto streamer = SIL.takeSILRemarkStreamer()) {
+    std::move(streamer)->intoLLVMContext(Module->getContext());
+  }
 }
 
 std::pair<IRGenerator *, IRGenModule *>
@@ -926,7 +932,7 @@ swift::irgen::createIRGenModule(SILModule *SILMod, StringRef OutputFilename,
       *irgen, std::move(targetMachine), nullptr, "", OutputFilename,
       MainInputFilenameForDebugInfo, PrivateDiscriminator);
 
-  initLLVMModule(*IGM, *SILMod->getSwiftModule());
+  initLLVMModule(*IGM, *SILMod);
 
   return std::pair<IRGenerator *, IRGenModule *>(irgen, IGM);
 }
@@ -969,7 +975,7 @@ performIRGeneration(const IRGenOptions &Opts, ModuleDecl *M,
                   PSPs.OutputFilename, PSPs.MainInputFilenameForDebugInfo,
                   PrivateDiscriminator);
 
-  initLLVMModule(IGM, *SILMod->getSwiftModule());
+  initLLVMModule(IGM, *SILMod);
 
   // Run SIL level IRGen preparation passes.
   runIRGenPreparePasses(*SILMod, IGM);
@@ -1217,7 +1223,7 @@ static void performParallelIRGeneration(
                         nextSF->getPrivateDiscriminator().str());
     IGMcreated = true;
 
-    initLLVMModule(*IGM, *SILMod->getSwiftModule());
+    initLLVMModule(*IGM, *SILMod);
     if (!DidRunSILCodeGenPreparePasses) {
       // Run SIL level IRGen preparation passes on the module the first time
       // around.
@@ -1431,7 +1437,7 @@ swift::createSwiftModuleObjectFile(SILModule &SILMod, StringRef Buffer,
 
   IRGenModule IGM(irgen, std::move(targetMachine), nullptr,
                   OutputPath, OutputPath, "", "");
-  initLLVMModule(IGM, *SILMod.getSwiftModule());
+  initLLVMModule(IGM, SILMod);
   auto *Ty = llvm::ArrayType::get(IGM.Int8Ty, Buffer.size());
   auto *Data =
       llvm::ConstantDataArray::getString(IGM.getLLVMContext(),

lib/SIL/Utils/SILRemarkStreamer.cpp

@@ -19,18 +19,20 @@ using namespace swift;
 SILRemarkStreamer::SILRemarkStreamer(
     std::unique_ptr<llvm::remarks::RemarkStreamer> &&streamer,
     std::unique_ptr<llvm::raw_fd_ostream> &&stream, const ASTContext &Ctx)
-    : streamerContext(std::make_unique<llvm::LLVMContext>()),
-      remarkStream(std::move(stream)), ctx(Ctx) {
-  streamerContext->setMainRemarkStreamer(std::move(streamer));
-}
+    : streamer(std::move(streamer)),
+      remarkStream(std::move(stream)), ctx(Ctx) { }
 
 llvm::remarks::RemarkStreamer &SILRemarkStreamer::getLLVMStreamer() {
-  return *streamerContext->getMainRemarkStreamer();
+  return *streamer.get();
 }
 
 const llvm::remarks::RemarkStreamer &
 SILRemarkStreamer::getLLVMStreamer() const {
-  return *streamerContext->getMainRemarkStreamer();
+  return *streamer.get();
+}
+
+void SILRemarkStreamer::intoLLVMContext(llvm::LLVMContext &Ctx) & {
+  Ctx.setMainRemarkStreamer(std::move(streamer));
 }
 
 std::unique_ptr<SILRemarkStreamer>

test/Driver/opt-record-bitstream.swift

@@ -1,9 +1,14 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-swiftc_driver -O -wmo -save-optimization-record=bitstream %s -module-name optrecordmod -o %t/opt-record 2>&1 | %FileCheck -allow-empty %s
+// RUN: %target-swift-frontend -c -O -wmo -save-optimization-record=bitstream -save-optimization-record-path %t/optrecordmod.opt.bitstream %s -module-name optrecordmod -o %t/opt-record.o 2>&1 | %FileCheck -allow-empty %s
 // RUN: llvm-bcanalyzer -dump %t/optrecordmod.opt.bitstream | %FileCheck -check-prefix=BITSTREAM %s
+// RUN: otool -l %t/opt-record.o | %FileCheck -check-prefix=OBJ %s
 
 // REQUIRES: OS=macosx || OS=ios || OS=tvos || OS=watchos
 
+// Ensure we emitted the appropriate section
+
+// OBJ: sectname __remarks
+
 // CHECK-NOT: remark
 
 var a: Int = 1