Merge pull request swiftlang#37784 from BradLarson/SR-14625

[AutoDiff] [SR-14625] Fix for subset parameters thunk over-consume

Author: BradLarson

lib/SILOptimizer/Differentiation/Thunk.cpp

@@ -610,7 +610,8 @@ getOrCreateSubsetParametersThunkForLinearMap(
   };
 
   // Build a `.zero` argument for the given `Differentiable`-conforming type.
-  auto buildZeroArgument = [&](SILType zeroSILType) {
+  auto buildZeroArgument = [&](SILParameterInfo zeroSILParameter) {
+    auto zeroSILType = zeroSILParameter.getSILStorageInterfaceType();
     auto zeroSILObjType = zeroSILType.getObjectType();
     auto zeroType = zeroSILType.getASTType();
     auto *swiftMod = parentThunk->getModule().getSwiftModule();
@@ -623,13 +624,17 @@ getOrCreateSubsetParametersThunkForLinearMap(
       localAllocations.push_back(buf);
       builder.emitZeroIntoBuffer(loc, buf, IsInitialization);
       if (zeroSILType.isAddress()) {
-        valuesToCleanup.push_back(buf);
         arguments.push_back(buf);
+        if (zeroSILParameter.isGuaranteed()) {
+          valuesToCleanup.push_back(buf);
+        }
       } else {
         auto arg = builder.emitLoadValueOperation(loc, buf,
                                                   LoadOwnershipQualifier::Take);
-        valuesToCleanup.push_back(arg);
         arguments.push_back(arg);
+        if (zeroSILParameter.isGuaranteed()) {
+          valuesToCleanup.push_back(arg);
+        }
       }
       break;
     }
@@ -687,10 +692,9 @@ getOrCreateSubsetParametersThunkForLinearMap(
       }
       // Otherwise, construct and use a zero argument.
       else {
-        auto zeroSILType =
-            linearMapType->getParameters()[mapOriginalParameterIndex(i)]
-                .getSILStorageInterfaceType();
-        buildZeroArgument(zeroSILType);
+        auto zeroSILParameter =
+            linearMapType->getParameters()[mapOriginalParameterIndex(i)];
+        buildZeroArgument(zeroSILParameter);
       }
     }
     break;

test/AutoDiff/compiler_crashers_fixed/sr14625-over-consume-in-subset-parameters-thunk.swift

@@ -0,0 +1,31 @@
+// RUN: %target-build-swift %s
+
+// SR-14625: An over-consume in a subset parameters thunk detected after
+// enabling OSSA.
+
+import _Differentiation
+
+struct Type2: Differentiable {
+  var test1: Double
+
+  @differentiable(reverse)
+  public init(test1: Double) {
+    self.test1 = test1
+  }
+}
+
+struct Type1: Differentiable {
+  var test1: Double
+  var test3: [Type2]
+
+  @differentiable(reverse)
+  public init(test1: Double, test3: [Type2]) {
+    self.test1 = test1
+    self.test3 = test3
+  }
+}
+
+@differentiable(reverse)
+func ingestValue(val1: Double, val2: Double) -> Type1 {
+  return Type1(test1: val1 * val2, test3: [])
+}