Python bindings

Overview

IREE offers Python bindings split into several packages, covering different components:

PIP package name	Description
`iree-compiler`	IREE's generic compiler tools and helpers
`iree-runtime`	IREE's runtime, including CPU and GPU backends
`iree-tools-tf`	Tools for importing from TensorFlow
`iree-tools-tflite`	Tools for importing from TensorFlow Lite
`iree-jax`	Tools for importing from JAX

Collectively, these packages allow for importing from frontends, compiling towards various targets, and executing compiled code on IREE's backends.

:octicons-download-16: Prerequisites

To use IREE's Python bindings, you will first need to install Python 3 and pip, as needed.

???+ Tip “Tip - Virtual environments” We recommend using virtual environments to manage python packages, such as through venv (about, tutorial):

=== ":fontawesome-brands-linux: Linux"

    ``` shell
    python -m venv .venv
    source .venv/bin/activate
    ```

=== ":fontawesome-brands-apple: macOS"

    ``` shell
    python -m venv .venv
    source .venv/bin/activate
    ```

=== ":fontawesome-brands-windows: Windows"

    ``` powershell
    python -m venv .venv
    .venv\Scripts\activate.bat
    ```

When done, run `deactivate`.

Installing IREE packages

:octicons-package-16: Prebuilt packages

=== “Stable releases”

Stable release packages are
[published to PyPI](https://pypi.org/user/google-iree-pypi-deploy/).

``` shell
python -m pip install \
  iree-compiler \
  iree-runtime
```

=== “:material-alert: Nightly releases”

Nightly releases are published on
[GitHub releases](https://github.com/openxla/iree/releases).

``` shell
python -m pip install \
  --find-links https://iree.dev/pip-release-links.html \
  --upgrade \
  iree-compiler \
  iree-runtime
```

:material-hammer-wrench: Building from source

See Building Python bindings page for instructions for building from source.

Usage

!!! info “Info - API reference pages”

API reference pages for IREE's runtime and compiler Python APIs are hosted on
[readthedocs](https://iree-python-api.readthedocs.io/en/latest/).

Documentation for the MLIR compiler Python APIs can be found at
[https://mlir.llvm.org/docs/Bindings/Python/](https://mlir.llvm.org/docs/Bindings/Python/).

Compile a program

from iree import compiler as ireec

# Compile a module.
INPUT_MLIR = """
module @arithmetic {
  func.func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
    %0 = arith.mulf %arg0, %arg1 : tensor<4xf32>
    return %0 : tensor<4xf32>
  }
}
"""

# Compile using the vmvx (reference) target:
compiled_flatbuffer = ireec.tools.compile_str(
    INPUT_MLIR,
    target_backends=["vmvx"])

Run a compiled program

from iree import runtime as ireert
import numpy as np

# Register the module with a runtime context.
# Use the "local-task" CPU driver, which can load the vmvx executable:
config = ireert.Config("local-task")
ctx = ireert.SystemContext(config=config)
vm_module = ireert.VmModule.copy_buffer(ctx.instance, compiled_flatbuffer)
ctx.add_vm_module(vm_module)

# Invoke the function and print the result.
print("INVOKE simple_mul")
arg0 = np.array([1., 2., 3., 4.], dtype=np.float32)
arg1 = np.array([4., 5., 6., 7.], dtype=np.float32)
f = ctx.modules.arithmetic["simple_mul"]
results = f(arg0, arg1).to_host()
print("Results:", results)

Samples

Check out the samples in IREE's samples/colab/ directory and the iree-samples repository for examples using the Python APIs.