[AutoDiff upstream] add more validation tests (#31190)

Author: marcrasi

test/AutoDiff/SILOptimizer/differentiation_subset_parameters_thunk.swift

@@ -0,0 +1,98 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+import _Differentiation
+
+func foo<T: Numeric>(_ x: T, _ y: T) -> T { x * y }
+
+@derivative(of: foo)
+func foo_vjp<T: Numeric & Differentiable>(_ x: T, _ y: T) -> (
+  value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)
+) {
+  (foo(x, y), { _ in (.zero, .zero) })
+}
+
+@differentiable
+func differentiate_foo_wrt_0(_ x: Float) -> Float {
+  foo(x, 1)
+}
+
+// CHECK-LABEL: sil hidden @{{.*}}differentiate_foo_wrt_0{{.*}}__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {
+// CHECK: bb0
+// CHECK:   [[FOO_ORIG:%.*]] = function_ref @{{.*}}foo{{.*}} : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> @out τ_0_0
+// CHECK:   [[FOO_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_ORIG]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> @out τ_0_0
+// CHECK:   [[FOO_JVP:%.*]] = differentiability_witness_function [jvp] [parameters 0 1] [results 0] <T where T : Numeric, T : Differentiable> @{{.*}}foo{{.*}} : $@convention(thin) <T where T : Numeric> (@in_guaranteed T, @in_guaranteed T) -> @out T
+// CHECK:   [[FOO_JVP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_JVP]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric, τ_0_0 : Differentiable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_1) -> @out τ_0_2 for <τ_0_0.TangentVector, τ_0_0.TangentVector, τ_0_0.TangentVector>)
+// CHECK:   [[FOO_JVP_SUBSET_THUNK_THIN:%.*]] = function_ref @AD__orig_{{.*}}foo{{.*}}_src_0_wrt_0_jvp_subset_parameters_thunk : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_JVP_SUBSET_THUNK:%.*]] = thin_to_thick_function [[FOO_JVP_SUBSET_THUNK_THIN]] : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float) to $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_VJP:%.*]] = differentiability_witness_function [vjp] [parameters 0 1] [results 0] <T where T : Numeric, T : Differentiable> @{{.*}}foo{{.*}} : $@convention(thin) <T where T : Numeric> (@in_guaranteed T, @in_guaranteed T) -> @out T
+// CHECK:   [[FOO_VJP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_VJP]]<Float>() : $@convention(thin) <τ_0_0 where τ_0_0 : Numeric, τ_0_0 : Differentiable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for <τ_0_0.TangentVector, τ_0_0.TangentVector, τ_0_0.TangentVector>)
+// CHECK:   [[FOO_VJP_SUBSET_THUNK_THIN:%.*]] = function_ref @AD__orig_{{.*}}foo{{.*}}_src_0_wrt_0_vjp_subset_parameters_thunk : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_VJP_SUBSET_THUNK:%.*]] = thin_to_thick_function [[FOO_VJP_SUBSET_THUNK_THIN]] : $@convention(thin) (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float) to $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
+// CHECK:   [[FOO_DIFF:%.*]] = differentiable_function [parameters 0] [[FOO_FLOAT]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> @out Float with_derivative {[[FOO_JVP_SUBSET_THUNK]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float), [[FOO_VJP_SUBSET_THUNK]] : $@callee_guaranteed (@in_guaranteed Float, @in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)}
+// CHECK: }
+
+func inoutIndirect<T: Differentiable, U: Differentiable, V: Differentiable>(
+  _ x: T, _ y: inout U, _ z: V
+) {}
+
+@derivative(of: inoutIndirect)
+func vjpInoutIndirect<T: Differentiable, U: Differentiable, V: Differentiable>(
+  _ x: T, _ y: inout U, _ z: V
+) -> (
+  value: Void, pullback: (inout U.TangentVector) -> (T.TangentVector, V.TangentVector)
+) {
+  return ((), { dy in
+    return (.zero, .zero)
+  })
+}
+
+@differentiable(wrt: x)
+@differentiable(wrt: y)
+@differentiable
+func inoutIndirectCaller<T: Differentiable, U: Differentiable, V: Differentiable>(
+  _ x: T, _ y: U, _ z: V
+) -> U {
+  var result = y
+  inoutIndirect(x, &result, z)
+  return result
+}
+
+@differentiable(wrt: (x, z))
+func concreteInoutIndirectCaller(
+  _ x: Float, _ y: Double, _ z: Float
+) -> Double {
+  return inoutIndirectCaller(x, y, z)
+}
+
+// CHECK-LABEL: sil shared [transparent] [serialized] [thunk] @AD__$sSdSfSdSfIegnrrr_SdS2fIegnrr_TR_src_0_wrt_0_2_pullback_index_subset_thunk : $@convention(thin) (@in_guaranteed Double, @guaranteed @callee_guaranteed (@in_guaranteed Double) -> (@out Float, @out Double, @out Float)) -> (@out Float, @out Float) {
+// CHECK: bb0(%0 : $*Float, %1 : $*Float, %2 : $*Double, %3 : $@callee_guaranteed (@in_guaranteed Double) -> (@out Float, @out Double, @out Float)):
+// CHECK:   %4 = alloc_stack $Double
+// CHECK:   %5 = apply %3(%0, %4, %1, %2) : $@callee_guaranteed (@in_guaranteed Double) -> (@out Float, @out Double, @out Float)
+// CHECK:   destroy_addr %4 : $*Double
+// CHECK:   dealloc_stack %4 : $*Double
+// CHECK:   %8 = tuple ()
+// CHECK:   return %8 : $()
+// CHECK: }
+
+// CHECK-LABEL: sil shared [transparent] [serialized] [thunk] @AD__$s13TangentVector16_Differentiation14DifferentiablePQy_AaDQzAaDQy0_Ieglrr_AeFIeglr_AbCRzAbCR_AbCR0_r1_lTR_src_0_wrt_0_1_pullback_index_subset_thunk : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2 where τ_0_0 : Differentiable, τ_0_1 : Differentiable, τ_0_2 : Differentiable> (@inout τ_0_1.TangentVector, @guaranteed @callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)) -> @out τ_0_0.TangentVector {
+// CHECK: bb0(%0 : $*τ_0_0.TangentVector, %1 : $*τ_0_1.TangentVector, %2 : $@callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)):
+// CHECK:   %3 = alloc_stack $τ_0_2.TangentVector
+// CHECK:   %4 = apply %2(%0, %3, %1) : $@callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)
+// CHECK:   destroy_addr %3 : $*τ_0_2.TangentVector
+// CHECK:   dealloc_stack %3 : $*τ_0_2.TangentVector
+// CHECK:   %7 = tuple ()
+// CHECK:   return %7 : $()
+// CHECK: }
+
+// CHECK-LABEL: sil shared [transparent] [serialized] [thunk] @AD__$s13TangentVector16_Differentiation14DifferentiablePQy_AaDQzAaDQy0_Ieglrr_AEIegl_AbCRzAbCR_AbCR0_r1_lTR_src_0_wrt_1_pullback_index_subset_thunk : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2 where τ_0_0 : Differentiable, τ_0_1 : Differentiable, τ_0_2 : Differentiable> (@inout τ_0_1.TangentVector, @guaranteed @callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)) -> () {
+// CHECK: bb0(%0 : $*τ_0_1.TangentVector, %1 : $@callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)):
+// CHECK:   %2 = alloc_stack $τ_0_0.TangentVector
+// CHECK:   %3 = alloc_stack $τ_0_2.TangentVector
+// CHECK:   %4 = apply %1(%2, %3, %0) : $@callee_guaranteed (@inout τ_0_1.TangentVector) -> (@out τ_0_0.TangentVector, @out τ_0_2.TangentVector)
+// CHECK:   destroy_addr %2 : $*τ_0_0.TangentVector
+// CHECK:   destroy_addr %3 : $*τ_0_2.TangentVector
+// CHECK:   dealloc_stack %3 : $*τ_0_2.TangentVector
+// CHECK:   dealloc_stack %2 : $*τ_0_0.TangentVector
+// CHECK:   %9 = tuple ()
+// CHECK:   return %9 : $()
+// CHECK: }

test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/main/main.swift

@@ -0,0 +1,13 @@
+import StdlibUnittest
+import _Differentiation
+
+import module1
+
+var Tests = TestSuite("CrossModuleDerivativeAttr")
+
+Tests.test("CrossFile") {
+  let grad = gradient(at: 0, in: fCrossFile)
+  expectEqual(10, grad)
+}
+
+runAllTests()

test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1.swift

@@ -0,0 +1 @@
+public func fCrossFile(_ x: Float) -> Float { x }

test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1_other_file.swift

@@ -0,0 +1,6 @@
+import _Differentiation
+
+@derivative(of: fCrossFile)
+public func vjpCrossFile(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 10 * $0 })
+}

test/AutoDiff/validation-test/cross_module_derivative_attr.swift

@@ -0,0 +1,5 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -working-directory %t -I%t -parse-as-library -emit-module -module-name module1 -emit-module-path %t/module1.swiftmodule -emit-library -static %S/Inputs/cross_module_derivative_attr/module1/module1.swift %S/Inputs/cross_module_derivative_attr/module1/module1_other_file.swift -Xfrontend -validate-tbd-against-ir=none
+// RUN: %target-build-swift -I%t -L%t %S/Inputs/cross_module_derivative_attr/main/main.swift -o %t/a.out -lm -lmodule1 -Xfrontend -validate-tbd-against-ir=none
+// RUN: %target-run %t/a.out
+// REQUIRES: executable_test

test/AutoDiff/validation-test/custom_derivatives.swift

@@ -0,0 +1,74 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+#if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
+import Darwin.C
+#else
+import Glibc
+#endif
+import DifferentiationUnittest
+
+var CustomDerivativesTests = TestSuite("CustomDerivatives")
+
+// Specify non-differentiable functions.
+// These will be wrapped in `differentiableFunction` and tested.
+
+func unary(_ x: Tracked<Float>) -> Tracked<Float> {
+  var x = x
+  x *= 2
+  return x
+}
+
+func binary(_ x: Tracked<Float>, _ y: Tracked<Float>) -> Tracked<Float> {
+  var x = x
+  x *= y
+  return x
+}
+
+CustomDerivativesTests.testWithLeakChecking("differentiableFunction-unary") {
+  let diffableUnary = differentiableFunction { x in
+    (value: unary(x), pullback: { v in v * x * 2 })
+  }
+  expectEqual(20, gradient(at: 10, in: diffableUnary))
+  // Test differentiation of @differentiable function.
+  expectEqual(20, gradient(at: 10, in: { diffableUnary($0) }))
+  expectEqual(40, gradient(at: 10, in: { diffableUnary($0) * 2 }))
+}
+
+CustomDerivativesTests.testWithLeakChecking("differentiableFunction-binary") {
+  let diffableBinary = differentiableFunction { (x, y) in
+    (value: binary(x, y), pullback: { v in (v * y, v * x) })
+  }
+  expectEqual((10, 5), gradient(at: 5, 10, in: diffableBinary))
+  // Test differentiation of @differentiable function.
+  expectEqual((10, 5), gradient(at: 5, 10, in: { diffableBinary($0, $1) }))
+  expectEqual((20, 10), gradient(at: 5, 10, in: { diffableBinary($0, $1) * 2 }))
+}
+
+CustomDerivativesTests.testWithLeakChecking("SumOfGradPieces") {
+  var grad: Tracked<Float> = 0
+  func addToGrad(_ x: inout Tracked<Float>) { grad += x }
+  _ = gradient(at: 4) { (x: Tracked<Float>) in
+    x.withDerivative(addToGrad)
+      * x.withDerivative(addToGrad)
+        * x.withDerivative(addToGrad)
+  }
+  expectEqual(48, grad)
+}
+
+CustomDerivativesTests.testWithLeakChecking("ModifyGradientOfSum") {
+  expectEqual(30, gradient(at: 4) { (x: Tracked<Float>) in
+    x.withDerivative { $0 *= 10 } + x.withDerivative { $0 *= 20 }
+  })
+}
+
+CustomDerivativesTests.testWithLeakChecking("WithoutDerivative") {
+  expectEqual(0, gradient(at: Tracked<Float>(4)) { x in
+    withoutDerivative(at: x) { x in
+      Tracked<Float>(sinf(x.value) + cosf(x.value))
+    }
+  })
+}
+
+runAllTests()

test/AutoDiff/validation-test/subset_parameters_thunk.swift

@@ -0,0 +1,77 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+import _Differentiation
+
+var SubsetParameterThunkTests = TestSuite("SubsetParameterThunks")
+
+func inoutDirect(_ x: Float, _ y: inout Double, _ z: Float) {}
+
+@derivative(of: inoutDirect)
+func vjpInoutDirect(_ x: Float, _ y: inout Double, _ z: Float) -> (
+  value: Void, pullback: (inout Double) -> (Float, Float)
+) {
+  return ((), { dy in
+    dy = 3
+    return (2, 4)
+  })
+}
+
+SubsetParameterThunkTests.test("InoutParametersDirect") {
+  @differentiable(wrt: x)
+  @differentiable(wrt: y)
+  @differentiable(wrt: z)
+  func inoutDirectCaller(_ x: Float, _ y: Double, _ z: Float) -> Double {
+    var result = y
+    inoutDirect(x, &result, z)
+    return result
+  }
+
+  let x: Float = 3
+  let y: Double = 4
+  let z: Float = 5
+  expectEqual((2, 3, 4), gradient(at: x, y, z, in: inoutDirectCaller))
+  expectEqual((3, 4), gradient(at: y, z, in: { y, z in inoutDirectCaller(x, y, z) }))
+  expectEqual((2, 4), gradient(at: x, z, in: { x, z in inoutDirectCaller(x, y, z) }))
+  expectEqual((2, 3), gradient(at: x, y, in: { x, y in inoutDirectCaller(x, y, z) }))
+}
+
+func inoutIndirect<T: Differentiable, U: Differentiable, V: Differentiable>(
+  _ x: T, _ y: inout U, _ z: V
+) {}
+
+@derivative(of: inoutIndirect)
+func vjpInoutIndirect<T: Differentiable, U: Differentiable, V: Differentiable>(
+  _ x: T, _ y: inout U, _ z: V
+) -> (
+  value: Void, pullback: (inout U.TangentVector) -> (T.TangentVector, V.TangentVector)
+) {
+  return ((), { dy in
+    return (.zero, .zero)
+  })
+}
+
+SubsetParameterThunkTests.test("InoutParametersIndirect") {
+  @differentiable(wrt: x)
+  @differentiable(wrt: y)
+  @differentiable(wrt: z)
+  @differentiable
+  func inoutIndirectCaller<T: Differentiable, U: Differentiable, V: Differentiable>(
+    _ x: T, _ y: U, _ z: V
+  ) -> U {
+    var result = y
+    inoutIndirect(x, &result, z)
+    return result
+  }
+
+  let x: Float = 3
+  let y: Double = 4
+  let z: Float = 5
+  expectEqual((0, 1, 0), gradient(at: x, y, z, in: inoutIndirectCaller))
+  expectEqual((1, 0), gradient(at: y, z, in: { y, z in inoutIndirectCaller(x, y, z) }))
+  expectEqual((0, 0), gradient(at: x, z, in: { x, z in inoutIndirectCaller(x, y, z) }))
+  expectEqual((0, 1), gradient(at: x, y, in: { x, y in inoutIndirectCaller(x, y, z) }))
+}
+
+runAllTests()