ProtocolDevirtualizer Part 2: Transformation Phase

Author: Raj Barik

include/swift/AST/SILOptions.h

@@ -64,6 +64,9 @@ class SILOptions {
   /// Remove all runtime assertions during optimizations.
   bool RemoveRuntimeAsserts = false;
 
+  /// Enable protocol devirtualization optimization.
+  bool ProtocolDevirtualizer = false;
+
   /// Controls whether the SIL ARC optimizations are run.
   bool EnableARCOptimizations = true;
 

include/swift/Option/FrontendOptions.td

@@ -390,6 +390,9 @@ def sil_unroll_threshold : Separate<["-"], "sil-unroll-threshold">,
   MetaVarName<"<250>">,
   HelpText<"Controls the aggressiveness of loop unrolling">;
 
+def sil_protocol_devirtualizer : Flag<["-"], "sil-protocol-devirtualizer">,
+  HelpText<"Enable SIL protocol devirtualizer optimization">;
+
 def sil_merge_partial_modules : Flag<["-"], "sil-merge-partial-modules">,
   HelpText<"Merge SIL from all partial swiftmodules into the final module">;
 

include/swift/SILOptimizer/PassManager/Passes.def

@@ -150,6 +150,8 @@ PASS(DeadStoreElimination, "dead-store-elim",
      "Dead Store Elimination")
 PASS(GenericSpecializer, "generic-specializer",
      "Generic Function Specialization on Static Types")
+PASS(ProtocolDevirtualizer, "protocol-devirtualizer",
+     "Protocol Devirtualizer")
 PASS(GlobalOpt, "global-opt",
      "SIL Global Optimization")
 PASS(GlobalPropertyOpt, "global-property-opt",

include/swift/SILOptimizer/Utils/Generics.h

@@ -20,7 +20,10 @@
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SILOptimizer/Analysis/ProtocolDevirtualizerAnalysis.h"
+#include "swift/SILOptimizer/Utils/FunctionSignatureOptUtils.h"
 #include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/Utils/SpecializationMangler.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -299,6 +302,61 @@ class GenericFuncSpecializer {
   }
 };
 
+/// ProtocolDevirtualizer transformation class that creates a protocol 
+/// constrained generic and a thunk.
+class ProtocolDevirtualizerTransform {
+  /// The original function to analyze and transform.
+  SILFunction *F;
+
+  /// The newly created inner function.
+  SILFunction *NewF;
+
+  /// The function signature mangler we are using.
+  Mangle::FunctionSignatureSpecializationMangler &Mangler;
+
+  /// Map Arguments to their ProtocolTypes.
+  llvm::SmallDenseMap<int, std::pair<ProtocolDecl *, ClassDecl *>> &Arg2DeclMap;
+
+  /// Argument to Generic Type Map for NewF.
+  llvm::SmallDenseMap<int, GenericTypeParamType *> Arg2GenericTypeMap;
+
+  // Allocate the argument descriptors.
+  llvm::SmallVector<ArgumentDescriptor, 4> &ArgumentDescList;
+
+  /// Create the Devirtualized Inner Function.
+  void createDevirtualizedProtocolFunction();
+
+  /// Create a name for the inner function.
+  std::string createDevirtualizedFunctionName();
+
+  /// Create the new devirtualized protocol function signature.
+  CanSILFunctionType createDevirtualizedFunctionType();
+
+  /// Populate the body of NewF.
+  void populateProtocolConstrainedGenericFunction(const SILDebugScope *DS);
+
+  /// Create the thunk.
+  void populateThunkBody();
+
+  public:
+  /// Constructor.
+  ProtocolDevirtualizerTransform(
+              SILFunction *F,
+              Mangle::FunctionSignatureSpecializationMangler &Mangler,
+              llvm::SmallVector<ArgumentDescriptor, 4> &ADL,
+              llvm::SmallDenseMap<int, std::pair<ProtocolDecl *, ClassDecl *>> &Arg2DeclMap
+            ) : F(F), NewF(nullptr), Mangler(Mangler), Arg2DeclMap(Arg2DeclMap) , ArgumentDescList(ADL) {}
+
+  /// Return the optimized iner function.
+  SILFunction *getDevirtualizedProtocolFunction() { return NewF; }
+
+  /// Run the optimization.
+  bool run() {
+    createDevirtualizedProtocolFunction();
+    return true;
+  }
+};
+
 // =============================================================================
 // Prespecialized symbol lookup.
 // =============================================================================

lib/Frontend/CompilerInvocation.cpp

@@ -568,6 +568,9 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
       return true;
     }
   }
+  if (Args.hasArg(OPT_sil_protocol_devirtualizer)) {
+    Opts.ProtocolDevirtualizer = true;
+  }
   if (const Arg *A = Args.getLastArg(OPT_num_threads)) {
     if (StringRef(A->getValue()).getAsInteger(10, Opts.NumThreads)) {
       Diags.diagnose(SourceLoc(), diag::error_invalid_arg_value,

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -331,6 +331,9 @@ static void addPerfEarlyModulePassPipeline(SILPassPipelinePlan &P) {
 
   // Add the outliner pass (Osize).
   P.addOutliner();
+
+  // Protocol Devirtualizer Pass.
+  P.addProtocolDevirtualizer();
 }
 
 static void addHighLevelEarlyLoopOptPipeline(SILPassPipelinePlan &P) {

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -24,9 +24,13 @@
 #include "swift/SILOptimizer/Utils/Local.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 
 using namespace swift;
 
+STATISTIC(NumFunctionsWithProtocolArgsDevirtualized, "Number of functions with protocol args devirtualized");
+using ArgIndexList = llvm::SmallVector<unsigned, 8>;
+
 namespace {
 
 class GenericSpecializer : public SILFunctionTransform {
@@ -45,6 +49,93 @@ class GenericSpecializer : public SILFunctionTransform {
 
 };
 
+/// ProtocolDevirtualizer class.
+class ProtocolDevirtualizer : public SILFunctionTransform {
+
+  /// determine if the current function is a target for protocol devirtualizer.
+  bool canDevirtualizeProtocolInFunction(ProtocolDevirtualizerAnalysis *PDA,
+          llvm::SmallDenseMap<int, std::pair<ProtocolDecl*, ClassDecl *>> &Arg2DeclMap);
+
+public:
+
+  void run() override {
+    auto *F = getFunction();
+
+    /// Don't run protocol devirtualizer at -Os.
+    if (F->optimizeForSize())
+      return;
+
+    /// Don't optimize functions that should not be optimized.
+    if (F->empty() || (!F->shouldOptimize())) {
+      return;
+    }
+
+    /// This is the function to optimize for protocol devirtualize.
+    DEBUG(llvm::dbgs() << "*** ProtocolDevirtualization Pass on function: " << F->getName() << " ***\n");
+
+    /// Use FunctionSignature Specialization to determine if this function
+    /// can be specialized, without call graph information.
+    if ( !canSpecializeFunction(F, nullptr, false)) {
+      DEBUG(llvm::dbgs() << "  cannot specialize function -> abort\n");
+      return;
+    }
+
+    CallerAnalysis *CA = PM->getAnalysis<CallerAnalysis>();
+    const CallerAnalysis::FunctionInfo &FuncInfo = CA->getCallerInfo(F);
+
+    /// Use FunctionSignature Specialization to determine if this function
+    /// can be specialized, based on call graph.
+    /// canSpecializeFunction does not consider generic methods -- TODO in future.
+    if (!canSpecializeFunction(F, &FuncInfo, true)) {
+      DEBUG(llvm::dbgs() << "  cannot specialize function -> abort\n");
+      return;
+    }
+
+    /// Get the protocol devirtualizer pass that contains which protocols can be devirtualized.
+    ProtocolDevirtualizerAnalysis *PDA = getAnalysis<ProtocolDevirtualizerAnalysis>();
+
+    /// Determine the arguments that can be devirtualized.
+    llvm::SmallDenseMap<int, std::pair<ProtocolDecl *, ClassDecl *>> Arg2DeclMap;
+    if(!canDevirtualizeProtocolInFunction(PDA, Arg2DeclMap)) {
+      DEBUG(llvm::dbgs() << "  cannot devirtualize function -> abort\n");
+      return;
+    }
+
+    DEBUG(llvm::dbgs() << "Function::" << F->getName() << " has a Protocol Argument and can be optimized via PDA\n");
+
+    /// Name Mangler for naming the protocol constrained generic method.
+    auto P = Demangle::SpecializationPass::GenericSpecializer;
+    Mangle::FunctionSignatureSpecializationMangler Mangler(P, F->isSerialized(), F);
+
+    /// Save the arguments in a descriptor.
+    llvm::SmallVector<ArgumentDescriptor, 4> ArgumentDescList;
+    auto Args = F->begin()->getFunctionArguments();
+    for (unsigned i = 0, e = Args.size(); i != e; ++i) {
+      ArgumentDescList.emplace_back(Args[i]);
+    }
+
+    /// Instantiate the ProtocolDevirtualizerTransform pass.
+    ProtocolDevirtualizerTransform PDT(F, Mangler, ArgumentDescList, Arg2DeclMap);
+
+    /// Run the protocol devirtualizer.
+    bool Changed = PDT.run();
+
+    if (Changed) {
+      /// Update statistics on the number of functions devirtualized.
+      ++ NumFunctionsWithProtocolArgsDevirtualized;
+
+      /// Invalidate Analysis as we introduce a new function.
+      invalidateAnalysis(SILAnalysis::InvalidationKind::Everything);
+
+      /// Make sure the PM knows about the new devirtualized inner function.
+      notifyAddFunction(PDT.getDevirtualizedProtocolFunction(), F);
+
+      /// should we restart pipeline? be conservative.
+      restartPassPipeline();
+    }
+  }
+};
+
 } // end anonymous namespace
 
 bool GenericSpecializer::specializeAppliesInFunction(SILFunction &F) {
@@ -129,6 +220,39 @@ bool GenericSpecializer::specializeAppliesInFunction(SILFunction &F) {
   return Changed;
 }
 
+/// Check if any argument to a function meet the criteria for devirtualization.
+bool ProtocolDevirtualizer::canDevirtualizeProtocolInFunction(ProtocolDevirtualizerAnalysis *PDA,
+        llvm::SmallDenseMap<int, std::pair<ProtocolDecl*, ClassDecl *>> &Arg2DeclMap) {
+  auto *F = getFunction();
+  auto Args = F->begin()->getFunctionArguments();
+  bool returnFlag = false;
+
+  /// Analyze the argument for protocol conformance.
+  for (unsigned i = 0, e = Args.size(); i != e; ++i) {
+    auto ArgType = Args[i]->getType();
+    /// Keep it simple for now. Do not handle Protocol constrained methods or Optionals.
+    auto SwiftArgType = ArgType.getSwiftRValueType();
+    if (ArgType && ArgType.isAnyExistentialType()) {
+      auto SwiftProtoDecl = SwiftArgType.getAnyNominal();
+      if (SwiftProtoDecl) {
+        auto ProtoDecl = dyn_cast<ProtocolDecl>(SwiftProtoDecl);
+        auto CD = PDA->getSoleClassImplementingProtocol(ProtoDecl);
+        if (CD) {
+          /// Save the mapping for transformation pass.
+          Arg2DeclMap[i] = std::make_pair(ProtoDecl, CD);
+          DEBUG(llvm::dbgs() << "Function: " << F->getName() << " has a singel class decl\n");
+          returnFlag |= true;
+        }
+      }
+    }
+  }
+  return returnFlag;
+}
+
 SILTransform *swift::createGenericSpecializer() {
   return new GenericSpecializer();
 }
+
+SILTransform *swift::createProtocolDevirtualizer() {
+  return new ProtocolDevirtualizer();
+}