[AutoDiff] Enables conditional differentiability on protocol requirements. (#25974)

This PR enables the functionality shown in this example:

```swift
public protocol Distribution {
  associatedtype Value
  func logProbability(of value: Value) -> Float
}

public protocol DifferentiableDistribution: Differentiable, Distribution {
  @differentiable(wrt: self)
  func logProbability(of value: Value) -> Float
}

struct Foo: DifferentiableDistribution {
  @differentiable(wrt: self)
  func logProbability(of value: Float) -> Float {
    .zero
  }
}

@differentiable
func blah<T: DifferentiableDistribution>(_ x: T) -> Float where T.Value: AdditiveArithmetic {
  x.logProbability(of: .zero)
}
```

The fix is based on the fact that the Swift compiler does not add entries to the witness tables of protocols for overridden functions, to avoid redundancy. However, the `@differentiable` attribute being added should not be interpreted as an override as it adds new functionality, and should result in entries being added to the witness table. This PR adds this check to the override checking code, thus enabling support for the aforementioned feature.

Author: eaplatanios

Diff for: include/swift/AST/DiagnosticsSema.def

@@ -2784,6 +2784,8 @@ ERROR(differentiable_attr_invalid_access,none,
 ERROR(differentiable_attr_result_not_differentiable,none,
       "can only differentiate functions with results that conform to "
       "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))
+ERROR(differentiable_attr_protocol_where_clause,none,
+      "'where' clauses cannot be used in a '@differentiable' attribute on a protocol requirement", ())
 ERROR(differentiable_attr_empty_where_clause,none,
       "empty 'where' clause in '@differentiable' attribute", ())
 ERROR(differentiable_attr_nongeneric_trailing_where,none,

Diff for: lib/AST/Attr.cpp

@@ -499,7 +499,16 @@ static void printDifferentiableAttrArguments(
     stream << "vjp: " << vjp->Name;
   }
   // Print 'where' clause, if any.
-  if (!attr->getRequirements().empty()) {
+  // First, filter out requirements satisfied by the original function's
+  // generic signature. They should not be printed.
+  auto requirementsToPrint =
+      makeFilterRange(attr->getRequirements(), [&](Requirement req) {
+        if (auto *originalGenSig = original->getGenericSignature())
+          if (originalGenSig->isRequirementSatisfied(req))
+            return false;
+        return true;
+      });
+  if (!requirementsToPrint.empty()) {
     if (!isLeadingClause)
       stream << ' ';
     stream << "where ";
@@ -515,15 +524,6 @@ static void printDifferentiableAttrArguments(
         return genericEnv->getSugaredType(Ty);
       };
     }
-    // Filter out requirements satisfied by original function's generic
-    // signature. They should not be printed.
-    auto requirementsToPrint =
-        makeFilterRange(attr->getRequirements(), [&](Requirement req) {
-          if (auto *originalGenSig = original->getGenericSignature())
-            if (originalGenSig->isRequirementSatisfied(req))
-              return false;
-          return true;
-        });
     interleave(requirementsToPrint, [&](Requirement req) {
       if (auto *originalGenSig = original->getGenericSignature())
         if (originalGenSig->isRequirementSatisfied(req))

Diff for: lib/Sema/TypeCheckAttr.cpp

@@ -3276,6 +3276,15 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   GenericSignature *whereClauseGenSig = nullptr;
   GenericEnvironment *whereClauseGenEnv = nullptr;
   if (auto whereClause = attr->getWhereClause()) {
+    // 'where' clauses in '@differentiable' attributes of protocol
+    // requirements are not supported.
+    if (isa<ProtocolDecl>(original->getDeclContext()) &&
+        original->isProtocolRequirement()) {
+      TC.diagnose(attr->getLocation(),
+                  diag::differentiable_attr_protocol_where_clause);
+      attr->setInvalid();
+      return;
+    }
     if (whereClause->getRequirements().empty()) {
       // Where clause must not be empty.
       TC.diagnose(attr->getLocation(),

Diff for: lib/Sema/TypeCheckDeclOverride.cpp

@@ -591,6 +591,92 @@ static bool parameterTypesMatch(const ValueDecl *derivedDecl,
   return true;
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+static bool overridesDifferentiableAttribute(ValueDecl *derivedDecl,
+                                             ValueDecl *baseDecl) {
+  ASTContext &ctx = derivedDecl->getASTContext();
+  auto &diags = ctx.Diags;
+
+  auto *derivedAFD = dyn_cast<AbstractFunctionDecl>(derivedDecl);
+  auto *baseAFD = dyn_cast<AbstractFunctionDecl>(baseDecl);
+
+  if (!derivedAFD || !baseAFD)
+    return false;
+
+  auto derivedDAs = derivedAFD->getAttrs().getAttributes<DifferentiableAttr>();
+  auto baseDAs = baseAFD->getAttrs().getAttributes<DifferentiableAttr>();
+
+  // Make sure all the differentiable attributes in `baseDecl` are
+  // also declared in `derivedDecl`.
+  for (auto baseDA : baseDAs) {
+    auto baseParameters = baseDA->getParameterIndices();
+    auto defined = false;
+    for (auto derivedDA : derivedDAs) {
+      auto derivedParameters = derivedDA->getParameterIndices();
+      if (derivedParameters &&
+          baseParameters &&
+          AutoDiffIndexSubset::get(
+              ctx, baseParameters->parameters)
+              ->isSubsetOf(AutoDiffIndexSubset::get(
+                  ctx, derivedParameters->parameters))) {
+        defined = true;
+        break;
+      }
+    }
+    if (!defined) {
+      // Omit printing wrt clause if attribute differentiation parameters match
+      // inferred differentiation parameters.
+      auto *inferredParameters = TypeChecker::inferDifferentiableParameters(
+          derivedAFD, nullptr);
+      bool omitWrtClause = !baseParameters ||
+          baseParameters->parameters.count() ==
+          inferredParameters->parameters.count();
+      // Get `@differentiable` attribute description.
+      std::string baseDAString;
+      llvm::raw_string_ostream stream(baseDAString);
+      baseDA->print(stream, derivedDecl, omitWrtClause);
+      diags.diagnose(
+          derivedDecl,
+          diag::protocol_witness_missing_differentiable_attr,
+          StringRef(stream.str()).trim());
+      return false;
+    }
+  }
+
+  // If there is no differentiable attribute in `derivedDecl`, then
+  // overriding is not allowed.
+  if (derivedDAs.empty())
+    return false;
+
+  // Finally, go through all differentiable attributes in
+  // `derivedDecl` and check if they subsume any of the
+  // differentiable attributes in `baseDecl`.
+  for (auto derivedDA : derivedDAs) {
+    auto derivedParameters = derivedDA->getParameterIndices();
+    auto overrides = true;
+    for (auto baseDA : baseDAs) {
+      auto baseParameters = baseDA->getParameterIndices();
+      // If the differentiable indices of `derivedDA` are a
+      // subset of those of `baseDA`, then `baseDA` subsumes
+      // `derivedDA` and the function is marked as overridden.
+      if (derivedParameters &&
+            baseParameters &&
+            AutoDiffIndexSubset::get(
+                ctx, derivedParameters->parameters)
+              ->isSubsetOf(AutoDiffIndexSubset::get(
+                  ctx, baseParameters->parameters))) {
+        overrides = false;
+        break;
+      }
+    }
+    if (overrides)
+      return true;
+  }
+
+  return false;
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 /// Returns true if the given declaration is for the `NSObject.hashValue`
 /// property.
 static bool isNSObjectHashValue(ValueDecl *baseDecl) {
@@ -746,6 +832,12 @@ SmallVector<OverrideMatch, 2> OverrideMatcher::match(
     if (!areOverrideCompatibleSimple(decl, parentDecl))
       continue;
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // Check whether the differentiable attribute allows overriding.
+    if (overridesDifferentiableAttribute(decl, parentDecl))
+      continue;
+    // SWIFT_ENABLE_TENSORFLOW END
+
     auto parentMethod = dyn_cast<AbstractFunctionDecl>(parentDecl);
     auto parentStorage = dyn_cast<AbstractStorageDecl>(parentDecl);
     assert(parentMethod || parentStorage);

Diff for: lib/Sema/TypeCheckProtocol.cpp

@@ -569,7 +569,11 @@ swift::matchWitness(
   // SWIFT_ENABLE_TENSORFLOW
   auto result = finalize(anyRenaming, optionalAdjustments);
   if (result.isViable()) {
-    // '@differentiable' attributes must match completely.
+    // '@differentiable' attributes must match completely. If there exists a
+    // '@differentiable' attribute with a superset of the "wrt" parameters of
+    // a requirement, then an '@differentiable' attribute is added
+    // automatically.
+    ASTContext &ctx = witness->getASTContext();
     for (auto *reqDiffAttr : reqAttrs.getAttributes<DifferentiableAttr>()) {
       auto witnessDiffAttrs = witnessAttrs
           .getAttributes<DifferentiableAttr, /*AllowInvalid*/ true>();
@@ -579,14 +583,44 @@ swift::matchWitness(
                    reqDiffAttr->getParameterIndices() &&
                    witnessDiffAttr->parametersMatch(*reqDiffAttr);
           });
+      bool reqDiffAttrSupersetMatch = llvm::any_of(
+          witnessDiffAttrs, [&](const DifferentiableAttr *witnessDiffAttr) {
+            return witnessDiffAttr->getParameterIndices() &&
+                   reqDiffAttr->getParameterIndices() &&
+                   AutoDiffIndexSubset::get(
+                       ctx, witnessDiffAttr->getParameterIndices()->parameters)
+                     ->isSupersetOf(AutoDiffIndexSubset::get(
+                         ctx, reqDiffAttr->getParameterIndices()->parameters));
+          });
       if (!reqDiffAttrMatch) {
-        if (auto *vdWitness = dyn_cast<VarDecl>(witness))
-          return RequirementMatch(
-              getStandinForAccessor(vdWitness, AccessorKind::Get),
-              MatchKind::DifferentiableConflict, reqDiffAttr);
-        else
-          return RequirementMatch(witness, MatchKind::DifferentiableConflict,
-                                  reqDiffAttr);
+        auto implicitDiffAttr = false;
+        if (reqDiffAttrSupersetMatch) {
+          auto *newAttr = DifferentiableAttr::create(
+              ctx, /*implicit*/ true, reqDiffAttr->AtLoc,
+              reqDiffAttr->getRange(), reqDiffAttr->isLinear(),
+              reqDiffAttr->getParameterIndices(), /*jvp*/ None,
+              /*vjp*/ None, reqDiffAttr->getRequirements());
+          auto insertion = ctx.DifferentiableAttrs.try_emplace(
+              {witness, newAttr->getParameterIndices()}, newAttr);
+          // Valid `@differentiable` attributes are uniqued by their parameter
+          // indices. Reject duplicate attributes for the same decl and parameter
+          // indices pair.
+          if (!insertion.second) {
+            newAttr->setInvalid();
+          } else {
+            witness->getAttrs().add(newAttr);
+            implicitDiffAttr = true;
+          }
+        }
+        if (!implicitDiffAttr) {
+          if (auto *vdWitness = dyn_cast<VarDecl>(witness))
+            return RequirementMatch(
+                getStandinForAccessor(vdWitness, AccessorKind::Get),
+                MatchKind::DifferentiableConflict, reqDiffAttr);
+          else
+            return RequirementMatch(witness, MatchKind::DifferentiableConflict,
+                                    reqDiffAttr);
+        }
       }
     }
   }

Diff for: stdlib/public/core/AutoDiff.swift

@@ -100,7 +100,7 @@ public extension VectorProtocol {
   }
 }
 
-/* Note: These default-implemented opreators will slow down type-checking
+/* Note: These default-implemented operators will slow down type-checking
    performance and break existing code.
 
 public extension VectorProtocol {

Diff for: test/AutoDiff/differentiable_attr_type_checking.swift

@@ -588,17 +588,6 @@ extension FloatingPoint {
   }
 }
 
-protocol MethodDiffReq {
-  @differentiable(wrt: self, vjp: vjpFoo where Self : Differentiable)
-  func foo() -> Self
-}
-
-extension MethodDiffReq where Self : Differentiable {
-  func vjpFoo(x: Self) -> (Self, (Self.TangentVector) -> Self.TangentVector) {
-    return (self, { $0 })
-  }
-}
-
 // expected-error @+1 {{'vjpNonvariadic' does not have expected type '(Float, Int32...) -> (Float, (Float.TangentVector) -> Float.TangentVector)' (aka '(Float, Int32...) -> (Float, (Float) -> Float)')}}
 @differentiable(wrt: x, vjp: vjpNonvariadic)
 func variadic(_ x: Float, indices: Int32...) -> Float {
@@ -647,10 +636,6 @@ protocol DifferentiableAttrRequirements : Differentiable {
   // expected-note @+2 {{protocol requires function 'f2'}}
   @differentiable(wrt: (self, x, y))
   func f2(_ x: Float, _ y: Float) -> Float
-
-  // expected-note @+2 {{protocol requires function 'generic'}}
-  @differentiable(where T : Differentiable)
-  func generic<T>(_ x: T) -> T
 }
 
 // expected-error @+1 {{does not conform to protocol 'DifferentiableAttrRequirements'}}
@@ -696,11 +681,6 @@ struct DiffAttrConformanceErrors : DifferentiableAttrRequirements {
   func f2(_ x: Float, _ y: Float) -> Float {
     return x + y
   }
-
-  // expected-note @+1 {{candidate is missing attribute '@differentiable(where T : Differentiable)'}}
-  func generic<T>(_ x: T) -> T {
-    return x
-  }
 }
 
 protocol NotRefiningDiffable {
@@ -865,3 +845,43 @@ func inout1(x: Float, y: inout Float) -> Void {
 func inout2(x: Float, y: inout Float) -> Float {
   let _ = x + y
 }
+
+
+// Missing `@differentiable` attribute, without printing the 'wrt' arguments.
+
+protocol DifferentiableWhereClause: Differentiable {
+  associatedtype Scalar
+
+  @differentiable(where Scalar: Differentiable) // expected-error {{'where' clauses cannot be used in a '@differentiable' attribute on a protocol requirement}}
+  func test(value: Scalar) -> Float
+}
+
+// Missing a `@differentiable` attribute.
+
+public protocol Distribution {
+  associatedtype Value
+  func logProbability(of value: Value) -> Float
+}
+
+public protocol DifferentiableDistribution: Differentiable, Distribution {
+  @differentiable(wrt: self)
+  func logProbability(of value: Value) -> Float
+}
+
+public protocol MissingDifferentiableDistribution: DifferentiableDistribution
+  where Value: Differentiable {
+  func logProbability(of value: Value) -> Float // expected-note {{candidate is missing attribute '@differentiable(wrt: self)'}}
+}
+
+// Missing `@differentiable` attribute, without printing the 'wrt' arguments.
+
+protocol Example: Differentiable {
+  associatedtype Scalar: Differentiable
+
+  @differentiable
+  func test(value: Scalar) -> Float
+}
+
+protocol MissingDifferentiableTest: Example {
+  func test(value: Scalar) -> Float // expected-note {{candidate is missing attribute '@differentiable'}}
+}

Diff for: test/AutoDiff/protocol_requirement_autodiff.swift

@@ -137,4 +137,60 @@ struct S : P {
   }
 }
 
+// MARK: - Overridden protocol method adding differentiable attribute.
+
+public protocol Distribution {
+  associatedtype Value
+  func logProbability(of value: Value) -> Float
+}
+
+public protocol DifferentiableDistribution: Differentiable, Distribution {
+  @differentiable(wrt: self)
+  func logProbability(of value: Value) -> Float
+}
+
+struct Foo: DifferentiableDistribution {
+  @differentiable(wrt: self)
+  func logProbability(of value: Float) -> Float {
+    .zero
+  }
+}
+
+@differentiable
+func blah<T: DifferentiableDistribution>(_ x: T) -> Float where T.Value: AdditiveArithmetic {
+  x.logProbability(of: .zero)
+}
+
+// Adding a more general `@differentiable` attribute.
+public protocol DoubleDifferentiableDistribution: DifferentiableDistribution
+  where Value: Differentiable {
+  @differentiable(wrt: self)
+  @differentiable(wrt: (self, value))
+  func logProbability(of value: Value) -> Float
+}
+
+@differentiable
+func blah2<T: DoubleDifferentiableDistribution>(_ x: T, _ value: T.Value) -> Float
+  where T.Value: AdditiveArithmetic {
+  x.logProbability(of: value)
+}
+
+protocol DifferentiableFoo {
+  associatedtype T: Differentiable
+  @differentiable(wrt: x)
+  func foo(_ x: T) -> Float
+}
+
+protocol MoreDifferentiableFoo: Differentiable, DifferentiableFoo {
+  @differentiable(wrt: (self, x))
+  func foo(_ x: T) -> Float
+}
+
+struct MoreDifferentiableFooStruct: MoreDifferentiableFoo {
+  @differentiable(wrt: (self, x))
+  func foo(_ x: Float) -> Float {
+    x
+  }
+}
+
 runAllTests()