[AutoDiff] Type-checking support for inout parameter differentiation. (swiftlang#29959)

Semantically, an `inout` parameter is both a parameter and a result.

`@differentiable` and `@derivative` attributes now support original functions
with one "semantic result": either a formal result or an `inout` parameter.

Derivative typing rules for functions with `inout` parameters are now defined.

The differential/pullback type of a function with `inout` differentiability
parameters also has `inout` parameters. This is ideal for performance.

Differential typing rules:
- Case 1: original function has no `inout` parameters.
  - Original:     `(T0, T1, ...) -> R`
  - Differential: `(T0.Tan, T1.Tan, ...) -> R.Tan`
- Case 2: original function has a non-wrt `inout` parameter.
  - Original:     `(T0, inout T1, ...) -> Void`
  - Differential: `(T0.Tan, ...) -> T1.Tan`
- Case 3: original function has a wrt `inout` parameter.
  - Original:     `(T0, inout T1, ...) -> Void`
  - Differential: `(T0.Tan, inout T1.Tan, ...) -> Void`

Pullback typing rules:
- Case 1: original function has no `inout` parameters.
  - Original: `(T0, T1, ...) -> R`
  - Pullback: `R.Tan -> (T0.Tan, T1.Tan, ...)`
- Case 2: original function has a non-wrt `inout` parameter.
  - Original: `(T0, inout T1, ...) -> Void`
  - Pullback: `(T1.Tan) -> (T0.Tan, ...)`
- Case 3: original function has a wrt `inout` parameter.
  - Original: `(T0, inout T1, ...) -> Void`
  - Pullback: `(inout T1.Tan) -> (T0.Tan, ...)`

Resolves TF-1164.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -114,6 +114,13 @@ struct AutoDiffConfig {
   SWIFT_DEBUG_DUMP;
 };
 
+/// A semantic function result type: either a formal function result type or
+/// an `inout` parameter type. Used in derivative function type calculation.
+struct AutoDiffSemanticFunctionResultType {
+  Type type;
+  bool isInout;
+};
+
 /// Key for caching SIL derivative function types.
 struct SILAutoDiffDerivativeFunctionKey {
   SILFunctionType *originalType;
@@ -263,11 +270,17 @@ using SILDifferentiabilityWitnessKey = std::pair<StringRef, AutoDiffConfig>;
 /// 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);
+/// Given a function type, collects its semantic result types in type order
+/// into `result`: first, the formal result type (if non-`Void`), followed by
+/// `inout` parameter types.
+///
+/// The function type may have at most two parameter lists.
+///
+/// Remaps the original semantic result using `genericEnv`, if specified.
+void getFunctionSemanticResultTypes(
+    AnyFunctionType *functionType,
+    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
+    GenericEnvironment *genericEnv = nullptr);
 
 /// "Constrained" derivative generic signatures require all differentiability
 /// parameters to conform to the `Differentiable` protocol.

include/swift/AST/DiagnosticsSema.def

@@ -2898,8 +2898,6 @@ ERROR(implements_attr_protocol_not_conformed_to,none,
       (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", ())
@@ -2996,6 +2994,11 @@ 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))
+ERROR(autodiff_attr_original_void_result,none,
+      "cannot differentiate void function %0", (DeclName))
+ERROR(autodiff_attr_original_multiple_semantic_results,none,
+      "cannot differentiate functions with both an 'inout' parameter and a "
+      "result", ())
 
 // differentiation `wrt` parameters clause
 ERROR(diff_function_no_parameters,none,

include/swift/AST/Types.h

@@ -3203,15 +3203,25 @@ class AnyFunctionType : public TypeBase {
     return getExtInfo().getRepresentation();
   }
 
+  /// Appends the parameters indicated by `parameterIndices` to `results`.
+  ///
+  /// For curried function types: if `reverseCurryLevels` is true, append
+  /// the `self` parameter last instead of first.
+  ///
+  /// TODO(TF-874): Simplify logic and remove the `reverseCurryLevels` flag.
+  void getSubsetParameters(IndexSubset *parameterIndices,
+                           SmallVectorImpl<AnyFunctionType::Param> &results,
+                           bool reverseCurryLevels = false);
+
   /// Returns the derivative function type for the given parameter indices,
   /// result index, derivative function kind, derivative function generic
   /// signature (optional), and other auxiliary parameters.
   ///
   /// Preconditions:
   /// - Parameters corresponding to parameter indices must conform to
   ///   `Differentiable`.
-  /// - The result corresponding to the result index must conform to
-  ///   `Differentiable`.
+  /// - There is one semantic function result type: either the formal original
+  ///   result or an `inout` parameter. It must conform to `Differentiable`.
   ///
   /// Typing rules, given:
   /// - Original function type. Three cases:
@@ -3257,6 +3267,11 @@ class AnyFunctionType : public TypeBase {
   ///              original result | deriv. wrt result | deriv. wrt params
   /// \endverbatim
   ///
+  /// The original type may have `inout` parameters. If so, the
+  /// differential/pullback typing rules are more nuanced: see documentation for
+  /// `getAutoDiffDerivativeFunctionLinearMapType` for details. Semantically,
+  /// `inout` parameters behave as both parameters and results.
+  ///
   /// By default, if the original type has a `self` parameter list and parameter
   /// indices include `self`, the computed derivative function type will return
   /// a linear map taking/returning self's tangent *last* instead of first, for
@@ -3267,12 +3282,54 @@ class AnyFunctionType : public TypeBase {
   /// derivative function types, e.g. when type-checking `@differentiable` and
   /// `@derivative` attributes.
   AnyFunctionType *getAutoDiffDerivativeFunctionType(
-      IndexSubset *parameterIndices, unsigned resultIndex,
-      AutoDiffDerivativeFunctionKind kind,
+      IndexSubset *parameterIndices, AutoDiffDerivativeFunctionKind kind,
       LookupConformanceFn lookupConformance,
       GenericSignature derivativeGenericSignature = GenericSignature(),
       bool makeSelfParamFirst = false);
 
+  /// Returns the corresponding linear map function type for the given parameter
+  /// indices, linear map function kind, and other auxiliary parameters.
+  ///
+  /// Preconditions:
+  /// - Parameters corresponding to parameter indices must conform to
+  ///   `Differentiable`.
+  /// - There is one semantic function result type: either the formal original
+  ///   result or an `inout` parameter. It must conform to `Differentiable`.
+  ///
+  /// Differential typing rules: takes "wrt" parameter derivatives and returns a
+  /// "wrt" result derivative.
+  ///
+  /// - Case 1: original function has no `inout` parameters.
+  ///   - Original:     `(T0, T1, ...) -> R`
+  ///   - Differential: `(T0.Tan, T1.Tan, ...) -> R.Tan`
+  /// - Case 2: original function has a non-wrt `inout` parameter.
+  ///   - Original:     `(T0, inout T1, ...) -> Void`
+  ///   - Differential: `(T0.Tan, ...) -> T1.Tan`
+  /// - Case 3: original function has a wrt `inout` parameter.
+  ///   - Original:     `(T0, inout T1, ...) -> Void`
+  ///   - Differential: `(T0.Tan, inout T1.Tan, ...) -> Void`
+  ///
+  /// Pullback typing rules: takes a "wrt" result derivative and returns "wrt"
+  /// parameter derivatives.
+  ///
+  /// - Case 1: original function has no `inout` parameters.
+  ///   - Original: `(T0, T1, ...) -> R`
+  ///   - Pullback: `R.Tan -> (T0.Tan, T1.Tan, ...)`
+  /// - Case 2: original function has a non-wrt `inout` parameter.
+  ///   - Original: `(T0, inout T1, ...) -> Void`
+  ///   - Pullback: `(T1.Tan) -> (T0.Tan, ...)`
+  /// - Case 3: original function has a wrt `inout` parameter.
+  ///   - Original: `(T0, inout T1, ...) -> Void`
+  ///   - Pullback: `(inout T1.Tan) -> (T0.Tan, ...)`
+  ///
+  /// If `makeSelfParamFirst` is true, `self`'s tangent is reordered to appear
+  /// first. `makeSelfParamFirst` should be true when working with user-facing
+  /// derivative function types, e.g. when type-checking `@differentiable` and
+  /// `@derivative` attributes.
+  AnyFunctionType *getAutoDiffDerivativeFunctionLinearMapType(
+      IndexSubset *parameterIndices, AutoDiffLinearMapKind kind,
+      LookupConformanceFn lookupConformance, bool makeSelfParamFirst = false);
+
   /// True if the parameter declaration it is attached to is guaranteed
   /// to not persist the closure for longer than the duration of the call.
   bool isNoEscape() const {
@@ -4404,6 +4461,28 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
     return getParameters().back();
   }
 
+  struct IndirectMutatingParameterFilter {
+    bool operator()(SILParameterInfo param) const {
+      return param.isIndirectMutating();
+    }
+  };
+  using IndirectMutatingParameterIter =
+      llvm::filter_iterator<const SILParameterInfo *,
+                            IndirectMutatingParameterFilter>;
+  using IndirectMutatingParameterRange =
+      iterator_range<IndirectMutatingParameterIter>;
+
+  /// A range of SILParameterInfo for all indirect mutating parameters.
+  IndirectMutatingParameterRange getIndirectMutatingParameters() const {
+    return llvm::make_filter_range(getParameters(),
+                                   IndirectMutatingParameterFilter());
+  }
+
+  /// Returns the number of indirect mutating parameters.
+  unsigned getNumIndirectMutatingParameters() const {
+    return llvm::count_if(getParameters(), IndirectMutatingParameterFilter());
+  }
+
   /// Get the generic signature used to apply the substitutions of a substituted function type
   CanGenericSignature getSubstGenericSignature() const {
     return GenericSigAndIsImplied.getPointer();
@@ -4486,18 +4565,27 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   /// - Returns original results, followed by a differential function, which
   ///   takes "wrt" parameter derivatives and returns a "wrt" result derivative.
   ///
+  /// \verbatim
   ///     $(T0, ...) -> (R0, ...,  (T0.Tan, T1.Tan, ...) -> R0.Tan)
   ///                    ^~~~~~~    ^~~~~~~~~~~~~~~~~~~     ^~~~~~
   ///          original results | derivatives wrt params | derivative wrt result
+  /// \endverbatim
   ///
   /// VJP derivative type:
   /// - Takes original parameters.
   /// - Returns original results, followed by a pullback function, which
   ///   takes a "wrt" result derivative and returns "wrt" parameter derivatives.
   ///
+  /// \verbatim
   ///     $(T0, ...) -> (R0, ...,       (R0.Tan)  ->     (T0.Tan, T1.Tan, ...))
   ///                    ^~~~~~~         ^~~~~~            ^~~~~~~~~~~~~~~~~~~
   ///          original results | derivative wrt result | derivatives wrt params
+  /// \endverbatim
+  ///
+  /// The original type may have `inout` parameters. If so, the
+  /// differential/pullback typing rules are more nuanced: see documentation for
+  /// `getAutoDiffDerivativeFunctionLinearMapType` for details. Semantically,
+  /// `inout` parameters behave as both parameters and results.
   ///
   /// A "constrained derivative generic signature" is computed from
   /// `derivativeFunctionGenericSignature`, if specified. Otherwise, it is

lib/AST/AutoDiff.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/AST/AutoDiff.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
@@ -72,13 +73,11 @@ static unsigned countNumFlattenedElementTypes(Type type) {
 }
 
 // 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) {
+void AnyFunctionType::getSubsetParameters(
+    IndexSubset *parameterIndices,
+    SmallVectorImpl<AnyFunctionType::Param> &results, bool reverseCurryLevels) {
   SmallVector<AnyFunctionType *, 2> curryLevels;
-  unwrapCurryLevels(type, curryLevels);
+  unwrapCurryLevels(this, curryLevels);
 
   SmallVector<unsigned, 2> curryLevelParameterIndexOffsets(curryLevels.size());
   unsigned currentOffset = 0;
@@ -99,8 +98,43 @@ void autodiff::getSubsetParameterTypes(IndexSubset *subset,
     unsigned parameterIndexOffset =
         curryLevelParameterIndexOffsets[curryLevelIndex];
     for (unsigned paramIndex : range(curryLevel->getNumParams()))
-      if (subset->contains(parameterIndexOffset + paramIndex))
-        results.push_back(curryLevel->getParams()[paramIndex].getOldType());
+      if (parameterIndices->contains(parameterIndexOffset + paramIndex))
+        results.push_back(curryLevel->getParams()[paramIndex]);
+  }
+}
+
+void autodiff::getFunctionSemanticResultTypes(
+    AnyFunctionType *functionType,
+    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
+    GenericEnvironment *genericEnv) {
+  auto &ctx = functionType->getASTContext();
+
+  // Remap type in `genericEnv`, if specified.
+  auto remap = [&](Type type) {
+    if (!genericEnv)
+      return type;
+    return genericEnv->mapTypeIntoContext(type);
+  };
+
+  // Collect formal result type as a semantic result, unless it is
+  // `Void`.
+  auto formalResultType = functionType->getResult();
+  if (auto *resultFunctionType =
+          functionType->getResult()->getAs<AnyFunctionType>()) {
+    formalResultType = resultFunctionType->getResult();
+  }
+  if (!formalResultType->isEqual(ctx.TheEmptyTupleType))
+    result.push_back({remap(formalResultType), /*isInout*/ false});
+
+  // Collect `inout` parameters as semantic results.
+  for (auto param : functionType->getParams())
+    if (param.isInOut())
+      result.push_back({remap(param.getPlainType()), /*isInout*/ true});
+  if (auto *resultFunctionType =
+          functionType->getResult()->getAs<AnyFunctionType>()) {
+    for (auto param : resultFunctionType->getParams())
+      if (param.isInOut())
+        result.push_back({remap(param.getPlainType()), /*isInout*/ true});
   }
 }