-
Notifications
You must be signed in to change notification settings - Fork 10.5k
[AutoDiff] First cut of coroutines differentiation #71461
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Tagging @BradLarson @dan-zheng |
@swift-ci please test |
@swift-ci please test windows platform |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should the changes to begin_apply
and end_apply
be in their own PR before the differentiation aspects?
@rxwei Rebased and ready for review. Please take a look when you will have a chance. |
@swift-ci please test |
@swift-ci please test |
@swift-ci please test |
@swift-ci please test macos platform |
@swift-ci please test |
@swift-ci please test macos platform |
1 similar comment
@swift-ci please test macos platform |
@swift-ci please test |
1 similar comment
@swift-ci please test |
Co-authored-by: Richard Wei <rxwei@apple.com>
@swift-ci please test |
@swift-ci please test macos platform |
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:
@yield_once
coroutines are naturally supported_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.begin_apply
site and pullback closure is available only fromend_apply
, so we need to track the order in which pullbacks are produced (and arrange consumption of the values accordingly – effectively delay them)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.Some diagnostics were contributed by @asavonic