Thoughts on PyTorch 2.0

Hi folks, I thought I’d write some notes about the recently-announced PyTorch 2.0.

tl;dr – due to our close collaboration with the PyTorch core devs, our roadmap already took all this into account, so nothing really changes :slight_smile:

In the Torch-MLIR project we were ecstatic to find TorchDynamo ‘Just Works’ when we plugged it into our system. TorchDynamo layers really nicely with PyTorch’s other compiler infra, so we also now have a way to deprecate large amounts of code we are maintaining ourselves.

Going into a few more specifics:

  • I’m really excited to hear about the direction towards compiler paths as first-class citizens.
  • The formalization of the PrimTorch and ATen op subsets as a fully-supported feature significantly reduces the burden on Torch-MLIR’s backend refactor from the long-term roadmap.
  • I’m really happy to hear the dynamic shape work as a headline feature. That is a critical feature for many of our users. Also functionalization.
  • Adding an interface into TorchInductor could be interesting. From what I hear from the grapevine those folks are possibly already interfacing via Triton’s MLIR code path. Adding a more direct MLIR dialect modeling TorchInductor’s IR is a possibility – please reach out of there is demand for this!
  • Since torch.compile is a wrapper around TorchDynamo, our recently-added support should Just Work (since backend is simply a TorchDynamo backend).
  • They mention “mobile” in fullgraph mode, but they are light on details about what tools are being provided for those use cases.
3 Likes

Sean, can you give us some insight into the integration etc ? Maybe a high level diagram etc.

The main integration point in PyTorch 2.0 is a TorchDynamo backend, which is just a Python function with the signature

def my_backend(gm: torch.fx.GraphModule,
               example_inputs: List[torch.Tensor]):

that return an arbitrary callable which has the same contract as gm when called with example_inputs.

To help with this, we provide a decorator make_simple_dynamo_backend that performs some basic transformations on the FX graph inside gm to prepare it for being passed to torch_mlir.compile.

You can find a longer description here: [torchdynamo] Initial TorchDynamo support · llvm/torch-mlir@28957ad · GitHub

A simple example backend using linalg-on-tensors and refbackend can be seen here

1 Like

Is there a path for bringing a code written In triton Python to MLIR? Will TorchDynamo handle triton python code? If I wanted to target triton Python code on a HW accelerator using MLIR, what are my options at present? Any guidance will be appreciated.

I doubt there is a Trition Python → MLIR lowering available today publicly if I am not mistaken. I know of some internal efforts but they may or may not go public in the coming days. So the best bet is probably to build your own for the time being if there is a time crunch. Maybe others who know more can comment.

Just to add to my previous comment - Trition Python has resemblances with Numba according to its documentation. As a result you can look at this discussion - MLIR based Numba backend which talks about a numba Python → MLIR effort

Triton was actually recently rewriten in MLIR.
I haven’t looked into it in a lot of detail yet, but the lowering goes something like Python → Triton IR → Triton GPU IR → LLVM IR → PTX.
Here Triton IR and Triton GPU IR are composed of custom Triton and Triton GPU dialects, with a mix of some of the standard MLIR dialects (arith, math, scf etc.). Triton IR models the Python program ~exactly, while Triton GPU IR adds layout information determining how the program is distributed on a GPU. This is then lowered directly to LLVM IR.

I’m not aware though of any efforts outside of Triton itself to mix this with other compiler tech.

Thanks @gflegar for the pointer. I heard the same from one of my colleagues yesterday.

  • Adding an interface into TorchInductor could be interesting. From what I hear from the grapevine those folks are possibly already interfacing via Triton’s MLIR code path. Adding a more direct MLIR dialect modeling TorchInductor’s IR is a possibility – please reach out of there is demand for this!

Hi @_sean_silva, do you have more information about the interface into TorchInductor from torch-mlir side?

There is currently not much more info other than what is written in the post. This is just speculation about future possibilities. Do you have a specific use case in mind?