Getting Started with Rust Extensions

Note

All the code in this tutorial is in the repository, as the Python package examples/extension.

This tutorial shows how to extend TIRx from Rust: define new IR types and operators, write a pass that lowers them, and use them from Python like built-in features.

In this chapter, we will first look at why extensions are written in Rust, then get ready for the rest of the tutorial: a Python package for the extension and the example program it will compile.

Why Rust

TIRx is built on a stable C++ library for its IR, analysis, transformations, and code generation. In most cases, you work with it from Python: tensor programs are written as Python functions in TVMScript, then compiled and run directly.

With agents, compiler development iterates quickly. Along the way, a project may need new types and operators to describe its own concepts, such as a new number format or a hardware feature, and passes that lower them to what the backend supports. To support this, TIRx provides Rust bindings, the tvm-tirx crate, which let you define new types, operators, and passes on the same IR objects as C++ and Python. An extension is a separate Rust library that builds and loads on its own, without changing the C++ library, so the C++ core stays stable while the extension iterates quickly. Once loaded, its types and operators can be written in TVMScript and its passes can run in Python compilation pipelines, just like built-in ones.

Setting Up the Project

An extension is shipped as a Python package that contains its Rust library, so it is installed and imported like any other package. The example package, my-ext, has this layout:

examples/extension/
├── Cargo.toml        # the Rust library, which depends on tvm-tirx
├── src/lib.rs        # the type, operators, and pass
├── CMakeLists.txt    # builds the library with cargo and installs it into my_ext
├── pyproject.toml    # the Python package, built with scikit-build-core
├── python/my_ext/    # loads the library and registers a compilation pipeline
└── run_example.py    # compiles and runs a function that uses the extension

The Rust library is an ordinary crate that builds a shared library. Its only dependency is the tvm-tirx crate, which re-exports TVM FFI as tvm_tirx::tvm_ffi:

[lib]
# The shared library that the Python package loads.
crate-type = ["cdylib"]

[dependencies]
# Use this checkout so the Rust API and native libraries share the same FFI revision.
tvm-tirx = { path = "../../rust" }
# Procedural macros resolve the FFI crate directly; use the same pinned sources.
tvm-ffi = { path = "../../3rdparty/tvm-ffi/rust/tvm-ffi" }

First build and install TIRx from this checkout. With Rust 1.84 or newer, install the extension from the repository root in the same Python environment:

uv pip install scikit-build-core
uv pip install --no-build-isolation --no-deps ./examples/extension

The example uses this checkout’s Rust crate and pinned TVM FFI sources. The package depends on apache-tvm-tirx at run time; disabling build isolation keeps its native build aligned with these sources instead of installing a released package in a temporary build environment.

Case Study: Complex Numbers

To show how to define new types, operators, and passes, we will use complex numbers as a running example throughout this tutorial: a complex type, operators that build and multiply complex values, and a pass that expands complex multiplication into real arithmetic:

\[(a + bi)(c + di) = (ac - bd) + (ad + bc)i.\]

The following function multiplies two arrays of complex numbers element by element. A and B hold the real and imaginary parts of the first operand, C and D those of the second:

N = 16


@T.function
def complex_product(
    A: T.Tensor((N,), "float32"),
    B: T.Tensor((N,), "float32"),
    C: T.Tensor((N,), "float32"),
    D: T.Tensor((N,), "float32"),
    R: T.Tensor((N,), "float32"),
    Im: T.Tensor((N,), "float32"),
):
    for i in range(N):
        z = my_ext.complex(A[i], B[i])
        w = my_ext.complex(C[i], D[i])
        product = my_ext.cmul(z, w)
        R[i] = my_ext.real(product)
        Im[i] = my_ext.imag(product)

my_ext is the extension’s Python package, which holds the builders of its type and operators; Using and Compiling the Extension shows how it provides them. z, w, and product are complex values that exist only in the IR. The function’s inputs and outputs are ordinary float32 arrays, so no runtime storage format for complex numbers is needed.

Next, Defining Types, Operators, and Passes defines the type, operators, and pass this function uses, and Using and Compiling the Extension compiles and runs it.