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.

Prerequisites

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

Tip - Virtual environments

We recommend using virtual environments to manage python packages, such as through venv (about, tutorial):

Linux macOS Windows

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

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

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

When done, run deactivate.

Installing IREE packages

Prebuilt packages

Stable releases Nightly releases

Stable release packages are published to PyPI.

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

Nightly releases are published on GitHub releases.

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

Building from source

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

Usage

Info - API reference pages

API reference pages for IREE's runtime and compiler Python APIs are hosted on readthedocs.

Documentation for the MLIR compiler Python APIs can be found at 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-experimental repository for examples using the Python APIs.

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-runtime
$ which iree-run-module

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

Profiling

The tools in the iree-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 - 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:

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