Add -internalize-at-link flag to allow dead stripping and VFE across modules at link time (#39214)

- Under -internalize-at-link, stop unconditionally marking all globals as used.
- Under -internalize-at-link, restrict visibility of vtables to linkage unit.
- Emit virtual method thunks for cross-module vcalls when VFE is enabled.
- Use thunks for vcalls across modules when VFE is enabled.
- Adjust TBDGen to account for virtual method thunks when VFE is enabled.
- Add an end-to-end test case for cross-module VFE.

Author: kubamracek

include/swift/AST/IRGenOptions.h

@@ -357,6 +357,8 @@ class IRGenOptions {
 
   unsigned WitnessMethodElimination : 1;
 
+  unsigned InternalizeAtLink : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -416,7 +418,7 @@ class IRGenOptions {
         DisableRoundTripDebugTypes(false), DisableDebuggerShadowCopies(false),
         DisableConcreteTypeMetadataMangledNameAccessors(false),
         EnableGlobalISel(false), VirtualFunctionElimination(false),
-        WitnessMethodElimination(false),
+        WitnessMethodElimination(false), InternalizeAtLink(false),
         CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All) {}

include/swift/Option/FrontendOptions.td

@@ -501,6 +501,12 @@ def enable_llvm_wme : Flag<["-"], "enable-llvm-wme">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
   HelpText<"Use LLVM IR Witness Method Elimination on Swift protocol witness tables">;
 
+def internalize_at_link : Flag<["-"], "internalize-at-link">,
+  Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
+  HelpText<"Allow internalizing public symbols and vtables at link time (assume"
+  " all client code of public types is part of the same link unit, or that"
+  " external symbols are explicitly requested via -exported_symbols_list)">;
+
 def disable_previous_implementation_calls_in_dynamic_replacements :
   Flag<["-"], "disable-previous-implementation-calls-in-dynamic-replacements">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,

include/swift/TBDGen/TBDGen.h

@@ -41,6 +41,9 @@ struct TBDGenOptions {
   /// Whether to include only symbols with public linkage.
   bool PublicSymbolsOnly = true;
 
+  /// Whether LLVM IR Virtual Function Elimination is enabled.
+  bool VirtualFunctionElimination = false;
+
   /// The install_name to use in the TBD file.
   std::string InstallName;
 
@@ -70,6 +73,7 @@ struct TBDGenOptions {
            lhs.IsInstallAPI == rhs.IsInstallAPI &&
            lhs.LinkerDirectivesOnly == rhs.LinkerDirectivesOnly &&
            lhs.PublicSymbolsOnly == rhs.PublicSymbolsOnly &&
+           lhs.VirtualFunctionElimination == rhs.VirtualFunctionElimination &&
            lhs.InstallName == rhs.InstallName &&
            lhs.ModuleLinkName == rhs.ModuleLinkName &&
            lhs.CurrentVersion == rhs.CurrentVersion &&
@@ -86,7 +90,8 @@ struct TBDGenOptions {
     using namespace llvm;
     return hash_combine(
         opts.HasMultipleIGMs, opts.IsInstallAPI, opts.LinkerDirectivesOnly,
-        opts.PublicSymbolsOnly, opts.InstallName, opts.ModuleLinkName,
+        opts.PublicSymbolsOnly, opts.VirtualFunctionElimination,
+        opts.InstallName, opts.ModuleLinkName,
         opts.CurrentVersion, opts.CompatibilityVersion,
         opts.ModuleInstallNameMapPath,
         hash_combine_range(opts.embedSymbolsFromModules.begin(),

lib/Frontend/CompilerInvocation.cpp

@@ -1528,6 +1528,8 @@ static bool ParseTBDGenArgs(TBDGenOptions &Opts, ArgList &Args,
 
   Opts.IsInstallAPI = Args.hasArg(OPT_tbd_is_installapi);
 
+  Opts.VirtualFunctionElimination = Args.hasArg(OPT_enable_llvm_vfe);
+
   if (const Arg *A = Args.getLastArg(OPT_tbd_compatibility_version)) {
     Opts.CompatibilityVersion = A->getValue();
   }
@@ -1912,6 +1914,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Opts.WitnessMethodElimination = true;
   }
 
+  if (Args.hasArg(OPT_internalize_at_link)) {
+    Opts.InternalizeAtLink = true;
+  }
+
   // Default to disabling swift async extended frame info on anything but
   // darwin. Other platforms are unlikely to have support for extended frame
   // pointer information.

lib/IRGen/GenDecl.cpp

@@ -880,6 +880,29 @@ IRGenModule::getAddrOfParentContextDescriptor(DeclContext *from,
                                              fromAnonymousContext);
 }
 
+static void markGlobalAsUsedBasedOnLinkage(IRGenModule &IGM, LinkInfo &link,
+    llvm::GlobalValue *global) {
+  // If we're internalizing public symbols at link time, don't make globals
+  // unconditionally externally visible.
+  if (IGM.getOptions().InternalizeAtLink)
+    return;
+
+  // Everything externally visible is considered used in Swift.
+  // That mostly means we need to be good at not marking things external.
+  if (link.isUsed())
+    IGM.addUsedGlobal(global);
+}
+
+bool LinkInfo::isUsed(IRLinkage IRL) {
+  // Everything externally visible is considered used in Swift.
+  // That mostly means we need to be good at not marking things external.
+  return IRL.Linkage == llvm::GlobalValue::ExternalLinkage &&
+         (IRL.Visibility == llvm::GlobalValue::DefaultVisibility ||
+          IRL.Visibility == llvm::GlobalValue::ProtectedVisibility) &&
+         (IRL.DLLStorage == llvm::GlobalValue::DefaultStorageClass ||
+          IRL.DLLStorage == llvm::GlobalValue::DLLExportStorageClass);
+}
+
 /// Add the given global value to @llvm.used.
 ///
 /// This value must have a definition by the time the module is finalized.
@@ -2041,10 +2064,8 @@ void irgen::updateLinkageForDefinition(IRGenModule &IGM,
                    ForDefinition, weakImported, isKnownLocal);
   ApplyIRLinkage(IRL).to(global);
 
-  // Everything externally visible is considered used in Swift.
-  // That mostly means we need to be good at not marking things external.
-  if (LinkInfo::isUsed(IRL))
-    IGM.addUsedGlobal(global);
+  LinkInfo link = LinkInfo::get(IGM, entity, ForDefinition);
+  markGlobalAsUsedBasedOnLinkage(IGM, link, global);
 }
 
 LinkInfo LinkInfo::get(IRGenModule &IGM, const LinkEntity &entity,
@@ -2143,25 +2164,11 @@ llvm::Function *irgen::createFunction(IRGenModule &IGM,
   if (!updatedAttrs.isEmpty())
     fn->setAttributes(updatedAttrs);
 
-  // Everything externally visible is considered used in Swift.
-  // That mostly means we need to be good at not marking things external.
-  if (linkInfo.isUsed()) {
-    IGM.addUsedGlobal(fn);
-  }
+  markGlobalAsUsedBasedOnLinkage(IGM, linkInfo, fn);
 
   return fn;
 }
 
-bool LinkInfo::isUsed(IRLinkage IRL) {
-  // Everything externally visible is considered used in Swift.
-  // That mostly means we need to be good at not marking things external.
-  return IRL.Linkage == llvm::GlobalValue::ExternalLinkage &&
-         (IRL.Visibility == llvm::GlobalValue::DefaultVisibility ||
-          IRL.Visibility == llvm::GlobalValue::ProtectedVisibility) &&
-         (IRL.DLLStorage == llvm::GlobalValue::DefaultStorageClass ||
-          IRL.DLLStorage == llvm::GlobalValue::DLLExportStorageClass);
-}
-
 /// Get or create an LLVM global variable with these linkage rules.
 llvm::GlobalVariable *swift::irgen::createVariable(
     IRGenModule &IGM, LinkInfo &linkInfo, llvm::Type *storageType,
@@ -2191,11 +2198,7 @@ llvm::GlobalVariable *swift::irgen::createVariable(
       .to(var, linkInfo.isForDefinition());
   var->setAlignment(llvm::MaybeAlign(alignment.getValue()));
 
-  // Everything externally visible is considered used in Swift.
-  // That mostly means we need to be good at not marking things external.
-  if (linkInfo.isUsed()) {
-    IGM.addUsedGlobal(var);
-  }
+  markGlobalAsUsedBasedOnLinkage(IGM, linkInfo, var);
 
   if (IGM.DebugInfo && !DbgTy.isNull() && linkInfo.isForDefinition())
     IGM.DebugInfo->emitGlobalVariableDeclaration(
@@ -4161,9 +4164,7 @@ llvm::GlobalValue *IRGenModule::defineAlias(LinkEntity entity,
   ApplyIRLinkage({link.getLinkage(), link.getVisibility(), link.getDLLStorage()})
       .to(alias);
 
-  if (link.isUsed()) {
-    addUsedGlobal(alias);
-  }
+  markGlobalAsUsedBasedOnLinkage(*this, link, alias);
 
   // Replace an existing external declaration for the address point.
   if (entry) {
@@ -4247,8 +4248,7 @@ llvm::GlobalValue *IRGenModule::defineTypeMetadata(
   }
 
   LinkInfo link = LinkInfo::get(*this, entity, ForDefinition);
-  if (link.isUsed())
-    addUsedGlobal(var);
+  markGlobalAsUsedBasedOnLinkage(*this, link, var);
 
   /// For concrete metadata, we want to use the initializer on the
   /// "full metadata", and define the "direct" address point as an alias.

lib/IRGen/GenMeta.cpp

@@ -350,6 +350,9 @@ void IRGenModule::addVTableTypeMetadata(
   } else if (AS.isPrivate() || AS.isInternal()) {
     var->setVCallVisibilityMetadata(
       llvm::GlobalObject::VCallVisibility::VCallVisibilityLinkageUnit);
+  } else if (getOptions().InternalizeAtLink) {
+    var->setVCallVisibilityMetadata(
+      llvm::GlobalObject::VCallVisibility::VCallVisibilityLinkageUnit);
   } else {
     var->setVCallVisibilityMetadata(
       llvm::GlobalObject::VCallVisibility::VCallVisibilityPublic);
@@ -1756,8 +1759,9 @@ namespace {
 
       // Emit method dispatch thunk if the class is resilient.
       auto *func = cast<AbstractFunctionDecl>(fn.getDecl());
-      if (Resilient &&
-          func->getEffectiveAccess() >= AccessLevel::Public) {
+
+      if ((Resilient && func->getEffectiveAccess() >= AccessLevel::Public) ||
+          IGM.getOptions().VirtualFunctionElimination) {
         IGM.emitDispatchThunk(fn);
       }
     }
@@ -1782,9 +1786,10 @@ namespace {
       // method for external clients.
 
       // Emit method dispatch thunk.
-      if (hasPublicVisibility(fn.getLinkage(NotForDefinition))) {
-        IGM.emitDispatchThunk(fn);
-      }
+     if (hasPublicVisibility(fn.getLinkage(NotForDefinition)) ||
+         IGM.getOptions().VirtualFunctionElimination) {
+       IGM.emitDispatchThunk(fn);
+     }
 
       if (IGM.getOptions().VirtualFunctionElimination) {
         auto offset = B.getNextOffsetFromGlobal() +

lib/IRGen/IRGenSIL.cpp

@@ -6776,8 +6776,26 @@ void IRGenSILFunction::visitClassMethodInst(swift::ClassMethodInst *i) {
   auto methodType = i->getType().castTo<SILFunctionType>();
 
   auto *classDecl = cast<ClassDecl>(method.getDecl()->getDeclContext());
-  if (IGM.hasResilientMetadata(classDecl,
-                               ResilienceExpansion::Maximal)) {
+  bool shouldUseDispatchThunk = false;
+  if (IGM.hasResilientMetadata(classDecl, ResilienceExpansion::Maximal)) {
+    shouldUseDispatchThunk = true;
+  } else if (IGM.getOptions().VirtualFunctionElimination) {
+    // For VFE, use a thunk if the target class is in another module. This
+    // enables VFE (which scans function bodies for used type identifiers) to
+    // work across modules by relying on:
+    //
+    // (1) virtual call sites are in thunks in the same module as the class,
+    //     therefore they are always visible to VFE,
+    // (2) if a thunk symbol is unused by any other module, we can safely
+    //     eliminate it.
+    //
+    // See the virtual-function-elimination-two-modules.swift testcase for an
+    // example of how cross-module VFE can be effectively used.
+    shouldUseDispatchThunk =
+        classDecl->getModuleContext() != IGM.getSwiftModule();
+  }
+
+  if (shouldUseDispatchThunk) {
     llvm::Constant *fnPtr = IGM.getAddrOfDispatchThunk(method, NotForDefinition);
 
     if (methodType->isAsync()) {

lib/TBDGen/TBDGen.cpp

@@ -934,7 +934,7 @@ void TBDGenVisitor::visitClassDecl(ClassDecl *CD) {
     void addMethod(SILDeclRef method) {
       assert(method.getDecl()->getDeclContext() == CD);
 
-      if (CD->hasResilientMetadata()) {
+      if (TBD.Opts.VirtualFunctionElimination || CD->hasResilientMetadata()) {
         if (FirstTime) {
           FirstTime = false;
 

test/IRGen/virtual-function-elimination-two-modules.swift

@@ -0,0 +1,71 @@
+// Tests that under -enable-llvm-vfe + -internalize-at-link, cross-module
+// virtual calls are done via thunks and LLVM GlobalDCE is able to remove unused
+// virtual 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-vfe \
+// 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-vfe \
+// 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-vfe -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 vtables.
+// REQUIRES: rdar81868930
+
+#if LIBRARY
+
+public class MyClass {
+  public init() {}
+  public func foo() { print("MyClass.foo") }
+  public func bar() { print("MyClass.bar") }
+}
+
+public class MyDerivedClass: MyClass {
+  override public func foo() { print("MyDerivedClass.foo") }
+  override public func bar() { print("MyDerivedClass.bar") }
+}
+
+// NM-NOT: $s7Library14MyDerivedClassC3baryyF
+// NM:     $s7Library14MyDerivedClassC3fooyyF
+// NM-NOT: $s7Library7MyClassC3baryyF
+// NM:     $s7Library7MyClassC3fooyyF
+
+#endif
+
+#if CLIENT
+
+import Library
+
+@_cdecl("main")
+func main() -> Int32 {
+   let o: MyClass = MyDerivedClass()
+   o.foo()
+   print("Done")
+   // CHECK: MyDerivedClass.foo
+   // CHECK-NEXT: Done
+   return 0
+}
+
+#endif