IRGen: Ensure that default witness thunks are emitted in the same thread as the protocol descriptor

lib/IRGen/GenDecl.cpp

@@ -984,7 +984,7 @@ static bool hasCodeCoverageInstrumentation(SILFunction &f, SILModule &m) {
   return f.getProfiler() && m.getOptions().EmitProfileCoverageMapping;
 }
 
-void IRGenerator::emitGlobalTopLevel(bool emitForParallelEmission) {
+void IRGenerator::emitGlobalTopLevel() {
   // Generate order numbers for the functions in the SIL module that
   // correspond to definitions in the LLVM module.
   unsigned nextOrderNumber = 0;
@@ -994,13 +994,6 @@ void IRGenerator::emitGlobalTopLevel(bool emitForParallelEmission) {
     FunctionOrder.insert(std::make_pair(&silFn, nextOrderNumber++));
   }
 
-  // Ensure that relative symbols are collocated in the same LLVM module.
-  for (SILWitnessTable &wt : PrimaryIGM->getSILModule().getWitnessTableList()) {
-    CurrentIGMPtr IGM = getGenModule(wt.getDeclContext());
-    if (emitForParallelEmission)
-      IGM->ensureRelativeSymbolCollocation(wt);
-  }
-
   for (SILGlobalVariable &v : PrimaryIGM->getSILModule().getSILGlobals()) {
     Decl *decl = v.getDecl();
     CurrentIGMPtr IGM = getGenModule(decl ? decl->getDeclContext() : nullptr);

lib/IRGen/GenMeta.cpp

@@ -4046,6 +4046,9 @@ void IRGenModule::emitProtocolDecl(ProtocolDecl *protocol) {
   if (isResilient(protocol, ResilienceExpansion::Minimal))
     defaultWitnesses = getSILModule().lookUpDefaultWitnessTable(protocol);
 
+  if (defaultWitnesses)
+    IRGen.ensureRelativeSymbolCollocation(*defaultWitnesses);
+
   {
     ProtocolDescriptorBuilder builder(*this, protocol, defaultWitnesses);
     builder.emit();

lib/IRGen/GenProto.cpp

@@ -38,8 +38,10 @@
 #include "swift/ClangImporter/ClangModule.h"
 #include "swift/IRGen/Linking.h"
 #include "swift/SIL/SILDeclRef.h"
+#include "swift/SIL/SILDefaultWitnessTable.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/SILValue.h"
+#include "swift/SIL/SILWitnessTable.h"
 #include "swift/SIL/SILWitnessVisitor.h"
 #include "swift/SIL/TypeLowering.h"
 #include "llvm/ADT/SmallString.h"
@@ -2079,7 +2081,10 @@ void IRGenModule::emitProtocolConformance(
   setTrueConstGlobal(var);
 }
 
-void IRGenModule::ensureRelativeSymbolCollocation(SILWitnessTable &wt) {
+void IRGenerator::ensureRelativeSymbolCollocation(SILWitnessTable &wt) {
+  if (!CurrentIGM)
+    return;
+
   // Only resilient conformances use relative pointers for witness methods.
   if (wt.isDeclaration() || isAvailableExternally(wt.getLinkage()) ||
       !isResilientConformance(wt.getConformance()))
@@ -2090,7 +2095,20 @@ void IRGenModule::ensureRelativeSymbolCollocation(SILWitnessTable &wt) {
       continue;
     auto *witness = entry.getMethodWitness().Witness;
     if (witness)
-      IRGen.forceLocalEmitOfLazyFunction(witness);
+      forceLocalEmitOfLazyFunction(witness);
+  }
+}
+
+void IRGenerator::ensureRelativeSymbolCollocation(SILDefaultWitnessTable &wt) {
+  if (!CurrentIGM)
+    return;
+
+  for (auto &entry : wt.getEntries()) {
+    if (entry.getKind() != SILWitnessTable::Method)
+      continue;
+    auto *witness = entry.getMethodWitness().Witness;
+    if (witness)
+      forceLocalEmitOfLazyFunction(witness);
   }
 }
 
@@ -2234,6 +2252,10 @@ void IRGenModule::emitSILWitnessTable(SILWitnessTable *wt) {
   if (isAvailableExternally(wt->getLinkage()))
     return;
 
+  // Ensure that relatively-referenced symbols for witness thunks are collocated
+  // in the same LLVM module.
+  IRGen.ensureRelativeSymbolCollocation(*wt);
+
   auto conf = wt->getConformance();
   PrettyStackTraceConformance _st(Context, "emitting witness table for", conf);
 

lib/IRGen/IRGen.cpp

@@ -964,11 +964,11 @@ static void performParallelIRGeneration(
   }
 
   // Emit the module contents.
-  irgen.emitGlobalTopLevel(true /*emitForParallelEmission*/);
+  irgen.emitGlobalTopLevel();
 
   for (auto *File : M->getFiles()) {
     if (auto *SF = dyn_cast<SourceFile>(File)) {
-      IRGenModule *IGM = irgen.getGenModule(SF);
+      CurrentIGMPtr IGM = irgen.getGenModule(SF);
       IGM->emitSourceFile(*SF);
     } else {
       File->collectLinkLibraries([&](LinkLibrary LinkLib) {

lib/IRGen/IRGenModule.h

@@ -93,6 +93,7 @@ namespace swift {
   class ProtocolCompositionType;
   class RootProtocolConformance;
   struct SILDeclRef;
+  class SILDefaultWitnessTable;
   class SILGlobalVariable;
   class SILModule;
   class SILProperty;
@@ -311,11 +312,7 @@ class IRGenerator {
 
   /// Emit functions, variables and tables which are needed anyway, e.g. because
   /// they are externally visible.
-  /// If \p emitForParallelEmission is true ensures that symbols that are
-  /// expressed as relative pointers are collocated in the same output module
-  /// with their base symbol. For example, witness methods need to be collocated
-  /// with the witness table in the same LLVM module.
-  void emitGlobalTopLevel(bool emitForParallelEmission = false);
+  void emitGlobalTopLevel();
 
   /// Emit references to each of the protocol descriptors defined in this
   /// IR module.
@@ -355,6 +352,10 @@ class IRGenerator {
     DefaultIGMForFunction[f] = CurrentIGM;
   }
 
+  void ensureRelativeSymbolCollocation(SILWitnessTable &wt);
+
+  void ensureRelativeSymbolCollocation(SILDefaultWitnessTable &wt);
+
   void noteUseOfTypeMetadata(NominalTypeDecl *type) {
     noteUseOfTypeGlobals(type, true, RequireMetadata);
   }
@@ -1386,8 +1387,6 @@ private:                            \
 
   void emitSharedContextDescriptor(DeclContext *dc);
 
-  void ensureRelativeSymbolCollocation(SILWitnessTable &wt);
-
   llvm::GlobalVariable *
   getGlobalForDynamicallyReplaceableThunk(LinkEntity &entity, llvm::Type *type,
                                           ForDefinition_t forDefinition);

test/IRGen/Inputs/multithread_resilience_other.swift

@@ -0,0 +1,3 @@
+extension Defaultable {
+  public func defaulted() {}
+}

test/IRGen/multithread_resilience.swift

@@ -0,0 +1,9 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -c %S/Inputs/multithread_resilience_other.swift %s -num-threads 2 -o %t/1.o -o %t/2.o -module-name multithread_resilience
+// RUN: %target-swift-frontend -c %S/Inputs/multithread_resilience_other.swift %s -num-threads 2 -o %t/1.o -o %t/2.o -module-name multithread_resilience -enable-testing
+// RUN: %target-swift-frontend -c %S/Inputs/multithread_resilience_other.swift %s -num-threads 2 -o %t/1.o -o %t/2.o -module-name multithread_resilience -enable-resilience
+// RUN: %target-swift-frontend -c %S/Inputs/multithread_resilience_other.swift %s -num-threads 2 -o %t/1.o -o %t/2.o -module-name multithread_resilience -enable-testing -enable-resilience
+
+public protocol Defaultable {
+  func defaulted()
+}