[AutoDiff] Fix differentiation for non-wrt inout parameters.

lib/SIL/IR/SILFunctionType.cpp

@@ -235,8 +235,23 @@ IndexSubset *SILFunctionType::getDifferentiabilityResultIndices() {
       resultIndices.push_back(resultAndIndex.index());
   // Check `inout` parameters.
   for (auto inoutParamAndIndex : enumerate(getIndirectMutatingParameters()))
-    if (inoutParamAndIndex.value().getDifferentiability() !=
-        SILParameterDifferentiability::NotDifferentiable)
+    // FIXME(TF-1305): The `getResults().empty()` condition is a hack.
+    //
+    // Currently, an `inout` parameter can either be:
+    // 1. Both a differentiability parameter and a differentiability result.
+    // 2. `@noDerivative`: neither a differentiability parameter nor a
+    //    differentiability result.
+    // However, there is no way to represent an `inout` parameter that:
+    // 3. Is a differentiability parameter but not a differentiability result.
+    // 4. Is a differentiability result but not a differentiability parameter.
+    //
+    // See TF-1305 for solution ideas. For now, `@noDerivative` `inout`
+    // parameters are not treated as differentiability results, unless the
+    // original function has no formal results, which case all `inout`
+    // parameters are treated as differentiability results.
+    if (getResults().empty() ||
+        inoutParamAndIndex.value().getDifferentiability() !=
+            SILParameterDifferentiability::NotDifferentiable)
       resultIndices.push_back(getNumResults() + inoutParamAndIndex.index());
   auto numSemanticResults =
       getNumResults() + getNumIndirectMutatingParameters();
@@ -428,8 +443,9 @@ static CanSILFunctionType getAutoDiffDifferentialType(
     }
   }
   SmallVector<SILResultInfo, 1> differentialResults;
-  if (inoutParamIndices->isEmpty()) {
-    for (auto resultIndex : resultIndices->getIndices()) {
+  for (auto resultIndex : resultIndices->getIndices()) {
+    // Handle formal original result.
+    if (resultIndex < originalFnTy->getNumResults()) {
       auto &result = originalResults[resultIndex];
       auto resultTan =
           result.getInterfaceType()->getAutoDiffTangentSpace(lookupConformance);
@@ -448,8 +464,27 @@ static CanSILFunctionType getAutoDiffDifferentialType(
         substReplacements.push_back(resultTanType);
         differentialResults.push_back({gpType, resultConv});
       }
+      continue;
     }
+    // Handle original `inout` parameter.
+    auto inoutParamIndex = resultIndex - originalFnTy->getNumResults();
+    auto inoutParamIt = std::next(
+        originalFnTy->getIndirectMutatingParameters().begin(), inoutParamIndex);
+    auto paramIndex =
+        std::distance(originalFnTy->getParameters().begin(), &*inoutParamIt);
+    // If the original `inout` parameter is a differentiability parameter, then
+    // it already has a corresponding differential parameter. Skip adding a
+    // corresponding differential result.
+    if (parameterIndices->contains(paramIndex))
+      continue;
+    auto inoutParam = originalFnTy->getParameters()[paramIndex];
+    auto paramTan = inoutParam.getInterfaceType()->getAutoDiffTangentSpace(
+        lookupConformance);
+    assert(paramTan && "Parameter type does not have a tangent space?");
+    differentialResults.push_back(
+        {paramTan->getCanonicalType(), ResultConvention::Indirect});
   }
+
   SubstitutionMap substitutions;
   if (!substGenericParams.empty()) {
     auto genericSig =
@@ -710,7 +745,9 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     CanGenericSignature derivativeFnInvocationGenSig,
     bool isReabstractionThunk) {
   assert(parameterIndices);
+  assert(!parameterIndices->isEmpty() && "Parameter indices must not be empty");
   assert(resultIndices);
+  assert(!resultIndices->isEmpty() && "Result indices must not be empty");
   auto &ctx = getASTContext();
 
   // Look up result in cache.

test/AutoDiff/compiler_crashers_fixed/sr13305-noderivative-inout-parameter.swift

@@ -0,0 +1,39 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+// REQUIRES: asserts
+
+import _Differentiation
+
+// SR-13305: Test protocol witness thunk for `@differentiable` protocol
+// requirement, where the required method has a non-wrt `inout` parameter
+// that should be treated as a differentiability result.
+
+protocol SR_13305_Protocol {
+  @differentiable(wrt: x)
+  func method(x: Float, y: inout Float)
+}
+
+struct SR_13305_Struct: SR_13305_Protocol {
+  @differentiable(wrt: x)
+  func method(x: Float, y: inout Float) {
+    y = y * x
+  }
+}
+
+// Original crash:
+// Assertion failed: (!array.empty() && "claiming next from empty array!"), function claimNext, file /Users/danielzheng/swift-build/swift/lib/SILGen/SILGenPoly.cpp, line 112.
+// Stack dump:
+// ...
+// 1.	Swift version 5.3-dev (LLVM f8bd914aadc2e7b, Swift ba9c433c81d51ea)
+// 2.	While evaluating request ASTLoweringRequest(Lowering AST to SIL for module main)
+// 3.	While generating SIL witness table protocol conformance to 'SR_13305_Protocol' (at sr-13305.swift:7:1) for type 'SR_13305_Struct' (declared at [sr-13305.swift:12:1 - line:17:1] RangeText="struct SR_13305_Struct: SR_13305_Protocol {
+//   @differentiable(wrt: x)
+//   func method(x: Float, y: inout Float) {
+//     y = y * x
+//   }
+// ")
+// 4.	While generating protocol witness thunk SIL function "@AD__$s4main15SR_13305_StructVAA0B15_13305_ProtocolA2aDP6method1x1yySf_SfztFTW_jvp_SUU".
+//  for 'method(x:y:)' (at sr-13305.swift:14:3)
+// 5.	While emitting reabstraction thunk in SIL function "@$sSfIegy_S2fIegyd_TR".
+// ...
+// 7  swift-frontend           0x0000000100fe80ad swift::SILResultInfo const& claimNext<swift::SILResultInfo>(llvm::ArrayRef<swift::SILResultInfo>&) + 93
+// 8  swift-frontend           0x0000000100fe6cc0 (anonymous namespace)::ResultPlanner::claimNextInnerResult((anonymous namespace)::ResultPlanner::PlanData&) + 32

test/AutoDiff/validation-test/inout_parameters.swift

@@ -0,0 +1,64 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import DifferentiationUnittest
+import StdlibUnittest
+
+var InoutParameterAutoDiffTests = TestSuite("InoutParameterDifferentiation")
+
+// SR-13305: Test function with non-wrt `inout` parameter, which should be
+// treated as a differentiability result.
+
+protocol SR_13305_Protocol {
+  @differentiable(wrt: x)
+  func method(_ x: Float, _ y: inout Float)
+
+  @differentiable(wrt: x)
+  func genericMethod<T: Differentiable>(_ x: T, _ y: inout T)
+}
+
+InoutParameterAutoDiffTests.test("non-wrt inout parameter") {
+  struct SR_13305_Struct: SR_13305_Protocol {
+    @differentiable(wrt: x)
+    func method(_ x: Float, _ y: inout Float) {
+      y = y * x
+    }
+
+    @differentiable(wrt: x)
+    func genericMethod<T: Differentiable>(_ x: T, _ y: inout T) {
+      y = x
+    }
+  }
+
+  @differentiable(wrt: x)
+  func foo(_ s: SR_13305_Struct, _ x: Float, _ y: Float) -> Float {
+    var y = y
+    s.method(x, &y)
+    return y
+  }
+
+  @differentiable(wrt: x)
+  func fooGeneric<T: SR_13305_Protocol>(_ s: T, _ x: Float, _ y: Float) -> Float {
+    var y = y
+    s.method(x, &y)
+    return x
+  }
+
+  let s = SR_13305_Struct()
+
+  do {
+    let (value, (dx, dy)) = valueWithGradient(at: 2, 3, in: { foo(s, $0, $1) })
+    expectEqual(6, value)
+    expectEqual((3, 2), (dx, dy))
+  }
+  expectEqual((value: 6, gradient: 3), valueWithGradient(at: 2, in: { foo(s, $0, 3) }))
+
+  do {
+    let (value, (dx, dy)) = valueWithGradient(at: 2, 3, in: { fooGeneric(s, $0, $1) })
+    expectEqual(2, value)
+    expectEqual((1, 0), (dx, dy))
+  }
+  expectEqual((value: 2, gradient: 1), valueWithGradient(at: 2, in: { fooGeneric(s, $0, 3) }))
+}
+
+runAllTests()