* [SR-13929][AutoDiff]: Enable [ossa] for Differentiation/Thunk.cpp:getOrCreateSubsetParametersThunkForLinearMap and promoteCurryThunkApplicationToDifferentiableFunction
Fixes:
[[SR-14290]: SIL verification fails when differentiating a function of [[Double]]](rdar://74876596)
[[SR-14298]: [AutoDiff] SIL verification failed: non-contiguous lexical scope at -Onone](rdar://75032457)
While it is very convenient to default the ExtInfo state when creating
new function types, it also make the intent unclear to those looking to
extend ExtInfo state. For example, did a given call site intend to have
the default ExtInfo state or does it just happen to work? This matters a
lot because function types are regularly unpacked and rebuilt and it's
really easy to accidentally drop ExtInfo state.
By changing the ExtInfo state to an optional, we can track when it is
actually needed.
Refactor SILGen's ApplyOptions into an OptionSet, add a
DoesNotAwait flag to go with DoesNotThrow, and sink it
all down into SILInstruction.h.
Then, replace the isNonThrowing() flag in ApplyInst and
BeginApplyInst with getApplyOptions(), and plumb it
through to TryApplyInst as well.
Set the flag when SILGen emits a sync call to a reasync
function.
When set, this disables the SIL verifier check against
calling async functions from sync functions.
Finally, this allows us to add end-to-end tests for
rdar://problem/71098795.
- 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.
Add the following new mangling rules.
```
global ::= from-type to-type 'TJO' AUTODIFF-FUNCTION-KIND // autodiff self-reordering reabstraction thunk
global ::= from-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff linear map subset parameters thunk
global ::= global to-type 'TJS' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' INDEX-SUBSET 'P' // autodiff derivative function subset parameters thunk
```
Example:
```console
$s13TangentVector16_Differentiation14DifferentiablePQzAaDQy_SdAFIegnnnr_TJSdSSSpSrSUSP ---> autodiff subset parameters thunk for differential from @escaping @callee_guaranteed (@in_guaranteed A._Differentiation.Differentiable.TangentVector, @in_guaranteed B._Differentiation.Differentiable.TangentVector, @in_guaranteed Swift.Double) -> (@out B._Differentiation.Differentiable.TangentVector) with respect to parameters {0, 1, 2} and results {0} to parameters {0, 2}
$sS2f8mangling3FooV13TangentVectorVIegydd_SfAESfIegydd_TJOp ---> autodiff self-reordering reabstraction thunk for pullback from @escaping @callee_guaranteed (@unowned Swift.Float) -> (@unowned Swift.Float, @unowned mangling.Foo.TangentVector) to @escaping @callee_guaranteed (@unowned Swift.Float) -> (@unowned mangling.Foo.TangentVector, @unowned Swift.Float)
```
Resolves rdar://72666310 / SR-13508.
Also fix a bug in `AutoDiffFunction` mangling where the original may be a global that contains more than 1 node (rdar://74151229 / SR-14106).
If we know that we have a FunctionRefInst (and not another variant of FunctionRefBaseInst), we know that getting the referenced function will not be null (in contrast to FunctionRefBaseInst::getReferencedFunctionOrNull).
NFC
- Set the builder's debug scope to the remapped debug scope in the overriden visiters.
- In trampoline blocks where we build predecessor enums, make the trampoline builder inherit the original terminator's remapped debug scope.
Resolves rdar://74087329.
- `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`.
`SILBuilder::createAllocStack` expects a debug variable when the location is a `VarDecl`. Since we are in pullbacks, there's no debug variables so we pass an empty one.
General support for debug-info-in-pullbacks will be added as part of SR-13535.
Also add a negative test for SR-13866.
Resolves SR-13865.
AD-generated data structures (linear map structs and branching trace enums) do not need to be resilient data structures. These decls ade missing a `@frozen` attribute.
Resolves rdar://71319547.
This makes it easier to understand conceptually why a ValueOwnershipKind with
Any ownership is invalid and also allowed me to explicitly document the lattice
that relates ownership constraints/value ownership kinds.
I have a need to have SwitchEnum{,Addr}Inst have different base classes
(TermInst, OwnershipForwardingTermInst). To do this I need to add a template to
SwitchEnumInstBase so I can switch that BaseTy. Sadly since we are using
SwitchEnumInstBase as an ADT type as well as an actual base type for
Instructions, this is impossible to do without introducing a template in a ton
of places.
Rather than doing that, I changed the code that was using SwitchEnumInstBase as
an ADT to instead use a proper ADT SwitchEnumBranch. I am happy to change the
name as possible see fit (maybe SwitchEnumTerm?).
We'll need this to get the right 'selfDC' when name lookup
finds a 'self' declaration in a capture list, eg
class C {
func bar() {}
func foo() {
_ = { [self] in bar() }
}
}
Previously, `LinearMapInfo::shouldDifferentiateInstruction` had a special case
for `copy_value`, returning true for `copy_value` instructions with an active
operand.
This is unexpected and led to "leaked owned value" ownership verification
failures due to unnecessarily cloned `copy_value` instructions during
differential generation.
Now, the special case is removed, fixing the failures.
`shouldDifferentiateInstruction` returns true for `copy_value` instructions
whose operand and result are both active.
Resolves SR-13530.
Fixes foward-mode crashed related to:
- Missing tangent buffers for non-wrt `inout` parameters.
- Tangent buffers not being initialized due to the corresponding original
buffer intialization instructions being non-active.
- Non-varied indirect results not being initialized.
- `emitDestroyValue` crashes due to `TangentVector` value category mismatch.
Resolves TF-984 and SR-13447.
Co-authored-by: Dan Zheng <danielzheng@google.com>
Make the differentiation transform generate transparent, ossa
reabstraction thunks.
This enables these thunks to be inlined into other ossa functions (e.g.
generated VJP and pullback functions) during mandatory inlining.
Resolves TF-989.
Unblocks further autodiff-related optimizations: SR-13390.
Fix SIL pullback function type calculation: remap original `unowned` results to
the correct pullback parameter convention depending on the `TangentVector` type
lowering.
Make `PullbackCloner` visitors for the following instructions check and handle
tangent value categories: `struct`, `struct_extract`, `tuple`, `destructure_tuple`.
Add differentiation tests for `Optional` struct and class stored properties,
exercising the instruction visitors above.
Resolves SR-13430.
Adds forward mode support for `apply` instruction with `inout` arguments.
Example of supported code:
```
func add(_ x: inout Float, _ y: inout Float) -> Float {
var result = x
result += y
return result
}
print(differential(at: 1, 1, in: add)(1, 1)) // prints "2"
```
Previously, JVP/VJP generation used a "return undef" hack when
differential/pullback values did not match the expected return type.
This was relevant before differentiation supported "loadable types with
address-only tangent types", which was diagnosed.
Now that support for the above exists, "return undef" should be removed and
replaced with an assertion.
