Merge pull request #31167 from CodaFi/NSStream-beans

Lazily Associate a SILRemarkStreamer with an LLVMContext at IRGen time

Author: CodaFi

include/swift/SIL/SILModule.h

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

include/swift/SIL/SILRemarkStreamer.h

@@ -26,8 +26,20 @@ namespace swift {
 
 class SILRemarkStreamer {
 private:
-  /// The \c LLVMContext the underlying streamer uses for scratch space.
-  std::unique_ptr<llvm::LLVMContext> streamerContext;
+  enum class Owner {
+    SILModule,
+    LLVM,
+  } owner;
+
+  /// The underlying LLVM streamer.
+  ///
+  /// If owned by a SILModule, this will be non-null.
+  std::unique_ptr<llvm::remarks::RemarkStreamer> streamer;
+  /// The owning LLVM context.
+  ///
+  /// If owned by LLVM, this will be non-null.
+  llvm::LLVMContext *context;
+
   /// The remark output stream used to record SIL remarks to a file.
   std::unique_ptr<llvm::raw_fd_ostream> remarkStream;
 
@@ -53,6 +65,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.getSILRemarkStreamer()) {
+    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,35 @@ 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));
-}
+    : owner(Owner::SILModule), streamer(std::move(streamer)), context(nullptr),
+      remarkStream(std::move(stream)), ctx(Ctx) { }
 
 llvm::remarks::RemarkStreamer &SILRemarkStreamer::getLLVMStreamer() {
-  return *streamerContext->getMainRemarkStreamer();
+  switch (owner) {
+  case Owner::SILModule:
+    return *streamer.get();
+  case Owner::LLVM:
+    return *context->getMainRemarkStreamer();
+  }
+  return *streamer.get();
 }
 
 const llvm::remarks::RemarkStreamer &
 SILRemarkStreamer::getLLVMStreamer() const {
-  return *streamerContext->getMainRemarkStreamer();
+  switch (owner) {
+  case Owner::SILModule:
+    return *streamer.get();
+  case Owner::LLVM:
+    return *context->getMainRemarkStreamer();
+  }
+  return *streamer.get();
+}
+
+void SILRemarkStreamer::intoLLVMContext(llvm::LLVMContext &Ctx) & {
+  assert(owner == Owner::SILModule);
+  Ctx.setMainRemarkStreamer(std::move(streamer));
+  context = &Ctx;
+  owner = Owner::LLVM;
 }
 
 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

tools/sil-opt/SILOpt.cpp

@@ -461,7 +461,7 @@ int main(int argc, char **argv) {
   if (CI.getSILModule())
     CI.getSILModule()->setSerializeSILAction([]{});
 
-  if (RemarksFilename != "") {
+  if (!RemarksFilename.empty()) {
     llvm::Expected<llvm::remarks::Format> formatOrErr =
         llvm::remarks::parseFormat(RemarksFormat);
     if (llvm::Error E = formatOrErr.takeError()) {