StableHLO + IREE Python compatibility

[Wheest]

The StableHLO Python bindings have been very handy.

This means I can run for example:

from mlir.dialects import stablehlo
from mlir.ir import Context, Module

# Define the MLIR string
mlir_source = """
module {
  func.func @gh_example() -> tensor<1xf32> {
    %cst = stablehlo.constant dense<4.200000e+01> : tensor<1xf32>
    %0 = stablehlo.add %cst, %cst : tensor<1xf32>
    return %0 : tensor<1xf32>
  }
}
"""

# Initialize context and register the StableHLO dialect
with Context() as ctx:
    stablehlo.register_dialect(ctx)

    # Parse the MLIR module from the string
    original_module = Module.parse(mlir_source)

    # Print the parsed module
    print("Parsed Module:")
    print(original_module)

And this will correctly parse the MLIR. I've been using this tooling for both parsing, trying out new passes (e.g., removing and updating ops), etc.

However, now I am using IREE more, which has its own Python packages (which bundles some of the core MLIR stuff).

This means I can do stuff like:

from iree.compiler import ir
from iree.compiler.dialects import flow, func, arith

This issue however is that since it bundles its own MLIR, its tricky to make this interoperable.

How would I register the StableHLO dialect with the IREE MLIR using the Python bindings?

Options:

1

from mlir.dialects import stablehlo
from iree.compiler import ir

with ir.Context() as ctx:
  stablehlo.register_dialect(ctx)

This crashes with:

    stablehlo.register_dialect(ctx)
TypeError: register_dialect(): incompatible function arguments. The following argument types are supported:
    1. (context: MlirContext, load: bool = True) -> None

Invoked with: <iree.compiler._mlir_libs._site_initialize.<locals>.Context object at 0x7c9d4cf6f770>

2

from iree.compiler.dialects import stablehlo

Fails with ImportError: cannot import name 'stablehlo' from 'iree.compiler.dialects'

I'm not sure if this is more of a StableHLO issue or an IREE issue. But I know there's a fair amount of dev overlap.

[GleasonK]

MLIR python bindings need to all be built together / refer to the same MLIR .so file to make sure all things interop well, see how JAX does this for example to expose StableHLO/arith/etc:
https://github.com/jax-ml/jax/blob/c1af808c8cc50fdd449544233e6038570ba8addb/jaxlib/mlir/_mlir_libs/BUILD.bazel#L340-L432

There are some global IDs that are instantiated by MLIR and if you import from two separate MLIR builds, they don't interop well (i.e. each instance has its own BuiltinDialect pointer used as a TypeID which doesn't work out well with isa, etc).

Likely your solution (2) is the path forward, which requires support from IREE folks, would recommend their discord server:
https://discord.com/channels/689900678990135345/1085666465061019728

[ScottTodd]

For option 1, if you can install both packages and import them independently, maybe you can pass IR between files, rather than stay in memory?

For option 2, IREE could expose the stablehlo dialect in its Python bindings. We'd need to somehow fit https://github.com/openxla/stablehlo/tree/main/stablehlo/integrations/python together with https://github.com/iree-org/iree/tree/main/compiler/bindings/python . A recent PR exposing a different (internal) dialect was here: iree-org/iree#18804 the author @makslevental may have suggestions.

This would be fairly complicated I think, because:

• There is a nontrivial amount of code in https://github.com/openxla/stablehlo/tree/main/stablehlo/integrations/python
• IREE currently pulls in the StableHLO CMake build but we may fork that to get more direct control over dependencies: Fork our own CMake/Bazel build files for StableHLO iree-org/iree#18844 . Notably there, IREE would only pull it parts of StableHLO that it uses (e.g. excluding the chlo dialect, excluding the reference interpreter), while you might want to make use of the full project in your own usage.
• StableHLO support in IREE is optional (just like Torch and TOSA support), so the build system would need some branches. The release that IREE publishes do include all optional components, at least.

[GleasonK]

For option 1, if you can install both packages and import them independently, maybe you can pass IR between files, rather than stay in memory?

Would recommend this if possible, IIUC wouldn't this require being able to get an MLIRContext in python with StableHLO registered so you can parse the StableHLO-in-a-file? If there's a way to do that then I think it's a good suggestion.

get_stablehlo_asm from https://github.com/openxla/stablehlo/blob/main/docs/tutorials/jax-export.ipynb is an example of how you could pass StableHLO between contexts using some methods exposed from JAX, not sure if there's an IREE equivalent to the get_mlir_context function which registers everything needed.

[makslevental]

StableHLO support in IREE is optional (just like Torch and TOSA support), so the build system would need some branches. The release that IREE publishes do include all optional components, at least.

I don't have the whole scope in my head right this moment but I'm sure we could have at least basic support (registering/building/parsing/printing/etc ie minus their own Python stuff) in IREE pretty easily with enough CMake branches and/or #ifdefs. Do we want to do that and then maintain it - I have no idea. I don't know how hipri StableHLO is as a frontend? Anyway @ScottTodd if you think it can land then I can put together a PR next week when I get back from break and we can adjudicate then/there (ie with something concrete to look at).

[stellaraccident]

not sure if there's an IREE equivalent to the get_mlir_context function which registers everything needed.

I'm pretty sure that just iree.compiler.ir.Context() registers everything that IREE knows about.

In this case, trying to directly register dialects from stablehlo with IREE is triggering type violations because TypeIDs internal to the libraries are distinct. If this were somehow not the case, it would probably just crash catastrophically unless if the two were built at exactly the same commit/flags because MLIR itself provides no C++ ABI compatibility.

Indeed, the most practical thing to do is to serialize in one and load in the other. If doing this as part of a real tool, vs just to play, I would use the bytecode serialization APIs. And there are a couple of other tricks for "hardening" stablehlo with respect to version, etc. This won't work if you're trying to stitch IR together that uses both stablehlo and IREE internal dialects (please say if this is the intent).

Note that if your intent is to invoke IREE's compiler vs just futzing with IR generation, then you either want to invoke its compiler binary (iree-compile) or use the Python compiler driver API, feeding it bytecode or assembly: iree.compiler.api (see tests). This gives you the control to manage sessions, sources, outputs, share thread pools, etc in a way that typically works for JIT-like cases.

I don't have a real problem with including the stablehlo Python bindings in IREE's overall python bindings, but if history is a guide, we need to do some work to better isolate the stablehlo build as it has been a maintenance burden in the past due to overly broad inclusion of sources outside of the core IR and that do not have full platform test coverage.

[makslevental]

we need to do some work to better isolate the stablehlo build as it has been a maintenance burden in the past due to overly broad inclusion of sources outside of the core IR and that do not have full platform test coverage.

I'm not familiar with StableHLO - is there more here than meets the eye? I.e., looks straightforward to build/link only that target and its deps and get minimal support for bindings but maybe there are weird transitive deps lurking?

[stellaraccident]

Yeah, there's transitive deps and a whole reference interpreter that sneaks in and has a bunch of code that doesn't even build on windows. Probably nothing insurmountable but just something that has bitten us many times during the integration grind and never got around to looking into more. We were discussing vendoring it in a similar way we do for torch-mlir in order to get a narrower build dep. This would go the opposite direction of that.

[ScottTodd]

FWIW, stablehlo/integrations/c/CMakeLists.txt may be simpler to take a dep on than the C++ compiler components... though I do see a few innocent-looking deps that could actually be problematic for downstream use like StablehloReferenceApi in there.

[makslevental]

[Wheest]

I would use the bytecode serialization APIs. And there are a couple of other tricks for "hardening" stablehlo with respect to version, etc. This won't work if you're trying to stitch IR together that uses both stablehlo and IREE internal dialects (please say if this is the intent).

The area I was playing around with was to use some IREE internal dialects (e.g., Flow for representing dataflow), and putting StableHLO ops inside of them (e.g., inside a Flow dispatch inside a ExecutableOp to get "nested functions").

Therefore in this case StableHLO parsing, op insertion. and op validation would be needed.

For what I was looking at, StableHLO passes wouldn't be strictly needed, but hey, canonicalisation is always nice to run.

My initial assumption as a user was since StableHLO is a frontend to IREE (and actually listed first here), that it would be relatively natural to mix them. Of course, it's MLIR and if I popped the hood on the C++ side I can do whatever I want. On the Python side, from the above discussion I can see how the different bundled versions of the MLIR builds could cause problems.

Is that a wider question for MLIR-based projects that expose Python bindings in future? I see two somewhat conflicting user-use cases, a) where one might want a bundled MLIR build inside the package to use that tool only (e.g., use IREE from Python); or 2) have multiple MLIR python packages that one wants to work together (e.g., use IREE+StableHLO+FooBarMLIR+etc).

[makslevental]

Is that a wider question for MLIR-based projects that expose Python bindings in future? I see two somewhat conflicting user-use cases, a) where one might want a bundled MLIR build inside the package to use that tool only (e.g., use IREE from Python); or 2) have multiple MLIR python packages that one wants to work together (e.g., use IREE+StableHLO+FooBarMLIR+etc).

This is a known, perennial point of friction. See this sub-discussion (of a much larger, broader scope discussion) here. It would be extreme to say it cannot be solved but it certainly not likely to be solved in the near or medium term.

[stellaraccident]

MLIR is fundamentally not modular across the C++ ABI boundary, and that means that, realistically, there is no such thing as "two MLIR compilers": there are just two compilers that happen to use some MLIR tech. I've been around for a long time, and I don't see this changing anytime soon. It's just one of those really hard problems. Further, IREE doesn't really believe in the view of "MLIR" as completely composable components: if IREE didn't wire it together, it's not part of IREE. Triton is the same way. Doesn't mean they can't work together, but it is more at the level that compiler components have always worked together and requires careful interface points to be defined.

With that said, we can make this work by including the stablehlo python bindings in IREE. Because each bit of python binding is forever maintenance and the compiler really is written in c++, we've historically only done this on demand as needed. I would say that the stablehlo python bindings are quite a bit more intricate than others, and this is a point of philosophy: IREE generally uses the simplest thing for IR generation at that level, even at the expense of the API being nice.

Generating torch and ONNX IR for example is done in python using only generic APIs and without additional, dedicated Python APIs. This was done because python APIs are generally slow, these are performance sensitive areas, and we anticipate needing to provide a faster path some day -- wanting a narrow API to replace vs a "nice" and highly ergonomic API that has a lot of touch points.

But the full python APIs are certainly better for hacking. I'm not opposed to seeing what it would take to put these together. It's probably a relatively easy patch that you could carry on a branch. Landing it will require us to do some of the deferred maintenance on the stablehlo side that has resulted in it being a fat dependency with some portability challenges.

[stellaraccident]

Scanning through the code, it's in better shape than I remember from a while back wrt deps. We still may want some cmake isolation. Would need to check with @ScottTodd on recent portability issues.