Using StableHLO for ONNX to HLO conversion

[AlexandreEichenberger]

In light of the presentation on OpenXLA where they touted StableHLO as the preferred representation for exchanges between compilers/tools, would it make sense to migrate our support from ONNX to StableHLO?

My understanding is that at this stage, StableHLO might only be starting, and initially it will be very close to MHLO, but as MHLO evolves to better correspond to optimizations, StableHLO will remain more stable and focused on providing compatibility.

Does it make sense to the ONNX-MHLO contributors?

@ZihengJiang @rsuderman @lipracer @raikonenfnu (apologies if I missed some MHLO contributors)

[lipracer]

I'm sorry I didn't participate in the meeting. I can't offer any constructive suggestions.

[jpienaar]

StableHLO will be a purely interchange format as positioned at the moment yes (e.g., no transformations on it). Given the goal of stable interchange, it should be easy to vendor, provides some backward and forward compatibility (at least in conjunction with MLIR bytecode format, not APIs yet, details TBD), and reduce the hard code linkage (e.g., it could enable decoupling some of the coordination of git rev bumping).

Targeting StableHLO rather than MHLO from ONNX has some advantages. And for first while, the switch could be at "sed" level. I would suggest waiting a little bit until there is a clearer signal here and then doing cost-benefit analysis. But goal is to reduce cost for frontends/interchange only interops. [I'm very excited about bytecode, I think it'll flush out some issues here too, but I'd let things bake a bit until "desired temperature" before incurring cost - while in same breath I should add that unstable interop is an ongoing cost]

[burmako]

Thank you, Jacques! I agree that it is prudent to experiment with StableHLO before making the decision either way, and I'd love to assist however I can.

In the next 2-3 weeks, we're planning to turn StableHLO into a drop-in replacement for MHLO by adding the currently missing Python bindings and Bazel build, introducing the StableHLO <=> MHLO conversion, and then landing all that in MLIR-HLO. (More details in https://github.com/openxla/stablehlo/issues).

Once that happens, experimenting with StableHLO will be as easy as changing the CMake/Bazel target you depend on (from MHLO to StableHLO) and making some textual replacements. By that time, I expect we'll also have the stability guarantees documented, so that it is clear what the users can rely upon.

As far as the current positioning goes (e.g. the "no transformations" part), this is a fairly new area for all of us, so there is still quite a bit of flexibility. E.g. if it becomes apparent that it would be beneficial for StableHLO to support transformations, then we could revisit this. I imagine that a good way of validating this would be experimental evaluation, and we'll be working hard to make it as easy as possible (as discussed above).

[AlexandreEichenberger]

Thanks @jpienaar and @burmako for keeping track of this new and exciting development. Stable interface (once stable :-) sounds like a good fit for this effort, I appreciate your careful evaluation on the matter.

[burmako]

@AlexandreEichenberger It took one week longer than initially expected, but the goals that were set in my previous comment have just been achieved. Earlier today we landed a bidirectional converter between StableHLO and MHLO: b9a0847, which concludes bootstrapping StableHLO as a software product.

At the moment, StableHLO includes all MHLO ops except the following 10 ops: AddDependencyOp, AsyncDoneOp, AsyncStartOp, AsyncUpdateOp, BitcastOp, CopyOp, DomainOp, FusionOp, PartitionIdOp, XlaRngGetAndUpdateStateOp. 9 of these ops aren't used by JAX, ONNX-MLIR, TensorFlow and Torch-MLIR, so we didn't include them in StableHLO. The 10th - CopyOp - immediately folds into its operand, so we didn't include it either.

In a nutshell, StableHLO can now be used as a drop-in replacement for MHLO, with the added benefit of guaranteed backward and forward compatibIlity - see the RFC about this: openxla/stablehlo#115.

In the cases when MHLO is still needed, e.g. to run an MHLO pass or to lower to Linalg, it's just one hop away thanks to the StableHLO <=> MHLO conversions. We designed these conversions to be super robust (3.5k+ loc of exhaustive FileCheck tests + continuous integration that makes sure that these tests stay in sync with StableHLO and MHLO evolution), and we'll be maintaining them as part of MLIR-HLO moving forward.

I'll be talking about this at the OpenXLA community meeting tomorrow 9/20 at 8am PT (see invite), and I could give a presentation at one of the upcoming ONNX-MLIR meetings too. Please let us know what you think!