Skip to content

iree-org/iree-turbine

This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Jan 13, 2025
b3d58f9 · Jan 13, 2025
Jan 13, 2025
Jan 10, 2025
Dec 5, 2024
Nov 8, 2024
Jan 8, 2025
Jan 8, 2025
Jan 10, 2025
Dec 5, 2024
Apr 22, 2024
Aug 9, 2024
Apr 22, 2024
Dec 16, 2024
Nov 27, 2024
Apr 22, 2024
Dec 18, 2024
Dec 18, 2024
Jan 10, 2025
Jan 10, 2025
Jan 10, 2025
Apr 22, 2024
Jan 10, 2025
Jan 8, 2025

Repository files navigation

IREE Turbine

PyPI version

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

  1. 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)

  2. 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:

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