[AutoDiff upstream] Upstream @derivative attribute type-checking. (#28738)

The `@derivative` attribute registers a function as a derivative of another
function-like declaration: a `func`, `init`, `subscript`, or `var` computed
property declaration.

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

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

`@derivative` attribute type-checking verifies that the type of the derivative
function declaration is consistent with the type of the referenced original
declaration and the differentiation parameters.

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

Resolves TF-829.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -62,6 +62,7 @@ namespace swift {
   class Decl;
   class DeclContext;
   class DefaultArgumentInitializer;
+  class DerivativeAttr;
   class ExtensionDecl;
   class ForeignRepresentationInfo;
   class FuncDecl;
@@ -282,6 +283,16 @@ class ASTContext final {
   /// across invocations of both the parser and the type-checker.
   unsigned NextAutoClosureDiscriminator = 0;
 
+  /// Cache of `@derivative` attributes keyed by parameter indices and
+  /// derivative function kind. Used to diagnose duplicate `@derivative`
+  /// attributes for the same key.
+  // TODO(TF-1042): remove `DerivativeAttrs` from `ASTContext`. Serialize
+  // derivative function configurations per original `AbstractFunctionDecl`.
+  llvm::DenseMap<
+      std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
+      DerivativeAttr *>
+      DerivativeAttrs;
+
 private:
   /// The current generation number, which reflects the number of
   /// times that external modules have been loaded.

include/swift/AST/Attr.h

@@ -1765,7 +1765,21 @@ class DifferentiableAttr final
   }
 };
 
-/// Attribute that registers a function as a derivative of another function.
+/// The `@derivative` attribute registers a function as a derivative of another
+/// function-like declaration: a 'func', 'init', 'subscript', or 'var' computed
+/// property declaration.
+///
+/// The `@derivative` attribute also has an optional `wrt:` clause specifying
+/// the parameters that are differentiated "with respect to", i.e. the
+/// differentiation parameters. The differentiation parameters must conform to
+/// the `Differentiable` protocol.
+///
+/// If the `wrt:` clause is unspecified, the differentiation parameters are
+/// inferred to be all parameters that conform to `Differentiable`.
+///
+/// `@derivative` attribute type-checking verifies that the type of the
+/// derivative function declaration is consistent with the type of the
+/// referenced original declaration and the differentiation parameters.
 ///
 /// Examples:
 ///   @derivative(of: sin(_:))

include/swift/AST/AutoDiff.h

@@ -1,4 +1,4 @@
-//===--- AutoDiff.h - Swift Automatic Differentiation ---------------------===//
+//===--- AutoDiff.h - Swift automatic differentiation utilities -----------===//
 //
 // This source file is part of the Swift.org open source project
 //
@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-//  This file defines AST support for automatic differentiation.
+//  This file defines utilities for automatic differentiation.
 //
 //===----------------------------------------------------------------------===//
 
@@ -21,11 +21,14 @@
 
 #include "swift/AST/Identifier.h"
 #include "swift/AST/IndexSubset.h"
-#include "swift/Basic/SourceLoc.h"
+#include "swift/AST/Type.h"
 #include "swift/Basic/Range.h"
+#include "swift/Basic/SourceLoc.h"
 
 namespace swift {
 
+class AnyFunctionType;
+
 /// A function type differentiability kind.
 enum class DifferentiabilityKind : uint8_t {
   NonDifferentiable = 0,
@@ -130,6 +133,46 @@ class ParsedAutoDiffParameter {
   }
 };
 
+/// Automatic differentiation utility namespace.
+namespace autodiff {
+
+/// Appends the subset's parameter's types to `results`, in the order in
+/// which they appear in the function type.
+void getSubsetParameterTypes(IndexSubset *indices, AnyFunctionType *type,
+                             SmallVectorImpl<Type> &results,
+                             bool reverseCurryLevels = false);
+
+} // end namespace autodiff
+
 } // end namespace swift
 
+namespace llvm {
+
+using swift::AutoDiffDerivativeFunctionKind;
+
+template <typename T> struct DenseMapInfo;
+
+template <> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
+  static AutoDiffDerivativeFunctionKind getEmptyKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getEmptyKey());
+  }
+
+  static AutoDiffDerivativeFunctionKind getTombstoneKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const AutoDiffDerivativeFunctionKind &Val) {
+    return DenseMapInfo<unsigned>::getHashValue(Val);
+  }
+
+  static bool isEqual(const AutoDiffDerivativeFunctionKind &LHS,
+                      const AutoDiffDerivativeFunctionKind &RHS) {
+    return LHS == RHS;
+  }
+};
+
+} // end namespace llvm
+
 #endif // SWIFT_AST_AUTODIFF_H

include/swift/AST/DiagnosticsSema.def

@@ -2756,6 +2756,9 @@ ERROR(attr_not_on_variadic_parameters,none,
 ERROR(attr_not_on_subscript_parameters,none,
       "'%0' must not be used on subscript parameters", (StringRef))
 
+ERROR(attr_ambiguous_reference_to_decl,none,
+      "ambiguous reference to %0 in '@%1' attribute", (DeclNameRef, StringRef))
+
 ERROR(override_final,none,
       "%0 overrides a 'final' %1", (DescriptiveDeclKind, DescriptiveDeclKind))
 
@@ -2891,6 +2894,68 @@ ERROR(implements_attr_protocol_not_conformed_to,none,
       "containing type %0 does not conform to protocol %1",
       (DeclName, DeclName))
 
+// @derivative
+ERROR(derivative_attr_expected_result_tuple,none,
+      "'@derivative(of:)' attribute requires function to return a two-element "
+      "tuple of type '(value: T..., pullback: (U.TangentVector) -> T.TangentVector...)' "
+      "or '(value: T..., differential: (T.TangentVector...) -> U.TangentVector)'", ())
+ERROR(derivative_attr_invalid_result_tuple_value_label,none,
+      "'@derivative(of:)' attribute requires function to return a two-element "
+      "tuple (first element must have label 'value:')", ())
+ERROR(derivative_attr_invalid_result_tuple_func_label,none,
+      "'@derivative(of:)' attribute requires function to return a two-element "
+      "tuple (second element must have label 'pullback:' or 'differential:')",
+      ())
+ERROR(derivative_attr_result_value_not_differentiable,none,
+      "'@derivative(of:)' attribute requires function to return a two-element "
+      "tuple (first element type %0 must conform to 'Differentiable')", (Type))
+ERROR(derivative_attr_result_func_type_mismatch,none,
+      "function result's %0 type does not match %1", (Identifier, DeclName))
+NOTE(derivative_attr_result_func_type_mismatch_note,none,
+     "%0 does not have expected type %1", (Identifier, Type))
+NOTE(derivative_attr_result_func_original_note,none,
+     "%0 defined here", (DeclName))
+ERROR(derivative_attr_not_in_same_file_as_original,none,
+      "derivative not in the same file as the original function", ())
+ERROR(derivative_attr_original_stored_property_unsupported,none,
+      "cannot register derivative for stored property %0", (DeclNameRef))
+ERROR(derivative_attr_original_already_has_derivative,none,
+      "a derivative already exists for %0", (DeclName))
+NOTE(derivative_attr_duplicate_note,none,
+     "other attribute declared here", ())
+
+// Automatic differentiation attributes
+ERROR(autodiff_attr_original_decl_invalid_kind,none,
+      "%0 is not a 'func', 'init', 'subscript', or 'var' computed property "
+      "declaration", (DeclNameRef))
+ERROR(autodiff_attr_original_decl_none_valid_found,none,
+      "could not find function %0 with expected type %1", (DeclNameRef, Type))
+ERROR(autodiff_attr_original_decl_not_same_type_context,none,
+      "%0 is not defined in the current type context", (DeclNameRef))
+
+// differentiation `wrt` parameters clause
+ERROR(diff_function_no_parameters,none,
+      "%0 has no parameters to differentiate with respect to", (DeclName))
+ERROR(diff_params_clause_param_name_unknown,none,
+      "unknown parameter name %0", (Identifier))
+ERROR(diff_params_clause_self_instance_method_only,none,
+      "'self' parameter is only applicable to instance methods", ())
+ERROR(diff_params_clause_self_must_be_first,none,
+      "'self' parameter must come first in the parameter list", ())
+ERROR(diff_params_clause_params_not_original_order,none,
+      "parameters must be specified in original order", ())
+ERROR(diff_params_clause_param_index_out_of_range,none,
+      "parameter index is larger than total number of parameters", ())
+ERROR(diff_params_clause_no_inferred_parameters,PointsToFirstBadToken,
+      "no differentiation parameters could be inferred; must differentiate "
+      "with respect to at least one parameter conforming to 'Differentiable'",
+      ())
+ERROR(diff_params_clause_cannot_diff_wrt_inout_parameter,none,
+      "cannot differentiate with respect to 'inout' parameter (%0)", (Type))
+ERROR(diff_params_clause_param_not_differentiable,none,
+      "can only differentiate with respect to parameters that conform to "
+      "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
+
 //------------------------------------------------------------------------------
 // MARK: Type Check Expressions
 //------------------------------------------------------------------------------

include/swift/AST/KnownIdentifiers.def

@@ -197,6 +197,9 @@ IDENTIFIER_(nsError)
 // Custom string interpolation type used by os log APIs.
 IDENTIFIER(OSLogMessage)
 
+// Differentiable programming
+IDENTIFIER(TangentVector)
+
 #undef IDENTIFIER
 #undef IDENTIFIER_
 #undef IDENTIFIER_WITH_NAME

lib/AST/AutoDiff.cpp

@@ -0,0 +1,68 @@
+//===--- AutoDiff.cpp - Swift automatic differentiation utilities ---------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/AST/AutoDiff.h"
+#include "swift/AST/Types.h"
+
+using namespace swift;
+
+// TODO(TF-874): This helper is inefficient and should be removed. Unwrapping at
+// most once (for curried method types) is sufficient.
+static void unwrapCurryLevels(AnyFunctionType *fnTy,
+                              SmallVectorImpl<AnyFunctionType *> &results) {
+  while (fnTy != nullptr) {
+    results.push_back(fnTy);
+    fnTy = fnTy->getResult()->getAs<AnyFunctionType>();
+  }
+}
+
+static unsigned countNumFlattenedElementTypes(Type type) {
+  if (auto *tupleTy = type->getCanonicalType()->getAs<TupleType>())
+    return accumulate(tupleTy->getElementTypes(), 0,
+                      [&](unsigned num, Type type) {
+                        return num + countNumFlattenedElementTypes(type);
+                      });
+  return 1;
+}
+
+// TODO(TF-874): Simplify this helper and remove the `reverseCurryLevels` flag.
+// See TF-874 for WIP.
+void autodiff::getSubsetParameterTypes(IndexSubset *subset,
+                                       AnyFunctionType *type,
+                                       SmallVectorImpl<Type> &results,
+                                       bool reverseCurryLevels) {
+  SmallVector<AnyFunctionType *, 2> curryLevels;
+  unwrapCurryLevels(type, curryLevels);
+
+  SmallVector<unsigned, 2> curryLevelParameterIndexOffsets(curryLevels.size());
+  unsigned currentOffset = 0;
+  for (unsigned curryLevelIndex : llvm::reverse(indices(curryLevels))) {
+    curryLevelParameterIndexOffsets[curryLevelIndex] = currentOffset;
+    currentOffset += curryLevels[curryLevelIndex]->getNumParams();
+  }
+
+  // If `reverseCurryLevels` is true, reverse the curry levels and offsets.
+  if (reverseCurryLevels) {
+    std::reverse(curryLevels.begin(), curryLevels.end());
+    std::reverse(curryLevelParameterIndexOffsets.begin(),
+                 curryLevelParameterIndexOffsets.end());
+  }
+
+  for (unsigned curryLevelIndex : indices(curryLevels)) {
+    auto *curryLevel = curryLevels[curryLevelIndex];
+    unsigned parameterIndexOffset =
+        curryLevelParameterIndexOffsets[curryLevelIndex];
+    for (unsigned paramIndex : range(curryLevel->getNumParams()))
+      if (subset->contains(parameterIndexOffset + paramIndex))
+        results.push_back(curryLevel->getParams()[paramIndex].getOldType());
+  }
+}

lib/AST/CMakeLists.txt

@@ -29,7 +29,7 @@ add_swift_host_library(swiftAST STATIC
   ASTVerifier.cpp
   ASTWalker.cpp
   Attr.cpp
-  IndexSubset.cpp
+  AutoDiff.cpp
   Availability.cpp
   AvailabilitySpec.cpp
   Builtins.cpp
@@ -55,6 +55,7 @@ add_swift_host_library(swiftAST STATIC
   Identifier.cpp
   ImportCache.cpp
   IncrementalRanges.cpp
+  IndexSubset.cpp
   InlinableText.cpp
   LayoutConstraint.cpp
   Module.cpp