Thoughts on PyTorch 2.0

Incubator torch-mlir
3

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