The PyTorch Frontend

This frontend enables the translation of PyTorch GraphModules into Structured SDFGs. It replaces the old approach of using torch-mlir to generate MLIR from PyTorch GraphModules and translating the MLIR to Structured SDFGs via an MLIR frontend.

How to Use

The PyTorch frontend is available as a backend for torch.compile. All you have to do to use it with your existing PyTorch model is add import docc.pytorch and extend your torch.compile call with backend="docc". You can set all additional options with the options={...} dictionary argument of torch.compile.

...
import docc.pytorch
...
... = torch.compile(..., backend="docc", ...)
...

Example

This example shows a simple image-classification task using ResNet18. The corresponding source file is located in pytorch/examples/resnet18.py.

import torch
import torchvision.models as models
import PIL.Image as Image
import requests

import docc.pytorch

image_url = "https://github.com/pytorch/hub/raw/master/images/dog.jpg"

weights = models.resnet.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)
model.eval()

image = Image.open(requests.get(image_url, stream=True).raw)
transforms = weights.transforms()
input_tensor = transforms(image).unsqueeze(0)

with torch.no_grad():
    program = torch.compile(model, backend="docc", options={"target": "openmp"})
    output_tensor = program(input_tensor)

probs = torch.softmax(output_tensor, dim=1)
print(f"Probability of Samoyed (dog): {(probs[0, 258] * 100):.3f}%")

The PyTorch model is parallelized on the CPU. If, for example, you want to execute this model on an NVIDIA GPU and enable remote tuning, adapt the options attribute as follows:

    program = torch.compile(model, backend="docc", options={"target": "cuda", "remote-tuning": True})

Installation From Source

1. Installing dependencies

Create a virtual Python environment and install the dependencies from the requirements (.txt) file that matches your needs.

python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip

Use requirements.txt to install PyTorch without GPU support.

pip install -r pytorch/requirements.txt

Use requirements-cuda.txt to install PyTorch with CUDA support and CuPy.

pip install -r pytorch/requirements-cuda.txt

Use requirements-rocm.txt to install PyTorch with ROCm support and CuPy (for ROCm).

pip install -r pytorch/requirements-rocm.txt

If you want to use your own PyTorch and CuPy versions because of different drivers, GPU runtimes, or other requirements, use requirements-base.txt to install only the base dependencies.

pip install -r pytorch/requirements-base.txt

2. Install Python frontend

Since the PyTorch frontend builds upon the Python frontend, we first need to install it. Because we are building the PyTorch frontend from source, it is recommended to also build the Python frontend from source.

pip install --no-build-isolation -e python/

3. Install the PyTorch frontend

pip install --no-build-isolation -e pytorch/

Testing

There are many integration tests for the functionality of the PyTorch frontend. They are written as pytests and located in pytorch/tests/. The following list gives an overview of the subfolders:

• graph_parser: Contains tests for specific graph parser functionality.
• ops: Contains a large number of tests for individual PyTorch operations.
• pipelines: Contains tests for multiple PyTorch operations working together.
• models: Contains tests for complete models.

Functionality

This frontend is written in Python. It uses the Python bindings from the Python frontend to construct Structured SDFGs.

The starting point and core of the frontend is PyTorchProgram, located in pytorch/docc/pytorch/pytorch_program.py. It is registered by a wrapper as a custom PyTorch Dynamo backend (torch.compile). When it receives a PyTorch graph module, it uses torch.export to lower it to PyTorch Core ATen IR. The resulting IR is then parsed by a custom graph parser, which translates it into a Structured SDFG and returns a CompiledSDFG (see the Python frontend for more details).

The GraphParser class is located in pytorch/docc/pytorch/graph_parser/init.py. It has a modular registry in which a GraphParserModule (located in pytorch/docc/pytorch/graph_parser/utils.py) can be registered for each Core ATen IR operation.