[AutoDiff] Fix @differentiable(linear) type-checking.

For all non-`@noDerivative` parameter and result types, `@differentiable(linear)`
function types should require and imply `T: Differentiable`, `T == T.TangentVector`
requirements instead of `T: Differentiable & AdditiveArithmetic`.

Update tests.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -4094,13 +4094,16 @@ ERROR(attr_only_on_parameters_of_differentiable,none,
       "'%0' may only be used on parameters of '@differentiable' function "
       "types", (StringRef))
 // SWIFT_ENABLE_TENSORFLOW
-ERROR(autodiff_attr_argument_not_differentiable,none,
-      "argument is not differentiable, but the enclosing function type is "
-      "marked '@differentiable'; did you want to add '@noDerivative' to this "
-      "argument?", ())
-ERROR(autodiff_attr_result_not_differentiable,none,
-      "result is not differentiable, but the function type is marked "
-      "'@differentiable'", ())
+ERROR(differentiable_function_type_invalid_parameter,none,
+      "parameter type '%0' does not conform to 'Differentiable'"
+      "%select{| and satisfy '%0 == %0.TangentVector'}1, but the enclosing "
+      "function type is '@differentiable%select{|(linear)}1'; did you want to "
+      "add '@noDerivative' to this parameter?", (StringRef, bool))
+ERROR(differentiable_function_type_invalid_result,none,
+      "result type '%0' does not conform to 'Differentiable'"
+      "%select{| and satisfy '%0 == %0.TangentVector'}1, but the enclosing "
+      "function type is '@differentiable%select{|(linear)}1'",
+      (StringRef, bool))
 ERROR(attr_differentiable_no_vjp_or_jvp_when_linear,none,
       "cannot specify 'vjp:' or 'jvp:' for linear functions; use "
       "'transpose:' instead", ())

lib/AST/GenericSignatureBuilder.cpp

@@ -5073,31 +5073,47 @@ class GenericSignatureBuilder::InferRequirementsWalker : public TypeWalker {
     }
 
     // SWIFT_ENABLE_TENSORFLOW
-    // Infer `Differentiable` or `Differentiable & AdditiveArithmetic` generic
-    // constraints from `@differentiable` or `@differentiable(linear)`.
+    // Infer requirements from `@differentiable` or `@differentiable(linear)`
+    // function types.
+    // For all non-`@noDerivative` parameter and result types:
+    // - `@differentiable`: add `T: Differentiable` requirement.
+    // - `@differentiable(linear)`: add
+    //   `T: Differentiable`, `T == T.TangentVector` requirements.
     if (auto *fnTy = ty->getAs<AnyFunctionType>()) {
       if (fnTy->isDifferentiable()) {
-        auto addConstraint = [&](Type typeToConstrain, ProtocolDecl *protocol) {
-          Requirement req(RequirementKind::Conformance, typeToConstrain,
+        auto addConformanceConstraint = [&](Type type, ProtocolDecl *protocol) {
+          Requirement req(RequirementKind::Conformance, type,
                           protocol->getDeclaredType());
           Builder.addRequirement(req, source, nullptr);
         };
-        auto constrainParametersAndResult = [&](ProtocolDecl *protocol) {
+        auto addSameTypeConstraint = [&](Type firstType,
+                                         AssociatedTypeDecl *assocType) {
+          auto *protocol = assocType->getProtocol();
+          auto conf = Builder.lookupConformance(CanType(), firstType, protocol);
+          auto secondType = conf.getAssociatedType(
+              firstType, assocType->getDeclaredInterfaceType());
+          Requirement req(RequirementKind::SameType, firstType, secondType);
+          Builder.addRequirement(req, source, nullptr);
+        };
+        auto &ctx = Builder.getASTContext();
+        auto *differentiableProtocol =
+            ctx.getProtocol(KnownProtocolKind::Differentiable);
+        auto *tangentVectorAssocType =
+            differentiableProtocol->getAssociatedType(ctx.Id_TangentVector);
+        auto addRequirements = [&](Type type, bool isLinear) {
+          addConformanceConstraint(type, differentiableProtocol);
+          if (isLinear)
+            addSameTypeConstraint(type, tangentVectorAssocType);
+        };
+        auto constrainParametersAndResult = [&](bool isLinear) {
           for (auto &param : fnTy->getParams())
             if (!param.isNoDerivative())
-              addConstraint(param.getPlainType(), protocol);
-          addConstraint(fnTy->getResult(), protocol);
+              addRequirements(param.getPlainType(), isLinear);
+          addRequirements(fnTy->getResult(), isLinear);
         };
-        // Add `Differentiable` constraints.
-        constrainParametersAndResult(
-            Builder.getASTContext()
-                .getProtocol(KnownProtocolKind::Differentiable));
-        // Add `AdditiveArithmetic` constraints if the function is linear.
-        if (fnTy->getDifferentiabilityKind() == DifferentiabilityKind::Linear) {
-          constrainParametersAndResult(
-              Builder.getASTContext()
-                  .getProtocol(KnownProtocolKind::AdditiveArithmetic));
-        }
+        // Add requirements.
+        constrainParametersAndResult(fnTy->getDifferentiabilityKind() ==
+                                     DifferentiabilityKind::Linear);
       }
     }
 

lib/Sema/TypeCheckType.cpp

@@ -1863,7 +1863,10 @@ namespace {
                                  TypeResolutionOptions options);
 
     // SWIFT_ENABLE_TENSORFLOW
-    bool isDifferentiableType(Type ty);
+    /// Returns true if the given type conforms to `Differentiable` in the
+    /// module of `DC`. If `tangentVectorEqualsSelf` is true, returns true iff
+    /// the given type additionally satisfies `Self == Self.TangentVector`.
+    bool isDifferentiable(Type type, bool tangentVectorEqualsSelf = false);
   };
 } // end anonymous namespace
 
@@ -2579,12 +2582,18 @@ bool TypeResolver::resolveASTFunctionTypeParams(
     }
 
     // SWIFT_ENABLE_TENSORFLOW
+    // All non-`@noDerivative` parameters of `@differentiable` and
+    // `@differentiable(linear)` function types must be differentiable.
     if (diffKind != DifferentiabilityKind::NonDifferentiable &&
         resolution.getStage() != TypeResolutionStage::Structural) {
-      if (!noDerivative && !isDifferentiableType(ty)) {
-        diagnose(eltTypeRepr->getLoc(),
-                 diag::autodiff_attr_argument_not_differentiable)
-            .fixItInsert(eltTypeRepr->getLoc(), "@noDerivative ");
+      if (!noDerivative) {
+        bool isLinear = diffKind == DifferentiabilityKind::Linear;
+        if (!isDifferentiable(ty, /*tangentVectorEqualsSelf*/ isLinear)) {
+          diagnose(eltTypeRepr->getLoc(),
+                   diag::differentiable_function_type_invalid_parameter,
+                   ty->getString(), isLinear)
+              .fixItInsert(eltTypeRepr->getLoc(), "@noDerivative ");
+        }
       }
     }
     // SWIFT_ENABLE_TENSORFLOW END
@@ -2717,30 +2726,38 @@ Type TypeResolver::resolveASTFunctionType(
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  // If the function is marked as `@differentiable`, the result must be a
-  // differentiable type.
+  // `@differentiable` and `@differentiable(linear)` function types must return
+  // a differentiable type.
   if (extInfo.isDifferentiable() &&
       resolution.getStage() != TypeResolutionStage::Structural) {
-    if (!isDifferentiableType(outputTy)) {
+    bool isLinear = diffKind == DifferentiabilityKind::Linear;
+    if (!isDifferentiable(outputTy, /*tangentVectorEqualsSelf*/ isLinear)) {
       diagnose(repr->getResultTypeRepr()->getLoc(),
-               diag::autodiff_attr_result_not_differentiable)
+               diag::differentiable_function_type_invalid_result,
+               outputTy->getString(), isLinear)
           .highlight(repr->getResultTypeRepr()->getSourceRange());
     }
   }
+// SWIFT_ENABLE_TENSORFLOW END
 
   return fnTy;
 }
 
 // SWIFT_ENABLE_TENSORFLOW
-bool TypeResolver::isDifferentiableType(Type ty) {
-  if (resolution.getStage() != TypeResolutionStage::Contextual) {
-    ty = DC->mapTypeIntoContext(ty);
-  }
-  return ty
-      ->getAutoDiffAssociatedTangentSpace(
-          LookUpConformanceInModule(DC->getParentModule()))
-      .hasValue();
+bool TypeResolver::isDifferentiable(Type type, bool tangentVectorEqualsSelf) {
+  if (resolution.getStage() != TypeResolutionStage::Contextual)
+    type = DC->mapTypeIntoContext(type);
+  auto tanSpace = type->getAutoDiffAssociatedTangentSpace(
+      LookUpConformanceInModule(DC->getParentModule()));
+  if (!tanSpace)
+    return false;
+  // If no `Self == Self.TangentVector` requirement, return true.
+  if (!tangentVectorEqualsSelf)
+    return true;
+  // Otherwise, return true if `Self == Self.TangentVector`.
+  return type->getCanonicalType() == tanSpace->getCanonicalType();
 }
+// SWIFT_ENABLE_TENSORFLOW END
 
 Type TypeResolver::resolveSILBoxType(SILBoxTypeRepr *repr,
                                      TypeResolutionOptions options) {

test/AutoDiff/differentiable_attr_type_checking.swift

@@ -776,7 +776,7 @@ extension TF_521 where T: Differentiable, T == T.TangentVector {
     return (TF_521(real: real, imaginary: imaginary), { ($0.real, $0.imaginary) })
   }
 }
-// expected-error @+1 {{result is not differentiable, but the function type is marked '@differentiable'}}
+// expected-error @+1 {{result type 'TF_521<Float>' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 let _: @differentiable(Float, Float) -> TF_521<Float> = { r, i in
   TF_521(real: r, imaginary: i)
 }

test/AutoDiff/differentiable_func_type_type_checking.swift

@@ -9,9 +9,9 @@ let _: @differentiable (Float) throws -> Float
 
 struct NonDiffType { var x: Int }
 // FIXME: Properly type-check parameters and the result's differentiability
-// expected-error @+1 {{argument is not differentiable, but the enclosing function type is marked '@differentiable'}} {{25-25=@noDerivative }}
+// expected-error @+1 {{parameter type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'; did you want to add '@noDerivative' to this parameter?}} {{25-25=@noDerivative }}
 let _: @differentiable (NonDiffType) -> Float
-// expected-error @+1 {{result is not differentiable, but the function type is marked '@differentiable'}}
+// expected-error @+1 {{result type 'NonDiffType' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 let _: @differentiable (Float) -> NonDiffType
 
 let _: @differentiable(linear) (Float) -> Float
@@ -80,10 +80,10 @@ func foo<T: Differentiable, U: Differentiable>(x: T) -> U {
 
 func test1<T: Differentiable, U: Differentiable>(_: @differentiable (T) -> @differentiable (U) -> Float) {}
 func test2<T: Differentiable, U: Differentiable>(_: @differentiable (T) -> (U) -> Float) {}
-// expected-error @+2 {{result is not differentiable, but the function type is marked '@differentiable'}}
-// expected-error @+1 {{result is not differentiable, but the function type is marked '@differentiable'}}
+// expected-error @+2 {{result type 'Int' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
+// expected-error @+1 {{result type '@differentiable (U) -> Int' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 func test3<T: Differentiable, U: Differentiable>(_: @differentiable (T) -> @differentiable (U) -> Int) {}
-// expected-error @+1 {{result is not differentiable, but the function type is marked '@differentiable'}}
+// expected-error @+1 {{result type '(U) -> Int' does not conform to 'Differentiable', but the enclosing function type is '@differentiable'}}
 func test4<T: Differentiable, U: Differentiable>(_: @differentiable (T) -> (U) -> Int) {}
 
 let diffFunc: @differentiable (Float) -> Float
@@ -107,23 +107,56 @@ struct Vector<T> {
   var x, y: T
 }
 extension Vector: Equatable where T: Equatable {}
-extension Vector: Differentiable where T: Differentiable {}
 extension Vector: AdditiveArithmetic where T: AdditiveArithmetic {}
+extension Vector: Differentiable where T: Differentiable {}
 
-// expected-note @+1 {{where 'T' = 'Int'}}
+// expected-note @+1 2 {{where 'T' = 'Int'}}
 func inferredConformancesGeneric<T, U>(_: @differentiable (Vector<T>) -> Vector<U>) {}
 
-// expected-note @+1 {{where 'T' = 'Int'}}
+// expected-error @+5 {{generic signature requires types 'Vector<T>' and 'Vector<T>.TangentVector' to be the same}}
+// expected-error @+4 {{generic signature requires types 'Vector<U>' and 'Vector<U>.TangentVector' to be the same}}
+// expected-error @+3 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'; did you want to add '@noDerivative' to this parameter?}}
+// expected-error @+2 {{result type 'Vector<U>' does not conform to 'Differentiable' and satisfy 'Vector<U> == Vector<U>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
+// expected-note @+1 2 {{where 'T' = 'Int'}}
 func inferredConformancesGenericLinear<T, U>(_: @differentiable(linear) (Vector<T>) -> Vector<U>) {}
 
-func nondiffVectorFunc(x: Vector<Int>) -> Vector<Int> {}
+func nondiff(x: Vector<Int>) -> Vector<Int> {}
 // expected-error @+1 {{global function 'inferredConformancesGeneric' requires that 'Int' conform to 'Differentiable}}
-inferredConformancesGeneric(nondiffVectorFunc)
+inferredConformancesGeneric(nondiff)
 // expected-error @+1 {{global function 'inferredConformancesGenericLinear' requires that 'Int' conform to 'Differentiable}}
-inferredConformancesGenericLinear(nondiffVectorFunc)
+inferredConformancesGenericLinear(nondiff)
 
-func diffVectorFunc(x: Vector<Float>) -> Vector<Float> {}
-inferredConformancesGeneric(diffVectorFunc) // okay!
+func diff(x: Vector<Float>) -> Vector<Float> {}
+inferredConformancesGeneric(diff) // okay!
 
 func inferredConformancesGenericResult<T, U>() -> @differentiable (Vector<T>) -> Vector<U> {}
+// expected-error @+4 {{generic signature requires types 'Vector<T>' and 'Vector<T>.TangentVector' to be the same}}
+// expected-error @+3 {{generic signature requires types 'Vector<U>' and 'Vector<U>.TangentVector' to be the same}}
+// expected-error @+2 {{parameter type 'Vector<T>' does not conform to 'Differentiable' and satisfy 'Vector<T> == Vector<T>.TangentVector', but the enclosing function type is '@differentiable(linear)'; did you want to add '@noDerivative' to this parameter?}}
+// expected-error @+1 {{result type 'Vector<U>' does not conform to 'Differentiable' and satisfy 'Vector<U> == Vector<U>.TangentVector', but the enclosing function type is '@differentiable(linear)'}}
 func inferredConformancesGenericResultLinear<T, U>() -> @differentiable(linear) (Vector<T>) -> Vector<U> {}
+
+struct Linear<T> {
+  var x, y: T
+}
+extension Linear: Equatable where T: Equatable {}
+extension Linear: AdditiveArithmetic where T: AdditiveArithmetic {}
+extension Linear: Differentiable where T: Differentiable, T == T.TangentVector {
+  typealias TangentVector = Self
+}
+
+func inferredConformancesGeneric<T, U>(_: @differentiable (Linear<T>) -> Linear<U>) {}
+
+func inferredConformancesGenericLinear<T, U>(_: @differentiable(linear) (Linear<T>) -> Linear<U>) {}
+
+func nondiff(x: Linear<Int>) -> Linear<Int> {}
+// expected-error @+1 {{global function 'inferredConformancesGeneric' requires that 'Int' conform to 'Differentiable}}
+inferredConformancesGeneric(nondiff)
+// expected-error @+1 {{global function 'inferredConformancesGenericLinear' requires that 'Int' conform to 'Differentiable}}
+inferredConformancesGenericLinear(nondiff)
+
+func diff(x: Linear<Float>) -> Linear<Float> {}
+inferredConformancesGeneric(diff) // okay!
+
+func inferredConformancesGenericResult<T, U>() -> @differentiable (Linear<T>) -> Linear<U> {}
+func inferredConformancesGenericResultLinear<T, U>() -> @differentiable(linear) (Linear<T>) -> Linear<U> {}