Consider an `@_alwaysEmitIntoClient` function and a custom derivative
defined
for it. Previously, such a combination resulted different errors under
different
circumstances.
Sometimes, there were linker errors due to missing derivative function
symbol -
these occurred when we tried to find the derivative in a module, while
it
should have been emitted into client's code (and it did not happen).
Sometimes, there were SIL verification failures like this:
```
SIL verification failed: internal/private function cannot be serialized or serializable: !F->isAnySerialized() || embedded
```
Linkage and serialization options for the derivative were not handled
properly,
and, instead of PublicNonABI linkage, we had Private one which is
unsupported
for serialization - but we need to serialize `@_alwaysEmitIntoClient`
functions
so the client's code is able to see them.
This patch resolves the issue and adds proper handling of custom
derivatives
of `@_alwaysEmitIntoClient` functions. Note that either both the
function and
its custom derivative or none of them should have
`@_alwaysEmitIntoClient`
attribute, mismatch in this attribute is not supported.
The following cases are handled (assume that in each case client's code
uses
the derivative).
1. Both the function and its derivative are defined in a single file in
one module.
2. Both the function and its derivative are defined in different files
which
are compiled to a single module.
3. The function is defined in one module, its derivative is defined in
another
module.
4. The function and the derivative are defined as members of a protocol
extension in two separate modules - one for the function and one for the
derivative. A struct conforming the protocol is defined in the third
module.
5. The function and the derivative are defined as members of a struct
extension in two separate modules - one for the function and one for the
derivative.
The changes allow to define derivatives for methods of `SIMD`.
Fixes#54445
<!--
If this pull request is targeting a release branch, please fill out the
following form:
https://github.com/swiftlang/.github/blob/main/PULL_REQUEST_TEMPLATE/release.md?plain=1
Otherwise, replace this comment with a description of your changes and
rationale. Provide links to external references/discussions if
appropriate.
If this pull request resolves any GitHub issues, link them like so:
Resolves <link to issue>, resolves <link to another issue>.
For more information about linking a pull request to an issue, see:
https://docs.github.com/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue
-->
<!--
Before merging this pull request, you must run the Swift continuous
integration tests.
For information about triggering CI builds via @swift-ci, see:
https://github.com/apple/swift/blob/main/docs/ContinuousIntegration.md#swift-ci
Thank you for your contribution to Swift!
-->
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)
getVarInfo() now always returns a variable with a location and scope.
To opt out of this change, getVarInfo(false) returns an incomplete variable.
This can be used to work around bugs, but should only really be used for
printing.
The complete var info will also contain the type, except for debug_values,
as its type depends on another instruction, which may be inconsistent if
called mid-pass.
All locations in debug variables are now also stripped of flags, to avoid
issues when comparing or hashing debug variables.
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>
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.
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.
`getValue` -> `value`
`getValueOr` -> `value_or`
`hasValue` -> `has_value`
`map` -> `transform`
The old API will be deprecated in the rebranch.
To avoid merge conflicts, use the new API already in the main branch.
rdar://102362022
Andy some time ago already created the new API but didn't go through and update
the old occurences. I did that in this PR and then deprecated the old API. The
tree is clean, so I could just remove it, but I decided to be nicer to
downstream people by deprecating it first.
In order to transfer debug info from `debug_value` to `alloc_stack` we need to drop `op_ref` debug expression, as the latter instruction represents the location and not the value itself. Also, fix the debug info verifier as variable type was deduced from `alloc_stack` / `alloc_box` improperly: pointer type of instruction itself was used instead of underlying object type.
The issue is only exposed when both optimizations (SIL mem2reg at least) and debug info are enabled
Resolves SR-15849
- Introduce a workaround while dealing with getLoweredParameterIndices() results:
it operates on AST and therefore does not take into account captured arguments
which are "on side" on AST and explicit in SIL
- Ensure captured arguments (@inout_aliased) are handled in a same way as ordinary
@inout arguments while generating a pullback type
Fixes SR-15205
- 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.
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 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.
This call only succeeds if the accessor has already been synthesized, which
makes it fragile. Replace it with calls to getOpaqueAccessor(),
which lazily evaluate the right requests ahead of time.
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.
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`).
Move differentiation-related SILOptimizer files to
{include/swift,lib}/SILOptimizer/Differentiation/.
This reduces directory nesting and gathers files together.
