[AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable. (swiftlang#27506)

This PR introduces `@differentiable` attribute to mark functions as differentiable. This PR only contains changes related to parsing the attribute. Type checking and other changes will be added in subsequent patches.

See https://github.com/apple/swift/pull/27506/files#diff-f3216f4188fd5ed34e1007e5a9c2490f for examples and tests for the new attribute.

Author: bgogul

include/swift/AST/Attr.def

@@ -502,6 +502,12 @@ SIMPLE_DECL_ATTR(_nonEphemeral, NonEphemeral,
   ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
   90)
 
+DECL_ATTR(differentiable, Differentiable,
+  OnAccessor | OnConstructor | OnFunc | OnVar | OnSubscript | LongAttribute |
+  AllowMultipleAttributes |
+  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove,
+  91)
+
 SIMPLE_DECL_ATTR(IBSegueAction, IBSegueAction,
   OnFunc |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,

include/swift/AST/Attr.h

@@ -27,6 +27,7 @@
 #include "swift/Basic/Version.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/AttrKind.h"
+#include "swift/AST/AutoDiff.h"
 #include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/DeclNameLoc.h"
 #include "swift/AST/KnownProtocols.h"
@@ -1724,6 +1725,148 @@ class DeclAttributes {
   SourceLoc getStartLoc(bool forModifiers = false) const;
 };
 
+/// A declaration name with location.
+struct DeclNameWithLoc {
+  DeclName Name;
+  DeclNameLoc Loc;
+};
+
+/// Attribute that marks a function as differentiable and optionally specifies
+/// custom associated derivative functions: 'jvp' and 'vjp'.
+///
+/// Examples:
+///   @differentiable(jvp: jvpFoo where T : FloatingPoint)
+///   @differentiable(wrt: (self, x, y), jvp: jvpFoo)
+class DifferentiableAttr final
+    : public DeclAttribute,
+      private llvm::TrailingObjects<DifferentiableAttr,
+                                    ParsedAutoDiffParameter> {
+  friend TrailingObjects;
+
+  /// Whether this function is linear (optional).
+  bool linear;
+  /// The number of parsed parameters specified in 'wrt:'.
+  unsigned NumParsedParameters = 0;
+  /// The JVP function.
+  Optional<DeclNameWithLoc> JVP;
+  /// The VJP function.
+  Optional<DeclNameWithLoc> VJP;
+  /// The JVP function (optional), resolved by the type checker if JVP name is
+  /// specified.
+  FuncDecl *JVPFunction = nullptr;
+  /// The VJP function (optional), resolved by the type checker if VJP name is
+  /// specified.
+  FuncDecl *VJPFunction = nullptr;
+  /// The differentiation parameters' indices, resolved by the type checker.
+  IndexSubset *ParameterIndices = nullptr;
+  /// The trailing where clause (optional).
+  TrailingWhereClause *WhereClause = nullptr;
+  /// The generic signature for autodiff associated functions. Resolved by the
+  /// type checker based on the original function's generic signature and the
+  /// attribute's where clause requirements. This is set only if the attribute
+  /// has a where clause.
+  GenericSignature DerivativeGenericSignature;
+
+  explicit DifferentiableAttr(ASTContext &context, bool implicit,
+                              SourceLoc atLoc, SourceRange baseRange,
+                              bool linear,
+                              ArrayRef<ParsedAutoDiffParameter> parameters,
+                              Optional<DeclNameWithLoc> jvp,
+                              Optional<DeclNameWithLoc> vjp,
+                              TrailingWhereClause *clause);
+
+  explicit DifferentiableAttr(ASTContext &context, bool implicit,
+                              SourceLoc atLoc, SourceRange baseRange,
+                              bool linear, IndexSubset *indices,
+                              Optional<DeclNameWithLoc> jvp,
+                              Optional<DeclNameWithLoc> vjp,
+                              GenericSignature derivativeGenericSignature);
+
+public:
+  static DifferentiableAttr *create(ASTContext &context, bool implicit,
+                                    SourceLoc atLoc, SourceRange baseRange,
+                                    bool linear,
+                                    ArrayRef<ParsedAutoDiffParameter> params,
+                                    Optional<DeclNameWithLoc> jvp,
+                                    Optional<DeclNameWithLoc> vjp,
+                                    TrailingWhereClause *clause);
+
+  static DifferentiableAttr *create(ASTContext &context, bool implicit,
+                                    SourceLoc atLoc, SourceRange baseRange,
+                                    bool linear, IndexSubset *indices,
+                                    Optional<DeclNameWithLoc> jvp,
+                                    Optional<DeclNameWithLoc> vjp,
+                                    GenericSignature derivativeGenSig);
+
+  /// Get the optional 'jvp:' function name and location.
+  /// Use this instead of `getJVPFunction` to check whether the attribute has a
+  /// registered JVP.
+  Optional<DeclNameWithLoc> getJVP() const { return JVP; }
+
+  /// Get the optional 'vjp:' function name and location.
+  /// Use this instead of `getVJPFunction` to check whether the attribute has a
+  /// registered VJP.
+  Optional<DeclNameWithLoc> getVJP() const { return VJP; }
+
+  IndexSubset *getParameterIndices() const {
+    return ParameterIndices;
+  }
+  void setParameterIndices(IndexSubset *pi) {
+    ParameterIndices = pi;
+  }
+
+  /// The parsed differentiation parameters, i.e. the list of parameters
+  /// specified in 'wrt:'.
+  ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  MutableArrayRef<ParsedAutoDiffParameter> getParsedParameters() {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  size_t numTrailingObjects(OverloadToken<ParsedAutoDiffParameter>) const {
+    return NumParsedParameters;
+  }
+
+  bool isLinear() const { return linear; }
+
+  TrailingWhereClause *getWhereClause() const { return WhereClause; }
+
+  GenericSignature getDerivativeGenericSignature() const {
+    return DerivativeGenericSignature;
+  }
+  void setDerivativeGenericSignature(ASTContext &context,
+                                     GenericSignature derivativeGenSig) {
+    DerivativeGenericSignature = derivativeGenSig;
+  }
+
+  FuncDecl *getJVPFunction() const { return JVPFunction; }
+  void setJVPFunction(FuncDecl *decl);
+  FuncDecl *getVJPFunction() const { return VJPFunction; }
+  void setVJPFunction(FuncDecl *decl);
+
+  bool parametersMatch(const DifferentiableAttr &other) const {
+    assert(ParameterIndices && other.ParameterIndices);
+    return ParameterIndices == other.ParameterIndices;
+  }
+
+  /// Get the derivative generic environment for the given `@differentiable`
+  /// attribute and original function.
+  GenericEnvironment *
+  getDerivativeGenericEnvironment(AbstractFunctionDecl *original) const;
+
+  // Print the attribute to the given stream.
+  // If `omitWrtClause` is true, omit printing the `wrt:` clause.
+  // If `omitAssociatedFunctions` is true, omit printing associated functions.
+  void print(llvm::raw_ostream &OS, const Decl *D,
+             bool omitWrtClause = false,
+             bool omitAssociatedFunctions = false) const;
+
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DAK_Differentiable;
+  }
+};
+
+
 void simple_display(llvm::raw_ostream &out, const DeclAttribute *attr);
 
 inline SourceLoc extractNearestSourceLoc(const DeclAttribute *attr) {

include/swift/AST/AutoDiff.h

@@ -0,0 +1,91 @@
+//===--- AutoDiff.h - Swift Automatic Differentiation ---------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+//  SWIFT_ENABLE_TENSORFLOW
+//  This file defines AST support for automatic differentiation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_AUTODIFF_H
+#define SWIFT_AST_AUTODIFF_H
+
+#include "swift/AST/IndexSubset.h"
+#include "swift/Basic/Range.h"
+
+namespace swift {
+
+class ParsedAutoDiffParameter {
+public:
+  enum class Kind { Named, Ordered, Self };
+
+private:
+  SourceLoc loc;
+  Kind kind;
+  union Value {
+    struct { Identifier name; } Named;
+    struct { unsigned index; } Ordered;
+    struct {} self;
+    Value(Identifier name) : Named({name}) {}
+    Value(unsigned index) : Ordered({index}) {}
+    Value() {}
+  } value;
+
+public:
+  ParsedAutoDiffParameter(SourceLoc loc, Kind kind, Value value)
+    : loc(loc), kind(kind), value(value) {}
+
+  ParsedAutoDiffParameter(SourceLoc loc, Kind kind, unsigned index)
+    : loc(loc), kind(kind), value(index) {}
+
+  static ParsedAutoDiffParameter getNamedParameter(SourceLoc loc,
+                                                   Identifier name) {
+    return { loc, Kind::Named, name };
+  }
+
+  static ParsedAutoDiffParameter getOrderedParameter(SourceLoc loc,
+                                                     unsigned index) {
+    return { loc, Kind::Ordered, index };
+  }
+
+  static ParsedAutoDiffParameter getSelfParameter(SourceLoc loc) {
+    return { loc, Kind::Self, {} };
+  }
+
+  Identifier getName() const {
+    assert(kind == Kind::Named);
+    return value.Named.name;
+  }
+
+  unsigned getIndex() const {
+    return value.Ordered.index;
+  }
+
+  Kind getKind() const {
+    return kind;
+  }
+
+  SourceLoc getLoc() const {
+    return loc;
+  }
+
+  bool isEqual(const ParsedAutoDiffParameter &other) const {
+    if (getKind() != other.getKind())
+      return false;
+    if (getKind() == Kind::Named)
+      return getName() == other.getName();
+    return getKind() == Kind::Self;
+  }
+};
+
+} // end namespace swift
+
+#endif // SWIFT_AST_AUTODIFF_H

include/swift/AST/DiagnosticsParse.def

@@ -1511,11 +1511,29 @@ ERROR(attr_implements_expected_member_name,PointsToFirstBadToken,
       "expected a member name as second parameter in '_implements' attribute", ())
 
 // differentiable
+ERROR(attr_differentiable_expected_function_name,PointsToFirstBadToken,
+      "expected a %0 function name", (StringRef))
+ERROR(attr_differentiable_expected_parameter_list,PointsToFirstBadToken,
+      "expected a list of parameters to differentiate with respect to", ())
+ERROR(attr_differentiable_use_wrt_not_withrespectto,none,
+      "use 'wrt:' to specify parameters to differentiate with respect to", ())
+ERROR(attr_differentiable_missing_label,PointsToFirstBadToken,
+      "missing label '%0:' in '@differentiable' attribute", (StringRef))
+ERROR(attr_differentiable_expected_label,none,
+      "expected either 'wrt:' or a function specifier label, e.g. 'jvp:', "
+      "or 'vjp:'", ())
 ERROR(differentiable_attribute_expected_rparen,none,
       "expected ')' in '@differentiable' attribute", ())
 ERROR(unexpected_argument_differentiable,none,
       "unexpected argument '%0' in '@differentiable' attribute", (StringRef))
 
+// differentiation `wrt` parameters clause
+ERROR(expected_colon_after_label,PointsToFirstBadToken,
+      "expected a colon ':' after '%0'", (StringRef))
+ERROR(diff_params_clause_expected_parameter,PointsToFirstBadToken,
+      "expected a parameter, which can be a function parameter name, "
+      "parameter index, or 'self'", ())
+
 //------------------------------------------------------------------------------
 // MARK: Generics parsing diagnostics
 //------------------------------------------------------------------------------

include/swift/Parse/Parser.h

@@ -711,6 +711,20 @@ class Parser {
   /// Check whether the current token starts with '>'.
   bool startsWithGreater(Token Tok) { return startsWithSymbol(Tok, '>'); }
 
+  /// Returns true if token is an identifier with the given value.
+  bool isIdentifier(Token Tok, StringRef value) {
+    return Tok.is(tok::identifier) && Tok.getText() == value;
+  }
+
+  /// Returns true if token is the identifier "wrt".
+  bool isWRTIdentifier(Token tok) { return isIdentifier(Tok, "wrt"); }
+
+  /// Returns true if token is the identifier "jvp".
+  bool isJVPIdentifier(Token Tok) { return isIdentifier(Tok, "jvp"); }
+
+  /// Returns true if token is the identifier "vjp".
+  bool isVJPIdentifier(Token Tok) { return isIdentifier(Tok, "vjp"); }
+
   /// Consume the starting '<' of the current token, which may either
   /// be a complete '<' token or some kind of operator token starting with '<',
   /// e.g., '<>'.
@@ -796,6 +810,12 @@ class Parser {
     return parseAnyIdentifier(Result, L, Diagnostic(ID, Args...));
   }
 
+  /// \brief Parse an unsigned integer and returns it in \p Result. On failure
+  /// emit the specified error diagnostic, and a note at the specified note
+  /// location.
+  bool parseUnsignedInteger(unsigned &Result, SourceLoc &Loc,
+                            const Diagnostic &D);
+
   /// The parser expects that \p K is next token in the input.  If so,
   /// it is consumed and false is returned.
   ///
@@ -973,6 +993,20 @@ class Parser {
   ParserResult<ImplementsAttr> parseImplementsAttribute(SourceLoc AtLoc,
                                                         SourceLoc Loc);
 
+  /// Parse the @differentiable attribute.
+  ParserResult<DifferentiableAttr> parseDifferentiableAttribute(SourceLoc AtLoc,
+                                                                SourceLoc Loc);
+
+  /// Parse the arguments inside the @differentiable attribute.
+  bool parseDifferentiableAttributeArguments(
+      bool &linear, SmallVectorImpl<ParsedAutoDiffParameter> &params,
+      Optional<DeclNameWithLoc> &jvpSpec, Optional<DeclNameWithLoc> &vjpSpec,
+      TrailingWhereClause *&whereClause);
+
+  /// Parse a differentiation parameters clause.
+  bool parseDifferentiationParametersClause(
+      SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName);
+
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc);