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

This was mistakenly reverted in an attempt to fix buildbots.
Unfortunately it's now smashed into one commit.

---
Introduce @_specialize(<type list>) internal attribute.

This attribute can be attached to generic functions. The attribute's
arguments must be a list of concrete types to be substituted in the
function's generic signature. Any number of specializations may be
associated with a generic function.

This attribute provides a hint to the compiler. At -O, the compiler
will generate the specified specializations and emit calls to the
specialized code in the original generic function guarded by type
checks.

The current attribute is designed to be an internal tool for
performance experimentation. It does not affect the language or
API. This work may be extended in the future to add user-visible
attributes that do provide API guarantees and/or direct dispatch to
specialized code.

This attribute works on any generic function: a freestanding function
with generic type parameters, a nongeneric method declared in a
generic class, a generic method in a nongeneric class or a generic
method in a generic class. A function's generic signature is a
concatenation of the generic context and the function's own generic
type parameters.

e.g.

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)

---
[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

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

@@ -1976,6 +1976,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, \