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
The changes are intentionally were made close to the original implementation w/o possible simplifications to ease the review
Fixes#63207, supersedes #63379 (and fixes#63234)
We had two notions of canonical types, one is the structural property
where it doesn't contain sugared types, the other one where it does
not contain reducible type parameters with respect to a generic
signature.
Rename the second one to a 'reduced type'.
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`.
Move differentiation-related SILOptimizer files to
{include/swift,lib}/SILOptimizer/Differentiation/.
This reduces directory nesting and gathers files together.
