MLIR Dialect Import for Mojo

Are there plans for out of tree dialects to be supported in Mojo? Interesting MLIR-based compilers for specific hardware, like TT-MLIR, are starting to pop up. That one in particular should only require turning on the RISC-V LLVM backend since most of the rest can be done in normal Mojo at this point.

1 Like

Hi Owen! There aren’t plans in the roadmap, but this is absolutely an aspirational goal for Mojo as a native frontend for MLIR.

In principle it’s not hard to do, the standard library uses tools to generate the dialects we use in our compilation stack for example. In practice there are a few big hurdles we’d need to clear first:

  • How do we tell the parser how to import your dialect? In principle this isn’t difficult, we just need to pick a way. In practice there’s some lurking complexities in MLIR where it’s not possible in general to link against two different libraries that both compile against MLIR and send objects between them. We would need to provide a general solution to this problem.
  • Once you are able to generate your new dialect, how do we make it practical for people to use?
    • Your code will likely always translate to a mixture of dialects. Say you use an if statement or a function call; that will be a native Mojo if or call and not one from a dialect you may want to be using. This may not be an issue in the specific case you mention, but limits its usefulness in general.
    • Our compiler stack doesn’t know how to lower the custom dialect, so there’s no way for us to translate it down. I think this is realistically the biggest unknown. One could imagine Mojo providing native MLIR metaprogramming tools where you could write your lowerings as a Mojo function and register it with the compiler. With less effort we could allow registering and linking against custom passes in our compiler.

Hopefully that gives a taste of some of the problems here. I’d love to get there, but one step at a time!

2 Likes