Added tuple result tests, extracting tuple elements as semantic result types.

Author: BradLarson

lib/AST/AutoDiff.cpp

@@ -196,8 +196,16 @@ void autodiff::getFunctionSemanticResultTypes(
           functionType->getResult()->getAs<AnyFunctionType>()) {
     formalResultType = resultFunctionType->getResult();
   }
-  if (!formalResultType->isEqual(ctx.TheEmptyTupleType))
-    result.push_back({remap(formalResultType), /*isInout*/ false});
+  if (!formalResultType->isEqual(ctx.TheEmptyTupleType)) {
+    // Separate tuple elements into individual results.
+    if (formalResultType->is<TupleType>()) {
+      for (auto elt : formalResultType->castTo<TupleType>()->getElements()) {
+        result.push_back({remap(elt.getType()), /*isInout*/ false});
+      }
+    } else {
+      result.push_back({remap(formalResultType), /*isInout*/ false});
+    }
+  }
 
   // Collect `inout` parameters as semantic results.
   for (auto param : functionType->getParams())

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -871,6 +871,22 @@ extension InoutParameters {
   ) { fatalError() }
 }
 
+// Test tuple results.
+
+extension InoutParameters {
+  func tupleResults(_ x: Float) -> (Float, Float) { (x, x) }
+  @derivative(of: tupleResults, wrt: x)
+  func vjpTupleResults(_ x: Float) -> (
+    value: (Float, Float), pullback: (Float, Float) -> Float
+  ) { fatalError() }
+
+  func tupleResultsInt(_ x: Float) -> (Int, Float) { (1, x) }
+  @derivative(of: tupleResultsInt, wrt: x)
+  func vjpTupleResults(_ x: Float) -> (
+    value: (Int, Float), pullback: (Float) -> Float
+  ) { fatalError() }
+}
+
 // Test original/derivative function `inout` parameter mismatches.
 
 extension InoutParameters {

test/AutoDiff/Sema/differentiable_attr_type_checking.swift

@@ -685,6 +685,19 @@ extension InoutParameters {
   mutating func mutatingMethod(_ other: Self) -> Self {}
 }
 
+// Test tuple results.
+
+extension InoutParameters {
+  @differentiable(reverse)
+  static func tupleResults(_ x: Self) -> (Self, Self) {}
+
+  @differentiable(reverse)
+  static func tupleResultsInt(_ x: Self) -> (Int, Self) {}
+
+  @differentiable(reverse)
+  static func tupleResultsInt2(_ x: Self) -> (Self, Int) {}
+}
+
 // Test accessors: `set`, `_read`, `_modify`.
 
 struct Accessors: Differentiable {

test/AutoDiff/validation-test/simple_math.swift

@@ -147,6 +147,43 @@ SimpleMathTests.test("MultipleResultsWithCustomPullback") {
   expectEqual((10, 5), gradient(at: 5, 10, of: multiply_swapCustom))
 }
 
+// Test functions returning tuples.
+@differentiable(reverse)
+func swapTuple(_ x: Float, _ y: Float) -> (Float, Float) {
+  return (y, x)
+}
+
+@differentiable(reverse)
+func swapTupleCustom(_ x: Float, _ y: Float) -> (Float, Float) {
+  return (y, x)
+}
+@derivative(of: swapTupleCustom)
+func vjpSwapTupleCustom(_ x: Float, _ y: Float) -> (
+  value: (Float, Float), pullback: (Float, Float) -> (Float, Float)
+) {
+  return (swapTupleCustom(x, y), {v1, v2 in
+    return (v2, v1)
+  })
+}
+
+SimpleMathTests.test("ReturningTuples") {
+  func multiply_swapTuple(_ x: Float, _ y: Float) -> Float {
+    let result = swapTuple(x, y)
+    return result.0 * result.1
+  }
+
+  expectEqual((4, 3), gradient(at: 3, 4, of: multiply_swapTuple))
+  expectEqual((10, 5), gradient(at: 5, 10, of: multiply_swapTuple))
+
+  func multiply_swapTupleCustom(_ x: Float, _ y: Float) -> Float {
+    let result = swapTupleCustom(x, y)
+    return result.0 * result.1
+  }
+
+  expectEqual((4, 3), gradient(at: 3, 4, of: multiply_swapTupleCustom))
+  expectEqual((10, 5), gradient(at: 5, 10, of: multiply_swapTupleCustom))
+}
+
 SimpleMathTests.test("CaptureLocal") {
   let z: Float = 10
   func foo(_ x: Float) -> Float {