Torch-mlir operation registry access

Hi folks – we’ve been doing a sprint to uplift torch-mlir to be fully based on torch.compile and torch.export, eliminating the need for the TorchScript era interop. The work is going well, and I’ve been really pleased to dig in and see the amazing progress the team has made on the PyTorch side in bringing all of these features together. (side note: I was elsewhere for a couple of years and it is rare to step away and come back to find so many things that were causing issues before to have been overhauled and vastly improved – kudos to everyone)

In any case, we’ve now got a path working without any C++ dependency on PyTorch, and it is doing pure Python bridging to the torch dialect in MLIR, using only PyTorch and MLIR’s Python APIs. It’s really nice to have the features (dynamic shapes being the big one I’m happy about) working this way for both eager and exported paths, and I’m looking at how to clean up the rest and sequester the old TorchScript paths to a legacy part of the project.

All that’s left that is load bearing is this getRegisteredOps API that we added ages ago. It is used by project maintainers to introspect the PyTorch registry of ops and do code generation for a majority of the torch dialect ops. It is only used by our utility to generate that file. I remember when I wrote this many years ago, there wasn’t an easy way to get at this information from Python. I figured I’d ask now if anyone had a suggestion… if we eliminate this piece, then the core of the torch-mlir project just becomes dialects and compiler infra with no build time dependency on PyTorch, which would be really nice.

Let me know if anyone has suggestions, and it’ll save me some effort of poking around and seeing what the right way to do this now is.


(and again: thanks so much for the great work, especially on the export path. it is such a vast improvement over what was there before)