// RUN: %target-sil-opt -sil-print-types -test-runner %s -o /dev/null 2>&1 | %FileCheck %s // REQUIRES: swift_in_compiler sil_stage canonical import Builtin import Swift import SwiftShims import _Differentiation //////////////////////////////////////////////////////////////// // Single closure call site where closure is passed as @owned // //////////////////////////////////////////////////////////////// sil @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) sil private @$pullback_f : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float { bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> (Float, Float)): %2 = apply %1(%0) : $@callee_guaranteed (Float) -> (Float, Float) strong_release %1 : $@callee_guaranteed (Float) -> (Float, Float) // id: %3 %4 = tuple_extract %2 : $(Float, Float), 0 %5 = tuple_extract %2 : $(Float, Float), 1 %6 = struct_extract %5 : $Float, #Float._value %7 = struct_extract %4 : $Float, #Float._value %8 = builtin "fadd_FPIEEE32"(%6 : $Builtin.FPIEEE32, %7 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %9 = struct $Float (%8 : $Builtin.FPIEEE32) debug_value %9 : $Float, let, name "x", argno 1 // id: %10 return %9 : $Float // id: %11 } // reverse-mode derivative of f(_:) sil hidden @$s4test1fyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { bb0(%0 : $Float): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s4test1fyS2fFTJrSpSr // CHECK: PartialApply call site: %8 = partial_apply [callee_guaranteed] %7(%6) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float // CHECK: Passed in closures: // CHECK: 1. %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s11$pullback_f12$vjpMultiplyS2fTf1nc_n // CHECK: sil private @$s11$pullback_f12$vjpMultiplyS2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float { // CHECK: bb0(%0 : $Float, %1 : $Float, %2 : $Float): // CHECK: %3 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %4 = partial_apply [callee_guaranteed] %3(%1, %2) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %5 = apply %4(%0) : $@callee_guaranteed (Float) -> (Float, Float) // CHECK: strong_release %4 : $@callee_guaranteed (Float) -> (Float, Float) // id: %6 // CHECK: return //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s4test1fyS2fFTJrSpSr // CHECK: sil hidden @$s4test1fyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { // CHECK: bb0(%0 : $Float): // CHECK: %2 = struct_extract %0 : $Float, #Float._value // CHECK: %3 = builtin "fmul_FPIEEE32"(%2 : $Builtin.FPIEEE32, %2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %4 = struct $Float (%3 : $Builtin.FPIEEE32) // CHECK: %5 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %8 = function_ref @$s11$pullback_f12$vjpMultiplyS2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float // CHECK: %9 = partial_apply [callee_guaranteed] %8(%0, %0) : $@convention(thin) (Float, Float, Float) -> Float // CHECK: release_value %6 : $@callee_guaranteed (Float) -> (Float, Float) // id: %10 // CHECK: %11 = tuple (%4 : $Float, %9 : $@callee_guaranteed (Float) -> Float) // CHECK: return %11 debug_value %0 : $Float, let, name "x", argno 1 // id: %1 %2 = struct_extract %0 : $Float, #Float._value %3 = builtin "fmul_FPIEEE32"(%2 : $Builtin.FPIEEE32, %2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %4 = struct $Float (%3 : $Builtin.FPIEEE32) // function_ref closure #1 in static Float._vjpMultiply(lhs:rhs:) %5 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // function_ref pullback of f(_:) %7 = function_ref @$pullback_f : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float %8 = partial_apply [callee_guaranteed] %7(%6) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float %9 = tuple (%4 : $Float, %8 : $@callee_guaranteed (Float) -> Float) return %9 : $(Float, @callee_guaranteed (Float) -> Float) // id: %10 } ///////////////////////////////////////////////////////////////////// // Single closure call site where closure is passed as @guaranteed // ///////////////////////////////////////////////////////////////////// sil private @$pullback_k : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> Float { // %0 // %1 bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> (Float, Float)): %2 = apply %1(%0) : $@callee_guaranteed (Float) -> (Float, Float) %3 = tuple_extract %2 : $(Float, Float), 0 %4 = tuple_extract %2 : $(Float, Float), 1 %5 = struct_extract %4 : $Float, #Float._value %6 = struct_extract %3 : $Float, #Float._value %7 = builtin "fadd_FPIEEE32"(%5 : $Builtin.FPIEEE32, %6 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %8 = struct $Float (%7 : $Builtin.FPIEEE32) debug_value %8 : $Float, let, name "x", argno 1 // id: %9 return %8 : $Float // id: %10 } // end sil function '$pullback_k' // reverse-mode derivative of k(_:) sil hidden @$s4test1kyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { bb0(%0 : $Float): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s4test1kyS2fFTJrSpSr // CHECK: PartialApply call site: %8 = partial_apply [callee_guaranteed] %7(%6) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> Float // CHECK: Passed in closures: // CHECK: 1. %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s11$pullback_k12$vjpMultiplyS2fTf1nc_n // CHECK: sil private @$s11$pullback_k12$vjpMultiplyS2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float { // CHECK: bb0(%0 : $Float, %1 : $Float, %2 : $Float): // CHECK: %3 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %4 = partial_apply [callee_guaranteed] %3(%1, %2) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %5 = apply %4(%0) : $@callee_guaranteed (Float) -> (Float, Float) // CHECK: release_value %4 : $@callee_guaranteed (Float) -> (Float, Float) // id: %13 // CHECK: return %11 //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s4test1kyS2fFTJrSpSr // CHECK: sil hidden @$s4test1kyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { // CHECK: bb0(%0 : $Float): // CHECK: %2 = struct_extract %0 : $Float, #Float._value // CHECK: %3 = builtin "fmul_FPIEEE32"(%2 : $Builtin.FPIEEE32, %2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %4 = struct $Float (%3 : $Builtin.FPIEEE32) // CHECK: %5 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %8 = function_ref @$s11$pullback_k12$vjpMultiplyS2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float // CHECK: %9 = partial_apply [callee_guaranteed] %8(%0, %0) : $@convention(thin) (Float, Float, Float) -> Float // CHECK: strong_release %6 : $@callee_guaranteed (Float) -> (Float, Float) // CHECK: %11 = tuple (%4 : $Float, %9 : $@callee_guaranteed (Float) -> Float) // CHECK: return %11 debug_value %0 : $Float, let, name "x", argno 1 // id: %1 %2 = struct_extract %0 : $Float, #Float._value %3 = builtin "fmul_FPIEEE32"(%2 : $Builtin.FPIEEE32, %2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %4 = struct $Float (%3 : $Builtin.FPIEEE32) // function_ref $vjpMultiply %5 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) %6 = partial_apply [callee_guaranteed] %5(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // function_ref $pullback_k %7 = function_ref @$pullback_k : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> Float %8 = partial_apply [callee_guaranteed] %7(%6) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> Float strong_release %6 : $@callee_guaranteed (Float) -> (Float, Float) // id: %9 %10 = tuple (%4 : $Float, %8 : $@callee_guaranteed (Float) -> Float) return %10 : $(Float, @callee_guaranteed (Float) -> Float) // id: %11 } // end sil function '$s4test1kyS2fFTJrSpSr' /////////////////////////////// // Multiple closure callsite // /////////////////////////////// sil @$vjpSin : $@convention(thin) (Float, Float) -> Float sil @$vjpCos : $@convention(thin) (Float, Float) -> Float // pullback of g(_:) sil private @$pullback_g : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float { bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> Float, %2 : $@callee_guaranteed (Float) -> Float, %3 : $@callee_guaranteed (Float) -> (Float, Float)): %4 = apply %3(%0) : $@callee_guaranteed (Float) -> (Float, Float) strong_release %3 : $@callee_guaranteed (Float) -> (Float, Float) // id: %5 %6 = tuple_extract %4 : $(Float, Float), 0 %7 = tuple_extract %4 : $(Float, Float), 1 %8 = apply %2(%7) : $@callee_guaranteed (Float) -> Float strong_release %2 : $@callee_guaranteed (Float) -> Float // id: %9 %10 = apply %1(%6) : $@callee_guaranteed (Float) -> Float strong_release %1 : $@callee_guaranteed (Float) -> Float // id: %11 %12 = struct_extract %8 : $Float, #Float._value %13 = struct_extract %10 : $Float, #Float._value %14 = builtin "fadd_FPIEEE32"(%13 : $Builtin.FPIEEE32, %12 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %15 = struct $Float (%14 : $Builtin.FPIEEE32) debug_value %15 : $Float, let, name "x", argno 1 // id: %16 return %15 : $Float // id: %17 } // reverse-mode derivative of g(_:) sil hidden @$s4test1gyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { bb0(%0 : $Float): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s4test1gyS2fFTJrSpSr // CHECK: PartialApply call site: %16 = partial_apply [callee_guaranteed] %15(%6, %10, %14) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float // CHECK: Passed in closures: // CHECK: 1. %6 = partial_apply [callee_guaranteed] %5(%0) : $@convention(thin) (Float, Float) -> Float // CHECK: 2. %10 = partial_apply [callee_guaranteed] %9(%0) : $@convention(thin) (Float, Float) -> Float // CHECK: 3. %14 = partial_apply [callee_guaranteed] %13(%8, %4) : $@convention(thin) (Float, Float, Float) -> (Float, Float) //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s11$pullback_g7$vjpSinSf0B3CosSf0B8MultiplyS2fTf1nccc_n // CHECK: sil private @$s11$pullback_g7$vjpSinSf0B3CosSf0B8MultiplyS2fTf1nccc_n : $@convention(thin) (Float, Float, Float, Float, Float) -> Float { // CHECK: bb0(%0 : $Float, %1 : $Float, %2 : $Float, %3 : $Float, %4 : $Float): // CHECK: %5 = function_ref @$vjpSin : $@convention(thin) (Float, Float) -> Float // CHECK: %6 = partial_apply [callee_guaranteed] %5(%1) : $@convention(thin) (Float, Float) -> Float // CHECK: %7 = function_ref @$vjpCos : $@convention(thin) (Float, Float) -> Float // CHECK: %8 = partial_apply [callee_guaranteed] %7(%2) : $@convention(thin) (Float, Float) -> Float // CHECK: %9 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %10 = partial_apply [callee_guaranteed] %9(%3, %4) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %11 = apply %10(%0) : $@callee_guaranteed (Float) -> (Float, Float) // CHECK: strong_release %10 : $@callee_guaranteed (Float) -> (Float, Float) // id: %12 // CHECK: %15 = apply %8(%14) : $@callee_guaranteed (Float) -> Float // CHECK: strong_release %8 : $@callee_guaranteed (Float) -> Float // id: %16 // CHECK: %17 = apply %6(%13) : $@callee_guaranteed (Float) -> Float // CHECK: strong_release %6 : $@callee_guaranteed (Float) -> Float // id: %18 // CHECK: return //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s4test1gyS2fFTJrSpSr // CHECK: sil hidden @$s4test1gyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { // CHECK: bb0(%0 : $Float): // CHECK: %2 = struct_extract %0 : $Float, #Float._value // CHECK: %3 = builtin "int_sin_FPIEEE32"(%2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %4 = struct $Float (%3 : $Builtin.FPIEEE32) // CHECK: %5 = function_ref @$vjpSin : $@convention(thin) (Float, Float) -> Float // CHECK: %6 = partial_apply [callee_guaranteed] %5(%0) : $@convention(thin) (Float, Float) -> Float // CHECK: %7 = builtin "int_cos_FPIEEE32"(%2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %8 = struct $Float (%7 : $Builtin.FPIEEE32) // CHECK: %9 = function_ref @$vjpCos : $@convention(thin) (Float, Float) -> Float // CHECK: %10 = partial_apply [callee_guaranteed] %9(%0) : $@convention(thin) (Float, Float) -> Float // CHECK: %11 = builtin "fmul_FPIEEE32"(%3 : $Builtin.FPIEEE32, %7 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %12 = struct $Float (%11 : $Builtin.FPIEEE32) // CHECK: %13 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %14 = partial_apply [callee_guaranteed] %13(%8, %4) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %16 = function_ref @$s11$pullback_g7$vjpSinSf0B3CosSf0B8MultiplyS2fTf1nccc_n : $@convention(thin) (Float, Float, Float, Float, Float) -> Float // CHECK: %17 = partial_apply [callee_guaranteed] %16(%0, %0, %8, %4) : $@convention(thin) (Float, Float, Float, Float, Float) -> Float // CHECK: release_value %6 : $@callee_guaranteed (Float) -> Float // id: %18 // CHECK: release_value %10 : $@callee_guaranteed (Float) -> Float // id: %19 // CHECK: release_value %14 : $@callee_guaranteed (Float) -> (Float, Float) // id: %20 // CHECK: %21 = tuple (%12 : $Float, %17 : $@callee_guaranteed (Float) -> Float) // CHECK: return %21 debug_value %0 : $Float, let, name "x", argno 1 // id: %1 %2 = struct_extract %0 : $Float, #Float._value %3 = builtin "int_sin_FPIEEE32"(%2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %4 = struct $Float (%3 : $Builtin.FPIEEE32) // function_ref closure #1 in _vjpSin(_:) %5 = function_ref @$vjpSin : $@convention(thin) (Float, Float) -> Float %6 = partial_apply [callee_guaranteed] %5(%0) : $@convention(thin) (Float, Float) -> Float %7 = builtin "int_cos_FPIEEE32"(%2 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %8 = struct $Float (%7 : $Builtin.FPIEEE32) // function_ref closure #1 in _vjpCos(_:) %9 = function_ref @$vjpCos : $@convention(thin) (Float, Float) -> Float %10 = partial_apply [callee_guaranteed] %9(%0) : $@convention(thin) (Float, Float) -> Float %11 = builtin "fmul_FPIEEE32"(%3 : $Builtin.FPIEEE32, %7 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 %12 = struct $Float (%11 : $Builtin.FPIEEE32) // function_ref closure #1 in static Float._vjpMultiply(lhs:rhs:) %13 = function_ref @$vjpMultiply : $@convention(thin) (Float, Float, Float) -> (Float, Float) %14 = partial_apply [callee_guaranteed] %13(%8, %4) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // function_ref pullback of g(_:) %15 = function_ref @$pullback_g : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float %16 = partial_apply [callee_guaranteed] %15(%6, %10, %14) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float %17 = tuple (%12 : $Float, %16 : $@callee_guaranteed (Float) -> Float) return %17 : $(Float, @callee_guaranteed (Float) -> Float) // id: %18 } /////////////////////////////// /// Parameter subset thunks /// /////////////////////////////// struct X : Differentiable { @_hasStorage var a: Float { get set } @_hasStorage var b: Double { get set } struct TangentVector : AdditiveArithmetic, Differentiable { @_hasStorage var a: Float { get set } @_hasStorage var b: Double { get set } static func + (lhs: X.TangentVector, rhs: X.TangentVector) -> X.TangentVector static func - (lhs: X.TangentVector, rhs: X.TangentVector) -> X.TangentVector @_implements(Equatable, ==(_:_:)) static func __derived_struct_equals(_ a: X.TangentVector, _ b: X.TangentVector) -> Bool typealias TangentVector = X.TangentVector init(a: Float, b: Double) static var zero: X.TangentVector { get } } init(a: Float, b: Double) mutating func move(by offset: X.TangentVector) } sil [transparent] [thunk] @subset_parameter_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector sil @pullback_f : $@convention(thin) (Float, Double) -> X.TangentVector sil shared @pullback_g : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) -> X.TangentVector { bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> X.TangentVector): %2 = apply %1(%0) : $@callee_guaranteed (Float) -> X.TangentVector strong_release %1 : $@callee_guaranteed (Float) -> X.TangentVector return %2 : $X.TangentVector } sil hidden @$s5test21g1xSfAA1XV_tFTJrSpSr : $@convention(thin) (X) -> (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) { bb0(%0 : $X): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s5test21g1xSfAA1XV_tFTJrSpSr // CHECK: PartialApply call site: %7 = partial_apply [callee_guaranteed] %6(%5) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) -> X.TangentVector // CHECK: Passed in closures: // CHECK: 1. %3 = thin_to_thick_function %2 : $@convention(thin) (Float, Double) -> X.TangentVector to $@callee_guaranteed (Float, Double) -> X.TangentVector //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s10pullback_g0A2_fTf1nc_n // CHECK: sil shared @$s10pullback_g0A2_fTf1nc_n : $@convention(thin) (Float) -> X.TangentVector { // CHECK: bb0(%0 : $Float): // CHECK: %1 = function_ref @pullback_f : $@convention(thin) (Float, Double) -> X.TangentVector // CHECK: %2 = thin_to_thick_function %1 : $@convention(thin) (Float, Double) -> X.TangentVector to $@callee_guaranteed (Float, Double) -> X.TangentVector // CHECK: %3 = function_ref @subset_parameter_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector // CHECK: %4 = partial_apply [callee_guaranteed] %3(%2) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector // CHECK: %5 = apply %4(%0) : $@callee_guaranteed (Float) -> X.TangentVector // CHECK: strong_release %4 : $@callee_guaranteed (Float) -> X.TangentVector // id: %6 // CHECK: return %5 : $X.TangentVector // id: %7 //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s5test21g1xSfAA1XV_tFTJrSpSr // CHECK: sil hidden @$s5test21g1xSfAA1XV_tFTJrSpSr : $@convention(thin) (X) -> (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) { // CHECK: bb0(%0 : $X): // CHECK: %1 = struct_extract %0 : $X, #X.a // CHECK: %2 = function_ref @pullback_f : $@convention(thin) (Float, Double) -> X.TangentVector // CHECK: %3 = thin_to_thick_function %2 : $@convention(thin) (Float, Double) -> X.TangentVector to $@callee_guaranteed (Float, Double) -> X.TangentVector // CHECK: %4 = function_ref @subset_parameter_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector // CHECK: %5 = partial_apply [callee_guaranteed] %4(%3) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector // CHECK: %7 = function_ref @$s10pullback_g0A2_fTf1nc_n : $@convention(thin) (Float) -> X.TangentVector // CHECK: %8 = partial_apply [callee_guaranteed] %7() : $@convention(thin) (Float) -> X.TangentVector // CHECK: release_value %3 : $@callee_guaranteed (Float, Double) -> X.TangentVector // id: %9 // CHECK: %10 = tuple (%1 : $Float, %8 : $@callee_guaranteed (Float) -> X.TangentVector) // CHECK: return %10 %1 = struct_extract %0 : $X, #X.a // function_ref pullback_f %2 = function_ref @pullback_f : $@convention(thin) (Float, Double) -> X.TangentVector %3 = thin_to_thick_function %2 : $@convention(thin) (Float, Double) -> X.TangentVector to $@callee_guaranteed (Float, Double) -> X.TangentVector // function_ref subset_parameter_thunk %4 = function_ref @subset_parameter_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector %5 = partial_apply [callee_guaranteed] %4(%3) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (Float, Double) -> X.TangentVector) -> X.TangentVector // function_ref pullback_g %6 = function_ref @pullback_g : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) -> X.TangentVector %7 = partial_apply [callee_guaranteed] %6(%5) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> X.TangentVector) -> X.TangentVector %8 = tuple (%1 : $Float, %7 : $@callee_guaranteed (Float) -> X.TangentVector) return %8 : $(Float, @callee_guaranteed (Float) -> X.TangentVector) // id: %9 } /////////////////////////////////////////////////////////////////////// ///////// Specialized generic closures - PartialApply Closure ///////// /////////////////////////////////////////////////////////////////////// // closure #1 in static Float._vjpMultiply(lhs:rhs:) sil @$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktSf_SftFZSf_SftSfcfU_ : $@convention(thin) (Float, Float, Float) -> (Float, Float) // thunk for @escaping @callee_guaranteed (@unowned Float) -> (@unowned Float, @unowned Float) sil [transparent] [reabstraction_thunk] @$sS3fIegydd_S3fIegnrr_TR : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) // function_ref specialized pullback of f(a:) sil [transparent] [thunk] @pullback_f_specialized : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // thunk for @escaping @callee_guaranteed (@in_guaranteed Float) -> (@out Float) sil [transparent] [reabstraction_thunk] @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float sil private [signature_optimized_thunk] [always_inline] @pullback_h : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float { bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> Float): %2 = apply %1(%0) : $@callee_guaranteed (Float) -> Float strong_release %1 : $@callee_guaranteed (Float) -> Float return %2 : $Float } // reverse-mode derivative of h(x:) sil hidden @$s5test21h1xS2f_tFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { bb0(%0 : $Float): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s5test21h1xS2f_tFTJrSpSr // CHECK: PartialApply call site: %14 = partial_apply [callee_guaranteed] %13(%11) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float // CHECK: Passed in closures: // CHECK: 1. %4 = partial_apply [callee_guaranteed] %3(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s10pullback_h073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbackti1_j5FZSf_J6SfcfU_S2fTf1nc_n // CHECK: sil private [signature_optimized_thunk] [always_inline] @$s10pullback_h073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbackti1_j5FZSf_J6SfcfU_S2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float { // CHECK: bb0(%0 : $Float, %1 : $Float, %2 : $Float): // CHECK: %3 = function_ref @$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktSf_SftFZSf_SftSfcfU_ : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %4 = partial_apply [callee_guaranteed] %3(%1, %2) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %5 = function_ref @$sS3fIegydd_S3fIegnrr_TR : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) // CHECK: %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) // CHECK: %7 = convert_function %6 : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @out Float) to $@callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for // CHECK: %8 = function_ref @pullback_f_specialized : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // CHECK: %9 = partial_apply [callee_guaranteed] %8(%7) : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // CHECK: %10 = function_ref @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %11 = partial_apply [callee_guaranteed] %10(%9) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %12 = apply %11(%0) : $@callee_guaranteed (Float) -> Float // CHECK: strong_release %11 : $@callee_guaranteed (Float) -> Float // id: %13 // CHECK: return %12 : $Float //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s5test21h1xS2f_tFTJrSpSr // CHECK:sil hidden @$s5test21h1xS2f_tFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { // CHECK:bb0(%0 : $Float): // CHECK: %1 = struct_extract %0 : $Float, #Float._value // CHECK: %2 = builtin "fmul_FPIEEE32"(%1 : $Builtin.FPIEEE32, %1 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // CHECK: %3 = function_ref @$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktSf_SftFZSf_SftSfcfU_ : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %4 = partial_apply [callee_guaranteed] %3(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // CHECK: %5 = function_ref @$sS3fIegydd_S3fIegnrr_TR : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) // CHECK: %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) // CHECK: %7 = convert_function %6 : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @out Float) to $@callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for // CHECK: %8 = function_ref @pullback_f_specialized : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // CHECK: %9 = partial_apply [callee_guaranteed] %8(%7) : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // CHECK: %10 = function_ref @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %11 = partial_apply [callee_guaranteed] %10(%9) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %12 = struct $Float (%2 : $Builtin.FPIEEE32) // CHECK: %14 = function_ref @$s10pullback_h073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbackti1_j5FZSf_J6SfcfU_S2fTf1nc_n : $@convention(thin) (Float, Float, Float) -> Float // CHECK: %15 = partial_apply [callee_guaranteed] %14(%0, %0) : $@convention(thin) (Float, Float, Float) -> Float // CHECK: release_value %4 : $@callee_guaranteed (Float) -> (Float, Float) // id: %16 // CHECK: %17 = tuple (%12 : $Float, %15 : $@callee_guaranteed (Float) -> Float) // CHECK: return %17 %1 = struct_extract %0 : $Float, #Float._value %2 = builtin "fmul_FPIEEE32"(%1 : $Builtin.FPIEEE32, %1 : $Builtin.FPIEEE32) : $Builtin.FPIEEE32 // function_ref closure #1 in static Float._vjpMultiply(lhs:rhs:) %3 = function_ref @$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktSf_SftFZSf_SftSfcfU_ : $@convention(thin) (Float, Float, Float) -> (Float, Float) %4 = partial_apply [callee_guaranteed] %3(%0, %0) : $@convention(thin) (Float, Float, Float) -> (Float, Float) // function_ref thunk for @escaping @callee_guaranteed (@unowned Float) -> (@unowned Float, @unowned Float) %5 = function_ref @$sS3fIegydd_S3fIegnrr_TR : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, Float)) -> (@out Float, @out Float) %7 = convert_function %6 : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @out Float) to $@callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for // function_ref pullback_f_specialized %8 = function_ref @pullback_f_specialized : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float %9 = partial_apply [callee_guaranteed] %8(%7) : $@convention(thin) (@in_guaranteed Float, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1, τ_0_2> (@in_guaranteed τ_0_0) -> (@out τ_0_1, @out τ_0_2) for ) -> @out Float // function_ref thunk for @escaping @callee_guaranteed (@in_guaranteed Float) -> (@out Float) %10 = function_ref @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float %11 = partial_apply [callee_guaranteed] %10(%9) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float %12 = struct $Float (%2 : $Builtin.FPIEEE32) // function_ref pullback_h %13 = function_ref @pullback_h : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float %14 = partial_apply [callee_guaranteed] %13(%11) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float %15 = tuple (%12 : $Float, %14 : $@callee_guaranteed (Float) -> Float) return %15 : $(Float, @callee_guaranteed (Float) -> Float) // id: %16 } ////////////////////////////////////////////////////////////////////////////// ///////// Specialized generic closures - ThinToThickFunction closure ///////// ////////////////////////////////////////////////////////////////////////////// sil [transparent] [thunk] @pullback_y_specialized : $@convention(thin) (@in_guaranteed Float) -> @out Float sil [transparent] [reabstraction_thunk] @reabstraction_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float sil private [signature_optimized_thunk] [always_inline] @pullback_z : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float { bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> Float): %2 = apply %1(%0) : $@callee_guaranteed (Float) -> Float strong_release %1 : $@callee_guaranteed (Float) -> Float return %2 : $Float } sil hidden @$s5test21z1xS2f_tFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { bb0(%0 : $Float): //=========== Test callsite and closure gathering logic ===========// specify_test "closure_specialize_gather_call_sites" // CHECK-LABEL: Specializing closures in function: $s5test21z1xS2f_tFTJrSpSr // CHECK: PartialApply call site: %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float // CHECK: Passed in closures: // CHECK: 1. %2 = thin_to_thick_function %1 : $@convention(thin) (@in_guaranteed Float) -> @out Float to $@callee_guaranteed (@in_guaranteed Float) -> @out Float //=========== Test specialized function signature and body ===========// specify_test "closure_specialize_specialized_function_signature_and_body" // CHECK-LABEL: Generated specialized function: $s10pullback_z0A14_y_specializedTf1nc_n // CHECK: sil private [signature_optimized_thunk] [always_inline] @$s10pullback_z0A14_y_specializedTf1nc_n : $@convention(thin) (Float) -> Float { // CHECK: bb0(%0 : $Float): // CHECK: %1 = function_ref @pullback_y_specialized : $@convention(thin) (@in_guaranteed Float) -> @out Float // CHECK: %2 = thin_to_thick_function %1 : $@convention(thin) (@in_guaranteed Float) -> @out Float to $@callee_guaranteed (@in_guaranteed Float) -> @out Float // CHECK: %3 = function_ref @reabstraction_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %4 = partial_apply [callee_guaranteed] %3(%2) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %5 = apply %4(%0) : $@callee_guaranteed (Float) -> Float // CHECK: strong_release %4 : $@callee_guaranteed (Float) -> Float // id: %6 // CHECK: return %5 : $Float //=========== Test rewritten body ===========// specify_test "closure_specialize_rewritten_caller_body" // CHECK-LABEL: Rewritten caller body for: $s5test21z1xS2f_tFTJrSpSr // CHECK: sil hidden @$s5test21z1xS2f_tFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { // CHECK: bb0(%0 : $Float): // CHECK: %1 = function_ref @pullback_y_specialized : $@convention(thin) (@in_guaranteed Float) -> @out Float // CHECK: %2 = thin_to_thick_function %1 : $@convention(thin) (@in_guaranteed Float) -> @out Float to $@callee_guaranteed (@in_guaranteed Float) -> @out Float // CHECK: %3 = function_ref @reabstraction_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %4 = partial_apply [callee_guaranteed] %3(%2) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // CHECK: %6 = function_ref @$s10pullback_z0A14_y_specializedTf1nc_n : $@convention(thin) (Float) -> Float // CHECK: %7 = partial_apply [callee_guaranteed] %6() : $@convention(thin) (Float) -> Float // CHECK: release_value %2 : $@callee_guaranteed (@in_guaranteed Float) -> @out Float // id: %8 // CHECK: %9 = tuple (%0 : $Float, %7 : $@callee_guaranteed (Float) -> Float) // CHECK: return %9 // function_ref pullback_y_specialized %1 = function_ref @pullback_y_specialized : $@convention(thin) (@in_guaranteed Float) -> @out Float %2 = thin_to_thick_function %1 : $@convention(thin) (@in_guaranteed Float) -> @out Float to $@callee_guaranteed (@in_guaranteed Float) -> @out Float // function_ref reabstraction_thunk %3 = function_ref @reabstraction_thunk : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float %4 = partial_apply [callee_guaranteed] %3(%2) : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float // function_ref pullback_z %5 = function_ref @pullback_z : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Float) -> Float %7 = tuple (%0 : $Float, %6 : $@callee_guaranteed (Float) -> Float) return %7 : $(Float, @callee_guaranteed (Float) -> Float) // id: %8 }