Pure-Python Torch FX to MLIR importer

Hello all,

Has any progress been made on Pure Python Torch FX Graph to MLIR importer?

I see that following post mentions about it:

In torch-mlir/__init__.py at main · llvm/torch-mlir · GitHub , I do not see any mention of support for Torch Fx graph. When using torch MLIR with PyTorch 2.0, am I still required to get a Torchscript graph before calling torch-mlir.compile?

Thanks for your help!

@makslevental Is currently working on it [WIP] Pure-Python FX importer. by makslevental · Pull Request #1996 · llvm/torch-mlir · GitHub

1 Like

It’s almost done - stay tuned.

1 Like

That is great to hear! Looking forward to it! So, for now, is there any way to import a Torch Fx Graph to Torch MLIR through some sort of Torch Fx to torchscript converter as torch-mlir.compile seems to expect torchscript graph ?

Passing an FX graph to torch_mlir.compile should “Just Work” – it will be scripted/traced as usual.

Thanks Sean! The Fx graph captured by TorchDynamo seems to have issue while lowering to torch MLIR for Transformer models from HuggingFace. The Fx graph captured by make_fx seems to do a lot better and lowers to torch MLIR without issues.