[RFC] Adding StableHLO transformations for CHLO, Shape, Canonicalizations

[Kevin Gleason]

Hello everyone! I’ve been noticing a bit of code duplication arise in the ecosystem lately, and figured it’s probably time to propose centralizing some of this code in the StableHLO repo. This is related to the recent openxla/discussions#106, I propose we keep discussion in this email thread since this proposal encompasses that discussion.

Out of scope: I am not proposing to make StableHLO the default transformation dialect, since in the current state of things this would introduce a large pain point of “some things happen in MHLO, others in StableHLO, we need to round trip to get what we want''. I believe the proposed changes have no negative impact on this story, and actually make for a more clear lowering story from SomeDialect→StableHLO→MHLO, as opposed to the current setup which relies on bi-directional conversions in a few places.

(P1) Migrate CHLO→MHLO lowerings to CHLO→StableHLO.

In the current setup, CHLO lives in openxla/stablehlo, but all lowerings live in openxla/xla and target MHLO. This has caused some users who need these lowerings, but don’t want an XLA/MHLO dependency to duplicate these passes (ex1, ex2), and other users to convert CHLO→MHLO→StableHLO (ex3). Migrating these to StableHLO and updating existing XLA uses of the pass to use CHLO→StableHLO→MHLO shouldn’t pose any additional risks, as lowering StableHLO→MHLO is well tested, and is a part of native JAX execution today.

(P2) Migrate Shape→MHLO to Shape→StableHLO

In its current state, this pass is primarily used after CHLO decompositions which may introduce shape operations (ex1, ex2), so it makes sense to migrate these passes together. This poses little risk for the same reason as CHLO, where lowering in this direction is natural and well supported. There is a use of this pass in MHLO→HLO (ex3) but replacing this with the new lowering path should be simple, I’m happy to prototype and de-risk if needed.

(P3) Add opt-in canonicalization patterns to StableHLO

Another growing ecosystem tech debt has been canonicalizations (ex1, ex2). While this will add some duplication with MHLO, I think it is worth it for the broader ecosystem to have centralized backend-agnostic patterns that can be used. Given StableHLO’s role as an input dialect, I propose we add these only as opt-in patterns, no transformations should be applied by default. As for which canonicalization patterns to add, I propose doing this on an as-needed basis for patterns deemed useful by StableHLO users, and think the above two example links are a good starting point.

Interested in everyone’s thoughts!

Best,

Kevin

[Mehdi AMINI]

This is a great direction! Very supportive :)

The problem of canonicalization is something that I have seen coming up frequently and that is deeply missing. 

I would be wary about what we're getting into about the "opt-in" aspect: first it does not make sense conceptually to talk about fine grain opt-in for canonicalization. It's really hard to keep invariants in mind when the number of possible configurations is unbounded. It's not clear to me also how you'd expose each possible individual pattern as an opt-in? 

You could say either "enable/disable canonicalization" globally as a single "opt-in switch", but then isn't it strictly equivalent to say "you can run the canonicalization pass or just don't"?

In MLIR the canonicalization patterns are just exposed through a single accessor on every operation: someone who would want to canonicalize only some operation can already do that from any MLIR dialect, by construction.

Note that MLIR Operations, on top of the canonicalizer patterns, can also implement a "folder" which turns them into a constant. The folder is invoked in many places, including while applying canonicalization. As such they can be disruptive and should be extremely conservative. This is an orthogonal concern from the canonicalization hook though: we don't need to define any folder to talk about canonicalization here.

--
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/2f226258-19ef-4cea-9d91-506d7a24ea95n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Stella Laurenzo]

Fwiw, I think that the term "canonicalizations" is used for a lot of things that it shouldn't be. True canonicalizations and folders reduce variants to a single canonical form. If that line can be held in a community dialect like this, it is useful to have such things be universal/automatic.

However, in ML, I see a lot more things that are actually opinionated compositions and decompositions. It is also useful to have a library of these things available for use, but which ones and how they compose is a decision that gets made by the user of the dialect in some context.

Even arithmetic simplification beyond some trivial stuff falls into this second category due to the variations on tolerance to rounding or size thresholds. Sometimes you apply different rules there even at different parts of the same compiler...

My two cents.

To view this discussion on the web visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/CANF-O%3DZSK-BzbeuRYxwrw0XtQrOSaPQTrtz6NhDL8Zzr6vfeWg%40mail.gmail.com.

[Jacques Pienaar]

I think that's saying: it should be actual canonicalizations and one should also have general simplifications accessible. Mehdi's point being: anyone that wants fine grained can already get it via an existing mechanism and so no need for an additional fine grained mechanism (while also agreeing that these should be scoped ones). And there is nothing automatic about canonicalizations, it is a separate pass, not something always injected already.

I think it would be a misnomer to call it a canonicalization pass if the patterns doesn't fit the concept even if separate 🙂 So SGTM to use the existing mechanisms for patterns that fit the mold.

-- Jacques 

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

[Kevin Gleason]

> If that line can be held in a community dialect like this, it is useful to have such things be universal/automatic.

[...]

> So SGTM to use the existing mechanisms for patterns that fit the mold.

I agree with the end vision that "true canonicalizers should be universally/automatically applied". I'm more hesitant that we know exactly which patterns are "canonical" today. I also don't think there's a good use for fine-grained control of these patterns, so I'm also not proposing that.

My thinking is that can be done in two steps:
1. Introduce a non-native pass similar to iree-stablehlo-canonicalize, a "general simplifications" pass as Jacques described it
2. Graduate simplification patterns to canonicalizers over time and use the native MLIR machinery for canonicalizations

I anticipate some breakages as we do (2) which should help identify non-universal patterns, or bad code that should be fixed along the way. Either way should help with integration pains.


Best,
Kevin

[Mehdi AMINI]

On Thu, Feb 1, 2024 at 12:41 PM Kevin Gleason <glea...@google.com> wrote:

> If that line can be held in a community dialect like this, it is useful to have such things be universal/automatic.

[...]

> So SGTM to use the existing mechanisms for patterns that fit the mold.

I agree with the end vision that "true canonicalizers should be universally/automatically applied". I'm more hesitant that we know exactly which patterns are "canonical" today. I also don't think there's a good use for fine-grained control of these patterns, so I'm also not proposing that.

My thinking is that can be done in two steps:
1. Introduce a non-native pass similar to iree-stablehlo-canonicalize, a "general simplifications" pass as Jacques described it
2. Graduate simplification patterns to canonicalizers over time and use the native MLIR machinery for canonicalizations

Makes sense to me!

I'd suggest not calling such pass "canonicalize" in any way to avoid confusion. Maybe "stablehlo-aggressive-simplification-patterns" or something like that, this could run both the canonicalization patterns + the extra ones. So that graduating some patterns to be canonicalization would change the behavior of this pass.

Thanks,

[Kevin Gleason]

CC @linuxlonelyeagle on GH who has proposed a few of the decompositions / canonicalizations mentioned above. Also they have opened https://github.com/openxla/stablehlo/pull/1984 which upstreams some of these proposed changes from IREE - Thanks!

My feedback on the PR matches the above discussion, i.e. name the canonicalization pass "stablehlo-aggressive-simplification-patterns" and slowly move patterns to native canonicalizers.

Best,

Kevin