IRGen: Add a frontend option to force single LLVM module emission in multithreaded mode

This allows to experiment with single module LLVM emission without
having to change drivers that expect multiple output files.

rdar://94744623

Author: aschwaighofer

include/swift/AST/IRGenOptions.h

@@ -284,6 +284,8 @@ class IRGenOptions {
   /// objects.
   unsigned EmitStackPromotionChecks : 1;
 
+  unsigned UseSingleModuleLLVMEmission : 1;
+
   /// Emit functions to separate sections.
   unsigned FunctionSections : 1;
 
@@ -447,13 +449,13 @@ class IRGenOptions {
         DebugInfoFormat(IRGenDebugInfoFormat::None),
         DisableClangModuleSkeletonCUs(false), UseJIT(false),
         DisableLLVMOptzns(false), DisableSwiftSpecificLLVMOptzns(false),
-        Playground(false),
-        EmitStackPromotionChecks(false), FunctionSections(false),
+        Playground(false), EmitStackPromotionChecks(false),
+        UseSingleModuleLLVMEmission(false), FunctionSections(false),
         PrintInlineTree(false), EmbedMode(IRGenEmbedMode::None),
         LLVMLTOKind(IRGenLLVMLTOKind::None),
         SwiftAsyncFramePointer(SwiftAsyncFramePointerKind::Auto),
-        HasValueNamesSetting(false),
-        ValueNames(false), ReflectionMetadata(ReflectionMetadataMode::Runtime),
+        HasValueNamesSetting(false), ValueNames(false),
+        ReflectionMetadata(ReflectionMetadataMode::Runtime),
         EnableReflectionNames(true), EnableAnonymousContextMangledNames(false),
         ForcePublicLinkage(false), LazyInitializeClassMetadata(false),
         LazyInitializeProtocolConformances(false),
@@ -468,8 +470,7 @@ class IRGenOptions {
         EnableGlobalISel(false), VirtualFunctionElimination(false),
         WitnessMethodElimination(false), ConditionalRuntimeRecords(false),
         InternalizeAtLink(false), InternalizeSymbols(false),
-        NoPreallocatedInstantiationCaches(false),
-        CmdArgs(),
+        NoPreallocatedInstantiationCaches(false), CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All) {
 #ifndef NDEBUG
@@ -530,8 +531,8 @@ class IRGenOptions {
     return llvm::hash_value(0);
   }
 
-  bool hasMultipleIRGenThreads() const { return NumThreads > 1; }
-  bool shouldPerformIRGenerationInParallel() const { return NumThreads != 0; }
+  bool hasMultipleIRGenThreads() const { return !UseSingleModuleLLVMEmission && NumThreads > 1; }
+  bool shouldPerformIRGenerationInParallel() const { return !UseSingleModuleLLVMEmission && NumThreads != 0; }
   bool hasMultipleIGMs() const { return hasMultipleIRGenThreads(); }
 };
 

include/swift/Option/FrontendOptions.td

@@ -505,6 +505,10 @@ def function_sections: Flag<["-"], "function-sections">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Emit functions to separate sections.">;
 
+def enable_single_module_llvm_emission: Flag<["-"], "enable-single-module-llvm-emission">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Emit LLVM IR into a single LLVM module in multithreaded mode.">;
+
 def stack_promotion_checks : Flag<["-"], "emit-stack-promotion-checks">,
   HelpText<"Emit runtime checks for correct stack promotion of objects.">;
 

include/swift/Subsystems.h

@@ -258,6 +258,11 @@ namespace swift {
                    llvm::Module *Module,
                    StringRef OutputFilename);
 
+  bool writeEmptyOutputFilesFor(
+    const ASTContext &Context,
+    std::vector<std::string> &ParallelOutputFilenames,
+    const IRGenOptions &IRGenOpts);
+
   /// Run the LLVM passes. In multi-threaded compilation this will be done for
   /// multiple LLVM modules in parallel.
   /// \param Diags The Diagnostic Engine.

lib/Frontend/CompilerInvocation.cpp

@@ -2299,6 +2299,9 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
                      "-num-threads");
     }
   }
+  Opts.UseSingleModuleLLVMEmission =
+      Opts.NumThreads != 0 &&
+      Args.hasArg(OPT_enable_single_module_llvm_emission);
 
   if (SWIFT_ENABLE_GLOBAL_ISEL_ARM64 &&
       Triple.getArch() == llvm::Triple::aarch64 &&

lib/FrontendTool/FrontendTool.cpp

@@ -1727,6 +1727,13 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
   if (validateTBDIfNeeded(Invocation, MSF, *IRModule.getModule()))
     return true;
 
+  if (IRGenOpts.UseSingleModuleLLVMEmission) {
+    // Pretend the other files that drivers/build systems expect exist by
+    // creating empty files.
+    if (writeEmptyOutputFilesFor(Context, ParallelOutputFilenames, IRGenOpts))
+      return true;
+  }
+
   return generateCode(Instance, OutputFilename, IRModule.getModule(),
                       HashGlobal);
 }

lib/IRGen/IRGen.cpp

@@ -1485,7 +1485,8 @@ GeneratedModule swift::performIRGeneration(
       outModuleHash);
 
   if (Opts.shouldPerformIRGenerationInParallel() &&
-      !parallelOutputFilenames.empty()) {
+      !parallelOutputFilenames.empty() &&
+      !Opts.UseSingleModuleLLVMEmission) {
     ::performParallelIRGeneration(desc);
     // TODO: Parallel LLVM compilation cannot be used if a (single) module is
     // needed as return value.

lib/IRGen/IRGenModule.cpp

@@ -28,6 +28,7 @@
 #include "swift/IRGen/Linking.h"
 #include "swift/Runtime/RuntimeFnWrappersGen.h"
 #include "swift/Runtime/Config.h"
+#include "swift/Subsystems.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/TargetInfo.h"
@@ -1636,12 +1637,7 @@ void IRGenModule::emitAutolinkInfo() {
   }
 }
 
-void IRGenModule::cleanupClangCodeGenMetadata() {
-  // Remove llvm.ident that ClangCodeGen might have left in the module.
-  auto *LLVMIdent = Module.getNamedMetadata("llvm.ident");
-  if (LLVMIdent)
-    Module.eraseNamedMetadata(LLVMIdent);
-
+void emitSwiftVersionNumberIntoModule(llvm::Module *Module) {
   // LLVM's object-file emission collects a fixed set of keys for the
   // image info.
   // Using "Objective-C Garbage Collection" as the key here is a hack,
@@ -1651,21 +1647,29 @@ void IRGenModule::cleanupClangCodeGenMetadata() {
   const char *ObjectiveCGarbageCollection = "Objective-C Garbage Collection";
   uint8_t Major, Minor;
   std::tie(Major, Minor) = version::getSwiftNumericVersion();
-  uint32_t Value = (Major << 24) | (Minor << 16) | (swiftVersion << 8);
-
-  if (Module.getModuleFlag(ObjectiveCGarbageCollection)) {
+  uint32_t Value =
+      (Major << 24) | (Minor << 16) | (IRGenModule::swiftVersion << 8);
+  auto &llvmContext = Module->getContext();
+  if (Module->getModuleFlag(ObjectiveCGarbageCollection)) {
     bool FoundOldEntry = replaceModuleFlagsEntry(
-        Module.getContext(), Module, ObjectiveCGarbageCollection,
+        llvmContext, *Module, ObjectiveCGarbageCollection,
         llvm::Module::Override,
-        llvm::ConstantAsMetadata::get(
-            llvm::ConstantInt::get(Int32Ty, Value)));
+        llvm::ConstantAsMetadata::get(llvm::ConstantInt::get(
+            llvm::Type::getInt32Ty(llvmContext), Value)));
 
     (void)FoundOldEntry;
     assert(FoundOldEntry && "Could not replace old module flag entry?");
   } else
-    Module.addModuleFlag(llvm::Module::Override,
-                         ObjectiveCGarbageCollection,
-                         Value);
+    Module->addModuleFlag(llvm::Module::Override, ObjectiveCGarbageCollection,
+                          Value);
+}
+
+void IRGenModule::cleanupClangCodeGenMetadata() {
+  // Remove llvm.ident that ClangCodeGen might have left in the module.
+  auto *LLVMIdent = Module.getNamedMetadata("llvm.ident");
+  if (LLVMIdent)
+    Module.eraseNamedMetadata(LLVMIdent);
+  emitSwiftVersionNumberIntoModule(&Module);
 }
 
 bool IRGenModule::finalize() {
@@ -1884,3 +1888,40 @@ bool IRGenModule::isConcurrencyAvailable() {
     AvailabilityContext::forDeploymentTarget(ctx);
   return deploymentAvailability.isContainedIn(ctx.getConcurrencyAvailability());
 }
+
+/// Pretend the other files that drivers/build systems expect exist by
+/// creating empty files. Used by UseSingleModuleLLVMEmission when
+/// num-threads > 0.
+bool swift::writeEmptyOutputFilesFor(
+  const ASTContext &Context,
+  std::vector<std::string>& ParallelOutputFilenames,
+  const IRGenOptions &IRGenOpts) {
+
+  for (auto fileName : ParallelOutputFilenames) {
+    // The first output file, was use for genuine output.
+    if (fileName == ParallelOutputFilenames[0])
+      continue;
+
+    std::unique_ptr<llvm::LLVMContext> llvmContext(new llvm::LLVMContext());
+    std::unique_ptr<clang::CodeGenerator> clangCodeGen(
+      createClangCodeGenerator(const_cast<ASTContext&>(Context),
+                               *llvmContext, IRGenOpts, fileName, ""));
+    auto *llvmModule = clangCodeGen->GetModule();
+
+    auto *clangImporter = static_cast<ClangImporter *>(
+      Context.getClangModuleLoader());
+    llvmModule->setTargetTriple(
+      clangImporter->getTargetInfo().getTargetOpts().Triple);
+
+    // Add LLVM module flags.
+    auto &clangASTContext = clangImporter->getClangASTContext();
+    clangCodeGen->HandleTranslationUnit(
+        const_cast<clang::ASTContext &>(clangASTContext));
+
+    emitSwiftVersionNumberIntoModule(llvmModule);
+
+    swift::performLLVM(IRGenOpts, const_cast<ASTContext&>(Context),
+                       llvmModule, fileName);
+  }
+  return false;
+}

test/IRGen/Inputs/single-module-num-threads-2.swift

@@ -0,0 +1,3 @@
+public func B() {
+  print("B")
+}

test/IRGen/single-module-num-threads.swift

@@ -0,0 +1,16 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name X -num-threads 1 -O -enable-single-module-llvm-emission -emit-ir %s %S/Inputs/single-module-num-threads-2.swift -o %t/single-module-num-threads.ll -o %t/single-module-num-threads-2.ll
+// RUN: %FileCheck %s < %t/single-module-num-threads.ll
+// RUN: %FileCheck %s --check-prefix=EMPTY < %t/single-module-num-threads-2.ll
+
+// CHECK: define{{.*}} swiftcc void @"$s1X1AyyF"()
+// CHECK: define{{.*}} swiftcc void @"$s1X1ByyF"()
+
+// EMPTY-NOT: s1X1AyyF
+// EMPTY-NOT: s1X1ByyF
+
+// Make sure that with enable-single-module-llvm-emission we emit all code into
+// one llvm module.
+public func A() {
+    print("A")
+}