Using Torch-MLIR as a front-end for MLIR files

[rengolin]

We want to use different front-ends to look at the IR they produce and see if we can feed them into our optimising pipeline. Torch-MLIR seems to support a lot of the features we want, mainly being able to handle different MLIR dialects for the same models.

The idea is to get a model in Python, run through Torch-MLIR and get the MLIR file in the end. By looking at the tests, I see that there's a python folder with the JIT IR exporting MLIR files via ModuleBuilder. But it seems it only outputs the Torch dialect.

Further looking at the tests, I see in the Conversion folder that there's a torch-mlir-opt that converts the Torch dialect into the others. While I can see MHLO and TOSA, I can also see Linalg, SCF and Std separate from each other. From the docs, I assumed there was one that used all three together.

So I have two questions:

1. Is there a way to use the ModuleBuilder to export directly to a particular dialect? If not, is torch-mlir-opt the best option?
2. Can I export a full model to Linalg with torch-mlir-opt? Will it use SCF and Std for the remaining ops automatically?

@silvasean

[silvasean]

I think you are looking for the torch_mlir.compile Python function.

Examples: https://github.com/llvm/torch-mlir/blob/main/examples/torchscript_resnet18_all_output_types.py
Tests: https://github.com/llvm/torch-mlir/tree/main/python/test/compile_api

We definitely need to surface this better in our docs I guess -- where would you recommend we talk about / mention this API?

[rengolin]

Perhaps on the development page, right after the PYTHONPATH section?

I'll try to write something, just so I remember better next time, and then send a PR. Feel free to correct as much as needed (wording is not my strongest suit) before merging.