Implement polymorphic inline caches.

Swift SVN r20740

include/swift/SILAnalysis/ClassHierarchyAnalysis.h

@@ -44,6 +44,10 @@ public:
     }
   }
 
+  /// \brief Collect all of the known direct subclasses of a class \p C in the
+  /// current module.
+  void collectSubClasses(ClassDecl *C, std::vector<ClassDecl*> &Sub);
+
   /// \returns True if the class is inherited by another class in this module.
   bool inheritedInModule(ClassDecl *CD) {
     return InheritedClasses.count(CD); 

lib/SILAnalysis/ClassHierarchyAnalysis.cpp

@@ -33,4 +33,22 @@ void ClassHierarchyAnalysis::init() {
   }
 }
 
+void ClassHierarchyAnalysis::collectSubClasses(ClassDecl *C,
+                                               std::vector<ClassDecl*> &Sub) {
+  // TODO: Cache this search.
+  assert(Sub.empty() && "Incoming list is not empty");
+  for (auto &VT : M->getVTableList()) {
+    ClassDecl *CD = VT.getClass();
+    // Ignore classes that are at the top of the class hierarchy:
+    if (!CD->hasSuperclass())
+      continue;
+
+    // Add the superclass to the list of inherited classes.
+    ClassDecl *Super = CD->getSuperclass()->getClassOrBoundGenericClass();
+    if (Super == C) {
+      Sub.push_back(CD);
+    }
+  }
+}
+
 ClassHierarchyAnalysis::~ClassHierarchyAnalysis() {}

lib/SILPasses/Devirtualizer.cpp

@@ -43,6 +43,9 @@ STATISTIC(NumDevirtualized, "Number of calls devirtualzied");
 STATISTIC(NumDynApply, "Number of dynamic apply devirtualzied");
 STATISTIC(NumAMI, "Number of witness_method devirtualzied");
 
+// The number of subclasses to allow when placing polymorphic inline caches.
+static const int MaxNumPolymorphicInlineCaches = 2;
+
 //===----------------------------------------------------------------------===//
 //                         Class Method Optimization
 //===----------------------------------------------------------------------===//
@@ -84,7 +87,7 @@ static ClassDecl *findClassDeclForOperand(SILValue S) {
   return instTy.getClassOrBoundGenericClass();
 }
 
-/// Optimize a class_method and alloc_ref pair into a direct function
+/// \brief Optimize a class_method and alloc_ref pair into a direct function
 /// reference:
 ///
 /// \code
@@ -104,19 +107,23 @@ static ClassDecl *findClassDeclForOperand(SILValue S) {
 /// %YY = function_ref @...
 /// apply %YY(...)
 ///
-static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
+/// \p AI is the apply to devirtualize.
+/// \p Member is the class member to devirtualize.
+/// \p ClassInstance is the operand for the ClassMethodInst or an alternative
+///    reference (duch as downcasted class reference).
+/// \p KnownClass (can be null) is a specific class type to devirtualize to.
+static bool optimizeClassMethod(ApplyInst *AI, SILDeclRef Member,
+                                SILValue ClassInstance,
                                 ClassDecl *KnownClass) {
-  DEBUG(llvm::dbgs() << "    Trying to optimize : " << *CMI);
+  DEBUG(llvm::dbgs() << "    Trying to devirtualize : " << *AI);
 
   // First attempt to lookup the origin for our class method. The origin should
   // either be a metatype or an alloc_ref.
-  SILValue Origin = CMI->getOperand().stripCasts();
+  DEBUG(llvm::dbgs() << "        Origin: " << ClassInstance);
-  DEBUG(llvm::dbgs() << "        Origin: " << Origin);
 
   // Then attempt to lookup the class for origin.
-  ClassDecl *Class = KnownClass;
+  ClassDecl *Class = KnownClass ? KnownClass :
-  if (!Class)
+                     findClassDeclForOperand(ClassInstance);
-    Class = findClassDeclForOperand(Origin);
 
   // If we are unable to lookup a class decl for origin there is nothing here
   // that we can deserialize.
@@ -128,8 +135,8 @@ static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
 
   // Otherwise lookup from the module the least derived implementing method from
   // the module vtables.
-  SILModule &Mod = CMI->getModule();
+  SILModule &Mod = AI->getModule();
-  SILFunction *F = Mod.lookUpSILFunctionFromVTable(Class, CMI->getMember());
+  SILFunction *F = Mod.lookUpSILFunctionFromVTable(Class, Member);
 
   // If we do not find any such function, we have no function to devirtualize
   // to... so bail.
@@ -165,7 +172,7 @@ static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
   // Grab the self type from the function ref and the self type from the class
   // method inst.
   SILType FuncSelfTy = paramTypes[paramTypes.size() - 1];
-  SILType OriginTy = Origin.getType();
+  SILType OriginTy = ClassInstance.getType();
   SILBuilder B(AI);
 
   // Then compare the two types and if they are unequal...
@@ -175,11 +182,11 @@ static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
            " class.");
 
     // Otherwise, upcast origin to the appropriate type.
-    Origin = B.createUpcast(CMI->getLoc(), Origin, FuncSelfTy);
+    ClassInstance = B.createUpcast(AI->getLoc(), ClassInstance, FuncSelfTy);
   }
 
   // Success! Perform the devirtualization.
-  FunctionRefInst *FRI = B.createFunctionRef(CMI->getLoc(), F);
+  FunctionRefInst *FRI = B.createFunctionRef(AI->getLoc(), F);
 
   // Construct a new arg list. First process all non-self operands, ref, addr
   // casting them to the appropriate types for F so that we allow for covariant
@@ -216,7 +223,7 @@ static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
     }
   }
   // Add in self to the end.
-  NewArgs.push_back(Origin);
+  NewArgs.push_back(ClassInstance);
 
   // If we have a direct return type, make sure we use the subst callee return
   // type. If we have an indirect return type, AI's return type of the empty
@@ -250,8 +257,6 @@ static bool optimizeClassMethod(ApplyInst *AI, ClassMethodInst *CMI,
   }
 
   AI->eraseFromParent();
-  if (CMI->use_empty())
-    CMI->eraseFromParent();
 
   DEBUG(llvm::dbgs() << "        SUCCESS: " << F->getName() << "\n");
   NumDevirtualized++;
@@ -781,7 +786,8 @@ static bool optimizeApplyInst(ApplyInst *AI) {
   ///
   /// %YY = function_ref @...
   if (auto *CMI = dyn_cast<ClassMethodInst>(AI->getCallee()))
-    return optimizeClassMethod(AI, CMI, nullptr);
+    return optimizeClassMethod(AI, CMI->getMember(),
+                               CMI->getOperand().stripCasts(), nullptr);
 
   // Devirtualize protocol_method + project_existential + init_existential
   // instructions.  For example:
@@ -858,10 +864,12 @@ static ApplyInst *CloneApply(ApplyInst *AI, SILBuilder &Builder) {
                              Ret, AI->isTransparent());
 }
 
-static bool insertMonomorphicInlineCaches(ApplyInst *AI, SILType InstanceType,
+/// Insert monomorphic inline caches for a specific class type \p SubClassTy.
-                                          SILValue ClassInstance) {
+static ApplyInst* insertMonomorphicInlineCaches(ApplyInst *AI,
+                                                SILType SubClassTy) {
   ClassMethodInst *CMI = cast<ClassMethodInst>(AI->getCallee());
-  ClassDecl *CD = InstanceType.getClassOrBoundGenericClass();
+  SILValue ClassInstance = CMI->getOperand();
+  ClassDecl *CD = SubClassTy.getClassOrBoundGenericClass();
 
   // Create a diamond shaped control flow and a checked_cast_branch
   // instruction that checks the exact type of the object.
@@ -875,15 +883,17 @@ static bool insertMonomorphicInlineCaches(ApplyInst *AI, SILType InstanceType,
   SILBasicBlock *Iden = F->createBasicBlock();
   // Virt is the block containing the slow virtual call.
   SILBasicBlock *Virt = F->createBasicBlock();
-  Iden->createArgument(InstanceType);
+  Iden->createArgument(SubClassTy);
 
   SILBasicBlock *Continue = Entry->splitBasicBlock(It);
 
   SILBuilder Builder(Entry);
   // Create the checked_cast_branch instruction that checks at runtime if the
   // class instance is identical to the SILType.
+  assert(SubClassTy.getClassOrBoundGenericClass() &&
+         "Dest type must be a class type");
   It = Builder.createCheckedCastBranch(AI->getLoc(), /*exact*/ true,
-                                       ClassInstance, InstanceType, Iden, Virt);
+                                       ClassInstance, SubClassTy, Iden, Virt);
 
   SILBuilder VirtBuilder(Virt);
   SILBuilder IdenBuilder(Iden);
@@ -919,9 +929,12 @@ static bool insertMonomorphicInlineCaches(ApplyInst *AI, SILType InstanceType,
   // Update the stats.
   NumInlineCaches++;
 
+  // This is the class reference downcasted into subclass SubClassTy.
+  SILValue DownCastedClassInstance = Iden->getBBArg(0);
+
   // Devirtualize the apply instruction on the identical path.
-  optimizeClassMethod(IdenAI, CMI, CD);
+  optimizeClassMethod(IdenAI, CMI->getMember(), DownCastedClassInstance, CD);
-  return true;
+  return VirtAI;
 }
 
 /// Specialize virtual dispatch.
@@ -939,12 +952,31 @@ static bool insertInlineCaches(ApplyInst *AI, ClassHierarchyAnalysis *CHA) {
   if (!CHA->inheritedInModule(CD)) {
     DEBUG(llvm::dbgs() << "Inserting monomorphic inline caches for class " <<
           CD->getName() << "\n");
-    return insertMonomorphicInlineCaches(AI, InstanceType, ClassInstance);
+    return insertMonomorphicInlineCaches(AI, InstanceType);
+  }
+
+  std::vector<ClassDecl*> Subs;
+  CHA->collectSubClasses(CD, Subs);
+
+  if (Subs.size() > MaxNumPolymorphicInlineCaches) {
+    DEBUG(llvm::dbgs() << "Class " << CD->getName() << " has " <<  Subs.size()
+          << " subclasses. Not inserting polymorphic inline caches.\n");
   }
 
   DEBUG(llvm::dbgs() << "Class " << CD->getName() << " is a superclass. "
-        " Not inserting monomorphic inline caches.\n");
+        "Inserting monomorphic inline caches.\n");
-  return false;
+
+  for (auto S : Subs) {
+    CanType CanClassType = S->getDeclaredType()->getCanonicalType();
+    SILType InstanceType = SILType::getPrimitiveObjectType(CanClassType);
+    if (!InstanceType.getClassOrBoundGenericClass())
+      continue;
+
+    AI = insertMonomorphicInlineCaches(AI, InstanceType);
+    assert(AI && "Unable to insert inline caches!");
+  }
+
+  return true;
 }
 
 namespace {

test/SILPasses/devirt_ctors.sil

@@ -22,7 +22,8 @@ sil @_TFC4main3ABC7member2fS0_FT_T_ : $@cc(method) @thin (@owned ABC) -> ()
 sil @_TFC4main3ABCCfMS0_FT1ySi_S0_ : $@thin (Int64, @thick ABC.Type) -> @owned ABC
 
 //CHECK:@_TF4main3fooFT_T_
-//CHECK-NOT: class_method
+//CHECK: function_ref @_TFC4main3ABC6memberfS0_FT_T_
+//CHECK: function_ref @_TFC4main3ABC7member2fS0_FT_T_
 //CHECK: return
 sil @_TF4main3fooFT_T_ : $@thin () -> () {
 bb0:

test/SILPasses/polymorphic_inline_caches.swift

@@ -0,0 +1,22 @@
+// RUN: %swift %s -O -emit-sil
+
+class A {
+  func ping() -> Int { return 1 }
+}
+class B : A {
+  override func ping() -> Int { return 2 }
+}
+class C : A {
+  override func ping() -> Int{ return 3 }
+}
+
+//CHECK-LABEL: @_TF4main3fooFCS_1ASi
+//CHECK: checked_cast_br [exact] %0 : $A to $B
+//CHECK: integer_literal $Builtin.Word, 2
+//CHECK: checked_cast_br [exact] %0 : $A to $C
+//CHECK: integer_literal $Builtin.Word, 3
+//CHECK: return
+func foo(x : A) -> Int {
+  return x.ping()
+
+}

test/SILPasses/sil_combine_devirt.sil

@@ -21,7 +21,7 @@ class Foo {
 }
 
 //CHECK: _TF4main4mainFT_T_
-//CHECK-NOT: class_method
+//CHECK: function_ref @_TFC4main3Foo3getfS0_FT_Si
 //CHECK: return
 sil @_TF4main4mainFT_T_ : $@thin () -> () {
 bb0: