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>
The names of the private witness table accessor thunks we generate for
an opaque return type mangle the concrete conformance of the underlying
type.
If a conformance requirement of the opaque return type was witnessed by
a conditional conformance of a variadic generic type, we would crash
because of an unimplemented case in the mangler.
Fixes rdar://problem/125668798.
This includes runtime support for instantiating transferring param/result in
function types. This is especially important since that is how we instantiate
function types like: typealias Fn = (transferring X) -> ().
rdar://123118061
Function body macros allow one to introduce a function body for a
particular function, either providing a body for a function that
doesn't have one, or wholesale replacing the body of a function that
was written with a new one.
Using symbolic references instead of a text based mangling avoids the
expensive type descriptor scan when objective c protocols are requested.
rdar://111536582
Extend the name mangling scheme for macro expansions to cover attached
macros, and use that scheme for the names of macro expansions buffers.
Finishes rdar://104038303, stabilizing file/buffer names for macro
expansion buffers.
Use the name mangling scheme we've devised for macro expansions to
back the implementation of the macro expansion context's
`getUniqueName` operation. This way, we guarantee that the names
provided by macro expansions don't conflict, as well as making them
demangleable so we can determine what introduced the names.
When a declaration has a structural opaque return type like:
func foo() -> Bar<some P>
then to mangle that return type `Bar<some P>`, we have to mangle the `some P`
part by referencing its defining declaration `foo()`, which in turn includes
its return type `Bar<some P>` again (this time using a special mangling for
`some P` that prevents infinite recursion). Since we mangle `Bar<some P>`
once as part of mangling the declaration, and we register substitutions for
bound generic types when they're complete, we end up registering the
substitution for `Bar<some P>` twice, once as the return type of the
declaration name, and again as the actual type. This would be fine, except
that the mangler doesn't check for key collisions, and it picks
substitution indexes based on the number of entries in its hash map, so
the duplicated substitution ends up corrupting the substitution sequence,
causing the mangler to produce an invalid mangled name.
Fixing that exposes us to another problem in the remangler: the AST
mangler keys substitutions by type identity, but the remangler
uses the value of the demangled nodes to recognize substitutions.
The mangling for `Bar<current declaration's opaque return type>` can
appear multiple times in a demangled tree, but referring to different
declarations' opaque return types, and the remangler would reconstruct
an incorrect mangled name when this happens. To avoid this, change the
way the demangler represents `OpaqueReturnType` nodes so that they
contain a backreference to the declaration they represent, so that
substitutions involving different declarations' opaque return types
don't get confused.
For performance annotations we need the generic specializer to trop non-generic metatype argumentrs
(which we don't do in general). For this we need a separate mangling.
Upgrade the old mangling from a list of argument types to a
list of requiremnets. For now, only same-type requirements
may actually be mangled since those are all that are available
to the surface language.
Reconstruction of existential types now consists of demangling (a list of)
base protocol(s), decoding the constraints, and converting the same-type
constraints back into a list of arguments.
rdar://96088707
The layout of constant static arrays differs from non-constant static arrays.
Therefore use a different mangling to get symbol mismatches if for some reason two modules don't agree on which version a static array is.
I wrote out this whole analysis of why different existential types
might have the same logical content, and then I turned around and
immediately uniqued existential shapes purely by logical content
rather than the (generalized) formal type. Oh well. At least it's
not too late to make ABI changes like this.
We now store a reference to a mangling of the generalized formal
type directly in the shape. This type alone is sufficient to unique
the shape:
- By the nature of the generalization algorithm, every type parameter
in the generalization signature should be mentioned in the
generalized formal type in a deterministic order.
- By the nature of the generalization algorithm, every other
requirement in the generalization signature should be implied
by the positions in which generalization type parameters appear
(e.g. because the formal type is C<T> & P, where C constrains
its type parameter for well-formedness).
- The requirement signature and type expression are extracted from
the existential type.
As a result, we no longer rely on computing a unique hash at
compile time.
Storing this separately from the requirement signature potentially
allows runtimes with general shape support to work with future
extensions to existential types even if they cannot demangle the
generalized formal type.
Storing the generalized formal type also allows us to easily and
reliably extract the formal type of the existential. Otherwise,
it's quite a heroic endeavor to match requirements back up with
primary associated types. Doing so would also only allows us to
extract *some* matching formal type, not necessarily the *right*
formal type. So there's some good synergy here.
