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>
d971f125d9 added a usage of std::variant and std::holds_alternative to
SILOptimizer/Differentiation/AdjointValue.h, but did not include variant
directly. This caused issues while building on Linux.
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.
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'.
We had two copies of this code that had drifted apart. Bring them back
together so there is just one place where we compute the type of a
reabstraction thunk.
* [SR-13929][AutoDiff]: Enable [ossa] for Differentiation/Thunk.cpp:getOrCreateSubsetParametersThunkForLinearMap and promoteCurryThunkApplicationToDifferentiableFunction
- Properly clone and use debug scopes for all instructions in pullback functions.
- Emit `debug_value` instructions for adjoint values.
- Add debug locations and variable info to adjoint buffer allocations.
- Add `TangentBuilder` (a `SILBuilder` subclass) to unify and simplify special emitter utilities for tangent vector code generation. More simplifications to come.
Pullback variable inspection example:
```console
(lldb) n
Process 50984 stopped
* thread #1, queue = 'com.apple.main-thread', stop reason = step over
frame #0: 0x0000000100003497 main`pullback of foo(x=0) at main.swift:12:11
9 import _Differentiation
10
11 func foo(_ x: Float) -> Float {
-> 12 let y = sin(x)
13 let z = cos(y)
14 let k = tanh(z) + cos(y)
15 return k
Target 0: (main) stopped.
(lldb) fr v
(Float) x = 0
(Float) k = 1
(Float) z = 0.495846391
(Float) y = -0.689988375
```
Resolves rdar://68616528 / SR-13535.
- `Mangle::ASTMangler::mangleAutoDiffDerivativeFunction()` and `Mangle::ASTMangler::mangleAutoDiffLinearMap()` accept original function declarations and return a mangled name for a derivative function or linear map. This is called during SILGen and TBDGen.
- `Mangle::DifferentiationMangler` handles differentiation function mangling in the differentiation transform. This part is necessary because we need to perform demangling on the original function and remangle it as part of a differentiation function mangling tree in order to get the correct substitutions in the mangled derivative generic signature.
A mangled differentiation function name includes:
- The original function.
- The differentiation function kind.
- The parameter indices for differentiation.
- The result indices for differentiation.
- The derivative generic signature.
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`.
Add differentiation support for non-active `try_apply` SIL instructions.
Notable pullback generation changes:
* Original basic blocks are now visited in a different order:
* starting from the original basic block, all its predecessors
* are visited in a breadth-first search order. This ensures that
* all successors of any block are visited before the block itself.
Resolves TF-433.
Start `linear_function` canonicalization skeleton copying from
`differentiable_function` canonicalization. For now, transpose function
operands are filled in with `undef`.
`JVPCloner.h` is now tiny: `JVPCloner` exposes only a `bool run()` entry point.
All of the implementation is moved to `JVPCloner::Implementation` in
`JVPCloner.cpp`. Methods can be defined directly in `JVPCloner.cpp` without
separate declarations.
`VJPCloner.h` is now tiny: `VJPCloner` exposes only a `bool run()` entry point.
All of the implementation is moved to `VJPCloner::Implementation` in
`VJPCloner.cpp`. Methods can be defined directly in `VJPCloner.cpp` without
separate declarations.
Clarify `llvm_unreachable` after exhaustive switch is needed to silence
MSVC C4715, so we don't accidentally remove it.
Standardize some assertion messages.
Reimplement `PullbackCloner` using the pointer-to-implementation pattern.
`PullbackCloner.h` is now tiny: `PullbackCloner` exposes only a `bool run()`
entry point.
All of the implementation is moved to `PullbackCloner::Implementation` in
`PullbackCloner.cpp`.
Benefits of this approach:
- A main benefit is that methods can be defined directly in `PullbackCloner.cpp`
without needing to separately declare them in `PullbackCloner.h`.
- There is now no code duplication between `PullbackCloner.h` and
`PullbackCloner.cpp`.
- Consequently, method documentation is easier to read because it appears
directly on method definitions, instead of on method declarations in a
separate file. This is important for documentation of `PullbackCloner`
instruction visitor methods, which explain pullback transformation rules.
- Incremental recompilation may be faster since `PullbackCloner.h` changes less
often.
Partially resolves SR-13182.
Add base type parameter to `TangentStoredPropertyRequest`.
Use `TypeBase::getTypeOfMember` instead of `VarDecl::getType` to correctly
compute the member type of original stored properties, using the base type.
Resolves SR-13134.
Reverse-mode differentiation now supports `apply` instructions with multiple
active "semantic results" (formal results or `inout` parameters).
The "cannot differentiate through multiple results" non-differentiability error
is lifted.
Resolves TF-983.
Previously, PullbackEmitter assumed that original values' value category
matches their `TangentVector` types' value category. This was problematic
for loadable types with address-only `TangentVector` types.
Now, PullbackEmitter starts to support differentiation of loadable types with
address-only `TangentVector` types. This patch focuses on supporting and testing
class types, more support can be added incrementally.
Resolves TF-1149.
Use TangentStoredPropertyRequest in differentiation transform.
Improve non-differentiability diagnostics regarding invalid stored
property projection instructions:
`struct_extract`, `struct_element_addr`, `ref_element_addr`.
Diagnose the following cases:
- Original property's type does not conform to `Differentiable`.
- Base type's `TangentVector` is not a struct.
- Tangent property not found: base type's `TangentVector` does not have a
stored property with the same name as the original property.
- Tangent property's type is not equal to the original property's
`TangentVector` type.
- Tangent property is not a stored property.
Resolves TF-969 and TF-970.
- Show SIL type when printing `AdjointValue`.
- Add utilities to print `PullbackEmitter` adjoint value and buffer mappings.
- Print generated VJP before printing generated pullback.
- This is useful because pullback generation may crash after VJP
generation succeeds.
Update differentiation to handle `array.finalize_intrinsic` applications.
`VJPEmitter::visitApplyInst` does standard cloning for these applications.
`PullbackEmitter::visitApplyInst` treats the intrinsic like an identity
function, accumulating result's adjoint into argument's adjoint.
This fixes array literal initialization differentiation.
`DifferentiableFunctionInst` now stores result indices.
`SILAutoDiffIndices` now stores result indices instead of a source index.
`@differentiable` SIL function types may now have multiple differentiability
result indices and `@noDerivative` resutls.
`@differentiable` AST function types do not have `@noDerivative` results (yet),
so this functionality is not exposed to users.
Resolves TF-689 and TF-1256.
Infrastructural support for TF-983: supporting differentiation of `apply`
instructions with multiple active semantic results.
Fix SIL locations and debug scopes in `VJPEmitter` for:
- Pullback struct `struct` instructions
- Predecessor enums `enum` instructions
These instructions are not directly cloned from the original function and should
have auto-generated locations.
Resolves SR-12887: debug scope error for `VJPEmitter`-generated function.
Support differentiation of `is` and `as?` operators.
These operators lower to branching cast SIL instructions, requiring control
flow differentiation support.
Resolves SR-12898.
Add a common helper function `VJPEmitter::createTrampolineBasicBlock`.
Change `VJPEmitter::buildPullbackValueStructValue` to take an original basic
block instead of a terminator instruction.
Add special-case VJP generation support for "semantic member accessors".
This is necessary to avoid activity analysis related diagnostics and simplifies
generated code.
Fix "wrapped property mutability" check in `Differentiable` derived conformnances.
This resolves SR-12642.
Add e2e test using real world property wrappers (`@Lazy` and `@Clamping`).
Support differentiation of property wrapper wrapped value getters and setters.
Create new pullback generation code path for "semantic member accessors".
"Semantic member accessors" are attached to member properties that have a
corresponding tangent stored property in the parent `TangentVector` type.
These accessors have special-case pullback generation based on their semantic
behavior. Currently, only getters and setters are supported.
This special-case pullback generation is currently used for stored property
accessors and property wrapper wrapped value accessors. In the future, it can
also be used to support `@differentiable(useInTangentVector)` computed
properties: SR-12636.
User-defined accesors cannot use this code path because they may use custom
logic that does not semantically perform a member access.
Resolves SR-12639.
This became necessary after recent function type changes that keep
substituted generic function types abstract even after substitution to
correctly handle automatic opaque result type substitution.
Instead of performing the opaque result type substitution as part of
substituting the generic args the underlying type will now be reified as
part of looking at the parameter/return types which happens as part of
the function convention apis.
rdar://62560867
Move differentiation-related SILOptimizer files to
{include/swift,lib}/SILOptimizer/Differentiation/.
This reduces directory nesting and gathers files together.
