[SILOptimizer] Introduce an eager-specializer pass.

This pass finds generic functions with @_specialized attributes and
generates specialized code for the attribute's concrete types. It
inserts type checks and guarded dispatch at the beginning of the
generic function for each specialization. Since we don't currently
expose this attribute as API and don't specialize vtables and witness
tables yet, the only way to reach the specialized code is by calling
the generic function which performs the guarded dispatch.

In the future, we can build on this work in several ways:
- cross module dispatch directly to specialized code
- dynamic dispatch directly to specialized code
- automated specialization based on less specific hints
- partial specialization
- and so on...

I reorganized and refactored the optimizer's generic utilities to
support direct function specialization as opposed to apply
specialization.

Author: atrick

include/swift/SILOptimizer/PassManager/Passes.def

@@ -88,6 +88,8 @@ PASS(DiagnoseUnreachable, "diagnose-unreachable",
      "Diagnose Unreachable Code")
 PASS(DiagnosticConstantPropagation, "diagnostic-constant-propagation",
      "Propagate constants and emit diagnostics")
+PASS(EagerSpecializer, "eager-specializer",
+     "Specialize speculatively and insert dispatch guarded by type checks")
 PASS(EarlyCodeMotion, "early-codemotion",
      "Code motion without release hoisting")
 PASS(EarlyInliner, "early-inline",

include/swift/SILOptimizer/Utils/GenericCloner.h

@@ -40,24 +40,25 @@ class GenericCloner : public TypeSubstCloner<GenericCloner> {
   GenericCloner(SILFunction *F,
                 const ReabstractionInfo &ReInfo,
                 TypeSubstitutionMap &ContextSubs,
+                ArrayRef<Substitution> ParamSubs,
                 StringRef NewName,
-                ArrayRef<Substitution> ApplySubs,
                 CloneCollector::CallbackType Callback)
   : TypeSubstCloner(*initCloned(F, ReInfo, NewName), *F, ContextSubs,
-                    ApplySubs), ReInfo(ReInfo), Callback(Callback) {
+                    ParamSubs), ReInfo(ReInfo), Callback(Callback) {
     assert(F->getDebugScope()->Parent != getCloned()->getDebugScope()->Parent);
   }
   /// Clone and remap the types in \p F according to the substitution
   /// list in \p Subs. Parameters are re-abstracted (changed from indirect to
   /// direct) according to \p ReInfo.
-  static SILFunction *cloneFunction(SILFunction *F,
-                                    const ReabstractionInfo &ReInfo,
-                                    TypeSubstitutionMap &ContextSubs,
-                                    StringRef NewName, ApplySite Caller,
-                            CloneCollector::CallbackType Callback =nullptr) {
+  static SILFunction *
+  cloneFunction(SILFunction *F,
+                const ReabstractionInfo &ReInfo,
+                TypeSubstitutionMap &ContextSubs,
+                ArrayRef<Substitution> ParamSubs,
+                StringRef NewName,
+                CloneCollector::CallbackType Callback =nullptr) {
     // Clone and specialize the function.
-    GenericCloner SC(F, ReInfo, ContextSubs, NewName,
-                     Caller.getSubstitutions(), Callback);
+    GenericCloner SC(F, ReInfo, ContextSubs, ParamSubs, NewName, Callback);
     SC.populateCloned();
     SC.cleanUp(SC.getCloned());
     return SC.getCloned();

include/swift/SILOptimizer/Utils/Generics.h

@@ -22,24 +22,51 @@
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Utils/Local.h"
-#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 namespace swift {
 
+/// Tries to specialize an \p Apply of a generic function. It can be a full
+/// apply site or a partial apply.
+/// Replaced and now dead instructions are returned in \p DeadApplies.
+/// New created functions, like the specialized callee and thunks, are returned
+/// in \p NewFunctions.
+///
+/// This is the top-level entry point for specializing an existing call site.
+void trySpecializeApplyOfGeneric(
+  ApplySite Apply,
+  llvm::SmallVectorImpl<SILInstruction *> &DeadApplies,
+  llvm::SmallVectorImpl<SILFunction *> &NewFunctions);
+
 /// Helper class to describe re-abstraction of function parameters done during
 /// specialization.
 ///
 /// Specifically, it contains information which parameters and returns are
 /// changed from indirect values to direct values.
 class ReabstractionInfo {
+  /// A 1-bit means that this parameter/return value is converted from indirect
+  /// to direct.
+  llvm::SmallBitVector Conversions;
+
+  /// The first NumResults bits in Conversions refer to indirect out-parameters.
+  unsigned NumResults;
+
+  /// The function type after applying the substitutions of the original
+  /// apply site.
+  CanSILFunctionType SubstitutedType;
+
+  /// The function type after applying the re-abstractions on the
+  /// SubstitutedType.
+  CanSILFunctionType SpecializedType;
+
 public:
-  /// Constructs the ReabstractionInfo for an apply site \p AI calling the
-  /// generic function \p Orig.
+  /// Constructs the ReabstractionInfo for generic function \p Orig with
+  /// substitutions \p ParamSubs.
   /// If specialization is not possible getSpecializedType() will return an
   /// invalid type.
-  ReabstractionInfo(SILFunction *Orig, ApplySite AI);
+  ReabstractionInfo(SILFunction *Orig, ArrayRef<Substitution> ParamSubs);
 
   /// Does the \p ArgIdx refer to an indirect out-parameter?
   bool isResultIndex(unsigned ArgIdx) const {
@@ -88,8 +115,8 @@ class ReabstractionInfo {
     return Conversions.size() - numArgs;
   }
 
-  /// Get the function type after applying the substitutions of the original
-  /// apply site.
+  /// Get the function type after applying the substitutions to the original
+  /// generic function.
   CanSILFunctionType getSubstitutedType() const { return SubstitutedType; }
 
   /// Get the function type after applying the re-abstractions on the
@@ -101,41 +128,61 @@ class ReabstractionInfo {
   /// SubstFTy by applying the re-abstractions.
   CanSILFunctionType createSpecializedType(CanSILFunctionType SubstFTy,
                                            SILModule &M) const;
-private:
-  /// A 1-bit means that this parameter/return value is converted from indirect
-  /// to direct.
-  llvm::BitVector Conversions;
-
-  /// The first NumResults bits in Conversions refer to indirect out-parameters.
-  unsigned NumResults;
+};
 
-  /// The function type after applying the substitutions of the original
-  /// apply site.
-  CanSILFunctionType SubstitutedType;
+/// Helper class for specializing a generic function given a list of
+/// substitutions.
+class GenericFuncSpecializer {
+  SILModule &M;
+  SILFunction *GenericFunc;
+  ArrayRef<Substitution> ParamSubs;
+  const ReabstractionInfo &ReInfo;
 
-  /// The function type after applying the re-abstractions on the
-  /// SubstitutedType.
-  CanSILFunctionType SpecializedType;
+  TypeSubstitutionMap ContextSubs;
+  std::string ClonedName;
+public:
+  GenericFuncSpecializer(SILFunction *GenericFunc,
+                         ArrayRef<Substitution> ParamSubs,
+                         const ReabstractionInfo &ReInfo);
+
+  /// If we already have this specialization, reuse it.
+  SILFunction *lookupSpecialization();
+
+  /// Return a newly created specialized function.
+  SILFunction *tryCreateSpecialization();
+  
+  /// Try to specialize GenericFunc given a list of ParamSubs.
+  /// Returns either a new or existing specialized function, or nullptr.
+  SILFunction *trySpecialization() {
+    if (!ReInfo.getSpecializedType())
+      return nullptr;
+
+    SILFunction *SpecializedF = lookupSpecialization();
+    if (!SpecializedF)
+      SpecializedF = tryCreateSpecialization();
+
+    return SpecializedF;
+  }
 };
 
-/// Tries to specialize an \p Apply of a generic function. It can be a full
-/// apply site or a partial apply.
-/// Replaced and now dead instructions are returned in \p DeadApplies.
-/// New created functions, like the specialized callee and thunks, are returned
-/// in \p NewFunctions.
-void trySpecializeApplyOfGeneric(ApplySite Apply,
-                        llvm::SmallVectorImpl<SILInstruction *> &DeadApplies,
-                        llvm::SmallVectorImpl<SILFunction *> &NewFunctions);
+// =============================================================================
+// Prespecialized symbol lookup.
+// =============================================================================
 
-/// Checks if a given mangled name could be a name of a whitelisted specialization.
+/// Checks if a given mangled name could be a name of a whitelisted
+/// specialization.
 bool isWhitelistedSpecialization(StringRef SpecName);
 
 /// Create a new apply based on an old one, but with a different
 /// function being applied.
 ApplySite replaceWithSpecializedFunction(ApplySite AI, SILFunction *NewF,
                                          const ReabstractionInfo &ReInfo);
 
-SILFunction *getExistingSpecialization(SILModule &M, StringRef FunctionName);
+/// Returns a SILFunction for the symbol specified by FunctioName if it is
+/// visible to the current SILModule. This is used to link call sites to
+/// externally defined specialization and should only be used when the function
+/// body is not required for further optimization or inlining (-Onone).
+SILFunction *lookupPrespecializedSymbol(SILModule &M, StringRef FunctionName);
 
 } // end namespace swift
 

lib/SILOptimizer/IPO/CMakeLists.txt

@@ -3,6 +3,7 @@ set(IPO_SOURCES
   IPO/CapturePropagation.cpp
   IPO/ClosureSpecializer.cpp
   IPO/DeadFunctionElimination.cpp
+  IPO/EagerSpecializer.cpp
   IPO/ExternalDefsToDecls.cpp
   IPO/FunctionSignatureOpts.cpp
   IPO/GlobalOpt.cpp