Merge pull request #77004 from eeckstein/fix-vtable-specialization

embedded: fix several issues with vtable specialization

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -33,11 +33,6 @@ let mandatoryPerformanceOptimizations = ModulePass(name: "mandatory-performance-
   // For embedded Swift, optimize all the functions (there cannot be any
   // generics, type metadata, etc.)
   if moduleContext.options.enableEmbeddedSwift {
-    // We need to specialize all vtables which are referenced from non-generic contexts. Beside
-    // `alloc_ref`s of generic classes in non-generic functions, we also need to specialize generic
-    // superclasses of non-generic classes. E.g. `class Derived : Base<Int> {}`
-    specializeVTablesOfSuperclasses(moduleContext)
-
     worklist.addAllNonGenericFunctions(of: moduleContext)
   } else {
     worklist.addAllPerformanceAnnotatedFunctions(of: moduleContext)
@@ -102,15 +97,18 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
       // Embedded Swift specific transformations
       case let alloc as AllocRefInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: alloc.type, in: function,
-                                                  errorLocation: alloc.location,
-                                                  moduleContext, &worklist)
+          specializeVTable(forClassType: alloc.type, errorLocation: alloc.location, moduleContext) {
+            worklist.pushIfNotVisited($0)
+          }
         }
       case let metatype as MetatypeInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: metatype.type, in: function,
-                                                  errorLocation: metatype.location,
-                                                  moduleContext, &worklist)
+          let instanceType = metatype.type.loweredInstanceTypeOfMetatype(in: function)
+          if instanceType.isClass {
+            specializeVTable(forClassType: instanceType, errorLocation: metatype.location, moduleContext) {
+              worklist.pushIfNotVisited($0)
+            }
+          }
         }
       case let classMethod as ClassMethodInst:
         if context.options.enableEmbeddedSwift {
@@ -166,29 +164,6 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
   }
 }
 
-private func specializeVTableAndAddEntriesToWorklist(for type: Type, in function: Function,
-                                                     errorLocation: Location,
-                                                     _ moduleContext: ModulePassContext,
-                                                     _ worklist: inout FunctionWorklist) {
-  let vTablesCountBefore = moduleContext.vTables.count
-
-  guard specializeVTable(forClassType: type, errorLocation: errorLocation, moduleContext) != nil else {
-    return
-  }
-
-  // More than one new vtable might have been created (superclasses), process them all
-  let vTables = moduleContext.vTables
-  for i in vTablesCountBefore ..< vTables.count {
-    for entry in vTables[i].entries
-      // A new vtable can still contain a generic function if the method couldn't be specialized for some reason
-      // and an error has been printed. Exclude generic functions to not run into an assert later.
-      where !entry.implementation.isGeneric
-    {
-      worklist.pushIfNotVisited(entry.implementation)
-    }
-  }
-}
-
 private func inlineAndDevirtualize(apply: FullApplySite, alreadyInlinedFunctions: inout Set<PathFunctionTuple>,
                                    _ context: FunctionPassContext, _ simplifyCtxt: SimplifyContext) {
   // De-virtualization and inlining in/into a "serialized" function might create function references to functions

SwiftCompilerSources/Sources/Optimizer/PassManager/ModulePassContext.swift

@@ -168,6 +168,14 @@ struct ModulePassContext : Context, CustomStringConvertible {
     return VTable(bridged: bridgedVTable)
   }
 
+  func replaceVTableEntries(of vTable: VTable, with entries: [VTable.Entry]) {
+    let bridgedEntries = entries.map { $0.bridged }
+    bridgedEntries.withBridgedArrayRef {
+      vTable.bridged.replaceEntries($0)
+    }
+    notifyFunctionTablesChanged()
+  }
+
   func createEmptyFunction(
     name: String,
     parameters: [ParameterInfo],
@@ -201,6 +209,10 @@ struct ModulePassContext : Context, CustomStringConvertible {
       }
     }
   }
+
+  func notifyFunctionTablesChanged() {
+    _bridged.asNotificationHandler().notifyChanges(.functionTablesChanged)
+  }
 }
 
 extension GlobalVariable {

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -13,65 +13,92 @@
 import AST
 import SIL
 
-@discardableResult
 func specializeVTable(forClassType classType: Type,
                       errorLocation: Location,
-                      _ context: ModulePassContext) -> VTable?
+                      _ context: ModulePassContext,
+                      notifyNewFunction: (Function) -> ())
 {
-  guard let nominal = classType.nominal,
-        let classDecl = nominal as? ClassDecl,
-        classType.isGenericAtAnyLevel else
-  {
-    return nil
-  }
+  var specializer = VTableSpecializer(errorLocation: errorLocation, context)
+  specializer.specializeVTable(forClassType: classType, notifyNewFunction)
+}
 
-  if context.lookupSpecializedVTable(for: classType) != nil {
-    return nil
-  }
+private struct VTableSpecializer {
+  let errorLocation: Location
+  let context: ModulePassContext
 
-  guard let origVTable = context.lookupVTable(for: classDecl) else {
-    context.diagnosticEngine.diagnose(errorLocation.sourceLoc, .cannot_specialize_class, classType)
-    return nil
-  }
+  // The type of the first class in the hierarchy which implements a method
+  private var baseTypesOfMethods = Dictionary<Function, Type>()
 
-  let classContextSubs = classType.contextSubstitutionMap
+  init(errorLocation: Location, _ context: ModulePassContext) {
+    self.errorLocation = errorLocation
+    self.context = context
+  }
 
-  let newEntries = origVTable.entries.map { origEntry in
-    if !origEntry.implementation.isGeneric {
-      return origEntry
+  mutating func specializeVTable(forClassType classType: Type, _ notifyNewFunction: (Function) -> ()) {
+    // First handle super classes.
+    // This is also required for non-generic classes - in case a superclass is generic, e.g.
+    // `class Derived : Base<Int> {}` - for two reasons:
+    // * A vtable of a derived class references the vtable of the super class. And of course the referenced
+    //   super-class vtable needs to be a specialized vtable.
+    // * Even a non-generic derived class can contain generic methods of the base class in case a base-class
+    //   method is not overridden.
+    //
+    if let superClassTy = classType.superClassType {
+      specializeVTable(forClassType: superClassTy, notifyNewFunction)
     }
-    let methodSubs = classContextSubs.getMethodSubstitutions(for: origEntry.implementation)
 
-    guard !methodSubs.conformances.contains(where: {!$0.isValid}),
-          let specializedMethod = context.specialize(function: origEntry.implementation, for: methodSubs) else
-    {
-      context.diagnosticEngine.diagnose(origEntry.methodDecl.location.sourceLoc, .non_final_generic_class_function)
-      return origEntry
+    let classDecl = classType.nominal! as! ClassDecl
+    guard let origVTable = context.lookupVTable(for: classDecl) else {
+      context.diagnosticEngine.diagnose(errorLocation.sourceLoc, .cannot_specialize_class, classType)
+      return
     }
 
-    context.deserializeAllCallees(of: specializedMethod, mode: .allFunctions)
-    specializedMethod.set(linkage: .public, context)
-    specializedMethod.set(isSerialized: false, context)
+    for entry in origVTable.entries {
+      if baseTypesOfMethods[entry.implementation] == nil {
+        baseTypesOfMethods[entry.implementation] = classType
+      }
+    }
 
-    return VTable.Entry(kind: origEntry.kind, isNonOverridden: origEntry.isNonOverridden,
-                        methodDecl: origEntry.methodDecl, implementation: specializedMethod)
+    if classType.isGenericAtAnyLevel {
+      if context.lookupSpecializedVTable(for: classType) != nil {
+        // We already specialized the vtable
+        return
+      }
+      let newEntries = specializeEntries(of: origVTable, notifyNewFunction)
+      context.createSpecializedVTable(entries: newEntries, for: classType, isSerialized: false)
+    } else {
+      if !origVTable.entries.contains(where: { $0.implementation.isGeneric }) {
+        // The vtable (of the non-generic class) doesn't contain any generic functions (from a generic base class).
+        return
+      }
+      let newEntries = specializeEntries(of: origVTable, notifyNewFunction)
+      context.replaceVTableEntries(of: origVTable, with: newEntries)
+    }
   }
 
-  let specializedVTable = context.createSpecializedVTable(entries: newEntries, for: classType, isSerialized: false)
-  if let superClassTy = classType.superClassType {
-    specializeVTable(forClassType: superClassTy, errorLocation: classDecl.location, context)
-  }
-  return specializedVTable
-}
+  private func specializeEntries(of vTable: VTable, _ notifyNewFunction: (Function) -> ()) -> [VTable.Entry] {
+    return vTable.entries.compactMap { entry in
+      if !entry.implementation.isGeneric {
+        return entry
+      }
+      let baseType = baseTypesOfMethods[entry.implementation]!
+      let classContextSubs = baseType.contextSubstitutionMap
+      let methodSubs = classContextSubs.getMethodSubstitutions(for: entry.implementation)
+
+      guard !methodSubs.conformances.contains(where: {!$0.isValid}),
+            let specializedMethod = context.specialize(function: entry.implementation, for: methodSubs) else
+      {
+        context.diagnosticEngine.diagnose(entry.methodDecl.location.sourceLoc, .non_final_generic_class_function)
+        return nil
+      }
+      notifyNewFunction(specializedMethod)
+
+      context.deserializeAllCallees(of: specializedMethod, mode: .allFunctions)
+      specializedMethod.set(linkage: .public, context)
+      specializedMethod.set(isSerialized: false, context)
 
-func specializeVTablesOfSuperclasses(_ moduleContext: ModulePassContext) {
-  for vtable in moduleContext.vTables {
-    if !vtable.isSpecialized,
-       !vtable.class.isGenericAtAnyLevel,
-       let superClassTy = vtable.class.superClassType,
-       superClassTy.isGenericAtAnyLevel
-    {
-      specializeVTable(forClassType: superClassTy, errorLocation: vtable.class.location, moduleContext)
+      return VTable.Entry(kind: entry.kind, isNonOverridden: entry.isNonOverridden,
+                          methodDecl: entry.methodDecl, implementation: specializedMethod)
     }
   }
 }

SwiftCompilerSources/Sources/SIL/VTable.swift

@@ -14,7 +14,7 @@ import AST
 import SILBridging
 
 public struct VTable : CustomStringConvertible, NoReflectionChildren {
-  let bridged: BridgedVTable
+  public let bridged: BridgedVTable
 
   public init(bridged: BridgedVTable) { self.bridged = bridged }
 

include/swift/SIL/SILBridging.h

@@ -1065,17 +1065,18 @@ struct BridgedVTableEntryArray {
 };
 
 struct BridgedVTable {
-  const swift::SILVTable * _Nonnull vTable;
+  swift::SILVTable * _Nonnull vTable;
 
   BridgedOwnedString getDebugDescription() const;
   BRIDGED_INLINE SwiftInt getNumEntries() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedVTableEntry getEntry(SwiftInt index) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj getClass() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getSpecializedClassType() const;
+  BRIDGED_INLINE void replaceEntries(BridgedArrayRef bridgedEntries) const;
 };
 
 struct OptionalBridgedVTable {
-  const swift::SILVTable * _Nullable table;
+  swift::SILVTable * _Nullable table;
 };
 
 struct BridgedWitnessTableEntry {
@@ -1328,7 +1329,8 @@ struct BridgedChangeNotificationHandler {
     instructionsChanged,
     callsChanged,
     branchesChanged,
-    effectsChanged
+    effectsChanged,
+    functionTablesChanged
   };
 
   void notifyChanges(Kind changeKind) const;

include/swift/SIL/SILBridgingImpl.h

@@ -1715,6 +1715,14 @@ BridgedType BridgedVTable::getSpecializedClassType() const {
   return {vTable->getClassType()};
 }
 
+void BridgedVTable::replaceEntries(BridgedArrayRef bridgedEntries) const {
+  llvm::SmallVector<swift::SILVTableEntry, 8> entries;
+  for (const BridgedVTableEntry &e : bridgedEntries.unbridged<BridgedVTableEntry>()) {
+    entries.push_back(e.unbridged());
+  }
+  vTable->replaceEntries(entries);
+}
+
 //===----------------------------------------------------------------------===//
 //               BridgedWitnessTable, BridgedDefaultWitnessTable
 //===----------------------------------------------------------------------===//

include/swift/SIL/SILVTable.h

@@ -205,6 +205,8 @@ class SILVTable final : public SILAllocated<SILVTable>,
     NumEntries = std::distance(Entries.begin(), end);
   }
 
+  void replaceEntries(ArrayRef<Entry> newEntries);
+
   /// Verify that the vtable is well-formed for the given class.
   void verify(const SILModule &M) const;
 

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -69,6 +69,8 @@ class SwiftPassInvocation {
   SILAnalysis::InvalidationKind changeNotifications =
       SILAnalysis::InvalidationKind::Nothing;
 
+  bool functionTablesChanged = false;
+
   /// All slabs, allocated by the pass.
   SILModule::SlabList allocatedSlabs;
 
@@ -141,6 +143,8 @@ class SwiftPassInvocation {
   /// Called by the pass when changes are made to the SIL.
   void notifyChanges(SILAnalysis::InvalidationKind invalidationKind);
 
+  void notifyFunctionTablesChanged();
+
   /// Called by the pass manager before the pass starts running.
   void startModulePassRun(SILModuleTransform *transform);
 
@@ -513,6 +517,9 @@ notifyChanges(SILAnalysis::InvalidationKind invalidationKind) {
     changeNotifications = (SILAnalysis::InvalidationKind)
         (changeNotifications | invalidationKind);
 }
+inline void SwiftPassInvocation::notifyFunctionTablesChanged() {
+  functionTablesChanged = true;
+}
 
 } // end namespace swift
 

lib/SIL/IR/SILVTable.cpp

@@ -75,6 +75,22 @@ void SILVTable::updateVTableCache(const Entry &entry) {
   M.VTableEntryCache[{this, entry.getMethod()}] = entry;
 }
 
+void SILVTable::replaceEntries(ArrayRef<Entry> newEntries) {
+  auto entries = getMutableEntries();
+  ASSERT(newEntries.size() <= entries.size());
+  for (unsigned i = 0; i < entries.size(); ++i) {
+    entries[i].getImplementation()->decrementRefCount();
+    if (i < newEntries.size()) {
+      entries[i] = newEntries[i];
+      entries[i].getImplementation()->incrementRefCount();
+      updateVTableCache(entries[i]);
+    } else {
+      removeFromVTableCache(entries[i]);
+    }
+  }
+  NumEntries = newEntries.size();
+}
+
 SILVTable::SILVTable(ClassDecl *c, SILType classType,
                      SerializedKind_t serialized, ArrayRef<Entry> entries)
     : Class(c), classType(classType), SerializedKind(serialized),

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1516,6 +1516,7 @@ void SwiftPassInvocation::finishedModulePassRun() {
   endPass();
   assert(!function && transform && "not running a pass");
   assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && !functionTablesChanged
          && "unhandled change notifications at end of module pass");
   transform = nullptr;
 }
@@ -1551,6 +1552,7 @@ void SwiftPassInvocation::endPass() {
 void SwiftPassInvocation::beginTransformFunction(SILFunction *function) {
   assert(!this->function && transform && "not running a pass");
   assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && !functionTablesChanged
          && "change notifications not cleared");
   this->function = function;
 }
@@ -1561,6 +1563,10 @@ void SwiftPassInvocation::endTransformFunction() {
     passManager->invalidateAnalysis(function, changeNotifications);
     changeNotifications = SILAnalysis::InvalidationKind::Nothing;
   }
+  if (functionTablesChanged) {
+    passManager->invalidateFunctionTables();
+    functionTablesChanged = false;
+  }
   function = nullptr;
   assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
   assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
@@ -1580,6 +1586,7 @@ void SwiftPassInvocation::endVerifyFunction() {
   assert(function);
   if (!transform) {
     assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing &&
+           !functionTablesChanged &&
            "verifyication must not change the SIL of a function");
     assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
     assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
@@ -1634,6 +1641,9 @@ void BridgedChangeNotificationHandler::notifyChanges(Kind changeKind) const {
   case Kind::effectsChanged:
     invocation->notifyChanges(SILAnalysis::InvalidationKind::Effects);
     break;
+  case Kind::functionTablesChanged:
+    invocation->notifyFunctionTablesChanged();
+    break;
   }
 }
 

lib/SILOptimizer/Utils/Generics.cpp

@@ -2324,6 +2324,9 @@ bool swift::specializeClassMethodInst(ClassMethodInst *cm) {
 
   SILValue instance = cm->getOperand();
   SILType classTy = instance->getType();
+  if (classTy.is<MetatypeType>())
+    classTy = classTy.getLoweredInstanceTypeOfMetatype(cm->getFunction());
+
   CanType astType = classTy.getASTType();
   if (!astType->isSpecialized())
     return false;