Awesome

IREE Turbine

Turbine is IREE's frontend for PyTorch.

Turbine provides a collection of tools:

AOT Export: For compiling one or more nn.Modules to compiled, deployment ready artifacts. This operates via both a simple one-shot export API (Already upstreamed to torch-mlir) for simple models and an underlying advanced API for complicated models and accessing the full features of the runtime.
Eager Execution: A torch.compile backend is provided and a Turbine Tensor/Device is available for more native, interactive use within a PyTorch session.
Custom Ops: Integration for defining custom PyTorch ops and implementing them in terms of IREE's backend IR or a Pythonic kernel language.

Documentation is available at https://iree.dev/guides/ml-frameworks/pytorch/.

Contact us

Turbine is under active development. Feel free to reach out on one of IREE's communication channels (specifically, we monitor the #pytorch and #turbine channels on the IREE Discord server).

Quick start for users

Install from PyPI:

Install a torch version that fulfills your needs:

# Fast installation of torch with just CPU support.
# See other options at https://pytorch.org/get-started/locally/
pip install torch --index-url https://download.pytorch.org/whl/test/cpu

Then install iree-turbine:

# Stable releases
pip install iree-turbine

# Nightly releases
pip install --find-links https://iree.dev/pip-release-links.html --upgrade --pre iree-turbine

(or follow the "Developers" instructions below)

Try one of the examples:
Generally, we use Turbine to produce valid, dynamic shaped Torch IR (from the torch-mlir torch dialect with various approaches to handling globals). Depending on the use-case and status of the compiler, these should be compilable via IREE with --iree-input-type=torch for end to end execution. Dynamic shape support in torch-mlir is a work in progress, and not everything works at head with release binaries at present.

Developers

Use this as a guide to get started developing the project using pinned, pre-release dependencies. You are welcome to deviate as you see fit, but these canonical directions mirror what the CI does.

Setup a venv

We recommend setting up a virtual environment (venv). The project is configured to ignore .venv directories, and editors like VSCode pick them up by default.

python -m venv --prompt iree-turbine .venv
source .venv/bin/activate

Install PyTorch for your system

You need to explicit install a PyTorch version that fulfills your needs. On Linux, install a variant by either following the official instructions or by using one of our requirements.txt files:

CPU: pytorch-cpu-requirements.txt

pip install -r pytorch-cpu-requirements.txt

ROCM: pytorch-rocm-requirements.txt

pip install -r pytorch-rocm-requirements.txt

Install development packages

# Install editable local projects.
pip install -r requirements.txt -e .

Running tests

# Python unit tests
pytest .

# Lit tests
lit lit_tests/ -v

Optional: Pre-commits and developer settings

This project is set up to use the pre-commit tooling. To install it in your local repo, run: pre-commit install. After this point, when making commits locally, hooks will run automatically.

Using a development compiler

If doing native development of the compiler, it can be useful to switch to source builds for the iree-base-compiler and iree-base-runtime packages.

In order to do this, check out IREE and follow the instructions to build from source, making sure to specify additional options for the Python bindings:

-DIREE_BUILD_PYTHON_BINDINGS=ON -DPython3_EXECUTABLE="$(which python)"

Configuring python

Uninstall existing packages (including any with the old package names):

pip uninstall iree-compiler iree-base-compiler iree-runtime iree-base-runtime

Copy the .env file from iree/ to this source directory to get IDE support and add to your path for use from your shell:

source .env && export PYTHONPATH