[AutoDiff upstream] Enable @differentiable on setters. (swiftlang#35133)

Enable `@differentiable` attribute on setters of properties and
subscripts in `Differentiable`-conforming types.

Add automatically-differentiated `@differentiable` setter test.

Resolves TF-1166.

Author: dan-zheng

lib/Sema/TypeCheckAttr.cpp

@@ -4325,11 +4325,10 @@ resolveDifferentiableAttrOriginalFunction(DifferentiableAttr *attr) {
     }
   }
   // Non-`get` accessors are not yet supported: `set`, `read`, and `modify`.
-  // TODO(TF-129): Enable `set` when differentiation supports inout parameters.
   // TODO(TF-1080): Enable `read` and `modify` when differentiation supports
   // coroutines.
   if (auto *accessor = dyn_cast_or_null<AccessorDecl>(original))
-    if (!accessor->isGetter())
+    if (!accessor->isGetter() && !accessor->isSetter())
       original = nullptr;
   // Diagnose if original `AbstractFunctionDecl` could not be resolved.
   if (!original) {

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -171,15 +171,13 @@ struct SubscriptMethod: Differentiable {
   subscript(explicit x: Float) -> Float {
     @differentiable // ok
     get { return x }
-    // expected-error @+1 {{'@differentiable' attribute cannot be applied to this declaration}}
     @differentiable
     set {}
   }
 
   subscript(x: Float, y: Float) -> Float {
     @differentiable // ok
     get { return x + y }
-    // expected-error @+1 {{'@differentiable' attribute cannot be applied to this declaration}}
     @differentiable
     set {}
   }
@@ -700,7 +698,6 @@ struct Accessors: Differentiable {
 
   var stored: Float
   var computed: Float {
-    // expected-error @+1 {{'@differentiable' attribute cannot be applied to this declaration}}
     @differentiable
     set { stored = newValue }
 

test/AutoDiff/TBD/derivative_symbols.swift

@@ -22,10 +22,12 @@ public func topLevelDerivative<T: Differentiable>(_ x: T) -> (
 public struct Struct: Differentiable {
   var stored: Float
 
-  // Test property.
-  @differentiable
+  // Test property: getter and setter.
   public var property: Float {
-    stored
+    @differentiable
+    get { stored }
+    @differentiable
+    set { stored = newValue }
   }
 
   // Test initializer.
@@ -35,24 +37,37 @@ public struct Struct: Differentiable {
   }
 
   // Test method.
-  public func method(x: Float, y: Float) -> Float { x }
+  public func method(_ x: Float, _ y: Float) -> Float { x }
 
   @derivative(of: method)
-  public func jvpMethod(x: Float, y: Float) -> (
+  public func jvpMethod(_ x: Float, _ y: Float) -> (
     value: Float, differential: (TangentVector, Float, Float) -> Float
   ) {
     fatalError()
   }
 
-  // Test subscript.
-  public subscript(x: Float) -> Float { x }
+  // Test subscript: getter and setter.
+  public subscript(_ x: Float) -> Float {
+    @differentiable
+    get { x }
+
+    @differentiable
+    set { stored = newValue }
+  }
 
   @derivative(of: subscript)
-  public func vjpSubscript(x: Float) -> (
+  public func vjpSubscript(_ x: Float) -> (
     value: Float, pullback: (Float) -> (TangentVector, Float)
   ) {
     fatalError()
   }
+
+  @derivative(of: subscript.set)
+  public mutating func vjpSubscriptSetter(_ x: Float, _ newValue: Float) -> (
+    value: (), pullback: (inout TangentVector) -> (Float, Float)
+  ) {
+    fatalError()
+  }
 }
 
 extension Array where Element == Struct {

test/AutoDiff/validation-test/derivative_registration.swift

@@ -94,8 +94,11 @@ DerivativeRegistrationTests.testWithLeakChecking("InstanceMethod") {
 
 extension Wrapper {
   subscript(_ x: Tracked<Float>) -> Tracked<Float> {
+    @differentiable
     @_semantics("autodiff.opaque")
     get { float * x }
+
+    @differentiable
     set {}
   }
 
@@ -117,7 +120,10 @@ DerivativeRegistrationTests.testWithLeakChecking("SubscriptGetter") {
 
 extension Wrapper {
   subscript() -> Tracked<Float> {
+    @differentiable
     get { float }
+
+    @differentiable
     set { float = newValue }
   }
 

test/AutoDiff/validation-test/inout_parameters.swift

@@ -117,6 +117,15 @@ InoutParameterAutoDiffTests.test("SetAccessor") {
       get { x }
       set { x = newValue }
     }
+
+    // Computed property with explicit `@differentiable` accessors.
+    var doubled: Float {
+      @differentiable
+      get { x + x }
+
+      @differentiable
+      set { x = newValue / 2 }
+    }
   }
 
   // `squared` implemented using a `set` accessor.
@@ -126,8 +135,19 @@ InoutParameterAutoDiffTests.test("SetAccessor") {
     s.computed *= x
     return s.x
   }
-  expectEqual(6, gradient(at: 3, in: squared))
-  expectEqual(8, gradient(at: 4, in: squared))
+  expectEqual((9, 6), valueWithGradient(at: 3, in: squared))
+  expectEqual((16, 8), valueWithGradient(at: 4, in: squared))
+
+  // `quadrupled` implemented using a `set` accessor.
+  func quadrupled(_ x: Float) -> Float {
+    var s = S(x: 1)
+    s.doubled *= 4 * x
+    return s.x
+  }
+  print(valueWithGradient(at: 3, in: quadrupled))
+  print(valueWithGradient(at: 4, in: quadrupled))
+  expectEqual((12, 4), valueWithGradient(at: 3, in: quadrupled))
+  expectEqual((16, 4), valueWithGradient(at: 4, in: quadrupled))
 }
 
 // Test differentiation wrt `inout` parameters that have a class type.