Implement LLVM IR Virtual Function Elimination for Swift classes. (#39128)

- Virtual calls are done via a @llvm.type.checked.load instrinsic call with a type identifier
- Type identifier of a vfunc is the base method's mangling
- Type descriptors and class metadata get !type markers that list offsets and type identifiers of all vfuncs
- The -enable-llvm-vfe frontend flag enables VFE
- Two added tests verify the behavior on IR and by executing a program

Author: kubamracek

include/swift/AST/IRGenOptions.h

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

include/swift/Option/FrontendOptions.td

@@ -489,6 +489,10 @@ def enable_implicit_dynamic : Flag<["-"], "enable-implicit-dynamic">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
   HelpText<"Add 'dynamic' to all declarations">;
 
+def enable_llvm_vfe : Flag<["-"], "enable-llvm-vfe">,
+  Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
+  HelpText<"Use LLVM Virtual Function Elimination on Swift class virtual tables">;
+
 def disable_previous_implementation_calls_in_dynamic_replacements :
   Flag<["-"], "disable-previous-implementation-calls-in-dynamic-replacements">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,

lib/Frontend/CompilerInvocation.cpp

@@ -1896,6 +1896,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Opts.EnableGlobalISel = true;
   }
 
+  if (Args.hasArg(OPT_enable_llvm_vfe)) {
+    Opts.VirtualFunctionElimination = true;
+  }
+
   return false;
 }
 

lib/IRGen/GenClass.cpp

@@ -2672,6 +2672,44 @@ irgen::emitClassResilientInstanceSizeAndAlignMask(IRGenFunction &IGF,
   return {size, alignMask};
 }
 
+llvm::MDString *irgen::typeIdForMethod(IRGenModule &IGM, SILDeclRef method) {
+  assert(!method.getOverridden() && "must always be base method");
+
+  auto entity = LinkEntity::forMethodDescriptor(method);
+  auto mangled = entity.mangleAsString();
+  auto typeId = llvm::MDString::get(*IGM.LLVMContext, mangled);
+  return typeId;
+}
+
+static llvm::Value *emitVTableSlotLoad(IRGenFunction &IGF, Address slot,
+                                       SILDeclRef method, Signature signature) {
+  if (IGF.IGM.getOptions().VirtualFunctionElimination) {
+    // For LLVM IR VFE, emit a @llvm.type.checked.load with the type of the
+    // method.
+    llvm::Function *checkedLoadIntrinsic = llvm::Intrinsic::getDeclaration(
+        &IGF.IGM.Module, llvm::Intrinsic::type_checked_load);
+    auto slotAsPointer = IGF.Builder.CreateBitCast(slot, IGF.IGM.Int8PtrTy);
+    auto typeId = typeIdForMethod(IGF.IGM, method);
+
+    // Arguments for @llvm.type.checked.load: 1) target address, 2) offset -
+    // always 0 because target address is directly pointing to the right slot,
+    // 3) type identifier, i.e. the mangled name of the *base* method.
+    SmallVector<llvm::Value *, 8> args;
+    args.push_back(slotAsPointer.getAddress());
+    args.push_back(llvm::ConstantInt::get(IGF.IGM.Int32Ty, 0));
+    args.push_back(llvm::MetadataAsValue::get(*IGF.IGM.LLVMContext, typeId));
+
+    llvm::Value *checkedLoad =
+        IGF.Builder.CreateCall(checkedLoadIntrinsic, args);
+    auto fnPtr = IGF.Builder.CreateExtractValue(checkedLoad, 0);
+    return IGF.Builder.CreateBitCast(fnPtr,
+                                     signature.getType()->getPointerTo());
+  }
+
+  // Not doing LLVM IR VFE, can just be a direct load.
+  return IGF.emitInvariantLoad(slot);
+}
+
 FunctionPointer irgen::emitVirtualMethodValue(IRGenFunction &IGF,
                                               llvm::Value *metadata,
                                               SILDeclRef method,
@@ -2690,7 +2728,7 @@ FunctionPointer irgen::emitVirtualMethodValue(IRGenFunction &IGF,
     auto slot = IGF.emitAddressAtOffset(metadata, offset,
                                         signature.getType()->getPointerTo(),
                                         IGF.IGM.getPointerAlignment());
-    auto fnPtr = IGF.emitInvariantLoad(slot);
+    auto fnPtr = emitVTableSlotLoad(IGF, slot, method, signature);
     auto &schema = methodType->isAsync()
                        ? IGF.getOptions().PointerAuth.AsyncSwiftClassMethods
                        : IGF.getOptions().PointerAuth.SwiftClassMethods;

lib/IRGen/GenClass.h

@@ -25,6 +25,7 @@ namespace llvm {
   class Constant;
   class Value;
   class Function;
+  class MDString;
 }
 
 namespace swift {
@@ -201,6 +202,9 @@ namespace irgen {
                                              ClassDecl *theClass,
                                              llvm::Value *metadata);
 
+  /// For VFE, returns a type identifier for the given base method on a class.
+  llvm::MDString *typeIdForMethod(IRGenModule &IGM, SILDeclRef method);
+
   /// Given a metadata pointer, emit the callee for the given method.
   FunctionPointer emitVirtualMethodValue(IRGenFunction &IGF,
                                          llvm::Value *metadata,

lib/IRGen/GenDecl.cpp

@@ -4202,11 +4202,10 @@ llvm::GlobalValue *IRGenModule::defineAlias(LinkEntity entity,
 /// public symbol for the metadata references. This function will rewrite any
 /// existing external declaration to the address point as an alias into the
 /// full metadata object.
-llvm::GlobalValue *IRGenModule::defineTypeMetadata(CanType concreteType,
-                                                   bool isPattern,
-                                                   bool isConstant,
-                                                   ConstantInitFuture init,
-                                                   llvm::StringRef section) {
+llvm::GlobalValue *IRGenModule::defineTypeMetadata(
+    CanType concreteType, bool isPattern, bool isConstant,
+    ConstantInitFuture init, llvm::StringRef section,
+    SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries) {
   assert(init);
 
   auto isPrespecialized = concreteType->getAnyGeneric() &&
@@ -4241,6 +4240,12 @@ llvm::GlobalValue *IRGenModule::defineTypeMetadata(CanType concreteType,
   if (!section.empty())
     var->setSection(section);
 
+  if (getOptions().VirtualFunctionElimination) {
+    if (auto classDecl = concreteType->getClassOrBoundGenericClass()) {
+      addVTableTypeMetadata(classDecl, var, vtableEntries);
+    }
+  }
+
   LinkInfo link = LinkInfo::get(*this, entity, ForDefinition);
   if (link.isUsed())
     addUsedGlobal(var);

lib/IRGen/GenMeta.cpp

@@ -334,6 +334,28 @@ void IRGenModule::emitNonoverriddenMethodDescriptor(const SILVTable *VTable,
   getAddrOfLLVMVariable(entity, init, DebugTypeInfo());
 }
 
+void IRGenModule::addVTableTypeMetadata(
+    ClassDecl *decl, llvm::GlobalVariable *var,
+    SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries) {
+  for (auto ventry : vtableEntries) {
+    auto method = ventry.second;
+    auto offset = ventry.first.getValue();
+    var->addTypeMetadata(offset, typeIdForMethod(*this, method));
+  }
+
+  auto AS = decl->getFormalAccessScope();
+  if (AS.isFileScope()) {
+    var->setVCallVisibilityMetadata(
+      llvm::GlobalObject::VCallVisibility::VCallVisibilityTranslationUnit);
+  } else if (AS.isPrivate() || AS.isInternal()) {
+    var->setVCallVisibilityMetadata(
+      llvm::GlobalObject::VCallVisibility::VCallVisibilityLinkageUnit);
+  } else {
+    var->setVCallVisibilityMetadata(
+      llvm::GlobalObject::VCallVisibility::VCallVisibilityPublic);
+  }
+}
+
 namespace {
   template<class Impl>
   class ContextDescriptorBuilderBase {
@@ -1232,7 +1254,11 @@ namespace {
       auto addr = IGM.getAddrOfTypeContextDescriptor(Type, HasMetadata,
                                                      B.finishAndCreateFuture());
       auto var = cast<llvm::GlobalVariable>(addr);
-      
+
+      if (IGM.getOptions().VirtualFunctionElimination) {
+        asImpl().addVTableTypeMetadata(var);
+      }
+
       var->setConstant(true);
       IGM.setTrueConstGlobal(var);
       return var;
@@ -1440,6 +1466,10 @@ namespace {
       B.addRelativeAddress(IGM.getAddrOfReflectionFieldDescriptor(
         getType()->getDeclaredType()->getCanonicalType()));
     }
+
+    void addVTableTypeMetadata(llvm::GlobalVariable *var) {
+      // Structs don't have vtables.
+    }
   };
 
   class EnumContextDescriptorBuilder
@@ -1522,6 +1552,10 @@ namespace {
       B.addRelativeAddress(IGM.getAddrOfReflectionFieldDescriptor(
         getType()->getDeclaredType()->getCanonicalType()));
     }
+
+    void addVTableTypeMetadata(llvm::GlobalVariable *var) {
+      // Enums don't have vtables.
+    }
   };
 
   class ClassContextDescriptorBuilder
@@ -1547,6 +1581,11 @@ namespace {
     SmallVector<SILDeclRef, 8> VTableEntries;
     SmallVector<std::pair<SILDeclRef, SILDeclRef>, 8> OverrideTableEntries;
 
+    // As we're constructing the vtable, VTableEntriesForVFE stores the offset
+    // (from the beginning of the global) for each vtable slot. The offsets are
+    // later turned into !type metadata attributes.
+    SmallVector<std::pair<Size, SILDeclRef>, 8> VTableEntriesForVFE;
+
   public:
     ClassContextDescriptorBuilder(IRGenModule &IGM, ClassDecl *Type,
                                   RequireMetadata_t requireMetadata)
@@ -1703,6 +1742,13 @@ namespace {
       IGM.defineMethodDescriptor(fn, Type,
                       B.getAddrOfCurrentPosition(IGM.MethodDescriptorStructTy));
 
+      if (IGM.getOptions().VirtualFunctionElimination) {
+        auto offset = B.getNextOffsetFromGlobal() +
+                      // 1st field of MethodDescriptorStructTy
+                      Size(IGM.DataLayout.getTypeAllocSize(IGM.Int32Ty));
+        VTableEntriesForVFE.push_back(std::pair<Size, SILDeclRef>(offset, fn));
+      }
+
       // Actually build the descriptor.
       auto descriptor = B.beginStruct(IGM.MethodDescriptorStructTy);
       buildMethodDescriptorFields(IGM, VTable, fn, descriptor);
@@ -1715,7 +1761,15 @@ namespace {
         IGM.emitDispatchThunk(fn);
       }
     }
-    
+
+    void addVTableTypeMetadata(llvm::GlobalVariable *var) {
+      if (!IGM.getOptions().VirtualFunctionElimination)
+        return;
+      assert(VTable && "no vtable?!");
+
+      IGM.addVTableTypeMetadata(getType(), var, VTableEntriesForVFE);
+    }
+
     void emitNonoverriddenMethod(SILDeclRef fn) {
       // TODO: Derivative functions do not distinguish themselves in the mangled
       // names of method descriptor symbols yet, causing symbol name collisions.
@@ -1731,7 +1785,14 @@ namespace {
       if (hasPublicVisibility(fn.getLinkage(NotForDefinition))) {
         IGM.emitDispatchThunk(fn);
       }
-      
+
+      if (IGM.getOptions().VirtualFunctionElimination) {
+        auto offset = B.getNextOffsetFromGlobal() +
+                      // 1st field of MethodDescriptorStructTy
+                      Size(IGM.DataLayout.getTypeAllocSize(IGM.Int32Ty));
+        VTableEntriesForVFE.push_back(std::pair<Size, SILDeclRef>(offset, fn));
+      }
+
       // Emit a freestanding method descriptor structure. This doesn't have to
       // exist in the table in the class's context descriptor since it isn't
       // in the vtable, but external clients need to be able to link against the
@@ -1761,6 +1822,17 @@ namespace {
     }
 
     void emitMethodOverrideDescriptor(SILDeclRef baseRef, SILDeclRef declRef) {
+      if (IGM.getOptions().VirtualFunctionElimination) {
+        auto offset =
+            B.getNextOffsetFromGlobal() +
+            // 1st field of MethodOverrideDescriptorStructTy
+            Size(IGM.DataLayout.getTypeAllocSize(IGM.RelativeAddressTy)) +
+            // 2nd field of MethodOverrideDescriptorStructTy
+            Size(IGM.DataLayout.getTypeAllocSize(IGM.RelativeAddressTy));
+        VTableEntriesForVFE.push_back(
+            std::pair<Size, SILDeclRef>(offset, baseRef));
+      }
+
       auto descriptor = B.beginStruct(IGM.MethodOverrideDescriptorStructTy);
 
       // The class containing the base method.
@@ -2900,6 +2972,11 @@ namespace {
 
     Size AddressPoint;
 
+    // As we're constructing the vtable, VTableEntriesForVFE stores the offset
+    // (from the beginning of the global) for each vtable slot. The offsets are
+    // later turned into !type metadata attributes.
+    SmallVector<std::pair<Size, SILDeclRef>, 8> VTableEntriesForVFE;
+
   public:
     ClassMetadataBuilderBase(IRGenModule &IGM, ClassDecl *theClass,
                              ConstantStructBuilder &builder,
@@ -3201,12 +3278,21 @@ namespace {
         }
       }
 
+      if (IGM.getOptions().VirtualFunctionElimination) {
+        auto offset = B.getNextOffsetFromGlobal();
+        VTableEntriesForVFE.push_back(std::pair<Size, SILDeclRef>(offset, fn));
+      }
+
       PointerAuthSchema schema =
           afd->hasAsync() ? IGM.getOptions().PointerAuth.AsyncSwiftClassMethods
                           : IGM.getOptions().PointerAuth.SwiftClassMethods;
       B.addSignedPointer(ptr, schema, fn);
     }
 
+    SmallVector<std::pair<Size, SILDeclRef>, 8> getVTableEntriesForVFE() {
+      return VTableEntriesForVFE;
+    }
+
     void addPlaceholder(MissingMemberDecl *m) {
       assert(m->getNumberOfVTableEntries() == 0
              && "cannot generate metadata with placeholders in it");
@@ -3805,6 +3891,7 @@ void irgen::emitClassMetadata(IRGenModule &IGM, ClassDecl *classDecl,
   bool canBeConstant;
 
   auto strategy = IGM.getClassMetadataStrategy(classDecl);
+  SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries;
 
   switch (strategy) {
   case ClassMetadataStrategy::Resilient: {
@@ -3846,6 +3933,9 @@ void irgen::emitClassMetadata(IRGenModule &IGM, ClassDecl *classDecl,
     canBeConstant = builder.canBeConstant();
 
     builder.createMetadataAccessFunction();
+    if (IGM.getOptions().VirtualFunctionElimination) {
+      vtableEntries = builder.getVTableEntriesForVFE();
+    }
     break;
   }
   }
@@ -3859,7 +3949,8 @@ void irgen::emitClassMetadata(IRGenModule &IGM, ClassDecl *classDecl,
 
   bool isPattern = (strategy == ClassMetadataStrategy::Resilient);
   auto var = IGM.defineTypeMetadata(declaredType, isPattern, canBeConstant,
-                                    init.finishAndCreateFuture(), section);
+                                    init.finishAndCreateFuture(), section,
+                                    vtableEntries);
 
   // If the class does not require dynamic initialization, or if it only
   // requires dynamic initialization on a newer Objective-C runtime, add it

lib/IRGen/IRGen.cpp

@@ -192,6 +192,10 @@ void setModuleFlags(IRGenModule &IGM) {
   // error during LTO if the user tries to combine files across ABIs.
   Module->addModuleFlag(llvm::Module::Error, "Swift Version",
                         IRGenModule::swiftVersion);
+
+  if (IGM.getOptions().VirtualFunctionElimination) {
+    Module->addModuleFlag(llvm::Module::Error, "Virtual Function Elim", 1);
+  }
 }
 
 void swift::performLLVMOptimizations(const IRGenOptions &Opts,

lib/IRGen/IRGenModule.h

@@ -1448,11 +1448,14 @@ private:                            \
                                                      bool isDestroyer,
                                                      bool isForeign,
                                                      ForDefinition_t forDefinition);
-  llvm::GlobalValue *defineTypeMetadata(CanType concreteType,
-                                        bool isPattern,
-                                        bool isConstant,
-                                        ConstantInitFuture init,
-                                        llvm::StringRef section = {});
+  void addVTableTypeMetadata(
+      ClassDecl *decl, llvm::GlobalVariable *var,
+      SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries);
+
+  llvm::GlobalValue *defineTypeMetadata(
+      CanType concreteType, bool isPattern, bool isConstant,
+      ConstantInitFuture init, llvm::StringRef section = {},
+      SmallVector<std::pair<Size, SILDeclRef>, 8> vtableEntries = {});
 
   TypeEntityReference getTypeEntityReference(GenericTypeDecl *D);