PyTorch to Calyx Lowering Path

[jiahanxie353]

Hi folks,

I'd like to bring up the conversation regarding lowering PyTorch to Calyx in the CIRCT compiler, as mentioned in Theme 1 of this discussion.

After some research, there are two options here:

1. Using the Torch-MLIR project;
2. Using allo (which is not open source for now);

Option 1: Using Torch-MLIR

If we were to take this route, I believe we would be involved in both:

1. The frontend, i.e., the path from PyTorch Torch-MLIR (since they are refactoring the frontend according to the roadmap, I'd imagine we would more or less try to address some frontend issues);
2. The backend, i.e., the path from lowering Torch-MLIR to Linalg, and then to Affine/SCF, etc. This path can then be integrated into our existing lowerings in the CIRCT project. All we need to do is add more support to the lowering paths to Calyx from all different dialects.

Option 2: Using Allo

Alternatively, we can use Allo as our frontend, which takes machine learning models in Python, and it can directly lower to MLIR dialects, such Linalg, Tensor, etc. Specifically, it can be lowered to the AMC dialect (which is also not open source), which can in turn be lowered to Calyx if we want.

My Thoughts on Different Options

Based on my current understanding of the big picture, I would suggest to use Allo as our frontend and support different dialect lowerings in CIRCT in the meantime.

Works involving PyTorch frontend could be complicated and a headache since they are in a big shift (and PyTorch 2.0 just got released), and since PyTorch has a complete compilation flow of its own, I don't think they are putting a lot of effort in Torch-MLIR. Using Allo has a great advantage over the former.
After talking with @chhzh123 , I think we can add our support for the PyTorch-to-MLIR path in the Allo frontend. Although they have very comprehensive support for the Python-to-MLIR path, the PyTorch-to-MLIR path is still under development.

Once the PyTorch-to-MLIR path is complete in the Allo frontend, we can work on lowering from different dialects to Calyx in CIRCT. I mean although they have passes to lower from different dialects to AMC and we can simply just create a lower-amc-to-calyx pass, it might be constraining ourselves to relying on AMC too much. On the other hand, we can add comprehensive and reusable support to lower from different dialects (like loopschedule, scf) to calyx in CIRCT as we've been working on.

Please let me know what you think! @sampsyo @rachitnigam @andrewb1999

[chhzh123]

Thanks @jiahanxie353 for summarizing the key ideas. Torch-mlir can be a good target in the PyTorch 1.0 era, but may not be that attractive after PyTorch 2.0 was released. Most of the facilities of torch-mlir are based on TorchScript, which has lots of limitations (as pointed out in this paper) and is no longer actively maintained by the PyTorch team. PyTorch 2.0 greatly enhances the compiler support by leveraging TorchDynamo as the frontend and torch.fx as the IR. Users can then plug in different backends using the torch.compile API to connect with the PyTorch ecosystem, and thus is not necessary to use torch-mlir to generate the MLIR assembly.

As a result, in Allo, we chose to embrace torch.fx and minimize the dependency of external libraries like torch-mlir. Basically, we take in a PyTorch module, lower it to torch.fx, generate the corresponding Allo representation, and finally generate MLIR programs represented in linalg/affine. We can already support lowering some basic models from PyTorch. For example, lowering a BERT layer to MLIR.

I think we can collaborate more on this lowering path, and we may have more flexibility on the backend selection. Either going through the AMC dialect or directly lowering from linalg/affine to Calyx can be fine.

[rachitnigam]

Thanks for the summary @chhzh123! @jiahanxie353 does the summary make sense? Do you want to take a shot at setting up allo and trying to see what the flow looks like when lowering a simple program?

[jiahanxie353]

Yes, totally! I also got in touch with Torch-MLIR maintainers and confirmed it's true that the development of Torch-MLIR with TorchDynamo and torch.fx is far from fully fledged; and TorchScript is no longer being actively maintained. To this end, we can use Allo and add support PyTorch-lowering to it.
Yes! I'm starting to set up Allo and will try to run some lowerings!

[rachitnigam]

Okay, awesome! Can you also start/continue using the lab notebook from the last semester? I think both @sampsyo and I would appreciate weekly updates on things and be more able to help with things!

[jiahanxie353]

Sure, will do!

[jiahanxie353]

Just had a meeting with @evanmwilliams, he's also interested in working on the PyTorch-Allo side.
And we'll try to flush out more details regarding the future work soon.

[chhzh123]

FYI, PyTorch 2.0 got accepted by this year's ASPLOS, and here is the paper:
https://pytorch.org/blog/pytorch-2-paper-tutorial/