Currently, AutoDiff Closure Specialization pass iterates over all VJP
instructions and checks each of them against a set of conditions which
`partial_apply` of pullback must satisfy.
This logic could be re-implemented w/o loop, checking conditions in
opposite direction, starting from `return` instruction and transitively
going to defining instructions of operands (`tuple` and `partial_apply`
for the desired pullback case).
* insert borrow scopes for specialized guaranteed arguments: This is needed if the guaranteed argument is replaced by a copied owned value. If the argument has any use which is not compatible with "owned" ownership, we need a borrow scope.
* prevent creating a specialized function with more than one "isolated" parameters
This new OSSA invariant simplifies many optimizations because they don't have to take care of the corner case of incomplete lifetimes in dead-end blocks.
The implementation basically consists of these changes:
* add the lifetime completion utility
* add a flag in SILFunction which tells optimization that they need to run the lifetime completion utility
* let all optimizations complete lifetimes if necessary
* enable the ownership verifier to check complete lifetimes
The conformances `Type: Comparable` and `EnumCase: Equatable` were
previously introduced in #85757 for implementing AutoDiff closure
specialization logic. This patch moves the latter directly to the SIL
module. The conformance `Type: Comparable` is deleted, and `sort(by:)`
is used instead of `sort` to mimic the same behavior in tests. These
changes address warnings mentioned in the comment:
https://github.com/swiftlang/swift/pull/85757#issuecomment-3681636278
Calling `cloneRecursively` from `SpecializationInfo.cloneClosures`
requires the callee having ownership info. Otherwise, the cloner uses
`recordFoldedValue` instead of `recordClonedInstruction`, and
`postProcess` hook is not called, which leads to an assertion failure in
`BridgedClonerImpl::cloneInst`.
Beside supporting OSSA, this change significantly simplifies the pass.
The main change is that instead of starting at a closure (e.g. `partial_apply`) and finding all call sites, we now start at a call site and look for closures for all arguments. This makes a lot of things much simpler, e.g. not so many intermediate data structures are required to track all the states.
I needed to remove the 3 unit tests because the things those tests were testing are not there anymore. However, the pass is tested with a lot of sil tests (and I added quite a few), which should give good test coverage.
The old ClosureSpecializer pass is still kept in place, because at that point in the pipeline we don't have OSSA, yet. Once we have that, we can replace the old pass withe the new one.
However, the autodiff closure specializer already runs in the OSSA pipeline and there the new changes take effect.
Previously, AutoDiff closure specialization pass was triggered only on
VJPs containing single basic block. However, the pass logic allows
running on arbitrary VJPs. This PR enables the pass for all VJPs
unconditionally. So, if the pullback corresponding to multiple-BB VJP
accepts some closures directly as arguments, these closures might become
specialized by the pass. Closures passed via payload of branch tracing
enum are not specialized - this is subject for future changes.
The PR contains several commits.
1. The thing named "call site" in the code is partial_apply of pullback
corresponding to the VJP. This might appear only once, so we drop
support for multiple "call sites".
2. Enhance existing SILOptimizer tests for the pass.
3. Add validation-tests for single basic block case.
4. The change itself - delete check against single basic block.
5. Add validation-tests for multiple basic block case.
6. Add SILOptimizer tests for multiple basic block case.
Originally this was a "private" utility for the ClosureSpecialization pass.
Now, make it a general utility which can be used for all kind of function specializations.
* let `SIL.Type` conform to `TypeProperties` to share the implementation of common type properties between the AST types and `SIL.Type`
* call references to an `AST.Type` `rawType` (instead of just `type`)
* remove unneeded stuff
* add comments
As the optimizer uses more and more AST stuff, it's now time to create an "AST" module.
Initially it defines following AST datastructures:
* declarations: `Decl` + derived classes
* `Conformance`
* `SubstitutionMap`
* `Type` and `CanonicalType`
Some of those were already defined in the SIL module and are now moved to the AST module.
This change also cleans up a few things:
* proper definition of `NominalTypeDecl`-related APIs in `SIL.Type`
* rename `ProtocolConformance` to `Conformance`
* use `AST.Type`/`AST.CanonicalType` instead of `BridgedASTType` in SIL and the Optimizer
Make SILLInkage available in SIL as `SIL.Linkage`.
Also, rename the misleading Function and GlobalVariable ABI `isAvailableExternally` to `isDefinedExternally`
The OSSA elimination pass has not yet been moved below all high level
function passes. Until that work has been completed the Autodiff
closure-spec optimization pass cannot solely support OSSA instructions.
Daniil Kovalev
Erik Eckstein
Jakub Florek
Anton Korobeynikov
Kshitij
Kshitij Jain
Nate Chandler