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.
48 lines
4.8 KiB
Swift
48 lines
4.8 KiB
Swift
// RUN: %target-swift-frontend -emit-sil -enable-experimental-forward-mode-differentiation -module-name=mangling -verify %s | %FileCheck %s
|
|
|
|
import _Differentiation
|
|
|
|
@differentiable(reverse)
|
|
func foo(_ x: Float) -> Float { x }
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @$s8mangling3fooyS2fF : $@convention(thin) (Float) -> Float {
|
|
// CHECK: jvp: @$s8mangling3fooyS2fFTJfSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK: vjp: @$s8mangling3fooyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
|
|
@differentiable(reverse)
|
|
func generic<T: Differentiable>(_ x: T) -> T { x }
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] <T where T : Differentiable> @$s8mangling7genericyxx16_Differentiation14DifferentiableRzlF : $@convention(thin) <T where T : Differentiable> (@in_guaranteed T) -> @out T {
|
|
// CHECK: jvp: @$s8mangling7genericyxx16_Differentiation14DifferentiableRzlFAcDRzlTJfSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable> (@in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <τ_0_0.TangentVector, τ_0_0.TangentVector>)
|
|
// CHECK: vjp: @$s8mangling7genericyxx16_Differentiation14DifferentiableRzlFAcDRzlTJrSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable> (@in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <τ_0_0.TangentVector, τ_0_0.TangentVector>)
|
|
|
|
@differentiable(reverse where T: Differentiable)
|
|
func genericConstrained<T>(_ x: T) -> T { x }
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] <T where T : Differentiable> @$s8mangling18genericConstrainedyxxlF : $@convention(thin) <T> (@in_guaranteed T) -> @out T {
|
|
// CHECK: jvp: @$s8mangling18genericConstrainedyxxlF16_Differentiation14DifferentiableRzlTJfSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable> (@in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <τ_0_0.TangentVector, τ_0_0.TangentVector>)
|
|
// CHECK: vjp: @$s8mangling18genericConstrainedyxxlF16_Differentiation14DifferentiableRzlTJrSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable> (@in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <τ_0_0.TangentVector, τ_0_0.TangentVector>)
|
|
|
|
class Class: Differentiable {
|
|
@differentiable(reverse, wrt: (self, x))
|
|
@differentiable(reverse, wrt: x)
|
|
func f(_ x: Float) -> Float { x }
|
|
}
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @$s8mangling5ClassC1fyS2fF : $@convention(method) (Float, @guaranteed Class) -> Float {
|
|
// CHECK: jvp: @$s8mangling5ClassC1fyS2fFTJfSUpSr : $@convention(method) (Float, @guaranteed Class) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK: vjp: @$s8mangling5ClassC1fyS2fFTJrSUpSr : $@convention(method) (Float, @guaranteed Class) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
|
|
func test<C: Class>(_ c: C, _ x: Float) {
|
|
_ = gradient(at: c, x) { c, x in c.f(x) }
|
|
_ = gradient(at: x) { x in c.f(x) }
|
|
}
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness private [reverse] [parameters 0] [results 0] <τ_0_0 where τ_0_0 : Class> @$s8mangling4testyyx_SftAA5ClassCRbzlFS2fcfU0_ : $@convention(thin) <C where C : Class> (Float, @guaranteed C) -> Float {
|
|
// CHECK: jvp: @$s8mangling4testyyx_SftAA5ClassCRbzlFS2fcfU0_ADRbzlTJfSUpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Class> (Float, @guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK: vjp: @$s8mangling4testyyx_SftAA5ClassCRbzlFS2fcfU0_ADRbzlTJrSUpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Class> (Float, @guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
|
|
// CHECK-LABEL: sil_differentiability_witness private [reverse] [parameters 0 1] [results 0] <τ_0_0 where τ_0_0 : Class> @$s8mangling4testyyx_SftAA5ClassCRbzlFSfx_SftcfU_ : $@convention(thin) <C where C : Class> (@guaranteed C, Float) -> Float {
|
|
// CHECK: jvp: @$s8mangling4testyyx_SftAA5ClassCRbzlFSfx_SftcfU_ADRbzlTJfSSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Class> (@guaranteed τ_0_0, Float) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (@guaranteed τ_0_0, Float) -> Float for <τ_0_0.TangentVector>)
|
|
// CHECK: vjp: @$s8mangling4testyyx_SftAA5ClassCRbzlFSfx_SftcfU_ADRbzlTJrSSpSr : $@convention(thin) <τ_0_0 where τ_0_0 : Class> (@guaranteed τ_0_0, Float) -> (Float, @owned @callee_guaranteed @substituted <τ_0_0> (Float) -> (τ_0_0, Float) for <τ_0_0.TangentVector>)
|