Merge pull request #1743 from atrick/reapply-specialize

Reapply "Merge pull request #1725 from atrick/specialize"

Author: atrick

docs/Generics.rst

@@ -713,6 +713,35 @@ the virtual dispatch, inline calls when appropriate, and eliminate the overhead
 of the generic system. Such optimizations can be performed based on heuristics,
 user direction, or profile-guided optimization.
 
+An internal @_specialize function attribute allows developers to force
+full specialization by listing concrete type names corresponding to the
+function's generic signature. A function's generic signature is a
+concatenation of its generic context and the function's own generic
+type parameters.::
+
+  struct S<T> {
+    var x: T
+    @_specialize(Int, Float)
+    mutating func exchangeSecond<U>(u: U, _ t: T) -> (U, T) {
+      x = t
+      return (u, x)
+    }
+  }
+
+  // Substitutes: <T, U> with <Int, Float> producing:
+  // S<Int>::exchangeSecond<Float>(u: Float, t: Int) -> (Float, Int)
+
+@_specialize currently acts as a hint to the optimizer, which
+generates type checks and code to dispatch to the specialized routine
+without affecting the signature of the generic function. The
+intention is to support efforts at evaluating the performance of
+specialized code. The performance impact is not guaranteed and is
+likely to change with the optimizer. This attribute should only be
+used in conjunction with rigorous performance analysis. Eventually,
+a similar attribute could be defined in the language, allowing it to be
+exposed as part of a function's API. That would allow direct dispatch
+to specialized code without type checks, even across modules.
+
 Existential Types and Generics
 ------------------------------
 

include/swift/AST/Attr.def

@@ -164,6 +164,10 @@ SIMPLE_DECL_ATTR(nonobjc, NonObjC,
 SIMPLE_DECL_ATTR(_fixed_layout, FixedLayout,
                  OnVar | OnClass | OnStruct | OnEnum | UserInaccessible, 31)
 
+DECL_ATTR(_specialize, Specialize,
+          OnConstructor | OnFunc | AllowMultipleAttributes | LongAttribute
+          | UserInaccessible, 32)
+
 // Non-serialized attributes.
 
 SIMPLE_DECL_ATTR(mutating, Mutating, OnFunc | DeclModifier | NotSerialized,

include/swift/AST/Attr.h

@@ -23,6 +23,7 @@
 #include "swift/Basic/Range.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/AttrKind.h"
+#include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/KnownProtocols.h"
 #include "swift/AST/Ownership.h"
 #include "swift/AST/PlatformKind.h"
@@ -38,6 +39,7 @@ class ASTContext;
 struct PrintOptions;
 class Decl;
 class ClassDecl;
+struct TypeLoc;
 
 class InfixData {
   unsigned Precedence : 8;
@@ -1136,6 +1138,36 @@ class Swift3MigrationAttr : public DeclAttribute {
   }
 };
 
+/// The @_specialize attribute, which forces specialization on the specified
+/// type list.
+class SpecializeAttr : public DeclAttribute {
+  unsigned numTypes;
+  ConcreteDeclRef specializedDecl;
+
+  TypeLoc *getTypeLocData() {
+    return reinterpret_cast<TypeLoc *>(this + 1);
+  }
+
+  SpecializeAttr(SourceLoc atLoc, SourceRange Range,
+                 ArrayRef<TypeLoc> typeLocs);
+
+public:
+  static SpecializeAttr *create(ASTContext &Ctx, SourceLoc atLoc,
+                                SourceRange Range, ArrayRef<TypeLoc> typeLocs);
+
+  ArrayRef<TypeLoc> getTypeLocs() const;
+
+  MutableArrayRef<TypeLoc> getTypeLocs();
+
+  ConcreteDeclRef getConcreteDecl() const { return specializedDecl; }
+
+  void setConcreteDecl(ConcreteDeclRef ref) { specializedDecl = ref; }
+
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DAK_Specialize;
+  }
+};
+
 /// \brief Attributes that may be applied to declarations.
 class DeclAttributes {
   /// Linked list of declaration attributes.

include/swift/AST/DiagnosticsSema.def

@@ -1978,6 +1978,8 @@ ERROR(generic_type_requires_arguments,none,
       "reference to generic type %0 requires arguments in <...>", (Type))
 NOTE(generic_type_declared_here,none,
      "generic type %0 declared here", (Identifier))
+ERROR(cannot_partially_specialize_generic_function,none,
+     "cannot partially specialize a generic function", ())
 
 // Ambiguities
 ERROR(ambiguous_decl_ref,none,

include/swift/SIL/SILFunction.h

@@ -38,6 +38,24 @@ enum IsTransparent_t { IsNotTransparent, IsTransparent };
 enum Inline_t { InlineDefault, NoInline, AlwaysInline };
 enum IsThunk_t { IsNotThunk, IsThunk, IsReabstractionThunk };
 
+class SILSpecializeAttr final :
+    private llvm::TrailingObjects<SILSpecializeAttr, Substitution> {
+  friend TrailingObjects;
+
+  unsigned numSubs;
+  
+  SILSpecializeAttr(ArrayRef<Substitution> subs);
+
+public:
+  static SILSpecializeAttr *create(SILModule &M, ArrayRef<Substitution> subs);
+
+  ArrayRef<Substitution> getSubstitutions() const {
+    return { getTrailingObjects<Substitution>(), numSubs };
+  }
+  
+  void print(llvm::raw_ostream &OS) const;
+};
+
 /// SILFunction - A function body that has been lowered to SIL. This consists of
 /// zero or more SIL SILBasicBlock objects that contain the SILInstruction
 /// objects making up the function.
@@ -141,6 +159,9 @@ class SILFunction
   /// StringRefs?
   llvm::SmallVector<std::string, 1> SemanticsAttrSet;
 
+  /// The function's remaining set of specialize attributes.
+  std::vector<SILSpecializeAttr*> SpecializeAttrSet;
+
   /// The function's effects attribute.
   EffectsKind EffectsKindAttr;
 
@@ -385,6 +406,18 @@ class SILFunction
     SemanticsAttrSet.erase(Iter);
   }
 
+  /// \returns the range of specialize attributes.
+  ArrayRef<SILSpecializeAttr*> getSpecializeAttrs() const {
+    return SpecializeAttrSet;
+  }
+
+  /// Removes all specialize attributes from this function.
+  void clearSpecializeAttrs() { SpecializeAttrSet.clear(); }
+
+  void addSpecializeAttr(SILSpecializeAttr *attr) {
+    SpecializeAttrSet.push_back(attr);
+  }
+
   /// \returns True if the function is optimizable (i.e. not marked as no-opt),
   ///          or is raw SIL (so that the mandatory passes still run).
   bool shouldOptimize() const;

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
 

include/swift/Serialization/ModuleFormat.h

@@ -53,7 +53,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 245; // Last change: re-number SIL stuff
+const uint16_t VERSION_MINOR = 246; // Last change: @_specialize attribute
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -1379,6 +1379,11 @@ namespace decls_block {
                // strings, separated by the prior index
   >;
 
+  using SpecializeDeclAttrLayout = BCRecordLayout<
+    Specialize_DECL_ATTR,
+    BCArray<TypeIDField> // concrete types
+  >;
+
 #define SIMPLE_DECL_ATTR(X, CLASS, ...) \
   using CLASS##DeclAttrLayout = BCRecordLayout< \
     CLASS##_DECL_ATTR, \