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.
115 lines
8.9 KiB
Plaintext
115 lines
8.9 KiB
Plaintext
// Round-trip parsing/printing test.
|
|
|
|
// RUN: %target-sil-opt %s | %FileCheck %s
|
|
|
|
// Round-trip serialization-deserialization test.
|
|
|
|
// RUN: %empty-directory(%t)
|
|
// RUN: %target-sil-opt %s -emit-sib -o %t/tmp.sib -module-name main
|
|
// RUN: %target-sil-opt %t/tmp.sib -o %t/tmp.sil -module-name main
|
|
// NOTE(SR-12090): Workaround because import declarations are not preserved in .sib files.
|
|
// RUN: sed -e 's/import Swift$/import Swift; import _Differentiation/' %t/tmp.sil > %t/tmp_fixed.sil
|
|
// RUN: %target-sil-opt %t/tmp_fixed.sil -module-name main -emit-sorted-sil | %FileCheck %s
|
|
|
|
// IRGen test.
|
|
|
|
// RUN: %target-swift-frontend -emit-ir %s | %FileCheck %s --check-prefix=IRGEN --check-prefix %target-cpu
|
|
// NOTE: `%target-cpu`-specific FileCheck lines exist because lowered function types in LLVM IR differ between architectures.
|
|
|
|
// NOTE(SR-12090): `shell` is required only to run `sed` as a SR-12090 workaround.
|
|
// REQUIRES: shell
|
|
|
|
sil_stage raw
|
|
|
|
import Swift
|
|
import Builtin
|
|
|
|
import _Differentiation
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0] [results 0] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0 1] [results 0 1] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0] [results 0] <T where T : Differentiable> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] <T where T : Differentiable> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] <T where T : Differentiable, T == T.TangentVector> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] <T where T : Differentiable, T: AdditiveArithmetic> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
|
|
sil @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
|
|
sil @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil @genericreq : $@convention(thin) <T : FloatingPoint> (@in_guaranteed T, Float) -> @out T
|
|
|
|
sil @test_derivative_witnesses : $@convention(thin) () -> () {
|
|
bb0:
|
|
%foo_jvp_wrt_0 = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
%foo_vjp_wrt_0_1 = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
|
|
// Test multiple results.
|
|
%bar_jvp_wrt_0_results_0 = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
%bar_vjp_wrt_0_1_results_0_1 = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0 1] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
|
|
// Test generic requirements.
|
|
%generic_jvp_wrt_0 = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] <T : Differentiable> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
%generic_vjp_wrt_0_1 = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] <T : Differentiable & AdditiveArithmetic> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
// Test "dependent" generic requirements: `T == T.TangentVector` depends on `T: Differentiable`.
|
|
%generic_vjp_wrt_0_1_dependent_req = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] <T where T: Differentiable, T == T.TangentVector> @generic : $@convention(thin) <T> (@in_guaranteed T, Float) -> @out T
|
|
|
|
return undef : $()
|
|
}
|
|
|
|
// CHECK-LABEL: sil @test_derivative_witnesses : $@convention(thin) () -> () {
|
|
// CHECK: bb0:
|
|
// CHECK: {{%.*}} = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
// CHECK: {{%.*}} = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] @foo : $@convention(thin) (Float, Float, Float) -> Float
|
|
// CHECK: {{%.*}} = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
// CHECK: {{%.*}} = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0 1] @bar : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
// CHECK: {{%.*}} = differentiability_witness_function [jvp] [reverse] [parameters 0] [results 0] <τ_0_0 where τ_0_0 : Differentiable> @generic : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0
|
|
// CHECK: {{%.*}} = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] <τ_0_0 where τ_0_0 : AdditiveArithmetic, τ_0_0 : Differentiable> @generic : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0
|
|
// CHECK: {{%.*}} = differentiability_witness_function [vjp] [reverse] [parameters 0 1] [results 0] <τ_0_0 where τ_0_0 : Differentiable, τ_0_0 == τ_0_0.TangentVector> @generic : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0
|
|
// CHECK: }
|
|
|
|
// IRGEN: @fooWJrSUUpSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT:[0-9]+]]
|
|
// IRGEN: @fooWJrSSUpSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
// IRGEN: @barWJrSUUpSUr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
// IRGEN: @barWJrSSUpSSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
// IRGEN: @generic16_Differentiation14DifferentiableRzlWJrSUpSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
// IRGEN: @generics18AdditiveArithmeticRz16_Differentiation14DifferentiableRzlWJrSSpSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
// IRGEN: @generic16_Differentiation14DifferentiableRz13TangentVector{{.*}}WJrSSpSr = external global %swift.differentiability_witness, align [[PTR_ALIGNMENT]]
|
|
|
|
// IRGEN-LABEL: define {{.*}} @test_derivative_witnesses()
|
|
|
|
// IRGEN: [[PTR1:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @fooWJrSUUpSr, i32 0, i32 0), align [[PTR_ALIGNMENT]]
|
|
// IRGEN: [[FNPTR1:%.*]] = bitcast i8* [[PTR1]] to { float, i8*, %swift.refcounted* } (float, float, float)*
|
|
|
|
// IRGEN: [[PTR2:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @fooWJrSSUpSr, i32 0, i32 1), align [[PTR_ALIGNMENT]]
|
|
// IRGEN: [[FNPTR2:%.*]] = bitcast i8* [[PTR2]] to { float, i8*, %swift.refcounted* } (float, float, float)*
|
|
|
|
// IRGEN: [[PTR3:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @barWJrSUUpSUr, i32 0, i32 0), align [[PTR_ALIGNMENT]]
|
|
// x86_64: [[FNPTR3:%.*]] = bitcast i8* [[PTR3]] to { float, float, i8*, %swift.refcounted* } (float, float, float)*
|
|
// i386: [[FNPTR3:%.*]] = bitcast i8* [[PTR3]] to void (<{ %TSf, %TSf, %swift.function }>*, float, float, float)*
|
|
|
|
// IRGEN: [[PTR4:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @barWJrSSUpSSr, i32 0, i32 1), align [[PTR_ALIGNMENT]]
|
|
// x86_64: [[FNPTR4:%.*]] = bitcast i8* [[PTR4]] to { float, float, i8*, %swift.refcounted* } (float, float, float)*
|
|
// i386: [[FNPTR4:%.*]] = bitcast i8* [[PTR4]] to void (<{ %TSf, %TSf, %swift.function }>*, float, float, float)*
|
|
|
|
// IRGEN: [[PTR5:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @generic16_Differentiation14DifferentiableRzlWJrSUpSr, i32 0, i32 0), align [[PTR_ALIGNMENT]]
|
|
// IRGEN: [[FNPTR5:%.*]] = bitcast i8* [[PTR5]] to { i8*, %swift.refcounted* } (%swift.opaque*, %swift.opaque*, float, %swift.type*, i8**)*
|
|
|
|
// IRGEN: [[PTR6:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @generics18AdditiveArithmeticRz16_Differentiation14DifferentiableRzlWJrSSpSr, i32 0, i32 1), align [[PTR_ALIGNMENT]]
|
|
// IRGEN: [[FNPTR6:%.*]] = bitcast i8* [[PTR6]] to { i8*, %swift.refcounted* } (%swift.opaque*, %swift.opaque*, float, %swift.type*, i8**, i8**)*
|
|
|
|
// IRGEN: [[PTR7:%.*]] = load i8*, i8** getelementptr inbounds (%swift.differentiability_witness, %swift.differentiability_witness* @generic16_Differentiation14DifferentiableRz13TangentVector{{.*}}WJrSSpSr, i32 0, i32 1), align [[PTR_ALIGNMENT]]
|
|
// IRGEN: [[FNPTR7:%.*]] = bitcast i8* [[PTR7]] to { i8*, %swift.refcounted* } (%swift.opaque*, %swift.opaque*, float, %swift.type*, i8**)*
|