| // RUN: %target-sil-opt -differentiation -differentiation-skip-folding-differentiable-function-extraction %s | %FileCheck %s |
| |
| sil_stage raw |
| |
| import Builtin |
| import Swift |
| import SwiftShims |
| |
| protocol DiffReq { |
| @differentiable(wrt: (x)) |
| func f(_ x: Float) -> Float |
| } |
| |
| sil @differentiateWitnessMethod : $@convention(thin) <T where T : DiffReq> (@in_guaranteed T) -> () { |
| bb0(%0 : $*T): |
| %1 = witness_method $T, #DiffReq.f!1 : <Self where Self : DiffReq> (Self) -> (Float) -> Float : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float |
| %2 = differentiable_function [parameters 0] %1 : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float |
| |
| %ret = tuple () |
| return %ret : $() |
| } |
| |
| // CHECK-LABEL: sil @differentiateWitnessMethod |
| // CHECK: [[ORIG_REF:%.*]] = witness_method $T, #DiffReq.f!1 |
| // CHECK: [[JVP_REF:%.*]] = witness_method $T, #DiffReq.f!1.jvp.SU |
| // CHECK: [[VJP_REF:%.*]] = witness_method $T, #DiffReq.f!1.vjp.SU |
| // CHECK: differentiable_function [parameters 0] [[ORIG_REF]] : {{.*}} with_derivative {[[JVP_REF]] : {{.*}}, [[VJP_REF]] : {{.*}}} |
| // CHECK: } // end sil function 'differentiateWitnessMethod' |
| |
| sil @differentiatePartiallyAppliedWitnessMethod : $@convention(thin) <T where T : DiffReq> (@in_guaranteed T) -> () { |
| bb0(%0 : $*T): |
| %1 = witness_method $T, #DiffReq.f!1 : <Self where Self : DiffReq> (Self) -> (Float) -> Float : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float |
| %2 = partial_apply [callee_guaranteed] %1<T>(%0) : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float |
| %3 = differentiable_function [parameters 0] %2 : $@callee_guaranteed (Float) -> Float |
| |
| %ret = tuple () |
| return %ret : $() |
| } |
| |
| // CHECK-LABEL: sil @differentiatePartiallyAppliedWitnessMethod |
| // CHECK: bb0([[ARG:%.*]] : $*T): |
| // CHECK: [[ORIG_REF:%.*]] = witness_method $T, #DiffReq.f!1 |
| // CHECK: [[ARGCOPY1:%.*]] = alloc_stack $T |
| // CHECK: copy_addr [[ARG]] to [initialization] [[ARGCOPY1]] : $*T |
| // CHECK: [[ARGCOPY2:%.*]] = alloc_stack $T |
| // CHECK: copy_addr [[ARG]] to [initialization] [[ARGCOPY2]] : $*T |
| // CHECK: [[ORIG_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[ORIG_REF]]<T>(%0) |
| // CHECK: [[JVP_REF:%.*]] = witness_method $T, #DiffReq.f!1.jvp.SU |
| // CHECK: [[JVP_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[JVP_REF]]<T>([[ARGCOPY1]]) |
| // CHECK: [[VJP_REF:%.*]] = witness_method $T, #DiffReq.f!1.vjp.SU |
| // CHECK: [[VJP_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[VJP_REF]]<T>([[ARGCOPY2]]) |
| // CHECK: dealloc_stack [[ARGCOPY2]] |
| // CHECK: dealloc_stack [[ARGCOPY1]] |
| // CHECK: differentiable_function [parameters 0] [[ORIG_REF_PARTIALLY_APPLIED]] : {{.*}} with_derivative {[[JVP_REF_PARTIALLY_APPLIED]] : {{.*}}, [[VJP_REF_PARTIALLY_APPLIED]] : {{.*}}} |
| // CHECK: } // end sil function 'differentiatePartiallyAppliedWitnessMethod' |