diff --git a/test/AutoDiff/SILOptimizer/differentiation_subset_parameters_thunk.swift b/test/AutoDiff/SILOptimizer/differentiation_subset_parameters_thunk.swift new file mode 100644 index 0000000000000..e89129f0dafb7 --- /dev/null +++ b/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(_ x: T, _ y: T) -> T { x * y } + +@derivative(of: foo) +func foo_vjp(_ 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]]() : $@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] @{{.*}}foo{{.*}} : $@convention(thin) (@in_guaranteed T, @in_guaranteed T) -> @out T +// CHECK: [[FOO_JVP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_JVP]]() : $@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] @{{.*}}foo{{.*}} : $@convention(thin) (@in_guaranteed T, @in_guaranteed T) -> @out T +// CHECK: [[FOO_VJP_FLOAT:%.*]] = partial_apply [callee_guaranteed] [[FOO_VJP]]() : $@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( + _ x: T, _ y: inout U, _ z: V +) {} + +@derivative(of: inoutIndirect) +func vjpInoutIndirect( + _ 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( + _ 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: } diff --git a/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/main/main.swift b/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/main/main.swift new file mode 100644 index 0000000000000..b82ebb2a7d2de --- /dev/null +++ b/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() diff --git a/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1.swift b/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1.swift new file mode 100644 index 0000000000000..c76929b440e65 --- /dev/null +++ b/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1.swift @@ -0,0 +1 @@ +public func fCrossFile(_ x: Float) -> Float { x } diff --git a/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1_other_file.swift b/test/AutoDiff/validation-test/Inputs/cross_module_derivative_attr/module1/module1_other_file.swift new file mode 100644 index 0000000000000..75eced660cb85 --- /dev/null +++ b/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 }) +} diff --git a/test/AutoDiff/validation-test/cross_module_derivative_attr.swift b/test/AutoDiff/validation-test/cross_module_derivative_attr.swift new file mode 100644 index 0000000000000..3d5011b07c144 --- /dev/null +++ b/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 diff --git a/test/AutoDiff/validation-test/custom_derivatives.swift b/test/AutoDiff/validation-test/custom_derivatives.swift new file mode 100644 index 0000000000000..44d7a21916bf0 --- /dev/null +++ b/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) -> Tracked { + var x = x + x *= 2 + return x +} + +func binary(_ x: Tracked, _ y: Tracked) -> Tracked { + 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 = 0 + func addToGrad(_ x: inout Tracked) { grad += x } + _ = gradient(at: 4) { (x: Tracked) in + x.withDerivative(addToGrad) + * x.withDerivative(addToGrad) + * x.withDerivative(addToGrad) + } + expectEqual(48, grad) +} + +CustomDerivativesTests.testWithLeakChecking("ModifyGradientOfSum") { + expectEqual(30, gradient(at: 4) { (x: Tracked) in + x.withDerivative { $0 *= 10 } + x.withDerivative { $0 *= 20 } + }) +} + +CustomDerivativesTests.testWithLeakChecking("WithoutDerivative") { + expectEqual(0, gradient(at: Tracked(4)) { x in + withoutDerivative(at: x) { x in + Tracked(sinf(x.value) + cosf(x.value)) + } + }) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/subset_parameters_thunk.swift b/test/AutoDiff/validation-test/subset_parameters_thunk.swift new file mode 100644 index 0000000000000..2b4dce4d9cca7 --- /dev/null +++ b/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( + _ x: T, _ y: inout U, _ z: V +) {} + +@derivative(of: inoutIndirect) +func vjpInoutIndirect( + _ 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( + _ 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()