XLA Computation to MLIR Module for METAL Support

[Sean Moriarity]

Hi everyone,

I am one of the maintainers of an XLA library in the Elixir ecosystem. We are currently transitioning from the old XlaBuilder pattern to depending on MLIR, because as we understand that's the route JAX has gone and also because we would like to support METAL.

For the time being, we're relying on `ConvertHloToMlirHlo` to go from XLA computation -> MLIR Module for METAL support. Everything works fine on the host platform; however, when attempting to run computations on the METAL platfrom we get legalization errors:

```

<unknown>:0: error: failed to legalize operation 'func.func'

<unknown>:0: note: see current operation: 

"func.func"() ({

^bb0(%arg0: tensor<1000000xf32>):

  %0 = "mhlo.exponential"(%arg0) : (tensor<1000000xf32>) -> tensor<1000000xf32>

  %1 = "mhlo.tuple"(%0) {xla_shape = "(f32[1000000]{0})"} : (tensor<1000000xf32>) -> tuple<tensor<1000000xf32>>

  "func.return"(%1) : (tuple<tensor<1000000xf32>>) -> ()

}) {function_type = (tensor<1000000xf32>) -> tuple<tensor<1000000xf32>>, sym_name = "main"} : () -> ()

```

This happens regardless of the operation. The same computation runs fine on host, just not on metal. Is there something I'm missing? 

Thank you,

Sean

[Jacques Pienaar]

Hey Sean,

Could you perhaps point to your pipeline?* The import seems to have worked to MLIR & MHLO, but fails post. I don't know what pipeline is being run there and with what input config etc.

Thanks,

Jacques 

* If pipeline not possible to point to in OSS, then perhaps just list of passes as well as the pass where the error is reported (MLIR pipelines should report this, if not run a debug build with print-before-all and threading disabled.

--
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/cad35345-8a27-4a0b-b84c-47e834460391n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Jan Pfeifer]

Not meaning to hijack the question, but is the XlaBuilder API also being deprecated ? 

I'm using it in GoMLX (github.com/gomlx/gomlx) a similar ML framework for Go in development.

[Jacques Pienaar]

There is a GitHub issue to provide API that produces StableHLO directly - whether that is a drop in for the current or not I don't know (we had a variant that was drop in for MHLO but got removed recently). Today one could of course just use the MLIR builders directly but they are more verbose than what one can achieve in a mostly closed op system (I'd say biggest part would be operator overloading cases - now some of those get folks in trouble with the C++ evaluation orders, but still more concise).

-- Jacques

[Jan Pfeifer]

Thanks Jacques, but I'm not sure I'm followed your comment. It's the XlaBuilder C++ API deprecated?

If yes, then indeed what would be the replacement? Would it provide shape result inference (like XlaBuilder),  or is it something one would have to reimplement themselves?

Cheers

[Jacques Pienaar]

I'm saying https://github.com/openxla/stablehlo/issues/30 is still open and I don't know more than that.

Shape inference is already provided by the MLIR builders and C interface (https://github.com/openxla/stablehlo/blob/556db3fcf2928a9fb8d2273b4dc1196716af0aad/stablehlo/dialect/StablehloOps.cpp#L207), so whichever way this goes one wouldn't need to do that yourself.

-- Jacques 

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

[Jan Pfeifer]

Oh, cool! Thanks for the pointers Jacques!