Status: Approved
Initial version: 5/28/2023
Last updated: 6/12/2023
Discussion thread: openxla-discuss.
Sunsetting the MLIR-HLO repository has announced that Google is planning to wind down the development effort and support of the MLIR-HLO repository. Over years, this repository has become an important community asset but it has also been hitting scalability issues - from technical, organizational and infrastructural perspectives.
In our experience, the primary usefulness of MLIR-HLO lies in providing direct access to: 1) the MHLO dialect that can serve as a portability layer between ML frameworks and ML compilers, 2) MHLO passes, e.g. hlo-legalize-to-linalg, which connect MHLO to the larger ecosystem. Both functions have been used successfully in multiple projects.
The first function has been successfully served by StableHLO over the last six months and an increasing number of projects have been migrating from MHLO to StableHLO as the portability layer of choice.
However, the second function - which boils down to MHLO => Linalg lowerings - is still an open question. As Mehdi pointed out in OpenXLA overall architecture & components, StableHLO was originally designed “decoupled for MHLO, which in turn is positioned as a compiler IR”. This design decision allowed to focus StableHLO on just the matters of portability (which was already plenty challenging), but as a result its transformation story has been left undeveloped.
Recently, Jakub from the IREE team has been exploring a potential solution to the second problem. openxla/iree#12678 goes into low-level details, but in a nutshell the idea was to take all MHLO-based passes that IREE uses, migrate them to use StableHLO and see what happens. This work is almost done, and it's been a success - now there is a proposal to sunset the MHLO-based input conversion pipeline in IREE.
Furthermore, folks from the Torch-MLIR project have recently been thinking about the second problem as well (llvm/torch-mlir#2177). Torch-MLIR is also using MHLO => Linalg passes from the MLIR-HLO repository, and while these passes are available in the XLA repository, depending on it is not very convenient logistically. Anush remarked “I would rather copy / fork if we have to than take an XLA dep because a few passes live there”.
I believe that this recent community exploration suggests that there is a need in another repository to fill in the role of MLIR-HLO and provide HLO => Linalg lowerings. This repository needs well-maintained CMake and Bazel builds, a process for regularly bumping LLVM revisions and a community-friendly development environment.
In principle, we can create a new repository for this purpose, decide on the scope, bootstrap the infrastructure and find folks who would be interested in maintaining all this. But we don't have to do any of that - we can just use openxla/stablehlo which satisfies all the requirements above, and this is what I would like to propose.
Q1: Doesn‘t this proposal contradict the mission of StableHLO?
A1: Indeed, the original mission for StableHLO was focused on just compiler input, so that MHLO can focus on being a compiler IR. However, this mission was formulated almost two years ago, and a lot has been learned during that time, e.g. that: 1) MHLO has a lot of cruft, to the extent that it’s unclear how to evolve it into an awesome compiler IR, 2) the idea of splitting StableHLO and MHLO comes at a significant maintenance cost, but it's unclear whether it really carries its weight.
Q2: StableHLO has compatibility guarantees, so how can it be a good transformation dialect?
A2: This hasn't been discussed much, but within the StableHLO project there are actually two different dialects - StableHLO which is the interface that producers and consumers are using, and VHLO which is where compatibility guarantees are provided. As Stella put it, “the stablehlo dialect is actually defined in terms of an evolution process that is much closer to, say, LLVM IR than it is to a serialization format (i.e. the vhlo dialect and corresponding passes/utilities for serialization are what arbitrate the wire-compatibility guarantees)”.
Q3: One of the benefits of splitting StableHLO and MHLO is that MHLO can contain additional operations which are only relevant to compiler pipelines. Does this proposal mean that we'll start polluting the portability layer with these operations?
A3: This kind of pollution would indeed be undesirable, but the current design of MHLO being a copy of StableHLO + a few ops is not the only way to avoid it. There is an alternative design where StableHLO is used together with satellite dialects which add functionality and not duplicate it. More specifically, ops which are currently in MHLO but not in StableHLO could go into a new dialect (could call it stablehlo_ext or xla, etc).
Q4: Perhaps we should first align on the overall OpenXLA architecture and only then decide on this proposal?
A4: Aligning on the overall architecture is very useful, but I don't think it has to be a blocker for making this particular decision. Providing StableHLO => Linalg lowerings in openxla/stablehlo will immediately resolve an acute issue that multiple projects are facing, so I would like to propose bias for action. If a better place for these lowerings materializes in the future, these lowerings can be moved there - since the API will stay the same, the migration would be easy even if the implementation changes.
Q5: What is the relationship between this proposal and TCP?
A5: To quote the TCP RFC, “TCP's mission is to be a mid-level dialect in MLIR that enables transformations that complement those in existing dialects, by supporting ops at a high-level of abstraction”. This is a promising project, which I think should be part of a long-term discussion about transformation dialects. However, its lowering to Linalg doesn't yet have feature parity with the MHLO => Linalg lowering, so it cannot yet provide a solution for the short-term problem.
Q6: What does this proposal mean for MHLO?
A6: MHLO dialect would remain as an implementation detail of the XLA compiler, which provides 1:1 parity with HLO and therefore a gateway to the wealth of functionality implemented with HLO. In the future, it may also dissolve into multiple dialects, but this is out of scope for this proposal. What happens to MHLO passes, e.g. the existing MHLO => Linalg lowering, would be up to their owners.