[DNM] [AutoDiff upstream] @differentiable attribute type-checking.

`@differentiable` attribute type-checking mega-patch.
Upstreams all necessary code from `tensorflow` branch.

This mega-patch is just for end-to-end testing.
Code will be upstreamed incrementally.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -64,6 +64,7 @@ namespace swift {
   class DeclContext;
   class DefaultArgumentInitializer;
   class DerivativeAttr;
+  class DifferentiableAttr;
   class ExtensionDecl;
   class ForeignRepresentationInfo;
   class FuncDecl;
@@ -284,6 +285,14 @@ class ASTContext final {
   /// across invocations of both the parser and the type-checker.
   unsigned NextAutoClosureDiscriminator = 0;
 
+  /// Cached mapping from types to their associated tangent spaces.
+  llvm::DenseMap<Type, Optional<TangentSpace>> AutoDiffTangentSpaces;
+
+  /// Cache of `@differentiable` attributes keyed by parameter indices. Used to
+  /// diagnose duplicate `@differentiable` attributes for the same key.
+  llvm::DenseMap<std::pair<Decl *, IndexSubset *>, DifferentiableAttr *>
+      DifferentiableAttrs;
+
   /// Cache of `@derivative` attributes keyed by parameter indices and
   /// derivative function kind. Used to diagnose duplicate `@derivative`
   /// attributes for the same key.

include/swift/AST/ASTScope.h

@@ -1568,6 +1568,41 @@ class SpecializeAttributeScope final : public ASTScopeImpl {
                              DeclConsumer) const override;
 };
 
+/// A `@differentiable` attribute scope.
+///
+/// This exists because `@differentiable` attribute may have a `where` clause
+/// referring to generic parameters from some generic context.
+class DifferentiableAttributeScope final : public ASTScopeImpl {
+public:
+  DifferentiableAttr *const differentiableAttr;
+  ValueDecl *const attributedDeclaration;
+
+  DifferentiableAttributeScope(DifferentiableAttr *diffAttr, ValueDecl *decl)
+      : differentiableAttr(diffAttr), attributedDeclaration(decl) {}
+  virtual ~DifferentiableAttributeScope() {}
+
+  std::string getClassName() const override;
+  SourceRange
+  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
+  NullablePtr<const void> addressForPrinting() const override {
+    return differentiableAttr;
+  }
+
+  NullablePtr<AbstractStorageDecl>
+  getEnclosingAbstractStorageDecl() const override;
+
+  NullablePtr<DeclAttribute> getDeclAttributeIfAny() const override {
+    return differentiableAttr;
+  }
+  NullablePtr<const void> getReferrent() const override;
+
+protected:
+  ASTScopeImpl *expandSpecifically(ScopeCreator &) override;
+  bool lookupLocalsOrMembers(ArrayRef<const ASTScopeImpl *>,
+                             DeclConsumer) const override;
+  bool doesContextMatchStartingContext(const DeclContext *) const override;
+};
+
 class SubscriptDeclScope final : public ASTScopeImpl {
 public:
   SubscriptDecl *const decl;

include/swift/AST/Attr.h

@@ -1647,6 +1647,7 @@ class DifferentiableAttr final
       private llvm::TrailingObjects<DifferentiableAttr,
                                     ParsedAutoDiffParameter> {
   friend TrailingObjects;
+  friend class DifferentiableAttributeTypeCheckRequest;
 
   /// The declaration on which the `@differentiable` attribute is declared.
   /// May not be a valid declaration for `@differentiable` attributes.
@@ -1667,7 +1668,8 @@ class DifferentiableAttr final
   /// specified.
   FuncDecl *VJPFunction = nullptr;
   /// The differentiability parameter indices, resolved by the type checker.
-  IndexSubset *ParameterIndices = nullptr;
+  /// The bit stores whether the parameter indices have been computed.
+  llvm::PointerIntPair<IndexSubset *, 1, bool> ParameterIndicesAndBit;
   /// The trailing where clause (optional).
   TrailingWhereClause *WhereClause = nullptr;
   /// The generic signature for autodiff associated functions. Resolved by the
@@ -1723,12 +1725,14 @@ class DifferentiableAttr final
   /// registered VJP.
   Optional<DeclNameRefWithLoc> getVJP() const { return VJP; }
 
-  IndexSubset *getParameterIndices() const {
-    return ParameterIndices;
-  }
-  void setParameterIndices(IndexSubset *parameterIndices) {
-    ParameterIndices = parameterIndices;
-  }
+private:
+  /// Returns true if the given `@differentiable` attribute has been
+  /// type-checked.
+  bool hasBeenTypeChecked() const;
+
+public:
+  IndexSubset *getParameterIndices() const;
+  void setParameterIndices(IndexSubset *parameterIndices);
 
   /// The parsed differentiability parameters, i.e. the list of parameters
   /// specified in 'wrt:'.

include/swift/AST/AutoDiff.h

@@ -19,6 +19,7 @@
 
 #include <cstdint>
 
+#include "swift/AST/GenericSignature.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/IndexSubset.h"
 #include "swift/AST/Type.h"
@@ -28,6 +29,7 @@
 namespace swift {
 
 class AnyFunctionType;
+class TupleType;
 
 /// A function type differentiability kind.
 enum class DifferentiabilityKind : uint8_t {
@@ -70,6 +72,26 @@ struct AutoDiffDerivativeFunctionKind {
   }
 };
 
+/// Identifies an autodiff derivative function configuration:
+/// - Parameter indices.
+/// - Result indices.
+/// - Derivative generic signature (optional).
+struct AutoDiffConfig {
+  IndexSubset *parameterIndices;
+  IndexSubset *resultIndices;
+  GenericSignature derivativeGenericSignature;
+
+  /*implicit*/ AutoDiffConfig(IndexSubset *parameterIndices,
+                              IndexSubset *resultIndices,
+                              GenericSignature derivativeGenericSignature)
+      : parameterIndices(parameterIndices), resultIndices(resultIndices),
+        derivativeGenericSignature(derivativeGenericSignature) {}
+
+  void print(llvm::raw_ostream &s = llvm::outs()) const;
+  LLVM_ATTRIBUTE_DEPRECATED(void dump() const LLVM_ATTRIBUTE_USED,
+                            "only for use within the debugger");
+};
+
 class ParsedAutoDiffParameter {
 public:
   enum class Kind { Named, Ordered, Self };
@@ -133,6 +155,91 @@ class ParsedAutoDiffParameter {
   }
 };
 
+/// The tangent space of a type.
+///
+/// For `Differentiable`-conforming types:
+/// - The tangent space is the `TangentVector` associated type.
+///
+/// For tuple types:
+/// - The tangent space is a tuple of the elements' tangent space types, for the
+///   elements that have a tangent space.
+///
+/// For function types whose innermost result type has a tangent space:
+/// - The tangent space is the same function type, replacing the innermost
+///   result type with its tangent space type.
+///
+/// Other types have no tangent space.
+class TangentSpace {
+public:
+  /// A tangent space kind.
+  enum class Kind {
+    /// The `TangentVector` associated type of a `Differentiable`-conforming
+    /// type.
+    TangentVector,
+    /// A product of tangent spaces as a tuple.
+    Tuple,
+    /// A function type whose innermost result is the `TangentVector`
+    /// associated type of a `Differentiable`-conforming type.
+    Function
+  };
+
+private:
+  Kind kind;
+  union Value {
+    // TangentVector
+    Type tangentVectorType;
+    // Tuple
+    TupleType *tupleType;
+    // Function
+    AnyFunctionType *functionType;
+
+    Value(Type tangentVectorType) : tangentVectorType(tangentVectorType) {}
+    Value(TupleType *tupleType) : tupleType(tupleType) {}
+    Value(AnyFunctionType *functionType) : functionType(functionType) {}
+  } value;
+
+  TangentSpace(Kind kind, Value value) : kind(kind), value(value) {}
+
+public:
+  TangentSpace() = delete;
+
+  static TangentSpace getTangentVector(Type tangentVectorType) {
+    return {Kind::TangentVector, tangentVectorType};
+  }
+  static TangentSpace getTuple(TupleType *tupleTy) {
+    return {Kind::Tuple, tupleTy};
+  }
+  static TangentSpace getFunction(AnyFunctionType *fnTy) {
+    return {Kind::Function, fnTy};
+  }
+
+  bool isTangentVector() const { return kind == Kind::TangentVector; }
+  bool isTuple() const { return kind == Kind::Tuple; }
+
+  Kind getKind() const { return kind; }
+  Type getTangentVector() const {
+    assert(kind == Kind::TangentVector);
+    return value.tangentVectorType;
+  }
+  TupleType *getTuple() const {
+    assert(kind == Kind::Tuple);
+    return value.tupleType;
+  }
+  AnyFunctionType *getFunction() const {
+    assert(kind == Kind::Function);
+    return value.functionType;
+  }
+
+  /// Get the tangent space type.
+  Type getType() const;
+
+  /// Get the tangent space canonical type.
+  CanType getCanonicalType() const;
+
+  /// Get the underlying
+  NominalTypeDecl *getNominal() const;
+};
+
 /// Automatic differentiation utility namespace.
 namespace autodiff {
 
@@ -148,10 +255,59 @@ void getSubsetParameterTypes(IndexSubset *indices, AnyFunctionType *type,
 
 namespace llvm {
 
+using swift::AutoDiffConfig;
 using swift::AutoDiffDerivativeFunctionKind;
+using swift::GenericSignature;
+using swift::IndexSubset;
 
 template <typename T> struct DenseMapInfo;
 
+template <> struct DenseMapInfo<AutoDiffConfig> {
+  static AutoDiffConfig getEmptyKey() {
+    auto *ptr = llvm::DenseMapInfo<void *>::getEmptyKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
+    return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
+            nullptr};
+  }
+
+  static AutoDiffConfig getTombstoneKey() {
+    auto *ptr = llvm::DenseMapInfo<void *>::getTombstoneKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
+    return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
+            nullptr};
+  }
+
+  static unsigned getHashValue(const AutoDiffConfig &Val) {
+    auto canGenSig =
+        Val.derivativeGenericSignature
+            ? Val.derivativeGenericSignature->getCanonicalSignature()
+            : nullptr;
+    unsigned combinedHash = hash_combine(
+        ~1U, DenseMapInfo<void *>::getHashValue(Val.parameterIndices),
+        DenseMapInfo<void *>::getHashValue(Val.resultIndices),
+        DenseMapInfo<GenericSignature>::getHashValue(canGenSig));
+    return combinedHash;
+  }
+
+  static bool isEqual(const AutoDiffConfig &LHS, const AutoDiffConfig &RHS) {
+    auto lhsCanGenSig =
+        LHS.derivativeGenericSignature
+            ? LHS.derivativeGenericSignature->getCanonicalSignature()
+            : nullptr;
+    auto rhsCanGenSig =
+        RHS.derivativeGenericSignature
+            ? RHS.derivativeGenericSignature->getCanonicalSignature()
+            : nullptr;
+    return LHS.parameterIndices == RHS.parameterIndices &&
+           LHS.resultIndices == RHS.resultIndices &&
+           DenseMapInfo<GenericSignature>::isEqual(lhsCanGenSig, rhsCanGenSig);
+  }
+};
+
 template <> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
   static AutoDiffDerivativeFunctionKind getEmptyKey() {
     return static_cast<AutoDiffDerivativeFunctionKind::innerty>(

include/swift/AST/Decl.h

@@ -5734,6 +5734,20 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
 private:
   ParameterList *Params;
 
+  /// The generation at which we last loaded derivative function configurations.
+  unsigned DerivativeFunctionConfigGeneration = 0;
+  /// Prepare to traverse the list of derivative function configurations.
+  void prepareDerivativeFunctionConfigurations();
+
+  /// A uniqued list of derivative function configurations.
+  /// - `@differentiable` and `@derivative` attribute type-checking is
+  ///   responsible for populating derivative function configurations specified
+  ///   in the current module.
+  /// - Module loading is responsible for populating derivative function
+  ///   configurations from imported modules.
+  struct DerivativeFunctionConfigurationList;
+  DerivativeFunctionConfigurationList *DerivativeFunctionConfigs = nullptr;
+
 protected:
   // If a function has a body at all, we have either a parsed body AST node or
   // we have saved the end location of the unparsed body.
@@ -6058,6 +6072,12 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// constructor.
   bool hasDynamicSelfResult() const;
 
+  /// Get all derivative function configurations.
+  ArrayRef<AutoDiffConfig> getDerivativeFunctionConfigurations();
+
+  /// Add the given derivative function configuration.
+  void addDerivativeFunctionConfiguration(AutoDiffConfig config);
+
   using DeclContext::operator new;
   using Decl::getASTContext;
 };

include/swift/AST/DiagnosticsSema.def

@@ -2904,6 +2904,62 @@ ERROR(implements_attr_protocol_not_conformed_to,none,
       "containing type %0 does not conform to protocol %1",
       (DeclName, DeclName))
 
+// @differentiable
+ERROR(differentiable_attr_void_result,none,
+      "cannot differentiate void function %0", (DeclName))
+ERROR(differentiable_attr_no_vjp_or_jvp_when_linear,none,
+      "cannot specify 'vjp:' or 'jvp:' for linear functions; use '@transpose' "
+      "attribute for transpose registration instead", ())
+ERROR(differentiable_attr_overload_not_found,none,
+      "%0 does not have expected type %1", (DeclNameRef, Type))
+// TODO(TF-482): Change duplicate `@differentiable` attribute diagnostic to also
+// mention "same generic requirements".
+ERROR(differentiable_attr_duplicate,none,
+      "duplicate '@differentiable' attribute with same parameters", ())
+NOTE(differentiable_attr_duplicate_note,none,
+     "other attribute declared here", ())
+ERROR(differentiable_attr_function_not_same_type_context,none,
+      "%0 is not defined in the current type context", (DeclNameRef))
+ERROR(differentiable_attr_derivative_not_function,none,
+      "registered derivative %0 must be a 'func' declaration", (DeclNameRef))
+ERROR(differentiable_attr_class_derivative_not_final,none,
+      "class member derivative must be final", ())
+ERROR(differentiable_attr_invalid_access,none,
+      "derivative function %0 is required to either be public or "
+      "'@usableFromInline' because the original function %1 is public or "
+      "'@usableFromInline'", (DeclNameRef, DeclName))
+ERROR(differentiable_attr_result_not_differentiable,none,
+      "can only differentiate functions with results that conform to "
+      "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
+ERROR(differentiable_attr_protocol_req_where_clause,none,
+      "'@differentiable' attribute on protocol requirement cannot specify "
+      "'where' clause", ())
+ERROR(differentiable_attr_protocol_req_assoc_func,none,
+      "'@differentiable' attribute on protocol requirement cannot specify "
+      "'jvp:' or 'vjp:'", ())
+ERROR(differentiable_attr_stored_property_variable_unsupported,none,
+      "'@differentiable' attribute on stored property cannot specify "
+      "'jvp:' or 'vjp:'", ())
+ERROR(differentiable_attr_class_member_no_dynamic_self,none,
+      "'@differentiable' attribute cannot be declared on class methods "
+      "returning 'Self'", ())
+// TODO(TF-654): Remove when differentiation supports class initializers.
+ERROR(differentiable_attr_class_init_not_yet_supported,none,
+      "'@differentiable' attribute does not yet support class initializers",
+      ())
+ERROR(differentiable_attr_empty_where_clause,none,
+      "empty 'where' clause in '@differentiable' attribute", ())
+ERROR(differentiable_attr_where_clause_for_nongeneric_original,none,
+      "'where' clause is valid only when original function is generic %0",
+      (DeclName))
+ERROR(differentiable_attr_layout_req_unsupported,none,
+      "'@differentiable' attribute does not yet support layout requirements",
+      ())
+ERROR(overriding_decl_missing_differentiable_attr,none,
+      "overriding declaration is missing attribute '%0'", (StringRef))
+NOTE(protocol_witness_missing_differentiable_attr,none,
+     "candidate is missing attribute '%0'", (StringRef))
+
 // @derivative
 ERROR(derivative_attr_expected_result_tuple,none,
       "'@derivative(of:)' attribute requires function to return a two-element "