Pure-Python Torch FX to MLIR importer

Incubator torch-mlir
1

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!
Kumar

2

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

1 Like
3

It’s almost done - stay tuned.

1 Like
4

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 ?

5

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

6

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.