[RFC]: Switching to StableHLO-based input conversion pipeline in IREE

[Jakub Kuderski]

Hi folks,

TL;DR: We propose a plan towards switching to the StableHLO-based input conversion pipeline in IREE.

IREE accepts multiple input formats, including StableHLO and MHLO. The input processing pipeline is decided based on the --iree-input-type flag, and can also be decided automatically using --iree-auto-input-conversion. Currently, the default pipeline for both StableHLO and MHLO is called mhlo, as it does lightweight preprocessing using the MHLO dialect, before being converted to Linalg on tensors + IREE Flow/Utils dialects.

With mlir-hlo winding down (https://discourse.llvm.org/t/sunsetting-the-mlir-hlo-repository/70536), we decided to port the input conversion pipeline to ingest StableHLO directly. This work is being tracked by this IREE issue: https://github.com/openxla/iree/issues/12678. The new input conversion pipeline, enabled with --iree-input-type=stablehlo is nearing feature-parity with the mhlo one and can already ingest all of the IREE TF benchmarks.

We plan to make the stablehlo input conversion pipeline be the default and would like to request community feedback on this plan. We expect the switch to be a transparent change, as with StableHLO as the input format, input conversion is an implementation detail. Ideally, we would like to avoid having to maintain both pipelines over a long time and propose to phase out the mhlo pipeline gradually, and give the community time to signal any issues discovered along the way. 

The proposed transition plan is:

1. In 1 week:

1. Allow for the stablehlo pipeline to ingest MHLO as an input format and immediately convert it to StableHLO.

2. In 2 weeks:

1. Switch all the compatible benchmarks to --iree-input-type=stablehlo.

2. Port the tuple flattening pass (--iree-input-type=xla). 

3. Make the stablehlo pipeline be the default input type for --iree-auto-input-conversion.

4. Move all feature development work (e.g., collectives) to the stablehlo pipeline.

3. In 3 weeks:

1. Rename the mhlo pipeline to mhlo_legacy.

4. In 4 weeks:

1. Remove the mhlo_legacy pipeline.

2. Remove support for MHLO as an input format.

3. Remove mlir-hlo from IREE submodules.

Let us know what your thoughts are. If you have some inputs or use cases that you care about that may not be covered by our test suite, feel free to send them my way.

Regards,
Jakub 

[Geoffrey Martin-Noble]

I think we should keep the stablehlo input pipeline only accepting stablehlo. Rather than putting this immediate conversion in that pipeline, why don't we provide a separate shim over the stablehlo pipeline for doing this immediate conversion from mhlo to stablehlo? We can introduce this pipeline as mhlo_shim (IDK, naming is hard) and we can preserve it for longer than the full mhlo pipeline. We can preserve the old pipeline as mhlo_legacy for a shorter time. That way migration is clearly tracked with the usage of a flag rather than needing to inspect the content of all inputs. Concretely

1. In 1 week

a. Introduce mhlo_shim input option to ingest MHLO as an input format and immediately convert it to StableHLO.

b. Rename mhlo pipeline input option to mhlo_legacy

c. Switch all the compatible benchmarks to --iree-input-type=stablehlo.
d. Port the tuple flattening pass (--iree-input-type=xla).

e. Move all feature development work (e.g., collectives) to the stablehlo pipeline.

2. In 2 weeks:
a. Make the stablehlo pipeline be the default for all stablehlo input passed with --iree-auto-input-conversion.

b. Make the mhlo_shim pipeline be the default for input containing mhlo passed with --iree-auto-input-conversion.

3. In 3 weeks:

a. Remove the mhlo_legacy pipeline.

4. In 4 weeks:

a. Remove the mhlo_shim pipeline.
b. Remove support for MHLO as an input format.
c. Remove mlir-hlo from IREE submodules.

So from the perspective of a downstream user updating IREE relatively frequently they will need to update as soon as they hit 1, but in a way that is trivial to fix (rename the flag) and makes it obvious that a more comprehensive fix is needed. At that point they can also switch directly to mhlo_shim point, which will give them a bit longer to fix and allow them to get any new feature work in the meantime. At 3, if they switched to mhlo_legacy before, they'll break again, but will still have the out of mhlo_shim to process MHLO. At 4, they'll really just have to switch.

If there are users reliant on MHLO input for whom the proposed timeline is too aggressive, we could consider extending the point at which we do 4, without having to keep around a parallel lowering path any longer.

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/CAG%2BB7V%3D3Cvut%3DMwLrr8HoBqaHB6kDFdZ6fy4vMSfs3nD5YvrLQ%40mail.gmail.com.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Jakub Kuderski]

Or as a third alternative, we could decide not to allow MHLO in stablehlo, rename the old pipeline to mhlo_legacy and not provide any shim. This way we would still clearly communicate that the MHLO pipeline is going away, and not risk introducing new bugs around the shim. If the shim is going to be short-live anyway (O(1-2 weeks)), would that provide much value?

WDYT?

On Thu, May 11, 2023 at 3:37 AM 'Jacques Pienaar' via iree-discuss <iree-d...@googlegroups.com> wrote:

Nice to see us getting close here!

If folks already switch to iree-auto-input-conversion rather than specifying a specific type (auto is new and may require a fix or too, in particular it doesn't detect stablehlo ops at all, and for error cases it would return different error) then they may not notice any change (e.g., perhaps in 1c let's just elide the type anywhere it doesn't matter). I'd also be pro emitting a warning when detecting the shim form in week 2.

Thanks,

Jacques 

(Typed on mobile, so it reads a bit like an airport conversation...)

You received this message because you are subscribed to the Google Groups "iree-discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to iree-discuss...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/iree-discuss/8b923b1f-f962-4ebe-a828-30d4b217ac8bn%40googlegroups.com.

[Geoffrey Martin-Noble]

My thought with the shim was to decouple "we accept MHLO as input" from "we maintain a separate conversion pipeline". Agreed that if the life of the shim is only a week, it's probably not worth it. Let's only go with a shim if we get feedback that people will need more time to transition off of MHLO as an input

To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/CAG%2BB7VnvsaNkh4FBbytHaXcwjUPRFyZkw_RFiYig9hJ11nNVyA%40mail.gmail.com.

[Jakub Kuderski]

Hello,

Here's a short progress update since I sent the RFC.

• Tuple flattening is supported and made available using `--iree-input-type=stablehlo_xla`.
• Missing canonicalization patterns that other conversion passes depend on were added.
• CHLO support landed earlier today.
• I'm now able to compile larger models like BERT_LARGE, RESNET, and T5_LARGE.
• (Related: Eugene posted an RFC to add the StableHLO -> Linalg+ conversion passes/patterns to the stablehlo repo: https://groups.google.com/a/openxla.org/g/openxla-discuss/c/KsRp9euuuB0.)

Since we haven't heard any pushback, I propose to continue with the initial plan but *without* any shims to ingest mhlo in the new pipeline. Concretely, I'd like to rename the old pipeline to mhlo_legacy and freeze mhlo feature development as soon as the remaining large benchmarks (listed above) are switched over to use the stablehlo input conversion pipeline.

Please give the stablehlo pipeline and report any issues or features/considerations that we might have missed.

Regards,

Jakub

[Stella Laurenzo]

+1 on proceeding. Thank you.

To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/CAG%2BB7V%3Dad7L-7vMyzrEFVh-n0wh39Dcq7vY9iN30ip3KOhKg8g%40mail.gmail.com.

[Jakub Kuderski]

Hi folks,

Short update on where we currently stand on the stablehlo migration timeline:

• Last Wednesday, we deprecated the mhlo input conversion pipeline and renamed it to `mhlo_legacy` and `xla_legacy`.
• Last Thursday, we made the `stablehlo` and `stablehlo_xla` pipelines the default for stablehlo inputs in the automatic input conversion pipeline selection (`auto`).
• Also last Thursday, we added a pass to convert mhlo inputs to stablehlo to eady the transition: `iree-opt --iree-convert-mhlo-to-stablehlo`.
• Over the last few days, we migrated all uses of mhlo in tests, samples, and documentation to stablehlo.

Based on this progress and general lack of known issues, we intend to move forward and drop mhlo input format support next Monday (June 12 2023). Please migrate any downstream use of mhlo and use stablehlo as the IREE input dialect instead -- the command above can help with the migration.

-Jakub

To view this discussion on the web visit https://groups.google.com/d/msgid/iree-discuss/CAEkedjhYzQfHEmX-i2zijvM0cdO%3DitHZbhZP6DA%3Dfh4BrunorQ%40mail.gmail.com.

[Jakub Kuderski]

Hi folks,

The migration is now complete. IREE no longer supports MHLO in favor of StableHLO, and we switched from the mlir-hlo to stablehlo repo as the submodule.

-Jakub

[Stella Laurenzo]

Fantastic! Thank you Jakub. Great follow through, and it is really good to have straightened this out.

[Geoffrey Martin-Noble]

Great work Jakub! I'm excited to see a cleaner dependency story and have us rebased on a community asset designed for stability.