Merge pull request #69678 from kubamracek/embedded-devirt

kubamracek · web-flow · commit 71612e6e5841 · 2023-11-07T09:08:56.000-08:00
[embedded] Fix class_method devirtualizer to consider specialized VTables
diff --git a/include/swift/SILOptimizer/Utils/Devirtualize.h b/include/swift/SILOptimizer/Utils/Devirtualize.h
@@ -73,21 +73,21 @@ tryDevirtualizeApply(ApplySite AI, ClassHierarchyAnalysis *CHA,
                      bool isMandatory = false);
 bool canDevirtualizeApply(FullApplySite AI, ClassHierarchyAnalysis *CHA);
 bool canDevirtualizeClassMethod(FullApplySite AI, ClassDecl *CD,
+                                CanType ClassType,
                                 OptRemark::Emitter *ORE = nullptr,
                                 bool isEffectivelyFinalMethod = false);
 SILFunction *getTargetClassMethod(SILModule &M, ClassDecl *CD,
-                                  MethodInst *MI);
+                                  CanType ClassType, MethodInst *MI);
 CanType getSelfInstanceType(CanType ClassOrMetatypeType);
 
 /// Devirtualize the given apply site, which is known to be devirtualizable.
 ///
 /// The caller must call deleteDevirtualizedApply on the original apply site.
 ///
 /// Return the new apply and true if the CFG was also modified.
-std::pair<FullApplySite, bool> devirtualizeClassMethod(FullApplySite AI,
-                                                       SILValue ClassInstance,
-                                                       ClassDecl *CD,
-                                                       OptRemark::Emitter *ORE);
+std::pair<FullApplySite, bool>
+devirtualizeClassMethod(FullApplySite AI, SILValue ClassInstance, ClassDecl *CD,
+                        CanType classType, OptRemark::Emitter *ORE);
 
 /// Attempt to devirtualize the given apply site, which is known to be
 /// of a class method.  If this fails, the returned FullApplySite will be null.
@@ -98,7 +98,8 @@ std::pair<FullApplySite, bool> devirtualizeClassMethod(FullApplySite AI,
 /// Return the new apply and true if the CFG was also modified.
 std::pair<FullApplySite, bool>
 tryDevirtualizeClassMethod(FullApplySite AI, SILValue ClassInstance,
-                           ClassDecl *CD, OptRemark::Emitter *ORE,
+                           ClassDecl *CD, CanType ClassType,
+                           OptRemark::Emitter *ORE,
                            bool isEffectivelyFinalMethod = false);
 
 /// Attempt to devirtualize the given apply site, which is known to be
diff --git a/lib/SILOptimizer/Mandatory/DiagnoseInfiniteRecursion.cpp b/lib/SILOptimizer/Mandatory/DiagnoseInfiniteRecursion.cpp
@@ -90,7 +90,7 @@ static bool isRecursiveCall(FullApplySite applySite) {
     if (classDecl && classDecl->getModuleContext() != module.getSwiftModule())
       return false;
 
-    SILFunction *method = getTargetClassMethod(module, classDecl, CMI);
+    SILFunction *method = getTargetClassMethod(module, classDecl, classType, CMI);
     if (method != parentFunc)
       return false;
 
diff --git a/lib/SILOptimizer/Transforms/SpeculativeDevirtualizer.cpp b/lib/SILOptimizer/Transforms/SpeculativeDevirtualizer.cpp
@@ -123,15 +123,15 @@ static FullApplySite CloneApply(FullApplySite AI, SILValue SelfArg,
 /// Insert monomorphic inline caches for a specific class or metatype
 /// type \p SubClassTy.
 static FullApplySite speculateMonomorphicTarget(FullApplySite AI,
-                                                CanType SubType,
-                                                ClassDecl *CD,
+                                                CanType SubType, ClassDecl *CD,
+                                                CanType ClassType,
                                                 CheckedCastBranchInst *&CCBI) {
   if (SubType->hasDynamicSelfType())
     return FullApplySite();
 
   CCBI = nullptr;
   // Bail if this class_method cannot be devirtualized.
-  if (!canDevirtualizeClassMethod(AI, CD))
+  if (!canDevirtualizeClassMethod(AI, CD, ClassType))
     return FullApplySite();
 
   // Can't speculate begin_apply yet.
@@ -228,9 +228,9 @@ static FullApplySite speculateMonomorphicTarget(FullApplySite AI,
   ++NumTargetsPredicted;
 
   // Devirtualize the apply instruction on the identical path.
-  auto NewInst =
-      devirtualizeClassMethod(IdenAI, DownCastedClassInstance, CD, nullptr)
-          .first;
+  auto NewInst = devirtualizeClassMethod(IdenAI, DownCastedClassInstance, CD,
+                                         ClassType, nullptr)
+                     .first;
   assert(NewInst && "Expected to be able to devirtualize apply!");
   (void)NewInst;
 
@@ -401,15 +401,16 @@ static bool tryToSpeculateTarget(FullApplySite AI, ClassHierarchyAnalysis *CHA,
     ClassHierarchyAnalysis::ClassList Subs;
     if (isDefaultCaseKnown(CHA, AI, CD, Subs)) {
       auto NewInst =
-          tryDevirtualizeClassMethod(AI, SubTypeValue, CD, &ORE).first;
+          tryDevirtualizeClassMethod(AI, SubTypeValue, CD, ClassType, &ORE)
+              .first;
       if (NewInst)
         deleteDevirtualizedApply(AI);
       return bool(NewInst);
     }
 
     LLVM_DEBUG(llvm::dbgs() << "Inserting monomorphic speculative call for "
                "class " << CD->getName() << "\n");
-    return !!speculateMonomorphicTarget(AI, SubType, CD, LastCCBI);
+    return !!speculateMonomorphicTarget(AI, SubType, CD, ClassType, LastCCBI);
   }
 
   // True if any instructions were changed or generated.
@@ -435,7 +436,7 @@ static bool tryToSpeculateTarget(FullApplySite AI, ClassHierarchyAnalysis *CHA,
 
   // Try to devirtualize the static class of instance
   // if it is possible.
-  if (auto F = getTargetClassMethod(M, CD, CMI)) {
+  if (auto F = getTargetClassMethod(M, CD, ClassType, CMI)) {
     // Do not devirtualize if a method in the base class is marked
     // as non-optimizable. This way it is easy to disable the
     // devirtualization of this method in the base class and
@@ -444,7 +445,8 @@ static bool tryToSpeculateTarget(FullApplySite AI, ClassHierarchyAnalysis *CHA,
       return false;
   }
 
-  auto FirstAI = speculateMonomorphicTarget(AI, SubType, CD, LastCCBI);
+  auto FirstAI =
+      speculateMonomorphicTarget(AI, SubType, CD, ClassType, LastCCBI);
   if (FirstAI) {
     Changed = true;
     AI = FirstAI;
@@ -506,7 +508,8 @@ static bool tryToSpeculateTarget(FullApplySite AI, ClassHierarchyAnalysis *CHA,
     }
 
     // Pass the metatype of the subclass.
-    auto NewAI = speculateMonomorphicTarget(AI, ClassOrMetatypeType, S, LastCCBI);
+    auto NewAI = speculateMonomorphicTarget(AI, ClassOrMetatypeType, S,
+                                            CanClassType, LastCCBI);
     if (!NewAI) {
       ++NotHandledSubsNum;
       continue;
@@ -562,7 +565,8 @@ static bool tryToSpeculateTarget(FullApplySite AI, ClassHierarchyAnalysis *CHA,
     return true;
   }
   auto NewInst =
-      tryDevirtualizeClassMethod(AI, SubTypeValue, CD, nullptr).first;
+      tryDevirtualizeClassMethod(AI, SubTypeValue, CD, ClassType, nullptr)
+          .first;
   if (NewInst) {
     ORE.emit(RB);
     deleteDevirtualizedApply(AI);
diff --git a/lib/SILOptimizer/Utils/Devirtualize.cpp b/lib/SILOptimizer/Utils/Devirtualize.cpp
@@ -646,11 +646,17 @@ void swift::deleteDevirtualizedApply(ApplySite old) {
 }
 
 SILFunction *swift::getTargetClassMethod(SILModule &module, ClassDecl *cd,
-                                         MethodInst *mi) {
-  assert((isa<ClassMethodInst>(mi) || isa<SuperMethodInst>(mi))
-         && "Only class_method and super_method instructions are supported");
+                                         CanType classType, MethodInst *mi) {
+  assert((isa<ClassMethodInst>(mi) || isa<SuperMethodInst>(mi)) &&
+         "Only class_method and super_method instructions are supported");
 
   SILDeclRef member = mi->getMember();
+
+  SILType silType = SILType::getPrimitiveObjectType(classType);
+  if (auto *vtable = module.lookUpSpecializedVTable(silType)) {
+    return vtable->getEntry(module, member)->getImplementation();
+  }
+
   return module.lookUpFunctionInVTable(cd, member);
 }
 
@@ -672,6 +678,7 @@ CanType swift::getSelfInstanceType(CanType classOrMetatypeType) {
 /// \p cd is the class declaration we are devirtualizing for.
 /// return true if it is possible to devirtualize, false - otherwise.
 bool swift::canDevirtualizeClassMethod(FullApplySite applySite, ClassDecl *cd,
+                                       CanType classType,
                                        OptRemark::Emitter *ore,
                                        bool isEffectivelyFinalMethod) {
 
@@ -683,7 +690,7 @@ bool swift::canDevirtualizeClassMethod(FullApplySite applySite, ClassDecl *cd,
   auto *mi = cast<MethodInst>(applySite.getCallee());
 
   // Find the implementation of the member which should be invoked.
-  auto *f = getTargetClassMethod(module, cd, mi);
+  auto *f = getTargetClassMethod(module, cd, classType, mi);
 
   // If we do not find any such function, we have no function to devirtualize
   // to... so bail.
@@ -742,15 +749,15 @@ bool swift::canDevirtualizeClassMethod(FullApplySite applySite, ClassDecl *cd,
 std::pair<FullApplySite, bool /* changedCFG */>
 swift::devirtualizeClassMethod(FullApplySite applySite,
                                SILValue classOrMetatype, ClassDecl *cd,
-                               OptRemark::Emitter *ore) {
+                               CanType classType, OptRemark::Emitter *ore) {
   bool changedCFG = false;
   LLVM_DEBUG(llvm::dbgs() << "    Trying to devirtualize : "
                           << *applySite.getInstruction());
 
   SILModule &module = applySite.getModule();
   auto *mi = cast<MethodInst>(applySite.getCallee());
 
-  auto *f = getTargetClassMethod(module, cd, mi);
+  auto *f = getTargetClassMethod(module, cd, classType, mi);
 
   CanSILFunctionType genCalleeType = f->getLoweredFunctionTypeInContext(
       TypeExpansionContext(*applySite.getFunction()));
@@ -835,10 +842,11 @@ swift::devirtualizeClassMethod(FullApplySite applySite,
 
 std::pair<FullApplySite, bool> swift::tryDevirtualizeClassMethod(
     FullApplySite applySite, SILValue classInstance, ClassDecl *cd,
-    OptRemark::Emitter *ore, bool isEffectivelyFinalMethod) {
-  if (!canDevirtualizeClassMethod(applySite, cd, ore, isEffectivelyFinalMethod))
+    CanType classType, OptRemark::Emitter *ore, bool isEffectivelyFinalMethod) {
+  if (!canDevirtualizeClassMethod(applySite, cd, classType, ore,
+                                  isEffectivelyFinalMethod))
     return {FullApplySite(), false};
-  return devirtualizeClassMethod(applySite, classInstance, cd, ore);
+  return devirtualizeClassMethod(applySite, classInstance, cd, classType, ore);
 }
 
 //===----------------------------------------------------------------------===//
@@ -1322,7 +1330,7 @@ swift::tryDevirtualizeApply(ApplySite applySite, ClassHierarchyAnalysis *cha,
     auto *cd = classType.getClassOrBoundGenericClass();
 
     if (isEffectivelyFinalMethod(fas, classType, cd, cha))
-      return tryDevirtualizeClassMethod(fas, instance, cd, ore,
+      return tryDevirtualizeClassMethod(fas, instance, cd, classType, ore,
                                         true /*isEffectivelyFinalMethod*/);
 
     // Try to check if the exact dynamic type of the instance is statically
@@ -1333,13 +1341,14 @@ swift::tryDevirtualizeApply(ApplySite applySite, ClassHierarchyAnalysis *cha,
       CanType classType = getSelfInstanceType(instance->getType().getASTType());
       // This should never be null - make the check just to be on the safe side.
       if (ClassDecl *cd = classType.getClassOrBoundGenericClass())
-        return tryDevirtualizeClassMethod(fas, instance, cd, ore);
+        return tryDevirtualizeClassMethod(fas, instance, cd, classType, ore);
       return {ApplySite(), false};
     }
 
     if (auto exactTy = getExactDynamicType(cmi->getOperand(), cha)) {
       if (exactTy == cmi->getOperand()->getType())
-        return tryDevirtualizeClassMethod(fas, cmi->getOperand(), cd, ore);
+        return tryDevirtualizeClassMethod(fas, cmi->getOperand(), cd, classType,
+                                          ore);
     }
   }
 
@@ -1348,7 +1357,7 @@ swift::tryDevirtualizeApply(ApplySite applySite, ClassHierarchyAnalysis *cha,
     auto classType = getSelfInstanceType(instance->getType().getASTType());
     auto *cd = classType.getClassOrBoundGenericClass();
 
-    return tryDevirtualizeClassMethod(fas, instance, cd, ore);
+    return tryDevirtualizeClassMethod(fas, instance, cd, classType, ore);
   }
 
   return {ApplySite(), false};
@@ -1389,20 +1398,21 @@ bool swift::canDevirtualizeApply(FullApplySite applySite,
     auto *cd = classType.getClassOrBoundGenericClass();
 
     if (isEffectivelyFinalMethod(applySite, classType, cd, cha))
-      return canDevirtualizeClassMethod(applySite, cd, nullptr /*ore*/,
+      return canDevirtualizeClassMethod(applySite, cd, classType,
+                                        nullptr /*ore*/,
                                         true /*isEffectivelyFinalMethod*/);
 
     // Try to check if the exact dynamic type of the instance is statically
     // known.
     if (auto instance = getInstanceWithExactDynamicType(cmi->getOperand(), cha)) {
       CanType classType = getSelfInstanceType(instance->getType().getASTType());
       ClassDecl *cd = classType.getClassOrBoundGenericClass();
-      return cd && canDevirtualizeClassMethod(applySite, cd);
+      return cd && canDevirtualizeClassMethod(applySite, cd, classType);
     }
 
     if (auto exactTy = getExactDynamicType(cmi->getOperand(), cha)) {
       if (exactTy == cmi->getOperand()->getType())
-        return canDevirtualizeClassMethod(applySite, cd);
+        return canDevirtualizeClassMethod(applySite, cd, classType);
     }
   }
 
@@ -1411,7 +1421,7 @@ bool swift::canDevirtualizeApply(FullApplySite applySite,
     auto classType = getSelfInstanceType(instance->getType().getASTType());
     auto *cd = classType.getClassOrBoundGenericClass();
 
-    return canDevirtualizeClassMethod(applySite, cd);
+    return canDevirtualizeClassMethod(applySite, cd, classType);
   }
 
   return false;
diff --git a/test/embedded/generic-classes.swift b/test/embedded/generic-classes.swift
@@ -1,4 +1,5 @@
 // RUN: %target-run-simple-swift(%S/Inputs/print.swift -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) | %FileCheck %s
+// RUN: %target-run-simple-swift(-Osize %S/Inputs/print.swift -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) | %FileCheck %s
 
 // REQUIRES: executable_test
 // REQUIRES: optimized_stdlib

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`// RUN: %target-run-simple-swift(%S/Inputs/print.swift -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) \| %FileCheck %s`
	`2`	`+// RUN: %target-run-simple-swift(-Osize %S/Inputs/print.swift -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) \| %FileCheck %s`
`2`	`3`
`3`	`4`	`// REQUIRES: executable_test`
`4`	`5`	`// REQUIRES: optimized_stdlib`