Skip to content
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

Merged
merged 10 commits into from
Dec 19, 2018

Conversation

jrevels
Copy link
Collaborator

@jrevels jrevels commented Dec 17, 2018

This PR implements a few changes requested by @MikeInnes to enable Zygote's Cassette-style transformation to replaced with actual Cassette:

  1. Adds a disablehooks function to turn off prehook/posthook injection
  2. Replace the usually constant-folded isa(..., OverdubInstead) branch with direct recursion.
  3. User passes are now provided with the full Cassette.Reflection argument, instead of just the signature and CodeInfo.
  4. User passes can return 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 to overdub and overdub to recurse. For 3, the provided Reflection object contains the CodeInfo 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.

@jrevels jrevels changed the title Refactor to reduce IR bloat for nested overdubbing Refactor to reduce IR bloat for nested overdubbing, plus other stuff for Zygote Dec 17, 2018
@oxinabox
Copy link
Contributor

It would be good to test this against
#91

To see if disabling the hooks decreases the overhead

@jrevels
Copy link
Collaborator Author

jrevels commented Dec 18, 2018

It would be good to test this against #91

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 disablehooks to affect run-time performance vs. no-op hook injection unless you're doing deeply nested overdubbing and the compiler is giving up or something. The main use case for disablehooks is to reduce the scaling of pre-inference CodeInfo size w.r.t. overdub nesting.

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).
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you

@jrevels
Copy link
Collaborator Author

jrevels commented Dec 18, 2018

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.

@jrevels jrevels merged commit cae698f into master Dec 19, 2018
@jrevels jrevels deleted the jr/cycle branch December 19, 2018 15:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants