Commit Graph

9 Commits

Author SHA1 Message Date
Anton Korobeynikov
c7a216058f [AutoDiff] First cut of coroutines differentiation (#71461)
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>
2024-04-04 17:24:55 -07:00
Richard Wei
8d8614058b [AudoDiff] NFC: Replace 'SILAutoDiffIndices' with 'AutoDiffConfig'. (#35079)
Resolve rdar://71678394 / SR-13889.
2020-12-14 14:32:40 -08:00
Dan Zheng
a88bd15815 [AutoDiff] Remove raw_ostream & argument from IndexSubset::dump. (#31675)
Remove `llvm::raw_ostream &` argument so that `IndexSubset::dump` can be called
in lldb without issue.
2020-05-09 11:18:40 -07:00
Dan Zheng
63ce61b9aa [AutoDiff] NFC: Use consistent headers for IndexSubset.{h,cpp}.
Left-align headers. Use consistent message. Use accurate year in copyright.
2019-12-12 14:23:28 -08:00
Gwen Mittertreiner
5f3a6af0e5 Explicitly Specify a 64bit Shift to Fix Warning
Explicitly doing a 64bit shift silences the otherwise noisy warning:
```
[1548/1735] Building CXX object
tools\swift\lib\Sema\CMakeFiles\swiftSema.dir\NameBinding.cpp.obj

S:\toolchain\swift\include\swift/AST/IndexSubset.h(96): warning C4334:
'<<': result of 32-bit shift implicitly converted to 64 bits (was 64-bit
shift intended?)

S:\toolchain\swift\include\swift/AST/IndexSubset.h(184): warning C4334:
'<<': result of 32-bit shift implicitly converted to 64 bits (was 64-bit
shift intended?)
```
2019-11-13 15:59:27 -08:00
Brent Royal-Gordon
99faa033fc [NFC] Standardize dump() methods in frontend
By convention, most structs and classes in the Swift compiler include a `dump()` method which prints debugging information. This method is meant to be called only from the debugger, but this means they’re often unused and may be eliminated from optimized binaries. On the other hand, some parts of the compiler call `dump()` methods directly despite them being intended as a pure debugging aid. clang supports attributes which can be used to avoid these problems, but they’re used very inconsistently across the compiler.

This commit adds `SWIFT_DEBUG_DUMP` and `SWIFT_DEBUG_DUMPER(<name>(<params>))` macros to declare `dump()` methods with the appropriate set of attributes and adopts this macro throughout the frontend. It does not pervasively adopt this macro in SILGen, SILOptimizer, or IRGen; these components use `dump()` methods in a different way where they’re frequently called from debugging code. Nor does it adopt it in runtime components like swiftRuntime and swiftReflection, because I’m a bit worried about size.

Despite the large number of files and lines affected, this change is NFC.
2019-10-31 18:37:42 -07:00
Gogul Balakrishnan
68185cf59e Fix memory error in the implementation of IndexSubset. 2019-10-21 23:47:17 -07:00
Dan Zheng
7b5dfd6cf6 [AutoDiff upstream] NFC: update IndexSubset print/dump utilities. (#27775)
Cherry-pick #27772 from tensorflow branch.
2019-10-18 14:01:28 -07:00
Richard Wei
9765858d31 Rename to 'IndexSubset' and move it to its own file. 2019-10-11 13:07:53 -07:00