[AutoDiff] Properly collect inout parameter adjoints (#41559)

Apparently, the parameter index calculation in #37861 was not always correct in presence of other pullback parameters (e.g. captures and non-differentiated args). Collect all inout parameters and collect inout parameter adjoints correctly.

Resolves SR-15891

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -2085,37 +2085,38 @@ bool PullbackCloner::Implementation::run() {
 
   // Collect differentiation parameter adjoints.
   // Do a first pass to collect non-inout values.
-  unsigned pullbackInoutArgumentIndex = 0;
   for (auto i : getConfig().parameterIndices->getIndices()) {
-    auto isParameterInout = conv.getParameters()[i].isIndirectMutating();
-    if (!isParameterInout) {
+    if (!conv.getParameters()[i].isIndirectMutating()) {
+       addRetElt(i);
+     }
+  }
+
+  // Do a second pass for all inout parameters, however this is only necessary
+  // for functions with multiple basic blocks.  For functions with a single
+  // basic block adjoint accumulation for those parameters is already done by
+  // per-instruction visitors.
+  if (getOriginal().size() > 1) {
+    const auto &pullbackConv = pullback.getConventions();
+    SmallVector<SILArgument *, 1> pullbackInOutArgs;
+    for (auto pullbackArg : enumerate(pullback.getArgumentsWithoutIndirectResults())) {
+      if (pullbackConv.getParameters()[pullbackArg.index()].isIndirectMutating())
+        pullbackInOutArgs.push_back(pullbackArg.value());
+    }
+
+    unsigned pullbackInoutArgumentIdx = 0;
+    for (auto i : getConfig().parameterIndices->getIndices()) {
+      // Skip non-inout parameters.
+      if (!conv.getParameters()[i].isIndirectMutating())
+        continue;
+
+      // For functions with multiple basic blocks, accumulation is needed
+      // for `inout` parameters because pullback basic blocks have different
+      // adjoint buffers.
+      pullbackIndirectResults.push_back(pullbackInOutArgs[pullbackInoutArgumentIdx++]);
       addRetElt(i);
     }
   }
 
-  // Do a second pass for all inout parameters.
-  for (auto i : getConfig().parameterIndices->getIndices()) {
-    // Skip non-inout parameters.
-    auto isParameterInout = conv.getParameters()[i].isIndirectMutating();
-    if (!isParameterInout)
-      continue;
-
-    // Skip `inout` parameters for functions with a single basic block:
-    // adjoint accumulation for those parameters is already done by
-    // per-instruction visitors.
-    if (getOriginal().size() == 1)
-      continue;
-
-    // For functions with multiple basic blocks, accumulation is needed
-    // for `inout` parameters because pullback basic blocks have different
-    // adjoint buffers.
-    auto pullbackInoutArgument =
-        getPullback()
-            .getArgumentsWithoutIndirectResults()[pullbackInoutArgumentIndex++];
-    pullbackIndirectResults.push_back(pullbackInoutArgument);
-    addRetElt(i);
-  }
-
   // Copy them to adjoint indirect results.
   assert(indParamAdjoints.size() == pullbackIndirectResults.size() &&
          "Indirect parameter adjoint count mismatch");

test/AutoDiff/compiler_crashers_fixed/sr15891-inout-adjoint-param-idx.swift

@@ -0,0 +1,42 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// SR-15891: The parameter indices used for copying
+// inout tangent vectors were calculated improperly
+// in presence of other pullback parameters (e.g. 
+// captures)
+
+import _Differentiation
+
+struct Foo {
+  var bar : Float
+  var baz : Float
+  var name : String?
+}
+
+func outerFunc(doIterations : Int, value: inout Float) -> (Float, (Float) -> Float) {
+  @differentiable(reverse, wrt: param)
+  func innerFunc1(param: Float, other: Foo) -> Float {
+    value += param * other.bar
+    return value * param * 2.0
+  }
+  
+  @differentiable(reverse, wrt: param1)
+  func loop(param1 : Float, other1: Foo) -> Float {
+    var res : Float;
+    res = 0.0
+    if (doIterations > 0) {
+      res = innerFunc1(param: param1, other: other1)
+    }
+    
+    return res
+  }
+  
+  @differentiable(reverse)
+  func curriedFunc(param: Float) -> Float {
+    let other = Foo(bar: 7, baz: 9)
+    return loop(param1: param, other1: other)
+  }
+  
+  let valAndPullback = valueWithPullback(at: value, of: curriedFunc)
+  return (value + valAndPullback.value, valAndPullback.pullback)
+}