Python bindings

Overview

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

PIP package name	Description
`iree-base-compiler`	IREE's generic compiler tools and helpers
`iree-base-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.

???+ Note “Note - iree-compiler and iree-runtime are deprecated” The Python packages iree-compiler and iree-runtime have been renamed to iree-base-compiler and iree-base-runtime respectively, thus effectively deprecating the old packages. This name change only affects the names of the packages but not the modules. To make sure the new packages are used, run pip uninstall iree-compiler iree-runtime before installing the new packages.

: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-base-compiler \
  iree-base-runtime
```

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

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

``` shell
python -m pip install \
  --find-links https://iree.dev/pip-release-links.html \
  --upgrade \
  iree-base-compiler \
  iree-base-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)

:octicons-code-16: Samples

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

:material-console: Console scripts

The Python packages include console scripts for most of IREE's native tools like iree-compile and iree-run-module. After installing a package from pip, these should be added to your path automatically:

$ python -m pip install iree-base-runtime
$ which iree-run-module

/projects/.venv/Scripts/iree-run-module

:material-chart-line: Profiling

The tools in the iree-base-runtime package support variants:

Variant name	Description
default	Standard runtime tools
tracy	Runtime tools instrumented using the Tracy profiler

Switch between variants of the installed tools using the IREE_PY_RUNTIME environment variable:

IREE_PY_RUNTIME=tracy iree-run-module ...

See the developer documentation page on Profiling with Tracy for information on using Tracy.

!!! tip - “Tip - flushing profile data”

When writing a Python-based program that you want to profile you may need to
insert IREE runtime calls to periodically flush the profile data:

```python
device = ... # HalDevice
device.flush_profiling()
```