[AutoDiff] Register VJPs for SIMD subscript(_: Int) setters. (swiftlang#32747)

Author: dan-zheng

stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb

@@ -45,17 +45,17 @@ where
 extension SIMD${n}
 where
   Scalar: Differentiable & BinaryFloatingPoint,
-  Scalar.TangentVector: BinaryFloatingPoint
+  Scalar.TangentVector == Scalar
 {
   // NOTE(TF-1094): serialized `@derivative` for `.swiftinterface` compilation.
   @inlinable
   @derivative(of: subscript(_:))
-  internal func _vjpSubscript(index: Int)
+  internal func _vjpSubscript(_ index: Int)
     -> (value: Scalar, pullback: (Scalar.TangentVector) -> TangentVector)
   {
     return (self[index], { v in
       var zeros = Self.zero
-      zeros[index] = Scalar(v)
+      zeros[index] = v
       return zeros
     })
   }
@@ -69,6 +69,19 @@ where
       return .init(v[index])
     })
   }
+
+  @inlinable
+  @derivative(of: subscript(_:).set)
+  internal mutating func _vjpSubscriptSetter(_ newValue: Scalar, _ index: Int)
+    -> (value: Void, pullback: (inout TangentVector) -> Scalar.TangentVector)
+  {
+    self[index] = newValue
+    return ((), { dSelf in
+      let dNewValue = dSelf[index]
+      dSelf[index] = 0
+      return dNewValue
+    })
+  }
 }
 
 %end
@@ -421,9 +434,9 @@ where
   @inlinable
   @derivative(of: sum)
   func _jvpSum() -> (
-    value: Scalar, differential (TangentVector) -> Scalar.TangentVector
+    value: Scalar, differential: (TangentVector) -> Scalar.TangentVector
   ) {
-    return (sum(), { v in v.sum() }
+    return (sum(), { v in Scalar.TangentVector(v.sum()) })
   }
 }
 */

test/AutoDiff/stdlib/simd.swift

@@ -56,7 +56,7 @@ SIMDTests.test("Negate") {
   expectEqual(-g, pb1(g))
 }
 
-SIMDTests.test("subscript") {
+SIMDTests.test("Subscript") {
   let a = SIMD4<Float>(1, 2, 3, 4)
 
   func foo1(x: SIMD4<Float>) -> Float {
@@ -68,6 +68,35 @@ SIMDTests.test("subscript") {
   expectEqual(SIMD4<Float>(0, 0, 0, 7), pb1(7))
 }
 
+SIMDTests.test("SubscriptSetter") {
+  let a = SIMD4<Float>(1, 2, 3, 4)
+  let ones = SIMD4<Float>(1, 1, 1, 1)
+
+  // A wrapper around `subscript(_: Int).set`.
+  func subscriptSet(
+    _ simd: SIMD4<Float>, index: Int, newScalar: Float
+  ) -> SIMD4<Float> {
+    var result = simd
+    result[index] = newScalar
+    return result
+  }
+
+  let (val1, pb1) = valueWithPullback(at: a, 5, in: { subscriptSet($0, index: 2, newScalar: $1) })
+  expectEqual(SIMD4<Float>(1, 2, 5, 4), val1)
+  expectEqual((SIMD4<Float>(1, 1, 0, 1), 1), pb1(ones))
+
+  func doubled(_ x: SIMD4<Float>) -> SIMD4<Float> {
+    var result = x
+    for i in withoutDerivative(at: x.indices) {
+      result[i] = x[i] * 2
+    }
+    return result
+  }
+  let (val2, pb2) = valueWithPullback(at: a, in: doubled)
+  expectEqual(SIMD4<Float>(2, 4, 6, 8), val2)
+  expectEqual(SIMD4<Float>(2, 2, 2, 2), pb2(ones))
+}
+
 SIMDTests.test("Addition") {
   let a = SIMD4<Float>(1, 2, 3, 4)
   let g = SIMD4<Float>(1, 1, 1, 1)