mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
e494df2ee6
Differentiability witnesses are now keyed by the original function name, the differentiability kind, and the autodiff config.
Updated SIL syntax:
```
differentiability-kind ::= 'forward' | 'reverse' | 'normal' | 'linear'
sil-differentiability-witness ::=
'sil_differentiability_witness'
sil-linkage?
'[' differentiability-kind ']'
'[' 'parameters' sil-differentiability-witness-function-index-list ']'
'[' 'results' sil-differentiability-witness-function-index-list ']'
generic-parameter-clause?
sil-function-name ':' sil-type
sil-differentiability-witness-body?
sil-instruction ::=
'differentiability_witness_function'
'[' sil-differentiability-witness-function-kind ']'
'[' differentiability-kind ']'
'[' 'parameters' sil-differentiability-witness-function-index-list ']'
'[' 'results' sil-differentiability-witness-function-index-list ']'
generic-parameter-clause?
sil-function-name ':' sil-type
```
```console
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] <T where T: Differentiable> @foo : <T> $(T) -> T
differentiability_witness_function [vjp] [reverse] [parameters 0] [results 0] <T where T: Differentiable> @foo : $(T) -> T
```
New mangling:
```swift
global ::= global generic-signature? 'WJ' DIFFERENTIABILITY-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' // differentiability witness
```
```console
$s13test_mangling3fooyS2f_S2ftFWJrSpSr ---> reverse differentiability witness for test_mangling.foo(Swift.Float, Swift.Float, Swift.Float) -> Swift.Float with respect to parameters {0} and results {0}
```
Resolves rdar://74380324.
142 lines
7.6 KiB
Plaintext
142 lines
7.6 KiB
Plaintext
// RUN: %target-sil-opt -differentiation %s | %FileCheck %s
|
|
|
|
sil_stage raw
|
|
|
|
import _Differentiation
|
|
import Swift
|
|
|
|
// Test `function_ref` instructions.
|
|
|
|
sil @basic : $@convention(thin) (Float) -> Float {
|
|
bb0(%0 : $Float):
|
|
return %0 : $Float
|
|
}
|
|
|
|
sil @test_function_ref : $@convention(thin) () -> () {
|
|
bb0:
|
|
%1 = function_ref @basic : $@convention(thin) (Float) -> Float
|
|
%2 = thin_to_thick_function %1 : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float
|
|
%3 = differentiable_function [parameters 0] [results 0] %2 : $@callee_guaranteed (Float) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_function_ref
|
|
// CHECK: [[ORIG_FN_REF:%.*]] = function_ref @basic : $@convention(thin) (Float) -> Float
|
|
// CHECK: [[ORIG_FN:%.*]] = thin_to_thick_function [[ORIG_FN_REF]]
|
|
// CHECK: [[JVP_FN_REF:%.*]] = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @basic
|
|
// CHECK: [[JVP_FN:%.*]] = thin_to_thick_function [[JVP_FN_REF]]
|
|
// CHECK: [[VJP_FN_REF:%.*]] = differentiability_witness_function [vjp] [reverse] [parameters 0] [results 0] @basic
|
|
// CHECK: [[VJP_FN:%.*]] = thin_to_thick_function [[VJP_FN_REF]]
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : $@callee_guaranteed (Float) -> Float with_derivative {[[JVP_FN]] : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[VJP_FN]] : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
// CHECK: }
|
|
|
|
// Test `differentiable_function_extract` instructions.
|
|
|
|
sil @test_differentiable_function_extract : $@convention(thin) (@differentiable(reverse) @callee_guaranteed (Float) -> Float) -> () {
|
|
bb0(%0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float):
|
|
%1 = differentiable_function_extract [original] %0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
|
|
%2 = differentiable_function [parameters 0] [results 0] %1 : $@callee_guaranteed (Float) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_differentiable_function_extract
|
|
// CHECK: bb0([[ARG:%.*]] : $@differentiable(reverse) @callee_guaranteed (Float) -> Float):
|
|
// CHECK: [[ORIG_FN:%.*]] = differentiable_function_extract [original] [[ARG]]
|
|
// CHECK: [[JVP_FN:%.*]] = differentiable_function_extract [jvp] [[ARG]]
|
|
// CHECK: [[VJP_FN:%.*]] = differentiable_function_extract [vjp] [[ARG]]
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : {{.*}} with_derivative {[[JVP_FN]] : {{.*}}, [[VJP_FN]] : {{.*}}}
|
|
// CHECK: }
|
|
|
|
// Test `witness_method` instructions.
|
|
|
|
protocol Protocol {
|
|
@differentiable(reverse, wrt: x)
|
|
func method(_ x: Float) -> Float
|
|
}
|
|
|
|
sil @test_witness_method : $@convention(thin) <T where T : Protocol> (@in_guaranteed T) -> () {
|
|
bb0(%0 : $*T):
|
|
%1 = witness_method $T, #Protocol.method : <Self where Self : Protocol> (Self) -> (Float) -> Float : $@convention(witness_method: Protocol) <τ_0_0 where τ_0_0 : Protocol> (Float, @in_guaranteed τ_0_0) -> Float
|
|
%2 = differentiable_function [parameters 0] [results 0] %1 : $@convention(witness_method: Protocol) <τ_0_0 where τ_0_0 : Protocol> (Float, @in_guaranteed τ_0_0) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_witness_method
|
|
// CHECK: [[ORIG_FN:%.*]] = witness_method $T, #Protocol.method
|
|
// CHECK: [[JVP_FN:%.*]] = witness_method $T, #Protocol.method!jvp.SU
|
|
// CHECK: [[VJP_FN:%.*]] = witness_method $T, #Protocol.method!vjp.SU
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : {{.*}} with_derivative {[[JVP_FN]] : {{.*}}, [[VJP_FN]] : {{.*}}}
|
|
// CHECK: }
|
|
|
|
sil @test_witness_method_partial_apply : $@convention(thin) <T where T : Protocol> (@in_guaranteed T) -> () {
|
|
bb0(%0 : $*T):
|
|
%1 = witness_method $T, #Protocol.method : <Self where Self : Protocol> (Self) -> (Float) -> Float : $@convention(witness_method: Protocol) <τ_0_0 where τ_0_0 : Protocol> (Float, @in_guaranteed τ_0_0) -> Float
|
|
%2 = partial_apply [callee_guaranteed] %1<T>(%0) : $@convention(witness_method: Protocol) <τ_0_0 where τ_0_0 : Protocol> (Float, @in_guaranteed τ_0_0) -> Float
|
|
%3 = differentiable_function [parameters 0] [results 0] %2 : $@callee_guaranteed (Float) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_witness_method_partial_apply
|
|
// CHECK: bb0([[ARG:%.*]] : $*T):
|
|
// CHECK: [[ORIG_FN:%.*]] = witness_method $T, #Protocol.method
|
|
// 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_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[ORIG_FN]]<T>([[ARG]])
|
|
// CHECK: [[JVP_FN:%.*]] = witness_method $T, #Protocol.method!jvp.SU
|
|
// CHECK: [[JVP_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[JVP_FN]]<T>([[ARGCOPY1]])
|
|
// CHECK: [[VJP_FN:%.*]] = witness_method $T, #Protocol.method!vjp.SU
|
|
// CHECK: [[VJP_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[VJP_FN]]<T>([[ARGCOPY2]])
|
|
// CHECK: dealloc_stack [[ARGCOPY2]]
|
|
// CHECK: dealloc_stack [[ARGCOPY1]]
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN_PARTIALLY_APPLIED]] : {{.*}} with_derivative {[[JVP_FN_PARTIALLY_APPLIED]] : {{.*}}, [[VJP_FN_PARTIALLY_APPLIED]] : {{.*}}}
|
|
// CHECK: }
|
|
|
|
// Test `class_method` instructions.
|
|
|
|
class Class {
|
|
@differentiable(reverse, wrt: x)
|
|
func method(_ x: Float) -> Float
|
|
}
|
|
|
|
sil @test_class_method : $@convention(thin) <T where T : Class> (@guaranteed Class) -> () {
|
|
bb0(%0 : $Class):
|
|
%1 = class_method %0 : $Class, #Class.method : (Class) -> (Float) -> Float, $@convention(method) (Float, @guaranteed Class) -> Float
|
|
%2 = differentiable_function [parameters 0] [results 0] %1 : $@convention(method) (Float, @guaranteed Class) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_class_method
|
|
// CHECK: bb0([[ARG:%.*]] : $Class):
|
|
// CHECK: [[ORIG_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method
|
|
// CHECK: [[JVP_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method!jvp.SU
|
|
// CHECK: [[VJP_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method!vjp.SU
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : {{.*}} with_derivative {[[JVP_FN]] : {{.*}}, [[VJP_FN]] : {{.*}}}
|
|
// CHECK: }
|
|
|
|
sil @test_class_method_partial_apply : $@convention(thin) <T where T : Class> (@guaranteed Class) -> () {
|
|
bb0(%0 : $Class):
|
|
%1 = class_method %0 : $Class, #Class.method : (Class) -> (Float) -> Float, $@convention(method) (Float, @guaranteed Class) -> Float
|
|
%2 = partial_apply [callee_guaranteed] %1(%0) : $@convention(method) (Float, @guaranteed Class) -> Float
|
|
%3 = differentiable_function [parameters 0] [results 0] %2 : $@callee_guaranteed (Float) -> Float
|
|
%void = tuple ()
|
|
return %void : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_class_method_partial_apply
|
|
// CHECK: bb0([[ARG:%.*]] : $Class):
|
|
// CHECK: [[ORIG_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method
|
|
// CHECK: [[ORIG_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[ORIG_FN]]([[ARG]])
|
|
// CHECK: [[JVP_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method!jvp.SU
|
|
// CHECK: [[JVP_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[JVP_FN]]([[ARG]])
|
|
// CHECK: [[VJP_FN:%.*]] = class_method [[ARG]] : $Class, #Class.method!vjp.SU
|
|
// CHECK: [[VJP_FN_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[VJP_FN]]([[ARG]])
|
|
// CHECK: differentiable_function [parameters 0] [results 0] [[ORIG_FN_PARTIALLY_APPLIED]] : {{.*}} with_derivative {[[JVP_FN_PARTIALLY_APPLIED]] : {{.*}}, [[VJP_FN_PARTIALLY_APPLIED]] : {{.*}}}
|
|
// CHECK: }
|