Migration from MLIR-HLO to StableHLO

[ashay]

Currently, torch-mlir contains a lowering from the torch dialect to the mhlo dialect. However, lowering to the StableHLO dialect seems beneficial for several reasons:

1. StableHLO supports versioning (thus forward and backward compatibility within major versions) and serialization.

2. StableHLO has a smaller codebase than MLIR-HLO, thus causing less churn in downstream projects like torch-mlir due to updates (e.g. reorganized directories, renamed targets).

3. Unlike the CI for the StableHLO repository, the CI for the MLIR-HLO repository does not use CMake (it uses Bazel), so the CMake build for MLIR-HLO breaks often. I fix these breakages in the MHLO green commits, but it's recurring work to fix them and submit patches upstream:

• [mhlo] build: fix missing dependencies on Chlo and Stablehlo tensorflow/tensorflow#57484
• [mhlo] add target_include_directories() for Stablehlo tensorflow/tensorflow#57504
• [mhlo] add missing dependency tensorflow/tensorflow#57655
• build: tell CMake about dependency on MhloToArithmeticConversion tensorflow/tensorflow#58880
• build: fix CMake rules for the MHLOTOSATransforms target tensorflow/tensorflow#59019
• mhlo: fix build errors and warnings in tosa pipeline tensorflow/tensorflow#59187
• [mhlo] fix build rules for the MhloPasses target tensorflow/tensorflow#59219
• mhlo: fix CMake build for tosa pipeline tensorflow/tensorflow#59495

So my questions for the TorchToMhlo maintainers are: (a) would there be any loss in functionality if we switched to using StableHLO and (b) would you (or someone familiar with that code) be willing to update the code to switch to StableHLO please? Thanks!

[ashay]

cc: @tanyokwok, @ZihengJiang, @Vremold

[burmako]

Hi everyone! Smaller footprint is one of the reasons why we started StableHLO in a separate repository, and it's great that this could potentially come in handy.

We have recently migrated JAX from producing MHLO to producing StableHLO (jax-ml/jax@a1480c4), and that was fairly straightforward thanks to how close StableHLO is to MHLO. Based on this experience, I think that StableHLO would be a good fit for Torch-MLIR as well.

We'll be happy to support the migration effort - please let us know if there's anything we can do to help.

[joker-eph]

Unlike the CI for the StableHLO repository, the CI for the MLIR-HLO repository does not use CMake

Actually I'm pretty it does, before Bazel files were even added to this repo. I filed tensorflow/mlir-hlo#58 to make it publicly visible.

so the CMake build for MLIR-HLO breaks often.

This may be a coverage problem with the cmake configuration used by the CI, so in the same ticket as above I mentioned that open-sourcing the script used by the Kokoro CI (even if the runs aren't visible) may already help figuring out why you found issues that weren't uncovered by the CI?

[ashay]

My apologies, I wasn't aware of the CMake build for MHLO. Thanks for creating the issue! That should help me understand why my builds fail (or what else needs to be added to the existing build).

[ashay]

Here (#1840) is an initial attempt at the migration.

[ashay]

Closing, since #1840 resolves this issue.