Custom JAX to MLIR translation

[NicolaLancellotti]

I am interested in translating JAX to MLIR so that it can be used as input to other tools such as IREE. In particular, I am interested in custom translation, where an operator is translated to custom StableHLO.
I am aware that it is possible to define new JAX primitives, but I wanted to know if there is a way to define a custom translation for existing JAX primitives.
Also, I would like to know if it is possible to generate operations of other MLIR dialects besides StableHLO.

[hawkinsp]

It's certainly possible. In some sense we already have two different MLIR lowerings of JAX right now:

• to StableHLO, which is in the somewhat misleadingly named https://github.com/google/jax/blob/main/jax/_src/interpreters/mlir.py, together with lowering rules in various locations around the JAX codebase. I believe IREE accepts stablehlo, fwiw.
• to a mixture of upstream dialects and some TPU and Triton -specific dialects as part of Pallas: https://github.com/google/jax/blob/main/jax/_src/pallas/mosaic/lowering.py

You could write a new lowering, starting from one of those as inspiration.

Note that these are internal APIs, so I won't promise stability, but for a research use or a prototype it should be fine.

[NicolaLancellotti]

Thank you very much. Do you think it is doable to extend https://github.com/google/jax/blob/main/jax/_src/interpreters/mlir.py and customise a subset of the primitives instead of writing a new lowering? Maybe replacing the MLIR definition associated with the primitives through mlir.register_lowering?

[hawkinsp]

If this is for a prototype, then yes, you could just override one or more of the lowering rules and emit different MLIR.

(I'd be hesitant doing that for anything other than research/prototyping, since that globally affects all users of JAX, notwithstanding issues like the lack of API stability.)