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.
44 lines
2.7 KiB
Swift
44 lines
2.7 KiB
Swift
// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s --check-prefix=CHECK-SILGEN
|
|
// RUN: %target-swift-frontend -enable-experimental-forward-mode-differentiation -emit-sil %s | %FileCheck %s --check-prefix=CHECK-SIL
|
|
|
|
// Simple differentiation transform test: check SIL before and after the transform.
|
|
|
|
import _Differentiation
|
|
|
|
@_silgen_name("basic")
|
|
@differentiable(reverse)
|
|
func basic(_ x: Float) -> Float { x }
|
|
|
|
// Test differentiability witnesses.
|
|
|
|
// CHECK-SILGEN-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @basic : $@convention(thin) (Float) -> Float {
|
|
// CHECK-SILGEN-NEXT: }
|
|
|
|
// CHECK-SIL-LABEL: sil_differentiability_witness hidden [reverse] [parameters 0] [results 0] @basic : $@convention(thin) (Float) -> Float {
|
|
// CHECK-SIL-NEXT: jvp: @basicTJfSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK-SIL-NEXT: vjp: @basicTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK-SIL-NEXT: }
|
|
|
|
// Test `differentiable_function` instructions.
|
|
|
|
@_silgen_name("test_differentiable_function")
|
|
func testDifferentiableFunction() {
|
|
let _: @differentiable(reverse) (Float) -> Float = basic
|
|
}
|
|
|
|
// CHECK-SILGEN-LABEL: sil hidden [ossa] @test_differentiable_function : $@convention(thin) () -> () {
|
|
// CHECK-SILGEN: [[ORIG_FN_REF:%.*]] = function_ref @basic : $@convention(thin) (Float) -> Float
|
|
// CHECK-SILGEN: [[ORIG_FN:%.*]] = thin_to_thick_function [[ORIG_FN_REF]] : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float
|
|
// CHECK-SILGEN: [[DIFF_FN:%.*]] = differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : $@callee_guaranteed (Float) -> Float
|
|
// CHECK-SILGEN: }
|
|
|
|
// CHECK-SIL-LABEL: sil hidden @test_differentiable_function : $@convention(thin) () -> () {
|
|
// CHECK-SIL: [[ORIG_FN_REF:%.*]] = function_ref @basic : $@convention(thin) (Float) -> Float
|
|
// CHECK-SIL: [[ORIG_FN:%.*]] = thin_to_thick_function [[ORIG_FN_REF]]
|
|
// CHECK-SIL: [[JVP_FN_REF:%.*]] = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @basic
|
|
// CHECK-SIL: [[JVP_FN:%.*]] = thin_to_thick_function [[JVP_FN_REF]]
|
|
// CHECK-SIL: [[VJP_FN_REF:%.*]] = differentiability_witness_function [vjp] [reverse] [parameters 0] [results 0] @basic
|
|
// CHECK-SIL: [[VJP_FN:%.*]] = thin_to_thick_function [[VJP_FN_REF]]
|
|
// CHECK-SIL: [[DIFF_FN:%.*]] = 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-SIL: }
|