Anton Korobeynikov 8990a12bee Fix use after free when pullback is used multiple times. (#64647) ...

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

2023-03-28 01:36:38 -07:00

8.2 KiB

Raw Blame History

View Raw