Merge pull request #41401 from asl/sr15205-fix

[AutoDiff] Fix several issues related to captured arguments

Author: rxwei

lib/SIL/IR/SILFunctionType.cpp

@@ -714,7 +714,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
   for (auto &param : diffParams) {
     // Skip `inout` parameters, which semantically behave as original results
     // and always appear as pullback parameters.
-    if (param.isIndirectInOut())
+    if (param.isIndirectMutating())
       continue;
     auto paramTanType = getAutoDiffTangentTypeForLinearMap(
         param.getInterfaceType(), lookupConformance,

lib/SILOptimizer/Differentiation/Common.cpp

@@ -452,6 +452,14 @@ findMinimalDerivativeConfiguration(AbstractFunctionDecl *original,
     auto *silParameterIndices = autodiff::getLoweredParameterIndices(
         config.parameterIndices,
         original->getInterfaceType()->castTo<AnyFunctionType>());
+
+    if (silParameterIndices->getCapacity() < parameterIndices->getCapacity()) {
+      assert(original->getCaptureInfo().hasLocalCaptures());
+      silParameterIndices =
+        silParameterIndices->extendingCapacity(original->getASTContext(),
+                                               parameterIndices->getCapacity());
+    }
+
     // If all indices in `parameterIndices` are in `daParameterIndices`, and
     // it has fewer indices than our current candidate and a primitive VJP,
     // then `attr` is our new candidate.

test/AutoDiff/compiler_crashers_fixed/sr15205-diff-capture.swift

@@ -0,0 +1,36 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// SR-15205: fix assertions related to captured arguments, they should
+// treated as constants
+
+import _Differentiation
+
+func outerFunc(value: inout Float) -> (Float, (Float) -> (Float, Float)) {
+  
+  @differentiable(reverse, wrt: param)
+  func innerFunc(param: Float, other: Float) -> Float {
+    value += param * other
+    return value * param * 2.0
+  }
+  
+  let valAndPullback = valueWithPullback(at: value, 2.0, of: innerFunc)
+  return (value + valAndPullback.value, valAndPullback.pullback)
+}
+
+func outerFunc2(value: inout Float) -> (Float, (Float) -> Float) {
+
+  @differentiable(reverse, wrt: param)
+  func innerFunc(param: Float, other: Float) -> Float {
+    value += param * other
+    return value * param * 2.0
+  }
+
+  @differentiable(reverse)
+  func curriedFunc(param: Float) -> Float {
+    return innerFunc(param: param, other: 3.0)
+  }
+
+  let valAndPullback = valueWithPullback(at: value, of: curriedFunc)
+  return (value + valAndPullback.value, valAndPullback.pullback)
+}
+