[AutoDiff upstream] Add @differentiable protocol requirement tests. (#32939)

Author: dan-zheng

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -544,6 +544,23 @@ public protocol DoubleDifferentiableDistribution: DifferentiableDistribution
   func logProbability(of value: Value) -> Float
 }
 
+// Test failure to satisfy protocol requirement's `@differentiable` attribute.
+
+public protocol HasRequirement {
+  @differentiable
+  // expected-note @+1 {{protocol requires function 'requirement' with type '<T> (T, T) -> T'; do you want to add a stub?}}
+  func requirement<T: Differentiable>(_ x: T, _ y: T) -> T
+}
+
+// expected-error @+1 {{type 'AttemptsToSatisfyRequirement' does not conform to protocol 'HasRequirement'}}
+public struct AttemptsToSatisfyRequirement: HasRequirement {
+  // This `@differentiable` attribute does not satisfy the requirement because
+  // it is mroe constrained than the requirement's `@differentiable` attribute.
+  @differentiable(where T: CustomStringConvertible)
+  // expected-note @+1 {{candidate is missing attribute '@differentiable(wrt: (x, y))'}}
+  public func requirement<T: Differentiable>(_ x: T, _ y: T) -> T { x }
+}
+
 // Test protocol requirement `@differentiable` attribute unsupported features.
 
 protocol ProtocolRequirementUnsupported: Differentiable {

test/AutoDiff/validation-test/differentiable_protocol_requirements.swift

@@ -0,0 +1,203 @@
+// RUN: %target-run-simple-swift
+
+import StdlibUnittest
+import DifferentiationUnittest
+
+// Test end-to-end differentiation of `@differentiable` protocol requirements.
+
+var ProtocolRequirementAutodiffTests = TestSuite("ProtocolRequirementDifferentiation")
+
+// MARK: - Method requirements.
+
+protocol DiffReq: Differentiable {
+  @differentiable(wrt: (self, x))
+  func f(_ x: Tracked<Float>) -> Tracked<Float>
+}
+
+extension DiffReq where TangentVector: AdditiveArithmetic {
+  @inline(never)  // Prevent specialization, to test all witness code.
+  func gradF(at x: Tracked<Float>) -> (Self.TangentVector, Tracked<Float>) {
+    return (valueWithPullback(at: self, x) { s, x in s.f(x) }).1(1)
+  }
+}
+
+struct Quadratic: DiffReq, AdditiveArithmetic {
+  typealias TangentVector = Quadratic
+
+  @differentiable
+  let a: Tracked<Float>
+
+  @differentiable
+  let b: Tracked<Float>
+
+  @differentiable
+  let c: Tracked<Float>
+
+  init(_ a: Tracked<Float>, _ b: Tracked<Float>, _ c: Tracked<Float>) {
+    self.a = a
+    self.b = b
+    self.c = c
+  }
+
+  @differentiable(wrt: (self, x))
+  func f(_ x: Tracked<Float>) -> Tracked<Float> {
+    return a * x * x + b * x + c
+  }
+}
+
+ProtocolRequirementAutodiffTests.testWithLeakChecking("func") {
+  expectEqual((Quadratic(0, 0, 1), 12), Quadratic(11, 12, 13).gradF(at: 0))
+  expectEqual((Quadratic(1, 1, 1), 2 * 11 + 12),
+              Quadratic(11, 12, 13).gradF(at: 1))
+  expectEqual((Quadratic(4, 2, 1), 2 * 11 * 2 + 12),
+              Quadratic(11, 12, 13).gradF(at: 2))
+}
+
+// MARK: - Constructor, accessor, and subscript requirements.
+
+protocol FunctionsOfX: Differentiable {
+  @differentiable
+  init(x: Tracked<Float>)
+
+  @differentiable
+  var x: Tracked<Float> { get }
+
+  @differentiable
+  var y: Tracked<Float> { get }
+
+  @differentiable
+  var z: Tracked<Float> { get }
+
+  @differentiable
+  subscript() -> Tracked<Float> { get }
+}
+
+struct TestFunctionsOfX: FunctionsOfX {
+  @differentiable
+  init(x: Tracked<Float>) {
+    self.x = x
+    self.y = x * x
+  }
+
+  /// x = x
+  var x: Tracked<Float>
+
+  /// y = x * x
+  var y: Tracked<Float>
+
+  /// z = x * x + x
+  var z: Tracked<Float> {
+    return y + x
+  }
+
+  @differentiable
+  subscript() -> Tracked<Float> {
+    return z
+  }
+}
+
+@inline(never)  // Prevent specialization, to test all witness code.
+func derivatives<F: FunctionsOfX>(at x: Tracked<Float>, in: F.Type)
+  -> (Tracked<Float>, Tracked<Float>, Tracked<Float>, Tracked<Float>)
+{
+  let dxdx = gradient(at: x) { x in F(x: x).x }
+  let dydx = gradient(at: x) { x in F(x: x).y }
+  let dzdx = gradient(at: x) { x in F(x: x).z }
+  let dsubscriptdx = gradient(at: x) { x in F(x: x)[] }
+  return (dxdx, dydx, dzdx, dsubscriptdx)
+}
+
+ProtocolRequirementAutodiffTests.testWithLeakChecking("constructor, accessor, subscript") {
+  expectEqual(
+    (1.0, 4.0, 5.0, 5.0),
+    derivatives(at: 2.0, in: TestFunctionsOfX.self))
+}
+
+// MARK: - Test witness method SIL type computation.
+
+protocol P: Differentiable {
+  @differentiable(wrt: (x, y))
+  func foo(_ x: Tracked<Float>, _ y: Double) -> Tracked<Float>
+}
+struct S: P {
+  @differentiable(wrt: (x, y))
+  func foo(_ x: Tracked<Float>, _ y: Double) -> Tracked<Float> {
+    return x
+  }
+}
+
+// MARK: - Overridding protocol method adding `@differentiable` attribute.
+
+public protocol Distribution {
+  associatedtype Value
+  func logProbability(of value: Value) -> Tracked<Float>
+}
+
+public protocol DifferentiableDistribution: Differentiable, Distribution {
+  @differentiable(wrt: self)
+  func logProbability(of value: Value) -> Tracked<Float>
+}
+
+struct Foo: DifferentiableDistribution {
+  @differentiable(wrt: self)
+  func logProbability(of value: Tracked<Float>) -> Tracked<Float> {
+    .zero
+  }
+}
+
+@differentiable
+func blah<T: DifferentiableDistribution>(_ x: T) -> Tracked<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) -> Tracked<Float>
+}
+
+@differentiable
+func blah2<T: DoubleDifferentiableDistribution>(_ x: T, _ value: T.Value) -> Tracked<Float>
+where T.Value: AdditiveArithmetic {
+  x.logProbability(of: value)
+}
+
+// Satisfying the requirement with more wrt parameter indices than are necessary.
+
+protocol DifferentiableFoo {
+  associatedtype T: Differentiable
+  @differentiable(wrt: x)
+  func foo(_ x: T) -> Tracked<Float>
+}
+
+protocol MoreDifferentiableFoo: Differentiable, DifferentiableFoo {
+  @differentiable(wrt: (self, x))
+  func foo(_ x: T) -> Tracked<Float>
+}
+
+struct MoreDifferentiableFooStruct: MoreDifferentiableFoo {
+  @differentiable(wrt: (self, x))
+  func foo(_ x: Tracked<Float>) -> Tracked<Float> {
+    x
+  }
+}
+
+// Satisfying the requirement with a less-constrained derivative than is necessary.
+
+protocol ExtraDerivativeConstraint {}
+
+protocol HasExtraConstrainedDerivative {
+  @differentiable
+  func requirement<T: Differentiable & ExtraDerivativeConstraint>(_ x: T) -> T
+}
+
+struct SatisfiesDerivativeWithLessConstraint: HasExtraConstrainedDerivative {
+  @differentiable
+  func requirement<T: Differentiable>(_ x: T) -> T { x }
+}
+
+runAllTests()