[AutoDiff] Fix return type of subset parameters thunk function (#67487)

The patch resolves #67402.

When the original function has a tuple result type, we should append
thunkedLinearMap as the last element of the tuple to match the function
declaration. Before this patch, the compiler used to wrap the original result
tuple and thunkedLinearMap into another tuple, and caused the verifier error.

Before the patch:

  return %{{.*}} : $((Float, Double), @callee_guaranteed (Float) -> X.TangentVector)

After the patch:

  return %{{.*}} : $(Float, Double, @callee_guaranteed (Float) -> X.TangentVector)

Author: asavonic

lib/SILOptimizer/Differentiation/Thunk.cpp

@@ -786,10 +786,8 @@ getOrCreateSubsetParametersThunkForDerivativeFunction(
   // Extract all direct results.
   SmallVector<SILValue, 8> directResults;
   extractAllElements(apply, builder, directResults);
-  auto originalDirectResults = ArrayRef<SILValue>(directResults).drop_back(1);
-  auto originalDirectResult =
-      joinElements(originalDirectResults, builder, apply->getLoc());
   auto linearMap = directResults.back();
+  directResults.pop_back();
 
   auto linearMapType = linearMap->getType().castTo<SILFunctionType>();
   auto linearMapTargetType = targetType->getResults()
@@ -830,8 +828,8 @@ getOrCreateSubsetParametersThunkForDerivativeFunction(
          0);
   if (origFnType->getNumResults() > 0 &&
       origFnType->getResults().front().isFormalDirect()) {
-    auto result =
-        joinElements({originalDirectResult, thunkedLinearMap}, builder, loc);
+    directResults.push_back(thunkedLinearMap);
+    auto result = joinElements(directResults, builder, loc);
     builder.createReturn(loc, result);
   } else {
     builder.createReturn(loc, thunkedLinearMap);

test/AutoDiff/SILOptimizer/param_thunk_tuple.swift

@@ -0,0 +1,33 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+// Verify the result type of a subset parameters thunk matches the declaration:
+//
+// CHECK: // autodiff subset parameters thunk for forward-mode derivative from f(x:)
+// CHECK-NEXT: sil shared [transparent] [thunk] @$s17param_thunk_tuple{{.*}} : $@convention(thin) (X)
+// CHECK-SAME: -> (Float, Double, @owned @callee_guaranteed (X.TangentVector) -> Float)
+// CHECK: return
+// CHECK-SAME: %{{.*}} : $(Float, Double, @callee_guaranteed (X.TangentVector) -> Float)
+//
+// CHECK: // autodiff subset parameters thunk for reverse-mode derivative from f(x:)
+// CHECK-NEXT: sil shared [transparent] [thunk] @$s17param_thunk_tuple{{.*}} : $@convention(thin) (X)
+// CHECK-SAME: -> (Float, Double, @owned @callee_guaranteed (Float) -> X.TangentVector)
+// CHECK: return
+// CHECK-SAME: %{{.*}} : $(Float, Double, @callee_guaranteed (Float) -> X.TangentVector)
+
+import _Differentiation
+
+struct X: Differentiable {
+    var a: Float
+    var b: Double
+}
+
+@differentiable(reverse)
+func f(x: X) -> (Float, Double) {
+    (x.a, x.b)
+}
+
+@differentiable(reverse)
+func g1(x: X) -> Float {
+    f(x: x).0
+}
+