No stablehlo-legalize-to-linalg pass in the python binding

mamanain commented 4 days ago

What happened?

Hello! We need to convert stablehlo to linalg in our pipeline. I know there is a stablehlo-legalize-to-linalg pass which is available through stablehlo-opt but it seems like it's not available in the python binding. Looks like only passes in the transformations directory are registered there.

Am I doing something wrong? Is there a reason for this? Could we also add conversion transformations to the binding? mhlo used to have them there. Thanks!

Steps to reproduce your issue

#func îs some jax function
lowered = func.lower(*args)
stablehlo_module = lowered.compiler_ir(dialect="stablehlo")

  with stablehlo_module.context as ctx:
      ctx.append_dialect_registry(mlir.upstream_dialects)
      stablehlo.register_stablehlo_passes()
      pipeline = PassManager.parse("builtin.module(stablehlo-legalize-to-linalg)") #other passes from the transformation directories like stablehlo-legalize-to-vhlo work
      pipeline.run(stablehlo_module.operation)

-->>

error: 'stablehlo-legalize-to-linalg' does not refer to a registered pass or pass pipeline
stablehlo-legalize-to-linalg

Running the pass through stablehlo-opt works.

Version information

1c0b606503aac4f8e01f5511b0a10418bf93e7a6

GleasonK commented 4 days ago

Dup of https://github.com/openxla/stablehlo/issues/2559, let's resolve this in the other ticket. I can try to get to this this week.

What python bindings are you using? The StableHLO pip wheel or the one via JAX?

mamanain commented 4 days ago

Dup of #2559, let's resolve this in the other ticket. I can try to get to this this week.

What python bindings are you using? The StableHLO pip wheel or the one via JAX?

Oh, sorry for missing it. We use the one provided by jax. Is it ok if I try doing it myself? Seems like the fix shouldn't be very hard.

openxla / stablehlo