[AutoDiff upstream] Add @differentiable declaration attribute type-checking. (swiftlang#29231)

The `@differentiable` attribute marks a function as differentiable.

Example:
```
@differentiable(wrt: x, jvp: derivativeFoo where T: Differentiable)
func id<T>(_ x: T) -> T { x }
```

The `@differentiable` attribute has an optional `wrt:` clause specifying the
parameters that are differentiated "with respect to", i.e. the differentiability
parameters. The differentiability parameters must conform to the
`Differentiable` protocol.

If the `wrt:` clause is unspecified, the differentiability parameters are
currently inferred to be all parameters that conform to `Differentiable`.

The `@differentiable` attribute also has optional `jvp:` and `vjp:` labels
for registering derivative functions. These labels are deprecated in favor of
the `@derivative` attribute and will be removed soon.

The `@differentiable` attribute also has an optional `where` clause, specifying
extra differentiability requirements for generic functions.

The `@differentiable` attribute is gated by the
`-enable-experimental-differentiable-programming` flag.

Code changes:
- Add `DifferentiableAttributeTypeCheckRequest`.
  - Currently, the request returns differentiability parameter indices, while
    also resolving `JVPFunction`, `VJPFunction`, and
    `DerivativeGenericSignature` and mutating them in-place in
    `DifferentiableAttr`. This was the simplest approach that worked without
    introducing request cycles.
- Add "is type-checked" bit to `DifferentiableAttr`.
  - Alternatively, I tried changing `DifferentiableAttributeTypeCheckRequest` to
    use `CacheKind::Cache` instead of `CacheKind::SeparatelyCached`, but it did
    not seem to work: `@differentiable` attributes in non-primary-files were
    left unchecked.

Type-checking rules (summary):
- `@differentiable` attribute must be declared on a function-like "original"
  declaration: `func`, `init`, `subscript`, `var` (computed properties only).
- Parsed differentiability parameters must be valid (if they exist).
- Parsed `where` clause must be valid (if it exists).
- Differentiability parameters must all conform to `Differentiable`.
- Original result must all conform to `Differentiable`.
- If JVP/VJP functions are specified, they must match the expected type.
  - `@differentiable(jvp:vjp:)` for derivative registration is deprecated in
    favor of `@derivative` attribute, and will be removed soon.
- Duplicate `@differentiable` attributes with the same differentiability
  parameters are invalid.
- For protocol requirements and class members with `@differentiable` attribute,
  conforming types and subclasses must have the same `@differentiable` attribute
  (or one with a superset of differentiability parameter indices) on
  implementing/overriding declarations.

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;
@@ -287,6 +288,11 @@ class ASTContext final {
   /// 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/Attr.h

@@ -1638,6 +1638,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.
@@ -1658,7 +1659,12 @@ 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.
+  ///
+  /// Note: it is necessary to use a bit instead of `nullptr` parameter indices
+  /// to represent "parameter indices not yet type-checked" because invalid
+  /// attributes have `nullptr` parameter indices but have been type-checked.
+  llvm::PointerIntPair<IndexSubset *, 1, bool> ParameterIndicesAndBit;
   /// The trailing where clause (optional).
   TrailingWhereClause *WhereClause = nullptr;
   /// The generic signature for autodiff associated functions. Resolved by the
@@ -1714,12 +1720,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/Decl.h

@@ -5724,6 +5724,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.
@@ -6048,6 +6062,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 "

include/swift/AST/TypeCheckRequests.h

@@ -375,7 +375,9 @@ struct WhereClauseOwner {
 
   /// The source of the where clause, which can be a generic parameter list
   /// or a declaration that can have a where clause.
-  llvm::PointerUnion<GenericParamList *, TrailingWhereClause *, SpecializeAttr *> source;
+  llvm::PointerUnion<GenericParamList *, TrailingWhereClause *,
+                     SpecializeAttr *, DifferentiableAttr *>
+      source;
 
   WhereClauseOwner(GenericContext *genCtx);
   WhereClauseOwner(AssociatedTypeDecl *atd);
@@ -386,6 +388,9 @@ struct WhereClauseOwner {
   WhereClauseOwner(DeclContext *dc, SpecializeAttr *attr)
       : dc(dc), source(attr) {}
 
+  WhereClauseOwner(DeclContext *dc, DifferentiableAttr *attr)
+      : dc(dc), source(attr) {}
+
   SourceLoc getLoc() const;
 
   friend hash_code hash_value(const WhereClauseOwner &owner) {
@@ -2022,6 +2027,35 @@ class PatternTypeRequest
   }
 };
 
+/// Type-checks a `@differentiable` attribute and returns the resolved parameter
+/// indices on success. On failure, emits diagnostics and returns `nullptr`.
+///
+/// Currently, this request resolves other `@differentiable` attribute
+/// components but mutates them in place:
+/// - `JVPFunction`
+/// - `VJPFunction`
+/// - `DerivativeGenericSignature`
+class DifferentiableAttributeTypeCheckRequest
+    : public SimpleRequest<DifferentiableAttributeTypeCheckRequest,
+                           IndexSubset *(DifferentiableAttr *),
+                           CacheKind::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<IndexSubset *> evaluate(Evaluator &evaluator,
+                                         DifferentiableAttr *attr) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  Optional<IndexSubset *> getCachedResult() const;
+  void cacheResult(IndexSubset *value) const;
+};
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -43,6 +43,9 @@ SWIFT_REQUEST(TypeChecker, DefaultDefinitionTypeRequest,
 SWIFT_REQUEST(TypeChecker, DefaultTypeRequest,
               Type(KnownProtocolKind, const DeclContext *), SeparatelyCached,
               NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, DifferentiableAttributeTypeCheckRequest,
+              IndexSubset *(DifferentiableAttr *),
+              SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DynamicallyReplacedDeclRequest,
               ValueDecl *(ValueDecl *),
               Cached, NoLocationInfo)

lib/AST/Attr.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/IndexSubset.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ParameterList.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/Defer.h"
@@ -1507,6 +1508,30 @@ void DifferentiableAttr::setOriginalDeclaration(Decl *originalDeclaration) {
   OriginalDeclaration = originalDeclaration;
 }
 
+bool DifferentiableAttr::hasBeenTypeChecked() const {
+  return ParameterIndicesAndBit.getInt();
+}
+
+IndexSubset *DifferentiableAttr::getParameterIndices() const {
+  assert(getOriginalDeclaration() &&
+         "Original declaration must have been resolved");
+  auto &ctx = getOriginalDeclaration()->getASTContext();
+  return evaluateOrDefault(ctx.evaluator,
+                           DifferentiableAttributeTypeCheckRequest{
+                               const_cast<DifferentiableAttr *>(this)},
+                           nullptr);
+}
+
+void DifferentiableAttr::setParameterIndices(IndexSubset *paramIndices) {
+  assert(getOriginalDeclaration() &&
+         "Original declaration must have been resolved");
+  auto &ctx = getOriginalDeclaration()->getASTContext();
+  ctx.evaluator.cacheOutput(
+      DifferentiableAttributeTypeCheckRequest{
+          const_cast<DifferentiableAttr *>(this)},
+      std::move(paramIndices));
+}
+
 void DifferentiableAttr::setJVPFunction(FuncDecl *decl) {
   JVPFunction = decl;
   if (decl && !JVP)

lib/AST/Decl.cpp

@@ -7007,6 +7007,45 @@ StringRef AbstractFunctionDecl::getInlinableBodyText(
   return extractInlinableText(getASTContext().SourceMgr, body, scratch);
 }
 
+/// A uniqued list of derivative function configurations.
+struct AbstractFunctionDecl::DerivativeFunctionConfigurationList
+    : public llvm::SetVector<AutoDiffConfig> {
+  // Necessary for `ASTContext` allocation.
+  void *operator new(
+      size_t bytes, ASTContext &ctx,
+      unsigned alignment = alignof(DerivativeFunctionConfigurationList)) {
+    return ctx.Allocate(bytes, alignment);
+  }
+};
+
+void AbstractFunctionDecl::prepareDerivativeFunctionConfigurations() {
+  if (DerivativeFunctionConfigs)
+    return;
+  auto &ctx = getASTContext();
+  DerivativeFunctionConfigs = new (ctx) DerivativeFunctionConfigurationList();
+  // Register an `ASTContext` cleanup calling the list destructor.
+  ctx.addCleanup([this]() {
+    this->DerivativeFunctionConfigs->~DerivativeFunctionConfigurationList();
+  });
+}
+
+ArrayRef<AutoDiffConfig>
+AbstractFunctionDecl::getDerivativeFunctionConfigurations() {
+  prepareDerivativeFunctionConfigurations();
+  auto &ctx = getASTContext();
+  if (ctx.getCurrentGeneration() > DerivativeFunctionConfigGeneration) {
+    // TODO(TF-1100): Upstream derivative function configuration serialization
+    // logic.
+  }
+  return DerivativeFunctionConfigs->getArrayRef();
+}
+
+void AbstractFunctionDecl::addDerivativeFunctionConfiguration(
+    AutoDiffConfig config) {
+  prepareDerivativeFunctionConfigurations();
+  DerivativeFunctionConfigs->insert(config);
+}
+
 FuncDecl *FuncDecl::createImpl(ASTContext &Context,
                                SourceLoc StaticLoc,
                                StaticSpellingKind StaticSpelling,

lib/AST/TypeCheckRequests.cpp

@@ -391,6 +391,9 @@ MutableArrayRef<RequirementRepr> WhereClauseOwner::getRequirements() const {
   } else if (const auto attr = source.dyn_cast<SpecializeAttr *>()) {
     if (auto whereClause = attr->getTrailingWhereClause())
       return whereClause->getRequirements();
+  } else if (const auto attr = source.dyn_cast<DifferentiableAttr *>()) {
+    if (auto whereClause = attr->getWhereClause())
+      return whereClause->getRequirements();
   } else if (const auto whereClause = source.get<TrailingWhereClause *>()) {
     return whereClause->getRequirements();
   }
@@ -1241,6 +1244,24 @@ void CallerSideDefaultArgExprRequest::cacheResult(Expr *expr) const {
   defaultExpr->ContextOrCallerSideExpr = expr;
 }
 
+//----------------------------------------------------------------------------//
+// DifferentiableAttributeTypeCheckRequest computation.
+//----------------------------------------------------------------------------//
+
+Optional<IndexSubset *>
+DifferentiableAttributeTypeCheckRequest::getCachedResult() const {
+  auto *attr = std::get<0>(getStorage());
+  if (attr->hasBeenTypeChecked())
+    return attr->ParameterIndicesAndBit.getPointer();
+  return None;
+}
+
+void DifferentiableAttributeTypeCheckRequest::cacheResult(
+    IndexSubset *parameterIndices) const {
+  auto *attr = std::get<0>(getStorage());
+  attr->ParameterIndicesAndBit.setPointerAndInt(parameterIndices, true);
+}
+
 //----------------------------------------------------------------------------//
 // TypeCheckSourceFileRequest computation.
 //----------------------------------------------------------------------------//