This RFC is a proposal to bring the MHLO and LMHLO dialects into the official MLIR repo and continue their further development upstream. The
lmhlo dialects provide a lowering path from the MLIR TensorFlow dialect, and there are conversions out of these dialects to lower level dialects already in MLIR; in addition, these dialects also enable getting into and out of XLA proper. Most of the motivation for having these ML dialects upstream was already spelt out here: RAMBLE: How to position new ML dialects in tree but that discussion became more general due to its focus on several other dialects. Here are a few things to note reg.
The Tensorflow repository uses these dialects: the MLIR TF dialect is converted to
mhlousing a few passes (mostly
-xla-legalize-tf-control-flow). There are also translators that go back and forth between XLA proper and
mhlo. Moving these dialects to the MLIR repo would mean that TF would depend on them through its existing LLVM/MLIR dependency. This frees TF from non-TF-specific MLIR dialects, moving them to MLIR where the latter’s organization, coding convention, and development style naturally fits. (Creating yet another third repository (with
lmhloas out of tree dialects) would be too chaotic and cumbersome due to the number of things that will be in play in getting TF to use MLIR: they will all have to align on the LLVM/MLIR version with mhlo/lmhlo sandwiched between TF at one end (one of its main users) and LLVM/MLIR on the other. This option I feel would be worse than what we have today giving people with too many things to look at and track.)
The current infrastructure in mhlo/lmhlo is mostly “dialects ops”, their canonicalizations, and dialect conversions into and out of them, without any elaborate analysis or transformations. In that sense, the code’s footprint is lower than those of dialects with many transforms. There are some transformations that work on or apply to them: eg. copy removal and buffererization, but nicely, these were at some point moved upstream and made to work with op interfaces (for eg. CopyOpInterface) - as such, these dialects transparently benefit from upstream passes.
One concern that came up here RAMBLE: How to position new ML dialects in tree was the impact on upstream maintenance. I’m happy to maintain or help maintain these dialects including aligning them with conventions used upstream. This of course doesn’t reduce change impact when the core changes, but again recall (2).
In terms of what tools should link in these dialects, I think
mlir-optshould link in these dialects the way it is now. Although it’s already 30 MB in size on a release build, in the future, these dialects along like
tosaand other things with ops and transforms on tensor types could be removed from
mlir-optand added to a new binary
Other thoughts, comments, and suggestions are welcome.