[AD] Support @differentiable where clauses with dependent member types. (#21731)

Previously, where clauses with dependent member type requirements crashed
during serialization. This was because where clause requirements were manually
computed. The solution is to use requirements computed by the type checker.

- Improve comments about where clause type-checking. There are two goals:
  - Compute the required generic signature for autodiff associated functions,
    which is formed based on the original function's generic signature and
    the attribute's where clause requirements.
  - Store resolved where clause requirements in the attribute for
    serialization.
- Add exercising tests.

This completes `@differentiable` attribute where clause type-checking.

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1316,8 +1316,7 @@ class ClangImporterSynthesizedTypeAttr : public DeclAttribute {
 ///   @differentiable(reverse, wrt: (self, .0, .1), adjoint: bar(_:_:_:seed:))
 class DifferentiableAttr final
     : public DeclAttribute,
-      private llvm::TrailingObjects<DifferentiableAttr, AutoDiffParameter,
-                                    Requirement> {
+      private llvm::TrailingObjects<DifferentiableAttr, AutoDiffParameter> {
 public:
   struct DeclNameWithLoc {
     DeclName Name;
@@ -1350,21 +1349,25 @@ class DifferentiableAttr final
   FuncDecl *VJPFunction = nullptr;
   /// Checked parameter indices, to be resolved by the type checker.
   AutoDiffParameterIndices *CheckedParameterIndices = nullptr;
-  /// The constraint clauses for generic types.
+  /// The trailing where clause, if it exists.
   TrailingWhereClause *WhereClause = nullptr;
-  /// The number of requirements in the trailing where clause, to be resolved
-  /// by the type checker.
-  unsigned NumRequirements = 0;
-
-  explicit DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
+  /// The requirements for autodiff associated functions. Resolved by the type
+  /// checker based on the original function's generic signature and the
+  /// attribute's where clause requirements. This is set only if the attribute's
+  /// where clause exists.
+  MutableArrayRef<Requirement> Requirements;
+
+  explicit DifferentiableAttr(ASTContext &context, SourceLoc atLoc,
+                              SourceRange baseRange,
                               ArrayRef<AutoDiffParameter> parameters,
                               Optional<DeclNameWithLoc> primal,
                               Optional<DeclNameWithLoc> adjoint,
                               Optional<DeclNameWithLoc> jvp,
                               Optional<DeclNameWithLoc> vjp,
                               TrailingWhereClause *clause);
 
-  explicit DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
+  explicit DifferentiableAttr(ASTContext &context, SourceLoc atLoc,
+                              SourceRange baseRange,
                               ArrayRef<AutoDiffParameter> parameters,
                               Optional<DeclNameWithLoc> primal,
                               Optional<DeclNameWithLoc> adjoint,
@@ -1417,16 +1420,9 @@ class DifferentiableAttr final
 
   TrailingWhereClause *getWhereClause() const { return WhereClause; }
 
-  ArrayRef<Requirement> getRequirements() const {
-    return { getTrailingObjects<Requirement>(), NumRequirements };
-  }
-  MutableArrayRef<Requirement> getRequirements() {
-    return { getTrailingObjects<Requirement>(), NumRequirements };
-  }
-  void setRequirements(ASTContext &ctx, ArrayRef<Requirement> requirements);
-  size_t numTrailingObjects(OverloadToken<Requirement>) const {
-    return NumRequirements;
-  }
+  ArrayRef<Requirement> getRequirements() const { return Requirements; }
+  MutableArrayRef<Requirement> getRequirements() { return Requirements; }
+  void setRequirements(ASTContext &context, ArrayRef<Requirement> requirements);
 
   FuncDecl *getPrimalFunction() const { return PrimalFunction; }
   void setPrimalFunction(FuncDecl *decl) { PrimalFunction = decl; }

lib/AST/Attr.cpp

@@ -1044,7 +1044,8 @@ SpecializeAttr *SpecializeAttr::create(ASTContext &Ctx, SourceLoc atLoc,
 }
 
 // SWIFT_ENABLE_TENSORFLOW
-DifferentiableAttr::DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
+DifferentiableAttr::DifferentiableAttr(ASTContext &context, SourceLoc atLoc,
+                                       SourceRange baseRange,
                                        ArrayRef<AutoDiffParameter> parameters,
                                        Optional<DeclNameWithLoc> primal,
                                        Optional<DeclNameWithLoc> adjoint,
@@ -1054,13 +1055,13 @@ DifferentiableAttr::DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
   : DeclAttribute(DAK_Differentiable, atLoc, baseRange, /*Implicit*/false),
     NumParameters(parameters.size()),
     Primal(std::move(primal)), Adjoint(std::move(adjoint)),
-    JVP(std::move(jvp)), VJP(std::move(vjp)), WhereClause(clause),
-    NumRequirements(clause ? clause->getRequirements().size() : 0) {
+    JVP(std::move(jvp)), VJP(std::move(vjp)), WhereClause(clause) {
   std::copy(parameters.begin(), parameters.end(),
             getTrailingObjects<AutoDiffParameter>());
 }
 
-DifferentiableAttr::DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
+DifferentiableAttr::DifferentiableAttr(ASTContext &context, SourceLoc atLoc,
+                                       SourceRange baseRange,
                                        ArrayRef<AutoDiffParameter> parameters,
                                        Optional<DeclNameWithLoc> primal,
                                        Optional<DeclNameWithLoc> adjoint,
@@ -1070,12 +1071,10 @@ DifferentiableAttr::DifferentiableAttr(SourceLoc atLoc, SourceRange baseRange,
   : DeclAttribute(DAK_Differentiable, atLoc, baseRange, /*Implicit*/false),
     NumParameters(parameters.size()),
     Primal(std::move(primal)), Adjoint(std::move(adjoint)),
-    JVP(std::move(jvp)), VJP(std::move(vjp)),
-    NumRequirements(requirements.size()) {
+    JVP(std::move(jvp)), VJP(std::move(vjp)) {
   std::copy(parameters.begin(), parameters.end(),
             getTrailingObjects<AutoDiffParameter>());
-  std::copy(requirements.begin(), requirements.end(),
-            getTrailingObjects<Requirement>());
+  setRequirements(context, requirements);
 }
 
 DifferentiableAttr *
@@ -1087,10 +1086,9 @@ DifferentiableAttr::create(ASTContext &context, SourceLoc atLoc,
                            Optional<DeclNameWithLoc> jvp,
                            Optional<DeclNameWithLoc> vjp,
                            TrailingWhereClause *clause) {
-  unsigned size = totalSizeToAlloc<AutoDiffParameter, Requirement>(
-      parameters.size(), clause ? clause->getRequirements().size() : 0);
+  unsigned size = totalSizeToAlloc<AutoDiffParameter>(parameters.size());
   void *mem = context.Allocate(size, alignof(DifferentiableAttr));
-  return new (mem) DifferentiableAttr(atLoc, baseRange, parameters,
+  return new (mem) DifferentiableAttr(context, atLoc, baseRange, parameters,
                                       std::move(primal), std::move(adjoint),
                                       std::move(jvp), std::move(vjp), clause);
 }
@@ -1104,22 +1102,19 @@ DifferentiableAttr::create(ASTContext &context, SourceLoc atLoc,
                            Optional<DeclNameWithLoc> jvp,
                            Optional<DeclNameWithLoc> vjp,
                            ArrayRef<Requirement> requirements) {
-  unsigned size = totalSizeToAlloc<AutoDiffParameter, Requirement>(
-      parameters.size(), requirements.size());
+  unsigned size = totalSizeToAlloc<AutoDiffParameter>(parameters.size());
   void *mem = context.Allocate(size, alignof(DifferentiableAttr));
-  return new (mem) DifferentiableAttr(atLoc, baseRange, parameters,
+  return new (mem) DifferentiableAttr(context, atLoc, baseRange, parameters,
                                       std::move(primal), std::move(adjoint),
                                       std::move(jvp), std::move(vjp),
                                       requirements);
 }
 
-void DifferentiableAttr::setRequirements(ASTContext &ctx,
+void DifferentiableAttr::setRequirements(ASTContext &context,
                                          ArrayRef<Requirement> requirements) {
-  assert(requirements.size() <= NumRequirements &&
-         "Requirements size must not exceed number of allocated requirements");
-  NumRequirements = requirements.size();
-  std::copy(requirements.begin(), requirements.end(),
-            getTrailingObjects<Requirement>());
+  Requirements =
+      context.AllocateUninitialized<Requirement>(requirements.size());
+  std::copy(requirements.begin(), requirements.end(), Requirements.data());
 }
 
 ImplementsAttr::ImplementsAttr(SourceLoc atLoc, SourceRange range,

lib/Sema/TypeCheckAttr.cpp

@@ -2289,21 +2289,27 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
     return;
   }
 
-  // Type-check 'where' clause.
+  // Handle 'where' clause, if it exists.
+  // - Resolve attribute where clause requirements and store in the attribute
+  //   for serialization.
+  // - Compute generic signature for autodiff associated functions based on
+  //   the original function's generate signature and the attribute's where
+  //   clause requirements.
   GenericSignature *whereClauseGenSig = nullptr;
   GenericEnvironment *whereClauseGenEnv = nullptr;
   if (auto whereClause = attr->getWhereClause()) {
     if (whereClause->getRequirements().empty()) {
-      // Report an empty where clause.
+      // Where clause must not be empty.
       TC.diagnose(attr->getLocation(),
                   diag::differentiable_attr_empty_where_clause);
       attr->setInvalid();
       return;
     }
 
-    auto *genericSig = original->getGenericSignature();
-    if (!genericSig) {
-      // Only generic functions can have trailing where clauses.
+    auto *originalGenSig = original->getGenericSignature();
+    if (!originalGenSig) {
+      // Attributes with where clauses can only be declared on
+      // generic functions.
       TC.diagnose(attr->getLocation(),
                   diag::differentiable_attr_nongeneric_trailing_where,
                   original->getFullName())
@@ -2312,34 +2318,16 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
       return;
     }
 
-    // Form a new generic signature.
+    // Build a new generic signature for autodiff associated functions.
     GenericSignatureBuilder builder(ctx);
-    // First, add the old generic signature.
-    builder.addGenericSignature(genericSig);
-    // Go over the set of requirements, adding them to the builder.
-    SmallVector<Requirement, 4> convertedRequirements;
-
-    // Set where clause owner.
-    // Default to using @differentiable attribute as owner.
-    // If original function is declared on a protocol, use protocol's generic
-    // parameters instead.
-    WhereClauseOwner owner(original, attr);
-    if (auto nominal = original->getDeclContext()->getSelfNominalTypeDecl()) {
-      if (auto proto = dyn_cast<ProtocolDecl>(nominal)) {
-        auto DC = original->getDeclContext();
-        auto genericParams = proto->createGenericParams(DC);
-        TC.prepareGenericParamList(genericParams, DC);
-        genericParams->addTrailingWhereClause(ctx, whereClause->getWhereLoc(),
-                                              whereClause->getRequirements());
-        owner = WhereClauseOwner(original, genericParams);
-      }
-    }
+    // Add the original function's generic signature.
+    builder.addGenericSignature(originalGenSig);
 
     using FloatingRequirementSource =
-    GenericSignatureBuilder::FloatingRequirementSource;
+        GenericSignatureBuilder::FloatingRequirementSource;
 
     RequirementRequest::visitRequirements(
-      owner, TypeResolutionStage::Structural,
+      WhereClauseOwner(original, attr), TypeResolutionStage::Structural,
       [&](const Requirement &req, RequirementRepr *reqRepr) {
         // Check additional constraints.
         // TODO: refine constraints.
@@ -2355,7 +2343,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
             .highlight(reqRepr->getSourceRange());
           return false;
 
-        // Conformance requirements are supported if:
+        // Conformance requirements are valid if:
         // - The first type is a generic type parameter type.
         // - The second type is a protocol type.
         case RequirementKind::Conformance:
@@ -2373,20 +2361,20 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
           break;
         }
 
-        // Add the requirement to the generic signature builder.
+        // Add requirement to generic signature builder.
         builder.addRequirement(req, reqRepr,
                                FloatingRequirementSource::forExplicit(reqRepr),
                                nullptr, original->getModuleContext());
-        convertedRequirements.push_back(getCanonicalRequirement(req));
         return false;
       });
 
-    // Store the converted requirements in the attribute.
-    attr->setRequirements(ctx, convertedRequirements);
+    // Compute generic signature and environment for autodiff associated
+    // functions.
     whereClauseGenSig = std::move(builder).computeGenericSignature(
         attr->getLocation(), /*allowConcreteGenericParams=*/true);
     whereClauseGenEnv = whereClauseGenSig->createGenericEnvironment();
-    whereClauseGenEnv->setOwningDeclContext(original->getDeclContext());
+    // Store the resolved requirements in the attribute.
+    attr->setRequirements(ctx, whereClauseGenSig->getRequirements());
   }
 
   // Resolve the primal declaration, if it exists.

test/Serialization/differentiable_attr.swift

@@ -98,7 +98,7 @@ func vjpSimpleVJP(x: Float) -> (Float, (Float) -> Float) {
   return (x, { v in v })
 }
 
-// CHECK-DAG: @differentiable(vjp: vjpTestWhereClause where T : Differentiable)
+// CHECK-DAG: @differentiable(vjp: vjpTestWhereClause where T : Differentiable, T : Numeric)
 // CHECK-DAG: func testWhereClause<T>(x: T) -> T where T : Numeric
 @differentiable(vjp: vjpTestWhereClause where T : Differentiable)
 func testWhereClause<T : Numeric>(x: T) -> T {
@@ -112,7 +112,7 @@ func vjpTestWhereClause<T>(x: T) -> (T, (T.CotangentVector) -> T.CotangentVector
 
 protocol P {}
 extension P {
-  // CHECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClause where Self : Differentiable)
+  // CHECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClause where Self : Differentiable, Self : P)
   // CHECK-DAG: func testWhereClause() -> Self
   @differentiable(wrt: (self), vjp: vjpTestWhereClause where Self : Differentiable)
   func testWhereClause() -> Self {
@@ -125,12 +125,11 @@ extension P where Self : Differentiable {
   }
 }
 
-/*
 // NOTE: The failing tests involve where clauses with member type constraints.
 // They pass type-checking but crash during serialization.
 
-// HECK-DAG: @differentiable(vjp: vjpTestWhereClauseMemberTypeConstraint where T : Differentiable, T == T.CotangentVector)
-// HECK-DAG: func testWhereClauseMemberTypeConstraint<T : Numeric>(x: T) -> T {
+// CHECK-DAG: @differentiable(vjp: vjpTestWhereClauseMemberTypeConstraint where T : Differentiable, T : Numeric, T == T.CotangentVector)
+// CHECK-DAG: func testWhereClauseMemberTypeConstraint<T>(x: T) -> T where T : Numeric
 @differentiable(vjp: vjpTestWhereClauseMemberTypeConstraint where T : Differentiable, T == T.CotangentVector)
 func testWhereClauseMemberTypeConstraint<T : Numeric>(x: T) -> T {
   return x
@@ -141,18 +140,16 @@ func vjpTestWhereClauseMemberTypeConstraint<T>(x: T) -> (T, (T) -> T)
   return (x, { v in v })
 }
 
-protocol P {}
 extension P {
-  // HECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClauseMemberTypeConstraint where Self.CotangentVector == Self, Self : Differentiable)
-  // HECK-DAG: func testWhereClauseMemberTypeConstraint() -> Self {
+  // CHECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClauseMemberTypeConstraint where Self : Differentiable, Self : P, Self == Self.CotangentVector)
+  // CHECK-DAG: func testWhereClauseMemberTypeConstraint() -> Self
   @differentiable(wrt: (self), vjp: vjpTestWhereClauseMemberTypeConstraint where Self.CotangentVector == Self, Self : Differentiable)
   func testWhereClauseMemberTypeConstraint() -> Self {
     return self
   }
 }
-extension P where Self : Differentiable {
+extension P where Self : Differentiable, Self == Self.CotangentVector {
   func vjpTestWhereClauseMemberTypeConstraint() -> (Self, (Self.CotangentVector) -> Self.CotangentVector) {
     return (self, { v in v })
   }
 }
-*/