mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
993f7c3ed9
The inlining benefit for VJPs and pullbacks seems to be a bit different on macos and linux. This difference seems to be arising due to certain function calls such as, `sin`, `cos` etc. always being inlined on macos but not on linux. Until I figure out a better way, I'm modifying these tests to avoid matching the exact value of the inlining benefit, which causes them to fail on Linux.
90 lines
3.2 KiB
Swift
90 lines
3.2 KiB
Swift
// VJP inlining tests.
|
|
|
|
// RUN: %target-swift-frontend -emit-sil -O -verify -Xllvm -debug-only=sil-inliner %s 2>&1 | %FileCheck %s
|
|
|
|
// REQUIRES: asserts
|
|
// REQUIRES: swift_in_compiler
|
|
// UNSUPPORTED: OS=windows-msvc
|
|
|
|
import _Differentiation
|
|
#if canImport(Glibc)
|
|
import Glibc
|
|
#else
|
|
import Foundation
|
|
#endif
|
|
|
|
// =============================================================== //
|
|
// VJPs with control-flow are inlined into non-VJP callers
|
|
// =============================================================== //
|
|
@differentiable(reverse)
|
|
func with_control_flow(_ x: Float) -> Float {
|
|
if (x > 0) {
|
|
return sin(x) * cos(x)
|
|
} else {
|
|
return sin(x) + cos(x)
|
|
}
|
|
}
|
|
|
|
@inline(never)
|
|
@_silgen_name("caller_of_with_control_flow")
|
|
func caller_of_with_control_flow(x: Float) -> Float {
|
|
gradient(at: x, of: with_control_flow)
|
|
}
|
|
// CHECK-LABEL: decision {{.*}} $s12vjp_inlining17with_control_flowyS2fFTJrSpSr
|
|
// CHECK-NEXT: "reverse-mode derivative of vjp_inlining.with_control_flow(_:)" inlined into "caller_of_with_control_flow"
|
|
|
|
// =============================================================== //
|
|
// VJPs with control-flow are inlined into VJP callers
|
|
// =============================================================== //
|
|
|
|
@differentiable(reverse)
|
|
func wrapperOnWithControlFlow(x: Float) -> Float {
|
|
return with_control_flow(x)
|
|
}
|
|
// CHECK-LABEL: decision {{.*}} $s12vjp_inlining17with_control_flowyS2fFTJrSpSr
|
|
// CHECK-NEXT: "reverse-mode derivative of vjp_inlining.with_control_flow(_:)" inlined into "reverse-mode derivative of vjp_inlining.wrapperOnWithControlFlow(x:)"
|
|
|
|
// =============================================================== //
|
|
// VJPs without control-flow are not inlined into non-VJP callers
|
|
// =============================================================== //
|
|
@differentiable(reverse)
|
|
func simple(x: Float) -> Float {
|
|
let a = x * x;
|
|
let b = x + x;
|
|
let c = x * a;
|
|
let d = a + b;
|
|
let e = b * c;
|
|
|
|
return a * b / c + d - e ;
|
|
}
|
|
|
|
@inline(never)
|
|
@_silgen_name("caller_of_simple")
|
|
func caller_of_simple(x: Float) -> Float {
|
|
gradient(at: Float(4), of: simple)
|
|
}
|
|
// CHECK-NOT: "reverse-mode derivative of vjp_inlining.simple(x:)" inlined into "caller_of_simple"
|
|
|
|
// ============================================================================ //
|
|
// VJPs without control-flow are not inlined into VJP callers with control-flow
|
|
// ============================================================================ //
|
|
@differentiable(reverse)
|
|
func wrapperWithControlFlowOnSimple(x: Float) -> Float {
|
|
if (x > 0) {
|
|
return simple(x: x) + simple(x: x)
|
|
} else {
|
|
return simple(x: x) * simple(x: x)
|
|
}
|
|
}
|
|
// CHECK-NOT: "reverse-mode derivative of vjp_inlining.simple(x:)" inlined into "reverse-mode derivative of vjp_inlining.wrapperWithControlFlowOnSimple(x:)"
|
|
|
|
// =============================================================== //
|
|
// VJPs without control-flow are inlined into VJP callers
|
|
// =============================================================== //
|
|
@differentiable(reverse)
|
|
func wrapperOnSimple(x: Float) -> Float {
|
|
return simple(x: x)
|
|
}
|
|
// CHECK-LABEL: decision {{.*}} $s12vjp_inlining6simple1xS2f_tFTJrSpSr
|
|
// CHECK-NEXT: "reverse-mode derivative of vjp_inlining.simple(x:)" inlined into "reverse-mode derivative of vjp_inlining.wrapperOnSimple(x:)"
|