Enable Array.subscript._modify differentiation

Author: asl

lib/Serialization/DeserializeSIL.cpp

@@ -5135,7 +5135,8 @@ SILDeserializer::readDifferentiabilityWitness(DeclID DId) {
                        ArrayRef<unsigned>(parameterAndResultIndices)
                            .take_front(numParameterIndices));
   auto numResults = originalFnType->getNumResults() +
-                    originalFnType->getNumIndirectMutatingParameters();
+                    originalFnType->getNumIndirectMutatingParameters() +
+                    originalFnType->getNumYields();
   auto *resultIndices =
       IndexSubset::get(MF->getContext(), numResults,
                        ArrayRef<unsigned>(parameterAndResultIndices)

lib/Serialization/SerializeSIL.cpp

@@ -3547,11 +3547,9 @@ void SILSerializer::writeSILDifferentiabilityWitness(
              dw.getParameterIndices()->getCapacity() &&
          "Original function parameter count should match differentiability "
          "witness parameter indices capacity");
-  unsigned numInoutParameters = llvm::count_if(
-      originalFnType->getParameters(), [](SILParameterInfo paramInfo) {
-        return paramInfo.isIndirectMutating();
-      });
-  assert(originalFnType->getNumResults() + numInoutParameters ==
+  assert(originalFnType->getNumResults() +
+         originalFnType->getNumIndirectMutatingParameters() +
+         originalFnType->getNumYields() ==
              dw.getResultIndices()->getCapacity() &&
          "Original function result count should match differentiability "
          "witness result indices capacity");

stdlib/public/Differentiation/ArrayDifferentiation.swift

@@ -170,10 +170,15 @@ where Element: AdditiveArithmetic & Differentiable {
 
   @inlinable
   public subscript(_ index: Int) -> Element {
-    if index < base.count {
-      return base[index]
-    } else {
-      return Element.zero
+    get {
+      if index < base.count {
+        return base[index]
+      } else {
+        return Element.zero
+      }
+    }
+    _modify {
+      yield &base[index]
     }
   }
 }
@@ -228,6 +233,21 @@ extension Array where Element: Differentiable {
     return (self[index], differential)
   }
 
+  @inlinable
+  @derivative(of: subscript._modify)
+  @yield_once
+  mutating func _vjpModify(index: Int) -> (
+    value: inout @yields Element, pullback: @yield_once (inout TangentVector) -> inout @yields Element.TangentVector
+  ) {
+    yield &self[index]
+
+    @yield_once
+    func pullback(_ v: inout TangentVector) -> inout @yields Element.TangentVector {
+      yield &v[index]
+    }
+    return pullback
+  }
+
   @inlinable
   @derivative(of: +)
   static func _vjpConcatenate(_ lhs: Self, _ rhs: Self) -> (

test/AutoDiff/SILOptimizer/activity_analysis.swift

@@ -641,8 +641,6 @@ func testAccessorCoroutinesModify(_ x: HasCoroutineModifyAccessors) -> Float {
 func testBeginApplyActiveInoutArgument(array: [Float], x: Float) -> Float {
   var array = array
   // Array subscript assignment below calls `Array.subscript.modify`.
-  // expected-error @+2 {{expression is not differentiable}}  
-  // expected-note @+1 {{cannot differentiate functions that have not been marked '@differentiable' and that are defined in other files}}
   array[0] = x
   return array[0]
 }
@@ -678,8 +676,6 @@ func testBeginApplyActiveButInitiallyNonactiveInoutArgument(x: Float) -> Float {
   // `var array` is initially non-active.
   var array: [Float] = [0]
   // Array subscript assignment below calls `Array.subscript.modify`.
-  // expected-error @+2 {{expression is not differentiable}}  
-  // expected-note @+1 {{cannot differentiate functions that have not been marked '@differentiable' and that are defined in other files}}
   array[0] = x
   return array[0]
 }

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -779,7 +779,7 @@ public func fragileDifferentiable(_ x: Float) -> Float {
 }
 
 //===----------------------------------------------------------------------===//
-// Coroutines (SIL function yields, `begin_apply`) (not yet supported)
+// Coroutines (SIL function yields, `begin_apply`)
 //===----------------------------------------------------------------------===//
 
 struct HasReadAccessors: Differentiable {
@@ -819,8 +819,6 @@ func testModifyAccessorCoroutines(_ x: HasModifyAccessors) -> Float {
 func TF_1078(array: [Float], x: Float) -> Float {
   var array = array
   // Array subscript assignment below calls `Array.subscript.modify`.
-  // expected-error @+2 {{expression is not differentiable}}
-  // expected-note @+1 {{cannot differentiate functions that have not been marked '@differentiable' and that are defined in other files}}
   array[0] = x
   return array[0]
 }
@@ -830,8 +828,6 @@ func TF_1078(array: [Float], x: Float) -> Float {
 func TF_1115(_ x: Float) -> Float {
   var array: [Float] = [0]
   // Array subscript assignment below calls `Array.subscript.modify`.
-  // expected-error @+2 {{expression is not differentiable}}
-  // expected-note @+1 {{cannot differentiate functions that have not been marked '@differentiable' and that are defined in other files}}
   array[0] = x
   return array[0]
 }