Xla: Any plans to support MLIR

[Arjuna197]

❓ Questions and Help

Do you have any plans to support pytorch->mlir?

[dlibenzi]

Our IR nodes have a Lower() virtual function, which today lowers to have.
When MLIR will be stable enough, the first integration step for it would be likely to plug behind the XLA builder (the thing we use to lower to XLA), and generate MLIR behind the scenes.
Eventually, assuming MLIR will reach stability at some point, we will likely convert the XLA builder lowering to the proper MLIR counter-part.

[cjolivier01]

Is it known whether anyone is actively working on this yet?

[dlibenzi]

We have no plan ATM of moving away from XLA until there will be a replacement which is providing an advantage over the current XLA framework.
When this will happen, will need to consider whether rewriting thousands and thousands of LOC will be worth the investment.

[bhack]

You could be interested in https://llvm.discourse.group/t/torch-dialects-gradients-and-bears-oh-my/1833

[cjolivier01]

I saw this recently, I did find it odd that pytorch/xla was conspiculously omitted in that blog, especially since it's the easiest path to MLIR. afaict, TF currently just converts normal xla hlo into MLIR via

mlir::OwningModuleRef HloToMlirHloTranslateFunction(
      llvm::StringRef input, mlir::MLIRContext* context)

so the lowered XLA could simply be converted in kind if that's desired. tbh,. I don't think you're going to come across a better tensor engine than this pytorch/xla even for an MLIR variant.
If sopmeone wanted to develop native MLIR lowering, I would be willing to add a pluginable custom lowering mechanism if they'd agree to merge it :)