Merge pull request #59344 from aschwaighofer/single_module_in_num_threads

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

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")
+}