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

Remove branching and nested Refs #64

Merged
merged 13 commits into from
Aug 1, 2023
Merged

Remove branching and nested Refs #64

merged 13 commits into from
Aug 1, 2023

Conversation

samestep
Copy link
Contributor

@samestep samestep commented Jul 24, 2023

Over the past couple weeks @ravenrothkopf and I worked through automatic differentiation on a bunch of examples, and we realized a couple things:

  • As the Dex paper notes, we need a Read kind of Ref to serve as the transpose for Accum.
  • In contrast to Dex, we actually don't need a constraint stating that a type is a Vector space; instead, when we accumulate, we can just take an existing value of any type to choose the topology, and then work inside the dynamically-defined topology of that value. This will become more important when we Support sum types #65 but for now it just means that Bool/Fin/etc stay as-is instead of being transformed into Unit.
  • We can't handle branching until we Support sum types #65, so instead for now we're just replacing If with a Select expression just like Wasm's.
  • Dex doesn't handle nested Refs, and we can't either; to deal with this, we need to allow functions to take multiple parameters.

This PR makes the necessary changes to our codebase to accommodate those IR changes, and also simplifes the IR in a couple other ways:

  • Every instance of Vec<T> is replaced with Box<[T]> to save space, because the former stores three pointers while the latter only uses two; the exception is Ty which we're leaving as Vec<T> for now because ts-rs does not support Box<[T]> or Rc<[T]>.
  • No more Func and Block. Each Call now directly encodes the generic types at that callsite, and each For, Read, and Accum now directly owns its body.
  • The Scope type now explicitly says what Constraint it satisfies, because otherwise typechecking the IR would be more annoying.

@samestep samestep mentioned this pull request Jul 24, 2023
@samestep samestep marked this pull request as ready for review July 24, 2023 21:59
@samestep samestep requested a review from ravenrothkopf July 24, 2023 21:59
Copy link
Contributor

@ravenrothkopf ravenrothkopf left a comment

Choose a reason for hiding this comment

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

LGTM! I need to go back in and change the pprint stuff to support select tho so don't review the PR yet!

@samestep
Copy link
Contributor Author

Thanks for reviewing @ravenrothkopf! Let's hold off on merging this for now, though; from our discussion earlier today it seems that we may not want to support nested Refs at all, in which case we need to adjust the IR a bit more.

@samestep samestep marked this pull request as draft July 28, 2023 16:59
@samestep samestep changed the title Remove branching and handle nested Refs Remove branching and nested Refs Jul 28, 2023
@samestep samestep marked this pull request as ready for review July 31, 2023 16:34
@samestep samestep requested a review from ravenrothkopf July 31, 2023 16:34
Copy link
Contributor

@ravenrothkopf ravenrothkopf left a comment

Choose a reason for hiding this comment

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

nice nice nice! thank you so much for this

@samestep samestep merged commit bfd6253 into main Aug 1, 2023
@samestep samestep deleted the new-ir branch August 1, 2023 15:17
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