mhlo: migrate conversion to stablehlo

[ashay]

This patch replaces all MHLO operations with their StableHLO
counterparts and adds a validation pass to ensure that no MHLO operations
remain before translating all Stablehlo operations to the MHLO dialect
for further lowering to the Linalg dialect.

This patch also updates all lit tests so that they refer to the
convert-torch-to-stablehlo pass and so that they check for StableHLO
operations.

[burmako]

@ashay Oh that's a good catch! When we forked StableHLO from MHLO, we did a survey of frontends and looked into which MHLO ops are used and which ops aren't. Only the former ops made the cut, and copy was in the latter category. Sounds like we have missed the usage of copy in Torch-MLIR!

mhlo.copy is an interesting op. It was added a while ago - I imagine it was done for the reasons of parity with kCopy in HLO - but it didn't model the fullness of HLO's semantics, and became kind of a no-op (MHLO semantics involves immutable tensors, so a copy is just a no-op). Then, at a later point, this was recognized, and a folder was added which unconditionally replaced mhlo.copy with its operand.

I see that at the moment mhlo.copy is only used in one place in Torch-MLIR. I wonder if it's possible to just use its operand instead of creating this op in the first place:

torch-mlir/lib/Conversion/TorchToMhlo/Basic.cpp

Line 1382 in 5057e88

INSERT_UNARY_PATTERN(AtenCloneOp, mhlo::CopyOp);

.

[silvasean]

Nit: possibly in a separate PR, we can rename all the Mhlo in file paths / directories to Stablehlo (note h is lowercase officially in the dialect name).

Also when possible, updating the VerifyStablehloBackendContract pass to reject Mhlo would be good (in some sense that is the true milestone of completion of this workstream).

And if we can get sign-off from @burmako that it is indeed checking the correct set of ops: https://github.com/llvm/torch-mlir/pull/1840/files#diff-d8eb5bebe72fb7e5c0d49bd9c2234d6ac74e44f250d38e7f0d7af0fee52530d7 -- two questions:

• Is CHLO part of the StableHLO set?
• Is shape::ShapeOfOp part of StableHLO set?

[ashay]

Nit: possibly in a separate PR, we can rename all the Mhlo in file paths / directories to Stablehlo (note h is lowercase officially in the dialect name).

Fixed in this PR.

Also when possible, updating the VerifyStablehloBackendContract pass to reject Mhlo would be good (in some sense that is the true milestone of completion of this workstream).

Since the TorchToMhlo pipeline ultimately lowers to Mhlo, I don't think we can get rid of the VerifyMhloBackendContract pass. So I added the VerifyStablehloBackendContract pass after the torch-to-stablehlo conversion while retaining the existing VerifyMhloBackendContract pass.,

[silvasean]

Since the TorchToMhlo pipeline ultimately lowers to Mhlo, I don't think we can get rid of the VerifyMhloBackendContract pass. So I added the VerifyStablehloBackendContract pass after the torch-to-stablehlo conversion while retaining the existing VerifyMhloBackendContract pass.,

It should lower to StableHLO and not Mhlo. The later lowering to MHLO is really just used as part of e2e testing since only mhlo has a path to linalg. Other than that incidental aspect MHLO is irrelevant to Torch-MLIR. (StableHLO has a path to MHLO, but that is not our project's responsibility)

Concretely we would put lowering stablehlo->mhlo here:

torch-mlir/python/torch_mlir_e2e_test/mhlo_backends/linalg_on_tensors.py

Line 39 in 9536174

"builtin.module(func.func(symbolic-shape-optimization),func.func(hlo-legalize-to-linalg),func.func(canonicalize))",

(and that whole directory structure also needs to be renamed to stablehlo, as well as the e2e test config)

[burmako]

As far as VerifyStablehloBackendContract goes, how about we populate it with the contents of the current VerifyMhloBackendContract with the only change being mhlo::MhloDialect replaced with stablehlo::StablehloDialect, and then over time work on shrinking this contract?

Here's the rationale for this proposal: StableHLO is ready to serve as a drop-in replacement for MHLO in the role of an interface between ML frameworks and ML compilers. It supports everything that MHLO supports at the interface level, and then some, but some future work still remains, namely:

1. We are planning to replace func.func and func.return in StableHLO programs with StableHLO-specific ops (Spec and implement Module/Func/Call/Return ops openxla/stablehlo#425).
2. There is no standard way of writing dynamically-shaped programs at the moment. The spec doesn't take a position on this at the moment of writing, because it's ongoing work (Write a dynamism RFC openxla/stablehlo#8).
3. CHLO is not part of the spec at the moment of writing, and I'm starting to think that maybe it shouldn't be (Create a plan for CHLO openxla/stablehlo#602).

Given that, the existing pieces of VerifyMhloBackendContract make sense:

• target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(opHasLegalTypes) corresponds to 1).
• target.addDynamicallyLegalOp<shape::ShapeOfOp>(opHasLegalTypes), target.addLegalDialect<tensor::TensorDialect>(), target.addLegalDialect<arith::ArithDialect>() correspond to 2).
• target.addLegalDialect<chlo::ChloDialect>() correspond to 3).

As we're working through these open questions, I expect that we'll be able to eventually shrink the backend contract. For example, one potential future could involve: a) keeping builtin.module, b) introducing stablehlo.func and using stablehlo.return for all returns, c) vendoring all shape computations into the StableHLO dialect, obviating the need for ops from arith, shape and tensor dialects in StableHLO programs, d) deconstructing CHLO to StableHLO. In that future, the contract would shrink to allowing just the ModuleOp and the StableHLO dialect.

What do you think?

[silvasean]

Thanks Eugene, this makes sense to me! Looking forward to collaborating to tighten this up!

[silvasean]

LGTM, but let's wait for one of the owners of StableHLO support in Torch-MLIR approve as well.

[tanyokwok]

LGTM!