Extend LLVM IR WME to use thunks for cross-module witness method calls. (#39528)

This enables optimizing / dead-stripping of witness methods across modules at
LTO time.

- Under -internalize-at-link, restrict visibility of wtables to linkage unit.
- Emit thunks for cross-module wcalls when WME is enabled.
- Use thunks for wcalls across modules when WME is enabled.
- Adjust TBDGen to account for witness method thunks when WME is enabled.
- Add an IR test to check that thunks are used when doing cross-module calls.
- Add an end-to-end test case for cross-module WME.

include/swift/TBDGen/TBDGen.h

@@ -44,6 +44,9 @@ struct TBDGenOptions {
   /// Whether LLVM IR Virtual Function Elimination is enabled.
   bool VirtualFunctionElimination = false;
 
+  /// Whether LLVM IR Witness Method Elimination is enabled.
+  bool WitnessMethodElimination = false;
+
   /// The install_name to use in the TBD file.
   std::string InstallName;
 
@@ -74,6 +77,7 @@ struct TBDGenOptions {
            lhs.LinkerDirectivesOnly == rhs.LinkerDirectivesOnly &&
            lhs.PublicSymbolsOnly == rhs.PublicSymbolsOnly &&
            lhs.VirtualFunctionElimination == rhs.VirtualFunctionElimination &&
+           lhs.WitnessMethodElimination == rhs.WitnessMethodElimination &&
            lhs.InstallName == rhs.InstallName &&
            lhs.ModuleLinkName == rhs.ModuleLinkName &&
            lhs.CurrentVersion == rhs.CurrentVersion &&
@@ -91,6 +95,7 @@ struct TBDGenOptions {
     return hash_combine(
         opts.HasMultipleIGMs, opts.IsInstallAPI, opts.LinkerDirectivesOnly,
         opts.PublicSymbolsOnly, opts.VirtualFunctionElimination,
+        opts.WitnessMethodElimination,
         opts.InstallName, opts.ModuleLinkName,
         opts.CurrentVersion, opts.CompatibilityVersion,
         opts.ModuleInstallNameMapPath,

lib/Frontend/CompilerInvocation.cpp

@@ -1542,6 +1542,7 @@ static bool ParseTBDGenArgs(TBDGenOptions &Opts, ArgList &Args,
   Opts.IsInstallAPI = Args.hasArg(OPT_tbd_is_installapi);
 
   Opts.VirtualFunctionElimination = Args.hasArg(OPT_enable_llvm_vfe);
+  Opts.WitnessMethodElimination = Args.hasArg(OPT_enable_llvm_wme);
 
   if (const Arg *A = Args.getLastArg(OPT_tbd_compatibility_version)) {
     Opts.CompatibilityVersion = A->getValue();

lib/IRGen/GenMeta.cpp

@@ -824,7 +824,7 @@ namespace {
       SILDeclRef func(entry.getFunction());
 
       // Emit the dispatch thunk.
-      if (Resilient)
+      if (Resilient || IGM.getOptions().WitnessMethodElimination)
         IGM.emitDispatchThunk(func);
 
       // Classify the function.

lib/IRGen/GenProto.cpp

@@ -2193,8 +2193,13 @@ static void addWTableTypeMetadata(IRGenModule &IGM,
     break;
   case SILLinkage::Public:
   default:
-    global->setVCallVisibilityMetadata(
-        llvm::GlobalObject::VCallVisibility::VCallVisibilityPublic);
+    if (IGM.getOptions().InternalizeAtLink) {
+      global->setVCallVisibilityMetadata(
+          llvm::GlobalObject::VCallVisibility::VCallVisibilityLinkageUnit);
+    } else {
+      global->setVCallVisibilityMetadata(
+          llvm::GlobalObject::VCallVisibility::VCallVisibilityPublic);
+    }
     break;
   }
 }

lib/IRGen/IRGenSIL.cpp

@@ -6568,8 +6568,18 @@ void IRGenSILFunction::visitWitnessMethodInst(swift::WitnessMethodInst *i) {
 
   assert(member.requiresNewWitnessTableEntry());
 
-  if (IGM.isResilient(conformance.getRequirement(),
-                      ResilienceExpansion::Maximal)) {
+  bool shouldUseDispatchThunk = false;
+  if (IGM.isResilient(conformance.getRequirement(), ResilienceExpansion::Maximal)) {
+    shouldUseDispatchThunk = true;
+  } else if (IGM.getOptions().WitnessMethodElimination) {
+    // For WME, use a thunk if the target protocol is defined in another module.
+    // This way, we guarantee all wmethod call sites are visible to the LLVM VFE
+    // optimization in GlobalDCE.
+    auto protoDecl = cast<ProtocolDecl>(member.getDecl()->getDeclContext());
+    shouldUseDispatchThunk = protoDecl->getModuleContext() != IGM.getSwiftModule();
+  }
+
+  if (shouldUseDispatchThunk) {
     llvm::Constant *fnPtr = IGM.getAddrOfDispatchThunk(member, NotForDefinition);
     llvm::Constant *secondaryValue = nullptr;
 

lib/TBDGen/TBDGen.cpp

@@ -1077,7 +1077,7 @@ void TBDGenVisitor::visitProtocolDecl(ProtocolDecl *PD) {
           : TBD(TBD), PD(PD), Resilient(PD->getParentModule()->isResilient()) {}
 
       void addMethod(SILDeclRef declRef) {
-        if (Resilient) {
+        if (Resilient || TBD.Opts.WitnessMethodElimination) {
           TBD.addDispatchThunk(declRef);
           TBD.addMethodDescriptor(declRef);
         }

test/IRGen/witness-method-elimination-ir-thunks.swift

@@ -0,0 +1,46 @@
+// Tests that under -enable-llvm-wme, protocol witness table calls to protocols
+// defined by other modules are using thunks (instead of direct wtable loads).
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -enable-llvm-wme -parse-as-library %s -DLIBRARY -module-name Library -emit-module -o %t/Library.swiftmodule
+// RUN: %target-build-swift -Xfrontend -enable-llvm-wme -parse-as-library %s -DCLIENT -module-name Main -I%t -emit-ir -o - | %FileCheck %s
+
+#if LIBRARY
+
+public protocol MyLibraryProtocol {
+  func library_req()
+}
+
+#endif
+
+#if CLIENT
+
+import Library
+
+public protocol MyLocalProtocol {
+  func local_req()
+}
+
+extension MyLocalProtocol {
+  func func1() {
+    // CHECK: define hidden swiftcc void @"$s4Main15MyLocalProtocolPAAE5func1yyF"
+    self.local_req()
+    // CHECK: [[SLOT:%.*]] = getelementptr inbounds i8*, i8** {{.*}}, i32 1
+    // CHECK: [[SLOTASPTR:%.*]] = bitcast i8** [[SLOT]] to i8*
+    // CHECK: call { i8*, i1 } @llvm.type.checked.load(i8* [[SLOTASPTR]], i32 0, metadata !"$s4Main15MyLocalProtocolP9local_reqyyFTq")
+    // CHECK: ret void
+  }
+}
+
+extension MyLibraryProtocol {
+  func func2() {
+    // CHECK: define hidden swiftcc void @"$s7Library02MyA8ProtocolP4MainE5func2yyF"
+    self.library_req()
+    // CHECK: call swiftcc void @"$s7Library02MyA8ProtocolP11library_reqyyFTj"
+
+    // CHECK-NOT: @llvm.type.checked.load
+    // CHECK: ret void
+  }
+}
+
+#endif

test/IRGen/witness-method-elimination-two-modules.swift

@@ -0,0 +1,72 @@
+// Tests that under -enable-llvm-wme + -internalize-at-link, cross-module
+// witness method calls are done via thunks and LLVM GlobalDCE is able to remove
+// unused witness methods from a library based on a list of used symbols by a
+// client.
+
+// RUN: %empty-directory(%t)
+
+// (1) Build library swiftmodule
+// RUN: %target-build-swift -parse-as-library -Xfrontend -enable-llvm-wme \
+// RUN:     %s -DLIBRARY -module-name Library \
+// RUN:     -emit-module -o %t/Library.swiftmodule \
+// RUN:     -emit-tbd -emit-tbd-path %t/libLibrary.tbd -Xfrontend -tbd-install_name=%t/libLibrary.dylib
+
+// (2) Build client
+// RUN: %target-build-swift -parse-as-library -Xfrontend -enable-llvm-wme \
+// RUN:     %s -DCLIENT -module-name Main -I%t -L%t -lLibrary -o %t/main
+
+// (3) Extract a list of used symbols by client from library
+// RUN: %llvm-nm --undefined-only -m %t/main | grep 'from libLibrary' | awk '{print $3}' > %t/used-symbols
+
+// (4) Now produce the .dylib with just the symbols needed by the client
+// RUN: %target-build-swift -parse-as-library -Xfrontend -enable-llvm-wme -Xfrontend -internalize-at-link \
+// RUN:     %s -DLIBRARY -lto=llvm-full %lto_flags -module-name Library \
+// RUN:     -emit-library -o %t/libLibrary.dylib \
+// RUN:     -Xlinker -exported_symbols_list -Xlinker %t/used-symbols -Xlinker -dead_strip
+
+// (5) Check list of symbols in library
+// RUN: %llvm-nm --defined-only %t/libLibrary.dylib | %FileCheck %s --check-prefix=NM
+
+// (6) Execution test
+// RUN: %target-run %t/main | %FileCheck %s
+
+// REQUIRES: executable_test
+
+// Test disabled until LLVM GlobalDCE supports Swift wtables.
+// REQUIRES: rdar81868930
+
+#if LIBRARY
+
+public protocol MyProtocol {
+  func func1_used()
+  func func2_unused()
+}
+
+public struct MyStruct : MyProtocol {
+  public init() {}
+  public func func1_used() { print("MyStruct.func1_used") }
+  public func func2_unused() { print("MyStruct.func2_unused") }
+}
+
+// NM:     $s7Library8MyStructV10func1_usedyyF
+// NM-NOT: $s7Library8MyStructV12func2_unusedyyF
+// NM:     $s7Library8MyStructVAA0B8ProtocolA2aDP10func1_usedyyFTW
+// NM-NOT: $s7Library8MyStructVAA0B8ProtocolA2aDP12func2_unusedyyFTW
+
+#endif
+
+#if CLIENT
+
+import Library
+
+@_cdecl("main")
+func main() -> Int32 {
+   let o: MyProtocol = MyStruct()
+   o.func1_used()
+   print("Done")
+   // CHECK: MyStruct.func1_used
+   // CHECK-NEXT: Done
+   return 0
+}
+
+#endif