[AutoDiff] Enable differentiation of generic functions.

include/swift/AST/DiagnosticsSIL.def

@@ -368,8 +368,9 @@ ERROR(autodiff_unsupported_type,none,
       "differentiating '%0' is not supported yet", (Type))
 ERROR(autodiff_function_not_differentiable,none,
       "function is not differentiable", ())
-NOTE(autodiff_function_generic_functions_unsupported,none,
-     "differentiating generic functions is not supported yet", ())
+NOTE(autodiff_function_indirect_params_or_result_unsupported,none,
+     "differentiating functions with parameters or result of unknown size "
+     "is not supported yet", ())
 NOTE(autodiff_external_nondifferentiable_function,none,
      "cannot differentiate an external function that has not been marked "
      "'@differentiable'", ())
@@ -386,6 +387,9 @@ NOTE(autodiff_protocol_member_subset_indices_not_differentiable,none,
 NOTE(autodiff_function_subset_indices_not_differentiable,none,
      "function is differentiable only with respect to a smaller subset of "
      "arguments", ())
+NOTE(autodiff_function_assoc_func_requirements_unmet,none,
+     "function call is not differentiate because generic requirements are not "
+     "met", ())
 NOTE(autodiff_opaque_function_not_differentiable,none,
      "opaque non-'@autodiff' function is not differentiable", ())
 NOTE(autodiff_property_not_differentiable,none,

lib/SIL/SILFunctionType.cpp

@@ -289,9 +289,6 @@ CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(
         /*genericSignature*/ nullptr, ExtInfo(), SILCoroutineKind::None,
         ParameterConvention::Direct_Guaranteed, tangentParams, {},
         tangentResults, None, ctx);
-    SmallVector<SILResultInfo, 8> jvpResults(
-        curryLevels.back()->getResults().begin(),
-        curryLevels.back()->getResults().end());
     break;
   }
   case AutoDiffAssociatedFunctionKind::VJP: {

lib/Sema/TypeCheckAttr.cpp

@@ -2188,23 +2188,6 @@ static FuncDecl *resolveAutoDiffAssociatedFunction(
   return candidate;
 }
 
-// SWIFT_ENABLE_TENSORFLOW
-/// Require that the given type either not involve type parameters or be
-/// a type parameter.
-// TODO: Generalize function to take a `Diagnostic` and merge with
-// `diagnoseIndirectGenericTypeParam`.
-static bool diagnoseDifferentiableAttrIndirectGenericType(SourceLoc loc,
-                                                          Type type,
-                                                          TypeRepr *typeRepr) {
-  if (type->hasTypeParameter() && !type->is<GenericTypeParamType>()) {
-    type->getASTContext()
-        .Diags.diagnose(loc, diag::differentiable_attr_only_generic_param_req)
-        .highlight(typeRepr->getSourceRange());
-    return true;
-  }
-  return false;
-}
-
 void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   auto &ctx = TC.Context;
   auto lookupConformance =
@@ -2303,11 +2286,10 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
     RequirementRequest::visitRequirements(
       WhereClauseOwner(original, attr), TypeResolutionStage::Structural,
       [&](const Requirement &req, RequirementRepr *reqRepr) {
-        // Check additional constraints.
-        // TODO: refine constraints.
         switch (req.getKind()) {
         case RequirementKind::SameType:
         case RequirementKind::Superclass:
+        case RequirementKind::Conformance:
           break;
 
         // Layout requirements are not supported.
@@ -2316,24 +2298,6 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
                       diag::differentiable_attr_unsupported_req_kind)
             .highlight(reqRepr->getSourceRange());
           return false;
-
-        // Conformance requirements are valid if:
-        // - The first type is a generic type parameter type.
-        // - The second type is a protocol type or protocol composition type.
-        case RequirementKind::Conformance:
-          if (diagnoseDifferentiableAttrIndirectGenericType(
-                  attr->getLocation(), req.getFirstType(),
-                  reqRepr->getSubjectRepr()))
-            return false;
-
-          if (!req.getSecondType()->is<ProtocolType>() &&
-              !req.getSecondType()->is<ProtocolCompositionType>()) {
-            TC.diagnose(attr->getLocation(),
-                     diag::differentiable_attr_non_protocol_type_constraint_req)
-              .highlight(reqRepr->getSourceRange());
-            return false;
-          }
-          break;
         }
 
         // Add requirement to generic signature builder.

test/AutoDiff/autodiff_diagnostics.swift

@@ -24,7 +24,7 @@ _ = gradient(at: 0, in: one_to_one_0) // okay!
 // Generics
 //===----------------------------------------------------------------------===//
 
-// expected-note @+3 {{differentiating generic functions is not supported yet}}
+// expected-note @+3 {{differentiating functions with parameters or result of unknown size is not supported yet}}
 // expected-error @+2 {{function is not differentiable}}
 @differentiable()
 func generic<T: Differentiable & FloatingPoint>(_ x: T) -> T {

test/AutoDiff/differentiable_attr_type_checking.swift

@@ -391,7 +391,7 @@ extension FloatingPoint {
   }
 }
 
-// expected-error @+2 {{only conformances to protocol types are supported by @differentiable attribute}}
+// expected-error @+2 {{type 'Scalar' constrained to non-protocol, non-class type 'Float'}}
 // expected-error @+1 {{can only differentiate with respect to parameters that conform to 'Differentiable', but 'Scalar' does not conform to 'Differentiable'}}
 @differentiable(where Scalar : Float)
 func invalidRequirementConformance<Scalar>(x: Scalar) -> Scalar {

test/AutoDiff/generics.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+struct Tensor<T : VectorNumeric> : VectorNumeric, Differentiable {
+  var value: Float
+  init(_ value: Float) { self.value = value }
+}
+
+func generic<T : FloatingPoint & Differentiable>(_ x: Tensor<T>) -> Float {
+  return x.value + x.value
+}
+print(pullback(at: Tensor<Float>(1), in: generic))
+print(pullback(at: Tensor<Float>(3), in: generic))
+
+// TODO: add more tests.

test/TensorFlowRuntime/tensor_autodiff_runtime.swift

@@ -22,6 +22,13 @@ TensorADTests.testAllBackends("TestSimpleGrad") {
   expectTrue(gradient(at: [[10], [20]], in: square) == [[20], [40]])
 }
 
+TensorADTests.testAllBackends("TestGenericGrad") {
+  func square<T : FloatingPoint & Differentiable>(_ x: Tensor<T>) -> Tensor<T> {
+    return x * x
+  }
+  expectEqual([0.2, 0.4, 0.6], gradient(at: Tensor([0.1, 0.2, 0.3]), in: square))
+}
+
 TensorADTests.testAllBackends("+") {
   let f = { (a: Tensor<Float>, b: Tensor<Float>) in a + b }
   expectTrue((Tensor(1), Tensor(1)) == gradient(at: Tensor(0), Tensor(0), in: f))