Skip to content

[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

Merged
merged 8 commits into from
Apr 5, 2024
Merged

[AutoDiff] First cut of coroutines differentiation #71461

merged 8 commits into from
Apr 5, 2024

Conversation

asl
Copy link
Contributor

@asl asl commented Feb 8, 2024

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.

Some diagnostics were contributed by @asavonic

@asl asl added the AutoDiff label Feb 8, 2024
@asl asl requested a review from rxwei February 8, 2024 02:46
@asl
Copy link
Contributor Author

asl commented Feb 8, 2024

Tagging @BradLarson @dan-zheng

@asl
Copy link
Contributor Author

asl commented Feb 8, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Feb 8, 2024

@swift-ci please test windows platform

Copy link
Contributor

@rxwei rxwei left a 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?

@asl
Copy link
Contributor Author

asl commented Feb 16, 2024

Should the changes to begin_apply and end_apply be in their own PR before the differentiation aspects?

They are and this is #71645. This PR will be rebased after #71645 will be merged.

To be reviewed are 3aba19b and everything after

@asl
Copy link
Contributor Author

asl commented Feb 16, 2024

Another pre-requisite is partial_apply patch for coroutines. It is submitted in #71653 As I mentioned, w/o #71653 we will generate proper SIL, but will be unable to proceed further down to real code generation and execution

@asl asl requested a review from rxwei March 27, 2024 21:07
@asl
Copy link
Contributor Author

asl commented Mar 27, 2024

@rxwei Rebased and ready for review. Please take a look when you will have a chance.

@asl
Copy link
Contributor Author

asl commented Mar 27, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Mar 29, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Apr 2, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Apr 2, 2024

@swift-ci please test macos platform

@asl
Copy link
Contributor Author

asl commented Apr 2, 2024

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Apr 3, 2024

@swift-ci please test macos platform

1 similar comment
@asl
Copy link
Contributor Author

asl commented Apr 3, 2024

@swift-ci please test macos platform

@asl
Copy link
Contributor Author

asl commented Apr 4, 2024

@swift-ci please test

1 similar comment
@asl
Copy link
Contributor Author

asl commented Apr 4, 2024

@swift-ci please test

Co-authored-by: Richard Wei <rxwei@apple.com>
@asl
Copy link
Contributor Author

asl commented Apr 4, 2024

@swift-ci please test

@asl asl enabled auto-merge (squash) April 4, 2024 16:16
@asl
Copy link
Contributor Author

asl commented Apr 4, 2024

@swift-ci please test macos platform

@asl asl merged commit c7a2160 into main Apr 5, 2024
4 of 5 checks passed
@asl asl deleted the yield-diff branch April 5, 2024 00:24
@asl asl linked an issue Dec 3, 2024 that may be closed by this pull request
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[SR-14115] SIL: derivative type calculation for coroutines
3 participants