Correctly handle multiple semantic results for autodiff subset differential thunks (#84366)

Fixes #84365

Author: asl

lib/SILOptimizer/Differentiation/Thunk.cpp

@@ -621,22 +621,33 @@ getOrCreateSubsetParametersThunkForLinearMap(
 
   // If differential thunk, deallocate local allocations and directly return
   // `apply` result (if it is desired).
+  // TODO: Unify with VJP code below
   if (kind == AutoDiffDerivativeFunctionKind::JVP) {
     SmallVector<SILValue, 8> differentialDirectResults;
     extractAllElements(ai, builder, differentialDirectResults);
     SmallVector<SILValue, 8> allResults;
     collectAllActualResultsInTypeOrder(ai, differentialDirectResults, allResults);
-    unsigned numResults = thunk->getConventions().getNumDirectSILResults() +
-     thunk->getConventions().getNumDirectSILResults();
     SmallVector<SILValue, 8> results;
-    for (unsigned idx : *actualConfig.resultIndices) {
-      if (idx >= numResults)
-        break;
 
-      auto result = allResults[idx];
-      if (desiredConfig.isWrtResult(idx))
-        results.push_back(result);
-      else {
+    unsigned firstSemanticParamResultIdx = origFnType->getNumResults();
+    for (unsigned resultIndex : *actualConfig.resultIndices) {
+      SILValue result;
+      if (resultIndex >= firstSemanticParamResultIdx) {
+        auto semanticResultArgIdx = resultIndex - firstSemanticParamResultIdx;
+        result =
+          *std::next(ai->getAutoDiffSemanticResultArguments().begin(),
+                     semanticResultArgIdx);
+      } else
+        result = allResults[resultIndex];
+
+      // If result is desired:
+      // - Do nothing if result is indirect.
+      //   (It was already forwarded to the `apply` instruction).
+      // - Push it to `results` if result is direct.
+      if (desiredConfig.isWrtResult(resultIndex)) {
+        if (result->getType().isObject())
+          results.push_back(result);
+      } else { // Otherwise, cleanup the unused results.
         if (result->getType().isAddress())
           builder.emitDestroyAddrAndFold(loc, result);
         else

test/AutoDiff/compiler_crashers_fixed/issue-84365-autodiff-subset-thunks-differential.swift

@@ -0,0 +1,49 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// https://github.com/swiftlang/swift/issues/84365
+// Ensure autodiff subset thunks for differential correctly
+// handle multiple semantic results and release unwanted
+// result values
+
+import _Differentiation
+
+@differentiable(reverse,wrt: logits)
+public func softSolveForwardWithQ(logits: [Float]) -> ([Float], [Float]) {
+    return ([Float](repeating: 0, count: 0), [])
+}
+
+@derivative(of: softSolveForwardWithQ, wrt: logits)
+public func vjpSoftSolveForwardWithQ(logits: [Float]) -> (value: ([Float], [Float]), pullback: ([Float].TangentVector, [Float].TangentVector) -> [Float].TangentVector) {
+    let n = logits.count
+    let q = [Float](repeating: 0, count: 0)
+    let y = [Float](repeating: 0, count: 0)
+
+    return (
+        value: (y, q),
+        pullback: { _, _ in 
+            return Array<Float>.DifferentiableView([Float](repeating: 0, count: n))
+        }
+    )
+}
+
+@differentiable(reverse,wrt: logits)
+public func forwardPredict(logits: [Float]) -> ([Float], [Float], [Float]) {
+    let (y, q) = softSolveForwardWithQ(logits: logits)
+    return (y, q, [0.0])
+}
+
+@derivative(of: forwardPredict, wrt: logits)
+public func vjpForwardPredict(logits: [Float]) -> (
+    value: ([Float], [Float], [Float]),
+    pullback: ([Float].TangentVector, [Float].TangentVector, [Float].TangentVector) -> [Float].TangentVector
+) {
+    let (valYQ, pb) = vjpSoftSolveForwardWithQ(logits: logits)
+    let (y, q) = valYQ
+    return ((y, q, [0.0]), { upY, upQ, _ in pb(upY, upQ) }) 
+}
+
+@differentiable(reverse,wrt: logits)
+public func crossEntropyFromForwardPredict(logits: [Float]) -> Float {
+    let (_, q, _) = forwardPredict(logits: logits)
+    return q[0] + 1e-8
+}