AD and other compiler transformations from within dex?

[zenna]

I just got around to looking at dex more closely after meeting with @froystig at LAFI, and I'm impressed!

Correct me if I'm wrong, but it looks as if the AD implementation is currently necessarily in Haskell, as opposed to dex itself?

The point of this post is to discuss the degree to which it is possible and desirable to support these kinds of transformations from within dex itself. One of the most exciting areas of Julia is this trend towards writing compiler transformations from within Julia. This started with Cassette.jl and has grown into a ecosystem of tools that are in the upcoming release of Julia getting first class status within the compiler itself. @oxinabox has written about it here.

From the perspective of someone who likes to implement inference methods through program transformations (for instance, my LAFI talk was on inverting programs), it would be great to be able to do this in a simpler, functional, language like dex.

[oxinabox]

Complier transforms are also the most exciting part of Jax too.
It would be really cool if they were first class in Dex.
In many ways it is the next level of homoiconic macros.

I know from playing with writing Arborist.jl recursive AST tranforms are, not a useful thing at least in julia (even though it is the natural extension of julia's macros).
but it might be on Dex code, if there are very few ways to express the same thing.
IRTools/low level Cassette is also not nice.
High level Cassette where it is just overdub to replace functions is nice; but limitted in its capacity (and runs into implemention related performance problems in Julia)

idk how Jax does its code tranforms user interface wise.

[froystig]

Roughly, jax transformations are written in python, conventionally as interpreters of jax's IR. There is some extra undocumented complexity for writing transformations that run "inline", without forming IR along the way.

An obstacle to drawing analogies here is that jax itself isn't really a language. One could maybe delineate a functional subset of python, plus our array types, as the language. The jax IR is a language, but it is definitely not one capable of expressing our transformations. It's roughly comparable to dex's internal program representations.

[apaszke]

Great question! Yes, all code transformations such as AD (as well as optimization passes) are implemented in Haskell as part of the compiler.

Now, we have briefly discussed that we are interested in surfacing this functionality, but it is not immediately clear how to do that in the short-term. Our focus is in supporting (mostly dense) flat collections of data with structure that is well visible to the compiler, because this is the scenario where generating code for accelerators and the such is the most feasible. Most metaprogramming facilities are on the complete other end of the spectrum, requiring recursive data structures and heavily branching logic, so it does diverge a bit from our original goal.

They do not have to be contradictory of course. We could allow a richer, and slower, language that makes it possible to implement macros and rewrites, with the eventual programs we emit being stripped away of all this inefficiency. But I think that it would require a fair amount of extensions to the language. Finally, my biggest worry would be with the fact that our internal program representations are somewhat unstable, and we do not want to commit to surfacing them as a user API, because that would require more stability. Dex is still a research project!

[zenna]

Thanks for the response.

Yeah, this makes sense. I haven’t looked at dex’s compilation model, yet, so I have no informed opinions on how feasible it is. You’re right you would at the very least need to able to represent and manipulate an AST (I see #331), but I doin’t think it would need to be particularly optimized in these areas.

As for the program representations, does it pose additional problems over any other interface? Julia has a few IRs, and in the cases when they’ve changed, people have updated accordingly.

@oxinabox’s raises a good point, in Julia we haven’t quite found the right interface yet. Either you have to work at the low level IR, or at high level method dispatch. In PL theory, semantics are often expressed compositionally in terms of rules in the grammar. I wonder if something like that could work — a bit like Cassette but able to handle control flow and constructs other than function application.

Anyway, it was just a thought! Definitely a nice-to-have feature, but at the same time, I think if addressed early could heavily shape the future of the language. Maybe I’ll take a look at it myself at some point.

[oxinabox]

We could allow a richer, and slower, language that makes it possible to implement macros and rewrites, with the eventual programs we emit being stripped away of all this inefficiency. But I think that it would require a fair amount of extensions to the language.

One possibility would be to have such compiler tranformed implemented by the user in Haskell -- kind of the opposite of homoiconic.
From a brief bit of reading it seems compiling and loading Haskell code at run-time is a thing that can be done.

Maybe something that takes in some form of IR or AST and outputs the same.
Idealy other things like optimization passes would also like this so you can do sthings like CSE -> MyPass -> CSE, since seems like we are finding that running bits of the optimizer before and after a pass is important (though kinda expensive).
and then expose an a way to make up a dynamically scoped section of code (like a function call?) to have this done to it.

Having a nice place in the compilation stack to put it, is also thing.
Julia ended up with it kind of by accident in @generated functions, (as per my the blog post that @zenna referenced at the start).
I maintain that @generated is not a great place to put it, as it runs too early, and interesting work by the compiler hasn't been done yet. I want it to run after the Control Flow Graph and dominance tree has been computed, since that is needed for interesting tranforms. Though unless you can return the updated structures they would need to be recomputed afterwards. idk enough about the Dex compiler pipeline to know how much these apply. This is right at the end of what i know about the julia pipeline...
I haven't been closely tracking the latest work on doing it properly in julia, and i don't know when that runs.

Finally, my biggest worry would be with the fact that our internal program representations are somewhat unstable, and we do not want to commit to surfacing them as a user API, because that would require more stability. Dex is still a research project!

I don't think we need to worry about this too much.
everything is somewhat unstable in Dex. Its v0.1.0-; some of the common keywords changed in the last 2 weeks.
I will note that the internal program represenentations that Julia has that get manipulated by Cassette/IRTools etc are unstable today, in Julia 1.6.
Cassette and IRTools have broken basically every minor release since they were created.
But everyone who uses them knows that is the case.
(One of the main reasons IRTools exists is to negate that, so that only IRTools needs to be fixed when the IR changes; because its API is stable because it works on the IR for the user rather than the user manipulating it directly)
Anyway, i think it would be reasonable to put something pointy and unstable infront of users as long as they know they are dealling with the deep magic and that it is unstable.

[oxinabox]

On further thought.
The API of Cassette might be perfect.
The problem with overdubbing functions is that it can't handle control flow.
Including the fact that things happen in order one after the other.

But Dex being functional doesn't have much control flow.
So maybe just overdubbing functions and something that allows similar for case and a few others might be enough.

It wouldn't let you write optimisation passes. But you could write ADs, I think, and debuggers in the style of MagneticReadHead.jl.