We are going to need to add more flags to the various checked cast
instructions. Generalize the CastingIsolatedConformances bit in all of
these SIL instructions to an "options" struct that's easier to extend.
Precursor to rdar://152335805.
When performing a dynamic cast to an existential type that satisfies
(Metatype)Sendable, it is unsafe to allow isolated conformances of any
kind to satisfy protocol requirements for the existential. Identify
these cases and mark the corresponding cast instructions with a new flag,
`[prohibit_isolated_conformances]` that will be used to indicate to the
runtime that isolated conformances need to be rejected.
Since `move_value` is a destroying operation, the adjoint of `y = move_value x` should be `adj[x] += adj[y]; adj[y] = 0` instead of just `adj[x] += adj[y]`.
As autodiff happens on function types it is not in general possible to determine the real expansion context of the function being differentiated. Use of minimal context is a conservative approach that should work even when libraty evolution mode is enabled.
Fixes#55179
Although I don't plan to bring over new assertions wholesale
into the current qualification branch, it's entirely possible
that various minor changes in main will use the new assertions;
having this basic support in the release branch will simplify that.
(This is why I'm adding the includes as a separate pass from
rewriting the individual assertions)
This PR implements first set of changes required to support autodiff for coroutines. It mostly targeted to `_modify` accessors in standard library (and beyond), but overall implementation is quite generic.
There are some specifics of implementation and known limitations:
- Only `@yield_once` coroutines are naturally supported
- VJP is a coroutine itself: it yields the results *and* returns a pullback closure as a normal return. This allows us to capture values produced in resume part of a coroutine (this is required for defers and other cleanups / commits)
- Pullback is a coroutine, we assume that coroutine cannot abort and therefore we execute the original coroutine in reverse from return via yield and then back to the entry
- It seems there is no semantically sane way to support `_read` coroutines (as we will need to "accept" adjoints via yields), therefore only coroutines with inout yields are supported (`_modify` accessors). Pullbacks of such coroutines take adjoint buffer as input argument, yield this buffer (to accumulate adjoint values in the caller) and finally return the adjoints indirectly.
- Coroutines (as opposed to normal functions) are not first-class values: there is no AST type for them, one cannot e.g. store them into tuples, etc. So, everywhere where AST type is required, we have to hack around.
- As there is no AST type for coroutines, there is no way one could register custom derivative for coroutines. So far only compiler-produced derivatives are supported
- There are lots of common things wrt normal function apply's, but still there are subtle but important differences. I tried to organize the code to enable code reuse, still it was not always possible, so some code duplication could be seen
- The order of how pullback closures are produced in VJP is a bit different: for normal apply's VJP produces both value and pullback closure via a single nested VJP apply. This is not so anymore with coroutine VJP's: yielded values are produced at `begin_apply` site and pullback closure is available only from `end_apply`, so we need to track the order in which pullbacks are produced (and arrange consumption of the values accordingly – effectively delay them)
- On the way some complementary changes were required in e.g. mangler / demangler
This patch covers the generation of derivatives up to SIL level, however, it is not enough as codegen of `partial_apply` of a coroutine is completely broken. The fix for this will be submitted separately as it is not directly autodiff-related.
---------
Co-authored-by: Andrew Savonichev <andrew.savonichev@gmail.com>
Co-authored-by: Richard Wei <rxwei@apple.com>
Unreachable blocks possess some challenges to autodiff since in reverse pass (pullback generation) we need to execute the function backwards, pushing the values from return BB back to entry block. As a result, unreachable blocks might become reachable from the return BB and this might cause all kind of issues.
Optional's `init_enum_data_addr` and `inject_enum_addr` instructions are generated in presence of non-loadable Optional values. The compiler used to treat these instructions as inactive, and this resulted in silent run-time
issues described in #64223.
The patch marks `init_enum_data_addr` as "active" if its Optional operand is also active, and in PullbackCloner we differentiate through it and the related `inject_enum_addr`.
However, we only determine this relation in simple cases when both instructions are in the same block. There is no def-use relation between them (both take the same Optional operand), so if there is more than one set of instructions
operating on the same Optional, or there is some control flow, we currently bail out.
In PullbackCloner, we walk over instructions in reverse order and start from `inject_enum_addr` and its `Optional<Wrapped>.TangentVector` operand. Assuming that is is already initialized, we emit an `unchecked_take_enum_data_addr` and set it as the adjoint buffer of `init_enum_data_addr`. The Optional value is
invalidated, and we have to destroy the enum data address later when we reach `init_enum_data_addr`.
It lowers let property accesses of classes.
Lowering consists of two tasks:
* In class initializers, insert `end_init_let_ref` instructions at places where all let-fields are initialized.
This strictly separates the life-range of the class into a region where let fields are still written during
initialization and a region where let fields are truly immutable.
* Add the `[immutable]` flag to all `ref_element_addr` instructions (for let-fields) which are in the "immutable"
region. This includes the region after an inserted `end_init_let_ref` in an class initializer, but also all
let-field accesses in other functions than the initializer and the destructor.
This pass should run after DefiniteInitialization but before RawSILInstLowering (because it relies on `mark_uninitialized` still present in the class initializer).
Note that it's not mandatory to run this pass. If it doesn't run, SIL is still correct.
Simplified example (after lowering):
bb0(%0 : @owned C): // = self of the class initializer
%1 = mark_uninitialized %0
%2 = ref_element_addr %1, #C.l // a let-field
store %init_value to %2
%3 = end_init_let_ref %1 // inserted by lowering
%4 = ref_element_addr [immutable] %3, #C.l // set to immutable by lowering
%5 = load %4
For some values we cannot compute types for differentiation (for example,
tangent vector type), so it is better to diagnose them earlier. Otherwise we hit
assertions when generating code for such invalid values.
The LIT test is a reduced reproducer from the issue #66996. Before the patch the
compiler crashed while trying to get a tangent vector type for the following
value (partial_apply):
%54 = function_ref @$s4null1o2ffSdAA1FV_tFSdyKXEfu0_ :
$@convention(thin) @substituted <τ_0_0> (@inout_aliasable Double)
-> (@out τ_0_0, @error any Error) for <Double>
%55 = partial_apply [callee_guaranteed] %54(%2) :
$@convention(thin) @substituted <τ_0_0> (@inout_aliasable Double)
-> (@out τ_0_0, @error any Error) for <Double>
Now we emit a diagnostic instead.
The patch resolves issues #66996 and #63331
This is a futile attempt to discourage future use of getType() by
giving it a "scary" name.
We want people to use getInterfaceType() like with the other decl kinds.
Introduce the notion of "semantic result parameter". Handle differentiation of inouts via semantic result parameter abstraction. Do not consider non-wrt semantic result parameters as semantic results
Fixes#67174
Translating more llvm::Optional::transform calls to swift::transform.
`llvm::Optional` had a transform function that ran a lambda on the
element stored in the optional if it existed. After migrating to
std::optional under the hood, that function went away. Replacing it with
calls to swift::optional in STLExtras.h.
llvm::SmallSetVector changed semantics
(https://reviews.llvm.org/D152497) resulting in build failures in Swift.
The old semantics allowed usage of types that did not have an
`operator==` because `SmallDenseSet` uses `DenseSetInfo<T>::isEqual` to
determine equality. The new implementation switched to using
`std::find`, which internally uses `operator==`. This type is used
pretty frequently with `swift::Type`, which intentionally deletes
`operator==` as it is not the canonical type and therefore cannot be
compared in normal circumstances.
This patch adds a new type-alias to the Swift namespace that provides
the old semantic behavior for `SmallSetVector`. I've also gone through
and replaced usages of `llvm::SmallSetVector` with the
`Swift::SmallSetVector` in places where we're storing a type that
doesn't implement or explicitly deletes `operator==`. The changes to
`llvm::SmallSetVector` should improve compile-time performance, so I
left the `llvm::SmallSetVector` where possible.
All value ownership instructions (begin_borrow, copy_value, move_value)
need to be handled in the same way. Factored out a common function for
that handling in preparation to visit move_value instructions.
Reformatting everything now that we have `llvm` namespaces. I've
separated this from the main commit to help manage merge-conflicts and
for making it a bit easier to read the mega-patch.
This is phase-1 of switching from llvm::Optional to std::optional in the
next rebranch. llvm::Optional was removed from upstream LLVM, so we need
to migrate off rather soon. On Darwin, std::optional, and llvm::Optional
have the same layout, so we don't need to be as concerned about ABI
beyond the name mangling. `llvm::Optional` is only returned from one
function in
```
getStandardTypeSubst(StringRef TypeName,
bool allowConcurrencyManglings);
```
It's the return value, so it should not impact the mangling of the
function, and the layout is the same as `std::optional`, so it should be
mostly okay. This function doesn't appear to have users, and the ABI was
already broken 2 years ago for concurrency and no one seemed to notice
so this should be "okay".
I'm doing the migration incrementally so that folks working on main can
cherry-pick back to the release/5.9 branch. Once 5.9 is done and locked
away, then we can go through and finish the replacement. Since `None`
and `Optional` show up in contexts where they are not `llvm::None` and
`llvm::Optional`, I'm preparing the work now by going through and
removing the namespace unwrapping and making the `llvm` namespace
explicit. This should make it fairly mechanical to go through and
replace llvm::Optional with std::optional, and llvm::None with
std::nullopt. It's also a change that can be brought onto the
release/5.9 with minimal impact. This should be an NFC change.
Adjoint buffers of projections (e.g. obtained via begin_access) are same as adjoint buffer of underlying struct value. As a result, when propagating adjoint values to pullback successor blocks we tend to produce lots of identical copies (essentially for every struct access and in every basic block) of adjoint buffers.
These copy_addrs instructions are then lowered down to plain loads and stores and while the redundant copies are usually optimized away by subsequent optimization passes, presence of such copies leads to elevated memory consumption and compilation time as one needs to track liveness of these values being copied.
Track the values being propagated and simply do not generate extra copies if the same value was already propagated.
One step towards #61773
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
