RFC: Rename torch_mlir.compile APIs and introduce FX based analogs

Incubator torch-mlir
1

Hi folks – I’ve been fielding numerous questions/confusion over the last few weeks related to torch_mlir.compile being the API for the project. Most people are (rightly) confused because this is central and called out in all of the docs as the one stop shop. However, it doesn’t give access to the current generation work that is being driven via the fx_importer. It is also hopelessly tied to the old TorchScript path.

I propose that:

  • We rename torch_mlir/__init__.py to torch_mlir/torchscript.py
  • Create a new torch_mlir/fx.py for user level APIs related to importing modules via FX.
  • Rename the long_term_roadmap.md to roadmap.md and update it to reflect these two paths and that the first is deprecated, while the second isn’t quite ready yet.
  • Update all project documentation to provide clarity on what each entrypoint provides.

I’d like to avoid the situation in the future where one generation of the technology squats on the whole project namespace so would like to avoid having an __init__.py at the top level. While we do service end-users, we mostly service integrators, and I would prefer to keep the APIs structured as such.

(edit: fix typo called out by Ramiro in torch_mlir.compile)

1 Like
2

I like the proposed changes. I’ve also felt the friction lately when trying to answer the question “how can I import this model into torch-mlir?” for people. The torch_mlir.compile tool is the most visible one (mentioned in docs and used in our examples), but it is limited and based on the old importer, while the new fx importer is a bit more hidden. The split into two Python files should make it easier for people to interact with the different tools available.

Also, I think by torch.compile you mean torch_mlir.compile (trying to avoid confusion with the PyTorch API).

3

Thanks, Ramio. I’ll let this sit for a few days and then send a patch later this week.

(thanks for the typo fix. edited the original)

4

+1 on proposed changes (having just [re]started the sparse work in torch-mlir, I too, first looked at all the -now almost- legacy compilation paths before being pointed to the new fx importer by Ramiro/Stella)

5

Someone on my team was nice enough to put a patch up: Rename torch_mlir.compile APIs and introduce FX based analogs by saienduri · Pull Request #2842 · llvm/torch-mlir · GitHub

Would be good to get a person or two here to review, given that it impacts multiple people.