-
Notifications
You must be signed in to change notification settings - Fork 35
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
Refactor to reduce IR bloat for nested overdubbing, plus other stuff for Zygote #97
Conversation
It would be good to test this against To see if disabling the hooks decreases the overhead |
It doesn't (just checked to make sure). It would be very surprising if no-op hooks caused issues there, though - the compiler is very good at DCEing those, and it's a very easy thing for the compiler to do. I general, I wouldn't expect |
…-dependency testing for new releases
returns, we also want to return a callback that executes all the `println` calls that we stripped | ||
out. How would you implement this with Cassette? | ||
Let's say you wanted to use Cassette to ["slice" various separable | ||
subcomputations out from an overall computation](https://en.wikipedia.org/wiki/Program_slicing). |
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.
Thank you
I think this is a good stopping point for this already way-overlarge PR. I'll probably merge tomorrow and cut a Cassette v0.2 release. |
This PR implements a few changes requested by @MikeInnes to enable Zygote's Cassette-style transformation to replaced with actual Cassette:
disablehooks
function to turn offprehook
/posthook
injectionisa(..., OverdubInstead)
branch with direct recursion.Cassette.Reflection
argument, instead of just the signature andCodeInfo
.Expr
objects, which Cassette will essentially just emit as fallback thunks. I actually don't think this is totally necessary, but I think it will ease the transition for WIP: Replace Zygote's Cassette-style transform with actual Cassette FluxML/Zygote.jl#45.Changes 1 and 2 allow users to configure their contexts such that lowered IR size (in statement count) doesn't increase exponentially for when nesting overdubs.
Note that changes 2 and 3 are fairly breaking, but the migration paths for each are pretty simple. For 2, most cases just involves switching
execute
tooverdub
andoverdub
torecurse
. For 3, the providedReflection
object contains theCodeInfo
and signature previously passed to user's transform function. After this is merged, feel free to ping me/open issues and I can help with migrating folks' Cassette usage/fixing any breakage that pops up.Additionally, this PR implements better handling of
Core._apply
calls (probably should've factored that out into a separate PR, but got a bit carried way).Barring any additional stuff as FluxML/Zygote.jl#45 develops, this PR is essentially complete modulo documentation/tests.
I'll try to update those ASAP as any new changes are made.Basically done for now.