Interpreter-as-transformer idiom for new languages?

[femtomc]

Hi — this is a bit of a meandering “open discussion” question, but it includes some technical questions about how I should think about the compilation model of JAX/XLA:

Short: have you considered moving the interpreter-as-code-transformer idiom closer to the next layer (e.g. MLIR or, at the very least, a custom dialect of MLIR — that then converts to XLA)?

The reason why I ask this is because I’ve been curious about exposing transformations as library code for new programming languages (for quite some time now -- starting with a fascination with the way that Julia performs type inference for optimization).

Obviously, JAX is a practical tool — it’s focused on accelerating programs which fit a specific picture (the one compatible with XLA) and focused on as wide an audience as possible (numerical Python writers). Ignoring this obvious focus on practicality, one thing I've often asked myself: what if you took a principled IR (some dialect of MLIR) -- potentially supporting higher-order features, etc -- and built the interpreter-as-transformer idiom expressed in autodidax and jax.core as the main way to write compiler passes?

In autodidax, the way that Jaxpr as a concept arises is to support jax.jit and jax.vjp -- so the direction I've discussed above doesn't quite follow that motivation. On the other hand, there exist compelling examples of non-standard interpretation interpreters which expose the design of custom lattices to users -- these interpreters require a Jaxpr (or IR-like) representation as well.

I think there are some interesting language design benefits if the compiler middleware above is constructed from first principles. For one, JAX is obviously focused on XLA -- but pulling out the interpreter idiom means your new toolchain doesn't necessarily need to be -- e.g. I could imagine lowering analyses which identify when code is XLA dialect compat, vs. must fall back on LLVM (or other runtime specific) toolchains. I'm also curious how much this solution "recovers Julia" modulo the incredible engineering required to support dynamic dispatch and specialization at runtime.

Also: how much of a benefit would there be to expressing the interpreter/transformer stack in a low-level language (vs. falling back on Python's interpreter to run the stack). Is tracing a bad bottleneck? Is XLA a significantly worse one?

[sharadmv]

the way that Jaxpr as a concept arises is to support jax.jit and jax.vjp -- so the direction I've discussed above doesn't quite follow that motivation.

Jaxpr has been repurposed for use-cases beyond AD and XLA. Pallas is a lowering from Jaxpr -> Triton IR, for example.

I think although the set of JAX primitives was designed without non-XLA lowerings in mind, it's easy enough to extend the set of JAX primitives to enable alternative lowerings.

[femtomc]

That’s a good point.

So then I guess the reason why JAX -> LLVM for Python programs with Python native control flow (e.g. to LLVM native control flow constructs like branching between basic blocks) is difficult is a limitation of tracing using interpretation to lower to Jaxpr?

Tracing handles both abstract type inference and Jaxpr construction — but it also removes native control flow, preventing its representation in the IR.

And — I don’t think I’m wrong here — tracing to get a Jaxpr requires all the XLA lowering information invariants to remain true. For example, I can’t have a primitive which doesn’t declare its shape under abstract interpretation. I need that information during tracing to actually make my Jaxpr.

So I can write things like JAX -> Triton — but if something like Triton supports a more expressive set of programs — there’s no analogous concept in JAX — so I’m restricting myself to that JAX/XLA compatible subset.

This may not have been clear from my message — but I’m curious about applying JAX’s idioms outside of that subset.

[sharadmv]

So then I guess the reason why JAX -> LLVM for Python programs with Python native control flow (e.g. to LLVM native control flow constructs like branching between basic blocks) is difficult is a limitation of tracing using interpretation to lower to Jaxpr?

Yea this limitation seems more fundamental since tracing in JAX is (currently) aware of Python control flow constructs. We would need to trace programs that use JAX control flow (lax.cond, lax.scan`, etc.).

tracing to get a Jaxpr requires all the XLA lowering information invariants to remain true.

This is not as fundamental. JAX traces with abstract values, which don't necessarily have to be shaped arrays as they usually are. In principle, we could extend JAX to have more abstract types that don't necessarily have shapes.