mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
af8942d940
Compiler: - Add `Forward` and `Reverse` to `DifferentiabilityKind`. - Expand `DifferentiabilityMask` in `ExtInfo` to 3 bits so that it now holds all 4 cases of `DifferentiabilityKind`. - Parse `@differentiable(reverse)` and `@differentiable(_forward)` declaration attributes and type attributes. - Emit a warning for `@differentiable` without `reverse`. - Emit an error for `@differentiable(_forward)`. - Rename `@differentiable(linear)` to `@differentiable(_linear)`. - Make `@differentiable(reverse)` type lowering go through today's `@differentiable` code path. We will specialize it to reverse-mode in a follow-up patch. ABI: - Add `Forward` and `Reverse` to `FunctionMetadataDifferentiabilityKind`. - Extend `TargetFunctionTypeFlags` by 1 bit to store the highest bit of differentiability kind (linear). Note that there is a 2-bit gap in `DifferentiabilityMask` which is reserved for `AsyncMask` and `ConcurrentMask`; `AsyncMask` is ABI-stable so we cannot change that. _Differentiation module: - Replace all occurrences of `@differentiable` with `@differentiable(reverse)`. - Delete `_transpose(of:)`. Resolves rdar://69980056.
97 lines
7.8 KiB
Plaintext
97 lines
7.8 KiB
Plaintext
// Round-trip parsing/printing test.
|
|
|
|
// RUN: %target-sil-opt %s -emit-sorted-sil | %FileCheck %s --check-prefix=CHECK-SIL
|
|
|
|
// 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 --check-prefix=CHECK-SIL
|
|
|
|
// IRGen test.
|
|
|
|
// RUN: %target-swift-frontend -emit-ir %s | %FileCheck %s --check-prefix=CHECK-IRGEN
|
|
|
|
// NOTE(SR-12090): `shell` is required only to run `sed` as a SR-12090 workaround.
|
|
// REQUIRES: shell
|
|
|
|
// Would need to update for arm64e.
|
|
// UNSUPPORTED: CPU=arm64e
|
|
|
|
sil_stage raw
|
|
|
|
import Swift
|
|
import Builtin
|
|
|
|
import _Differentiation
|
|
|
|
sil @function : $@convention(thin) (Float) -> Float
|
|
sil @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
|
|
sil @foo : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
sil @foo_vjp : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)
|
|
|
|
sil @make_differentiable_func : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float) {
|
|
// sil @make_differentiable_func : $@convention(thin) () -> () {
|
|
bb0:
|
|
%orig_fn = function_ref @foo : $@convention(thin) (Float, Float, Float) -> (Float, Float)
|
|
%vjp_fn = function_ref @foo_vjp : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)
|
|
%diff_fn = differentiable_function [parameters 0 1] [results 0] %orig_fn : $@convention(thin) (Float, Float, Float) -> (Float, Float) with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float), %vjp_fn : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
%extracted_vjp = differentiable_function_extract [vjp] %diff_fn : $@differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float)
|
|
return %diff_fn : $@differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float)
|
|
}
|
|
|
|
sil @make_differentiable_function : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float) -> Float {
|
|
bb0:
|
|
%orig_fn = function_ref @function : $@convention(thin) (Float) -> Float
|
|
%vjp_fn = function_ref @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
%diff_fn = differentiable_function [parameters 0] [results 0] %orig_fn : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp_fn : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
%extracted_vjp = differentiable_function_extract [vjp] %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
%extracted_original = differentiable_function_extract [original] %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
return %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
}
|
|
|
|
// CHECK-SIL-LABEL: @make_differentiable_function : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float) -> Float {
|
|
// CHECK-SIL: [[ORIG_FN:%.*]] = function_ref @function : $@convention(thin) (Float) -> Float
|
|
// CHECK-SIL: [[VJP_FN:%.*]] = function_ref @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
|
|
// CHECK-SIL: [[DIFF_FN:%.*]] = differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[VJP_FN]] : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
// CHECK-SIL: [[EXTRACTED_VJP_FN:%.*]] = differentiable_function_extract [vjp] [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
// CHECK-SIL: [[EXTRACTED_ORIG_FN:%.*]] = differentiable_function_extract [original] [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
// CHECK-SIL: return [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float
|
|
|
|
// CHECK-IRGEN-LABEL: define {{.*}}swiftcc { i8*, i8*, i8* } @make_differentiable_function()
|
|
// CHECK-IRGEN-NEXT: entry:
|
|
// CHECK-IRGEN-NEXT: ret { i8*, i8*, i8* } { i8* bitcast (float (float)* @function to i8*), i8* undef, i8* bitcast ({ float, i8*, %swift.refcounted* } (float)* @function_vjp to i8*) }
|
|
|
|
sil @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
|
|
sil @examplemethod : $@convention(method) (Float, Float, Float) -> Float
|
|
|
|
// CHECK-SIL-LABEL: sil @test_roundtrip_parse
|
|
sil @test_roundtrip_parse : $@convention(thin) () -> () {
|
|
bb0:
|
|
%0 = function_ref @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
|
|
%1 = differentiable_function [parameters 0 1 2] [results 0] %0 : $@convention(thin) (Float, Float, Float) -> Float with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float, Float, Float) -> Float), undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float, Float))}
|
|
|
|
// CHECK-SIL: %2 = differentiable_function_extract [vjp] %1 : $@differentiable(reverse) @convention(thin) (Float, Float, Float) -> Float
|
|
%2 = differentiable_function_extract [vjp] %1 : $@differentiable(reverse) @convention(thin) (Float, Float, Float) -> Float
|
|
%3 = differentiable_function [parameters 0] [results 0] %0 : $@convention(thin) (Float, Float, Float) -> Float with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
|
|
// CHECK-SIL: %4 = differentiable_function_extract [vjp] %3 : $@differentiable(reverse) @convention(thin) (Float, @noDerivative Float, @noDerivative Float) -> Float
|
|
%4 = differentiable_function_extract [vjp] %3 : $@differentiable(reverse) @convention(thin) (Float, @noDerivative Float, @noDerivative Float) -> Float
|
|
%5 = function_ref @examplemethod : $@convention(method) (Float, Float, Float) -> Float
|
|
%6 = differentiable_function [parameters 0 1 2] [results 0] %5 : $@convention(method) (Float, Float, Float) -> Float with_derivative {undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float, Float, Float) -> Float), undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float, Float))}
|
|
|
|
// CHECK-SIL: %7 = differentiable_function_extract [vjp] %6 : $@differentiable(reverse) @convention(method) (Float, Float, Float) -> Float
|
|
%7 = differentiable_function_extract [vjp] %6 : $@differentiable(reverse) @convention(method) (Float, Float, Float) -> Float
|
|
%8 = differentiable_function [parameters 0] [results 0] %5 : $@convention(method) (Float, Float, Float) -> Float with_derivative {undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
|
|
|
|
// CHECK-SIL: %9 = differentiable_function_extract [vjp] %8 : $@differentiable(reverse) @convention(method) (Float, @noDerivative Float, @noDerivative Float) -> Float
|
|
%9 = differentiable_function_extract [vjp] %8 : $@differentiable(reverse) @convention(method) (Float, @noDerivative Float, @noDerivative Float) -> Float
|
|
|
|
%ret = tuple ()
|
|
return %ret : $()
|
|
}
|