mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
de2dbe57ed
In derivatives of loops, no longer allocate boxes for indirect case payloads. Instead, use a custom pullback context in the runtime which contains a bump-pointer allocator. When a function contains a differentiated loop, the closure context is a `Builtin.NativeObject`, which contains a `swift::AutoDiffLinearMapContext` and a tail-allocated top-level linear map struct (which represents the linear map struct that was previously directly partial-applied into the pullback). In branching trace enums, the payloads of previously indirect cases will be allocated by `swift::AutoDiffLinearMapContext::allocate` and stored as a `Builtin.RawPointer`.
25 lines
1.2 KiB
Swift
25 lines
1.2 KiB
Swift
// RUN: %target-swift-frontend -parse-stdlib %s -emit-ir | %FileCheck %s
|
|
|
|
import Swift
|
|
import _Differentiation
|
|
|
|
struct ExamplePullbackStruct<T: Differentiable> {
|
|
var pb0: (T.TangentVector) -> T.TangentVector
|
|
}
|
|
|
|
@_silgen_name("test_context_builtins")
|
|
func test_context_builtins() {
|
|
let pbStruct = ExamplePullbackStruct<Float>(pb0: { $0 })
|
|
let context = Builtin.autoDiffCreateLinearMapContext(Builtin.sizeof(type(of: pbStruct)))
|
|
let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context)
|
|
UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: pbStruct, as: type(of: pbStruct))
|
|
let newBuffer = Builtin.autoDiffAllocateSubcontext(context, Builtin.sizeof(type(of: pbStruct)))
|
|
UnsafeMutableRawPointer(newBuffer).storeBytes(of: pbStruct, as: type(of: pbStruct))
|
|
}
|
|
|
|
// CHECK-LABEL: define{{.*}}@test_context_builtins()
|
|
// CHECK: entry:
|
|
// CHECK: [[CTX:%.*]] = call swiftcc %swift.refcounted* @swift_autoDiffCreateLinearMapContext({{i[0-9]+}} {{.*}})
|
|
// CEHCK: call swiftcc i8* @swift_autoDiffProjectTopLevelSubcontext(%swift.refcounted* [[CTX]])
|
|
// CHECK: [[BUF:%.*]] = call swiftcc i8* @swift_autoDiffAllocateSubcontext(%swift.refcounted* [[CTX]], {{i[0-9]+}} {{.*}})
|