MHLO Extraction from XLA Compiler about xla

[cheshire]

Pytorch/xla is exclusively for TPU

It isn't, we use it for GPUs as well.

[cheshire]

1. LMHLO doesn't exist anymore I believe.

2. Technically, roundtripping back to MHLO should work, did you try to use the available translate tool to convert the intermediate HLO result back to MHLO?

@GleasonK should be able to provide more details.

[GleasonK]

Hello!

Frameworks/compilers interested in an MLIR interface for XLA should use StableHLO instead of MHLO, which should be possible so long as the HLO passes/optimizations used are not hardware specific. As an added benefit this would provide an easy entrypoint from JAX/TF/PT -- see tutorials for details on this.

For exporting HLO to MHLO/StableHLO: PyTorch/XLA's export to StableHLO goes via HLO->MHLO->StableHLO, so can provide a good code of how to do this programmatically: stablehlo_helper.cc

As for APIs to "access intermediate results", I'm not sure what you have in mind, a custom pipeline can be built which runs the HLO passes of interest and then export to MHLO/StableHLO (see APIs above). Alternatively you could look into dumping HLO at a specific point and using xla-translate (example) to convert from HLO to MHLO, and mlir-hlo-opt (example) if you want to go to StableHLO via available tooling.

[GrandChariot]

@GleasonK Thank you! Your answer was indeed insightful. I will proceed to use the converted StableHLO from Pytorch/xla. I do have a few more inquiries, however:

1. Pytorch/xla is exclusively for TPU, but the hardware I intend to use in my compiler is not TPUs. Will there be any issues concerning optimization or compatibility if I utilize the StableHLO generated in this pytorch/xla?
2. To my understanding, the operations of MHLO and StableHLO do not correspond on a 1:1 basis. Does this mean we can still leverage all the optimizations offered by the XLA Compiler even if we use StableHLO instead of MHLO?