Files
swift-mirror/test/AutoDiff/IRGen/differentiable_function.sil
Anthony Latsis 55e5618eab [test] Match nocapture to succeed both on main and rebranch
Both the syntax and relative order of the LLVM `nocapture` parameter
attribute changed upstream in 29441e4f5fa5f5c7709f7cf180815ba97f611297.
To reduce conflicts with rebranch, adjust FileCheck patterns to expect
both syntaxes and orders anywhere the presence of the attribute is not
critical to the test. These changes are temporary and will be cleaned
up once rebranch is merged into main.
2025-05-08 23:52:43 +01:00

117 lines
7.7 KiB
Plaintext

// RUN: %target-swift-frontend -emit-ir %s | %FileCheck %s
sil_stage raw
import Swift
import Builtin
import _Differentiation
sil @f : $@convention(thin) (Float) -> Float {
bb0(%0 : $Float):
return undef : $Float
}
sil @f_jvp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {
bb0(%0 : $Float):
return undef : $(Float, @callee_guaranteed (Float) -> Float)
}
sil @f_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {
bb0(%0 : $Float):
return undef : $(Float, @callee_guaranteed (Float) -> Float)
}
sil @test_form_diff_func : $@convention(thin) () -> @owned @differentiable(reverse) @callee_guaranteed (Float) -> Float {
bb0:
%orig = function_ref @f : $@convention(thin) (Float) -> Float
%origThick = thin_to_thick_function %orig : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float
%jvp = function_ref @f_jvp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
%jvpThick = thin_to_thick_function %jvp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
%vjp = function_ref @f_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
%vjpThick = thin_to_thick_function %vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) to $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
%result = differentiable_function [parameters 0] [results 0] %origThick : $@callee_guaranteed (Float) -> Float with_derivative {%jvpThick : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjpThick : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
return %result : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
}
// CHECK-LABEL: define{{.*}}test_form_diff_func(ptr noalias{{( nocapture)?}} sret(<{ %swift.function, %swift.function, %swift.function }>)
// CHECK-SAME: [[OUT:%.*]])
// CHECK: [[OUT_ORIG:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 0
// CHECK: [[OUT_ORIG_FN:%.*]] = getelementptr{{.*}}[[OUT_ORIG]], i32 0, i32 0
// CHECK: store{{.*}}@f{{.*}}[[OUT_ORIG_FN]]
// CHECK: [[OUT_ORIG_DATA:%.*]] = getelementptr{{.*}} [[OUT_ORIG]], i32 0, i32 1
// CHECK: store{{.*}}null{{.*}} [[OUT_ORIG_DATA]]
// CHECK: [[OUT_JVP:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 1
// CHECK: [[OUT_JVP_FN:%.*]] = getelementptr{{.*}}[[OUT_JVP]], i32 0, i32 0
// CHECK: store{{.*}}@f_jvp{{.*}}[[OUT_JVP_FN]]
// CHECK: [[OUT_JVP_DATA:%.*]] = getelementptr{{.*}}[[OUT_JVP]], i32 0, i32 1
// CHECK: store{{.*}} null{{.*}}[[OUT_JVP_DATA]]
// CHECK: [[OUT_VJP:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 2
// CHECK: [[OUT_VJP_FN:%.*]] = getelementptr{{.*}}[[OUT_VJP]], i32 0, i32 0
// CHECK: store{{.*}}@f_vjp{{.*}}[[OUT_VJP_FN]]
// CHECK: [[OUT_VJP_DATA:%.*]] = getelementptr{{.*}}[[OUT_VJP]], i32 0, i32 1
// CHECK: store{{.*}}null{{.*}}[[OUT_VJP_DATA]]
sil @test_extract_components : $@convention(thin) (@guaranteed @differentiable(reverse) @callee_guaranteed (Float) -> Float) -> (@owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), @owned @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)) {
bb0(%0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float):
%orig = differentiable_function_extract [original] %0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
%jvp = differentiable_function_extract [jvp] %0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
%vjp = differentiable_function_extract [vjp] %0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
%result = tuple (%orig : $@callee_guaranteed (Float) -> Float, %jvp : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float))
return %result : $(@callee_guaranteed (Float) -> Float, @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float))
}
// CHECK-LABEL: define{{.*}}@test_extract_components(ptr noalias{{( nocapture)?}} sret(<{ %swift.function, %swift.function, %swift.function }>)
// CHECK-SAME: [[OUT:%.*]], ptr{{.*}}[[IN:%.*]])
// CHECK: [[ORIG:%.*]] = getelementptr{{.*}}[[IN]], i32 0, i32 0
// CHECK: [[ORIG_FN_ADDR:%.*]] = getelementptr{{.*}}[[ORIG]], i32 0, i32 0
// CHECK: [[ORIG_FN:%.*]] = load{{.*}}[[ORIG_FN_ADDR]]
// CHECK: [[ORIG_DATA_ADDR:%.*]] = getelementptr{{.*}}[[ORIG]], i32 0, i32 1
// CHECK: [[ORIG_DATA:%.*]] = load{{.*}}[[ORIG_DATA_ADDR]]
// CHECK: [[JVP:%.*]] = getelementptr{{.*}}[[IN]], i32 0, i32 1
// CHECK: [[JVP_FN_ADDR:%.*]] = getelementptr{{.*}}[[JVP]], i32 0, i32 0
// CHECK: [[JVP_FN:%.*]] = load{{.*}}[[JVP_FN_ADDR]]
// CHECK: [[JVP_DATA_ADDR:%.*]] = getelementptr{{.*}}[[JVP]], i32 0, i32 1
// CHECK: [[JVP_DATA:%.*]] = load{{.*}}[[JVP_DATA_ADDR]]
// CHECK: [[VJP:%.*]] = getelementptr{{.*}}[[IN]], i32 0, i32 2
// CHECK: [[VJP_FN_ADDR:%.*]] = getelementptr{{.*}}[[VJP]], i32 0, i32 0
// CHECK: [[VJP_FN:%.*]] = load{{.*}}[[VJP_FN_ADDR]]
// CHECK: [[VJP_DATA_ADDR:%.*]] = getelementptr{{.*}}[[VJP]], i32 0, i32 1
// CHECK: [[VJP_DATA:%.*]] = load{{.*}}[[VJP_DATA_ADDR]]
// CHECK: [[OUT_0:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 0
// CHECK: [[OUT_0_FN:%.*]] = getelementptr{{.*}}[[OUT_0]], i32 0, i32 0
// CHECK: store{{.*}}[[ORIG_FN]]{{.*}}[[OUT_0_FN]]
// CHECK: [[OUT_0_DATA:%.*]] = getelementptr{{.*}}[[OUT_0]], i32 0, i32 1
// CHECK: store{{.*}}[[ORIG_DATA]]{{.*}}[[OUT_0_DATA]]
// CHECK: [[OUT_1:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 1
// CHECK: [[OUT_1_FN:%.*]] = getelementptr{{.*}}[[OUT_1]], i32 0, i32 0
// CHECK: store{{.*}}[[JVP_FN]]{{.*}}[[OUT_1_FN]]
// CHECK: [[OUT_1_DATA:%.*]] = getelementptr{{.*}}[[OUT_1]], i32 0, i32 1
// CHECK: store{{.*}}[[JVP_DATA]]{{.*}}[[OUT_1_DATA]]
// CHECK: [[OUT_2:%.*]] = getelementptr{{.*}}[[OUT]], i32 0, i32 2
// CHECK: [[OUT_2_FN:%.*]] = getelementptr{{.*}}[[OUT_2]], i32 0, i32 0
// CHECK: store{{.*}}[[VJP_FN]]{{.*}}[[OUT_2_FN]]
// CHECK: [[OUT_2_DATA:%.*]] = getelementptr{{.*}}[[OUT_2]], i32 0, i32 1
// CHECK: store{{.*}}[[VJP_DATA]]{{.*}}[[OUT_2_DATA]]
sil @test_call_diff_fn : $@convention(thin) (@guaranteed @differentiable(reverse) @callee_guaranteed (Float) -> Float, Float) -> Float {
bb0(%0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float, %1 : $Float):
%2 = apply %0(%1) : $@differentiable(reverse) @callee_guaranteed (Float) -> Float
return %2 : $Float
}
// CHECK-LABEL: define{{.*}}@test_call_diff_fn(ptr
// CHECK-SAME: [[DIFF_FN:%.*]], float [[INPUT_FLOAT:%.*]])
// CHECK: [[ORIG:%.*]] = getelementptr{{.*}}[[DIFF_FN]], i32 0, i32 0
// CHECK: [[ORIG_FN_ADDR:%.*]] = getelementptr{{.*}}[[ORIG]], i32 0, i32 0
// CHECK: [[ORIG_FN:%.*]] = load{{.*}}[[ORIG_FN_ADDR]]
// CHECK: [[ORIG_DATA_ADDR:%.*]] = getelementptr{{.*}}[[ORIG]], i32 0, i32 1
// CHECK: [[ORIG_DATA:%.*]] = load{{.*}}[[ORIG_DATA_ADDR]]
// CHECK: [[RESULT:%.*]] = call swiftcc float [[ORIG_FN]](float [[INPUT_FLOAT]], ptr swiftself [[ORIG_DATA]])
// CHECK: ret float [[RESULT]]
sil @test_convert_escape_to_noescape : $@convention(thin) (@guaranteed @differentiable(reverse) @callee_guaranteed (Float) -> Float) -> () {
bb0(%0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float):
%1 = convert_escape_to_noescape %0 : $@differentiable(reverse) @callee_guaranteed (Float) -> Float to $@noescape @differentiable(reverse) @callee_guaranteed (Float) -> Float
return undef : $()
}