mirror of
https://github.com/apple/swift.git
synced 2025-12-14 20:36:38 +01:00
8990a12bee
Linear maps are captured in vjp routine via callee-guaranteed partial apply and are passed as @owned references to the enclosing pullback that finally consumes them. Necessary retains are inserted by a partial apply forwarder. However, this is not the case when the function being differentiated contains loops as heap-allocated context is used and bare pointer is captured by the pullback partial apply. As a result, partial apply forwarder does not retain the linear maps that are owned by a heap-allocated context, however, they are still treated as @owned references and therefore are released in the pullback after the first call. As a result, subsequent pullback calls release linear maps and we'd end with possible use-after-free. Ensure we retain values when we load values from the context. Reproducible only when: * Loops (so, heap-allocated context) * Pullbacks of thick functions (so context is non-zero) * Multiple pullback calls * Some cleanup while there Fixes #64257
51 lines
1.5 KiB
Swift
51 lines
1.5 KiB
Swift
// RUN: %target-run-simple-swift
|
|
// REQUIRES: executable_test
|
|
|
|
import _Differentiation
|
|
public func simWithPullback(params: MinParams) -> (value: Output, pullback: (Output.TangentVector) -> (MinParams.TangentVector)){
|
|
let simulationValueAndPullback = valueWithPullback(at: params, of: run)
|
|
return (value: simulationValueAndPullback.value, pullback: simulationValueAndPullback.pullback)
|
|
}
|
|
|
|
@differentiable(reverse)
|
|
public func run(params: MinParams) -> Output {
|
|
for t in 0 ... 1 {
|
|
}
|
|
|
|
let res = MiniLoop(other: params._twoDArray).results
|
|
return Output(results: res)
|
|
}
|
|
|
|
struct MiniLoop: Differentiable {
|
|
var results: Float
|
|
var twoDArray: Float
|
|
|
|
@differentiable(reverse)
|
|
init(results: Float = 146, other: Float = 123) {self.results = results; self.twoDArray = other}
|
|
}
|
|
|
|
public struct Output: Differentiable {
|
|
public var results: Float
|
|
@differentiable(reverse)
|
|
public init(results: Float) {self.results = results}
|
|
}
|
|
|
|
public struct MinParams: Differentiable {
|
|
public var _twoDArray: Float
|
|
public init(foo: Float = 42) { _twoDArray = foo }
|
|
}
|
|
|
|
func test() {
|
|
let valueAndPullback = simWithPullback(params: MinParams())
|
|
let output = valueAndPullback.value
|
|
let resultOnes = Float(1.0)
|
|
var grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
|
|
print(grad)
|
|
grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
|
|
print(grad)
|
|
grad = valueAndPullback.pullback(Output.TangentVector(results: resultOnes))
|
|
print(grad)
|
|
}
|
|
|
|
test()
|