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)

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

@@ -37,6 +37,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.
@@ -140,6 +158,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;
 
@@ -384,6 +405,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/Serialization/ModuleFormat.h

@@ -53,8 +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.
-/// describe what change you made.
-const uint16_t VERSION_MINOR = 243; // Last change: re-number SIL stuff
+const uint16_t VERSION_MINOR = 244; // Last change: @_specialize attribute
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -1380,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, \

lib/AST/Attr.cpp

@@ -442,11 +442,21 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options)
     break;
   }
 
-  default:
-    llvm_unreachable("handled before this switch");
+  case DAK_Specialize: {
+    Printer << "@" << getAttrName() << "(";
+    auto *attr = cast<SpecializeAttr>(this);
+    interleave(attr->getTypeLocs(),
+               [&](TypeLoc tyLoc){ tyLoc.getType().print(Printer, Options); },
+               [&]{ Printer << ", "; });
+    Printer << ")";
+    break;
+  }
 
   case DAK_Count:
     llvm_unreachable("exceed declaration attribute kinds");
+
+  default:
+    llvm_unreachable("handled before this switch");
   }
 
   return true;
@@ -552,6 +562,8 @@ StringRef DeclAttribute::getAttrName() const {
     return "swift3_migration";
   case DAK_WarnUnusedResult:
     return "warn_unused_result";
+  case DAK_Specialize:
+    return "_specialize";
   }
   llvm_unreachable("bad DeclAttrKind");
 }
@@ -722,4 +734,26 @@ const AvailableAttr *AvailableAttr::isUnavailable(const Decl *D) {
   return D->getAttrs().getUnavailable(ctx);
 }
 
+SpecializeAttr::SpecializeAttr(SourceLoc atLoc, SourceRange range,
+                               ArrayRef<TypeLoc> typeLocs)
+    : DeclAttribute(DAK_Specialize, atLoc, range, /*Implicit=*/false),
+      numTypes(typeLocs.size())
+{
+  std::copy(typeLocs.begin(), typeLocs.end(), getTypeLocData());
+}
+
+ArrayRef<TypeLoc> SpecializeAttr::getTypeLocs() const {
+  return const_cast<SpecializeAttr*>(this)->getTypeLocs();
+}
 
+MutableArrayRef<TypeLoc> SpecializeAttr::getTypeLocs() {
+  return { this->getTypeLocData(), numTypes };
+}
+
+SpecializeAttr *SpecializeAttr::create(ASTContext &Ctx, SourceLoc atLoc,
+                                       SourceRange range,
+                                       ArrayRef<TypeLoc> typeLocs) {
+  unsigned size = sizeof(SpecializeAttr) + (typeLocs.size() * sizeof(TypeLoc));
+  void *mem = Ctx.Allocate(size, alignof(SpecializeAttr));
+  return new (mem) SpecializeAttr(atLoc, range, typeLocs);
+}

lib/Parse/ParseDecl.cpp

@@ -1291,6 +1291,42 @@ bool Parser::parseNewDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc,
                                                      rParenLoc, false));
     break;
     }
+  case DAK_Specialize: {
+    if (Tok.isNot(tok::l_paren)) {
+      diagnose(Loc, diag::attr_expected_lparen, AttrName,
+               DeclAttribute::isDeclModifier(DK));
+      return false;
+    }
+
+    SourceLoc lParenLoc = consumeToken();
+
+    SmallVector<TypeLoc, 8> TypeList;
+    do {
+      // Parse the concrete nominal type.
+      ParserResult<TypeRepr> Ty = parseTypeIdentifier();
+      if (Ty.isNull()) {
+        skipUntil(tok::r_paren);
+        return false;
+      }  
+
+      // Record the type.
+      TypeList.push_back(Ty.get());
+  
+      // Check for a ',', which indicates that there are more protocols coming.
+    } while (consumeIf(tok::comma));
+
+    // Parse the closing ')'.
+    SourceLoc rParenLoc;
+    if (!consumeIf(tok::r_paren, rParenLoc)) {
+      diagnose(lParenLoc, diag::attr_expected_rparen, AttrName,
+               /*DeclModifier=*/false);
+      return false;
+    }
+    Attributes.add(
+      SpecializeAttr::create(Context, AtLoc, SourceRange(Loc, rParenLoc),
+                             TypeList));
+    break;
+    }
   }
 
   if (DuplicateAttribute) {