The inlining benefit for VJPs and pullbacks seems to be a bit different on
macos and linux. This difference seems to be arising due to certain
function calls such as, `sin`, `cos` etc. always being inlined on macos
but not on linux.
Until I figure out a better way, I'm modifying these tests to avoid matching
the exact value of the inlining benefit, which causes them to fail on
Linux.
Changes in this CR add part of the, Swift based, Autodiff specific
closure specialization optimization pass. The pass does not modify any
code nor does it even exist in any of the optimization pipelines. The
rationale for pushing this partially complete optimization pass upstream
is to keep up with the breaking changes in the underlying Swift based
compiler infrastructure.
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>
Previously, the lexical attribute on allock_stack instructions was used.
This doesn't work for values without lexical lifetimes which are
consumed, e.g. stdlib CoW types. Here, the new var_decl attribute on
alloc_stack is keyed off of instead. This flag encodes exactly that a
value corresponds to a source-level VarDecl, which is the condition
under which checking needs to run.
Add a new mandatory BooleanLiteralFolding pass which constant folds conditional branches with boolean literals as operands.
```
%1 = integer_literal -1
%2 = apply %bool_init(%1) // Bool.init(_builtinBooleanLiteral:)
%3 = struct_extract %2, #Bool._value
cond_br %3, bb1, bb2
```
->
```
...
br bb1
```
This pass is intended to run before DefiniteInitialization, where mandatory inlining and constant folding didn't run, yet (which would perform this kind of optimization).
This optimization is required to let DefiniteInitialization handle boolean literals correctly.
For example in infinite loops:
```
init() {
while true { // DI need to know that there is no loop exit from this while-statement
if some_condition {
member_field = init_value
break
}
}
}
```
For linear maps containing control-flow, closures (representing the pullbacks
of intermediate values) may be passed as arguments, however, they may be
hidden behind a branch-tracing enum (tracing execution flow of the original
function).
Such linear maps did not use to get inlining benefits as the compiler
could not see that the intermediate pullback closures were actually part
of the input.
This change modifies the inliner logic to correctly award inlining
benefits to linear maps containing control-flow, by checking if a
"callee" in the linear map actually traces back to an input closure that
was received as part of a branch-tracing enum input argument.
Fixes#68945
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`.
We need a lowered type for branch trace enum in order to compute linear map tuple type. However, the lowering of branch trace enum type depends on the types of its elements (the payloads are linear map tuples of predecessor BB).
As lowered types are cached, we cannot populate branch trace enum entries in the end as we did before: we already used wrong lowered types for linear map tuples.
Traverse basic blocks in reverse post-order traverse order building linear map tuples and branch tracing enums in one go, ensuring that we've done with predecessor BBs before processing the BB itself.
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
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
It is necessary for opaque values where for casts that will newly start
out as checked_cast_brs and be lowered to checked_cast_addr_brs, since
the latter has the source formal type, IRGen relies on being able to
access it, and there's no way in general to obtain the source formal
type from the source lowered type.
The patch resolves#67402.
When the original function has a tuple result type, we should append
thunkedLinearMap as the last element of the tuple to match the function
declaration. Before this patch, the compiler used to wrap the original result
tuple and thunkedLinearMap into another tuple, and caused the verifier error.
Before the patch:
return %{{.*}} : $((Float, Double), @callee_guaranteed (Float) -> X.TangentVector)
After the patch:
return %{{.*}} : $(Float, Double, @callee_guaranteed (Float) -> X.TangentVector)
As described in the issue #62922, the compiler should not allow to discard @noDerivative attribute and keep @differentiable. The patch adds a diagnostic for this case.
Resolves#62922.
Now that `InferredGenericSignatureRequest` creates
`StructuralRequirement`s from of the generic signature with valid source
locations, additional redundancy warnings are produced. Update tests
with the new warnings.
* move the apply of partial_apply transformation from simplify-apply to simplify-partial_apply
* delete dead partial_apply instructions
* devirtualize apply, try_apply and begin_apply
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
Types that have "value semantics" should not have lexical lifetimes.
Value types are not expected to have custom deinits. Are not expected to
expose unsafe interior pointers. And cannot have weak references because
they are structs. Therefore, deinitialization barriers are irrelevant.
rdar://107076869
This essentially passes the members of a linear map tuple as individual arguments. It yields few nice simplifications:
* No linear map tuples at all for getters / setters
* No tuple formation / deconstruction around pullbacks
* Pullbacks with loops still use heap-allocated tuples
The Swift Simplification pass can do more than the old MandatoryCombine pass: simplification of more instruction types and dead code elimination.
The result is a better -Onone performance while still keeping debug info consistent.
Currently following code patterns are simplified:
* `struct` -> `struct_extract`
* `enum` -> `unchecked_enum_data`
* `partial_apply` -> `apply`
* `br` to a 1:1 related block
* `cond_br` with a constant condition
* `isConcrete` and `is_same_metadata` builtins
More simplifications can be added in the future.
rdar://96708429
rdar://104562580
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)
Fix the common error of using underscores instead of dashes.
In the rebranch this is an error (lit got more picky), but it also makes sense to fix the tests in the main branch
