Working with IR in Rust

Working with IR in Rust#

The tvm-tirx crate lets a Rust function read, build, and rewrite TIRx IR. The function can receive a program written in TVMScript and return the result, and nothing is converted or copied along the way. This guide covers:

  • TVM-FFI object system: how IR nodes are represented and used in Rust.

  • Structural utilities: walk, map, and compare programs.

  • Visitors and mutators: keep state and control the traversal.

Note

The runnable code in this guide is under examples/ir_guides in the repository.

Setting Up#

The examples are packaged as tirx_ir_guide_examples under examples/ir_guides. When cargo is available, installing the package builds the Rust crate in rust/ into a library; CMakeLists.txt and pyproject.toml in the same directory show how. The crate uses the bindings from this checkout, so install TIRx from the same source tree. We install without build isolation to build against that installation instead of fetching one from PyPI:

pip install scikit-build-core
pip install --no-build-isolation -e examples/ir_guides

A static_init_block! at the end of the crate registers each example as a global function under tirx_ir_guide_examples.rust:

tvm_ffi::Function::register_global_typed(
    "tirx_ir_guide_examples.rust.construct",
    construct,
    "Build x * 2 + 1.",
)?;

The package’s tirx_ir_guide_examples.rust module loads the library and exposes these functions, so Python calls them by name. The driver working_with_ir_in_rust.py calls each example this way and prints the result.

The object system examples take no arguments: each builds its own IR in Rust. The later sections receive the add_one function from Quick Start, written in TVMScript:

from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[i] = A[i] + 1

import tirx_ir_guide_examples.rust

tirx_ir_guide_examples.rust.construct()
tirx_ir_guide_examples.rust.written_tensors(add_one)

TVM-FFI Object System#

Every TIRx IR node is a TVM-FFI object. Rust IRs are typically built as enums and matched with match, which works well when every node type is known in advance. TIRx is designed for dynamic extension instead: a library can add a new node type, for example a new statement for a hardware feature, and existing code keeps working on programs that contain it. Generic utilities such as walking, rewriting, comparing, and printing work on every node through its registered fields.

TVM-FFI objects also work across languages. Python, C++, and Rust share the same native, immutable, reference-counted objects, so a program moves between them without conversion.

We designed the TVM-FFI object system to be ergonomic in Rust. The key constructs are:

Construct

Purpose

Constructors, expression operators

Build expressions with ordinary Rust operators

FooObj / Foo

Underlying object and its named managed reference: read fields directly, clone cheaply

.into()

Convert to a parent type

.try_cast()

Convert to a child type, with a checked Result

match_any!

Branch on the node type

#[dispatch]

Process each node type in its own method, with shared state

The subsections below introduce each construct.

Constructing Objects#

We build IR with constructors and ordinary Rust operators. Writing x * 2 + 1 as nested constructor calls would bury the expression, so the crate lets us write it as we would in Rust. Let’s build it:

fn x_times_2_plus_1() -> Result<(Var, Expr)> {
    let x = Var::new("x", PrimType::int(32), UNKNOWN_LOC);
    let expr = (&x * 2 + 1).try_into()?;
    Ok((x, expr))
}

fn construct() -> Result<tvm_ffi::String> {
    let (_, expr) = x_times_2_plus_1()?;
    script(&expr)
}

The output is:

x * 2 + 1

Var::new takes a name, a type, and a source location; UNKNOWN_LOC means there is none. The operators * and + build IR directly from x, and the literals 2 and 1 become constants of its type.

An operator can fail, for example on mismatched types, so ExprResult carries the first error along, much like a chain of Results. The standard .try_into() conversion extracts an Expr, and ? passes an error on. script prints the node as TVMScript.

We now have an expression; next, let’s see what kind of Rust value it is.

Objects and Managed References#

Every node type comes as a pair: an underlying object and a named managed reference. The top of x * 2 + 1 is an addition, so let’s convert the expression to Add with try_cast, which the next subsection explains, and read its two operands:

fn objects() -> Result<Strings> {
    let (_, expr) = x_times_2_plus_1()?;
    let add: Add = expr.try_cast()?;
    let alias = add.clone();
    Ok(lines([
        format!("a = {}", script(&add.a)?.as_str()),
        format!("b = {}", script(&add.b)?.as_str()),
        format!("same object: {}", alias.same_as(&add)),
    ]))
}

The output is:

a = x * 2
b = 1
same object: true

alias is the same object as add: clone copies only the managed reference, and same_as compares identity, like Arc::ptr_eq.

AddObj is the underlying object: the struct that holds the fields. Add is a named managed reference to it: an owning, reference-counted smart pointer, like Arc<AddObj>, so cloning it is cheap. Functions take and return managed references, and a managed reference dereferences to its object, so fields read as ordinary struct fields: add.a.

We define a separate named type Add instead of using a plain pointer to AddObj because the named type carries the interface: constructors, methods, conversions to parent types, and the traits that let it pass through TVM-FFI functions. It can also impose additional constraints on the underlying object that the object alone does not express.

This is the newtype pattern. Simplified from the crate, the pair is:

#[repr(C)]
#[derive(Object)]
#[type_key = "prim.Add"]
pub struct AddObj {
    base: ExprObj,
    pub a: Expr,
    pub b: Expr,
}

#[repr(C)]
#[derive(ObjectRef, Clone)]
pub struct Add {
    base: ObjectArc<AddObj>,
}

Rust has no inheritance; instead, #[repr(C)] puts the parent’s object, ExprObj, first in AddObj, so every addition is also an expression. Add implements Deref to AddObj, which is why add.a works. AddObj has the same memory layout as its C++ counterpart, so add.a is a direct field read, as efficient as in C++.

Nodes are immutable. As with data behind an Arc, a managed reference gives read access only; to change a node, we build a new one.

We reached Add with a cast; next, let’s look at casts on their own.

Casts#

Casts move a managed reference up and down the type hierarchy. A function that takes any expression wants an Expr, while reading the operands needs an Add. .into() converts to a parent type through From/Into and cannot fail. .try_cast() converts to a child type: like TryFrom, it checks the node’s actual type and returns a Result.

fn casts() -> Result<Strings> {
    let (_, expr) = x_times_2_plus_1()?;
    let add: Add = expr.try_cast()?;
    let value: Expr = add.into();
    let again: Add = value.clone().try_cast()?;
    let mut out = vec![format!("back to Add: {}", script(&again)?.as_str())];
    if let Err(err) = value.try_cast::<Mul>() {
        out.push(format!("as Mul: {}", err.message()));
    }
    Ok(lines(out))
}

The output is:

back to Add: x * 2 + 1
as Mul: Cannot convert from type `prim.Add` to `prim.Mul`

value is an Expr again, but the node behind it is still an addition, so the cast back to Add succeeds.

What if we guess wrong? The cast to Mul returns an Err, and its message names both type keys: the actual type prim.Add and the requested prim.Mul.

Note

try_cast consumes the managed reference it converts, so we clone value before the first cast; the clone costs one reference-count increment. For read-only inspection, value.downcast_ref::<AddObj>() returns Option<&AddObj> without consuming or cloning the reference. Any and AnyView provide the same borrowed downcast.

Casting to one type at a time gets tedious when a node can be one of many; that is what matching is for.

Matching Node Types#

match_any! branches on the actual type of a node, as match does on an enum. Let’s print an expression with explicit parentheses by matching each node and recursing into its operands:

fn describe(expr: &Expr) -> Result<String> {
    let text = match_any! {
        AnyView::from(expr) {
            AddObj(ref add) => format!("({} + {})", describe(&add.a)?, describe(&add.b)?),
            MulObj(ref mul) => format!("({} * {})", describe(&mul.a)?, describe(&mul.b)?),
            VarObj(ref var) => var.name.as_str().to_owned(),
            IntImmObj(ref imm) => imm.value.to_string(),
            _ => "<other>".to_owned(),
        }
    };
    Ok(text)
}

We call it on x * 2 + 1 and on x - 1:

fn matching() -> Result<Strings> {
    let (x, expr) = x_times_2_plus_1()?;
    let difference = (&x - 1).try_into()?;
    Ok(lines([describe(&expr)?, describe(&difference)?]))
}

The output is:

((x * 2) + 1)
<other>

Each arm names an object type and borrows the matching node with ref; the first matching arm runs. The AddObj and MulObj arms recurse into the operands, the VarObj and IntImmObj arms end the recursion, and x - 1 falls to the _ arm because describe has no arm for subtraction.

The catch-all arm keeps existing code working as node types are added. The set of node types is open, like a #[non_exhaustive] enum from another crate, so the match needs a _ arm. A function handles the types it knows, and _ lets programs with new types pass through.

Note

Obj(ref node) borrows the underlying object without incrementing its reference count. describe converts expr to an AnyView. An owning arm such as Add(add) still produces a managed reference. match_any! also works on Any and AnyView values, including tensors, shapes, and typed arrays. Plain values such as integers take the _ arm in this example.

A match looks at one node at a time; to process every node in a program while keeping state, we dispatch.

Dispatching on Node Types#

#[dispatch] lets one struct process each node type in its own method, with shared state. A compiler pass often needs one piece of logic per node type and state shared across them; #[dispatch] puts both on one struct. It works like a trait whose methods all have default bodies: we write methods for the node types we care about, and every other type gets the default behavior.

#[dispatch(walk)] on an impl block turns its walk_* methods into callbacks for structural_walk, which visits every node. Each method borrows the nodes of its parameter type, and the methods are tried in source order. Walk handlers share the Result<WalkResult> return type; one handler declares the trailing DefRegionKind argument used by the walker, and handlers that do not need it may omit it. Let’s count the variables and additions:

#[derive(Default)]
struct Counter {
    vars: usize,
    adds: usize,
}

#[dispatch(walk)]
impl Counter {
    fn walk_var(&mut self, _var: &VarObj, _kind: DefRegionKind) -> Result<WalkResult> {
        self.vars += 1;
        Ok(WalkResult::Advance)
    }

    fn walk_add(&mut self, _add: &AddObj) -> Result<WalkResult> {
        self.adds += 1;
        Ok(WalkResult::Advance)
    }
}

We walk x * 2 + 1 with it. WalkResult::Advance continues into the node’s children, and WalkOrder::PreOrder visits a node before its children:

fn count() -> Result<tvm_ffi::String> {
    let (_, expr) = x_times_2_plus_1()?;
    let mut counter = Counter::default();
    tvm_ffi::structural_walk(&expr, &mut counter, WalkOrder::PreOrder)?;
    Ok(format!("variables: {}, additions: {}", counter.vars, counter.adds).into())
}

The output is:

variables: 1, additions: 1

Node types without a method fall back to the default traversal: the walk continues into their children, so new types pass through. Here the multiplication and the constants have no method, and the walk still reaches x inside the multiplication.

Structural Utilities#

The structural utilities are part of TVM-FFI, and the tvm_ffi crate, re-exported as tvm_tirx::tvm_ffi, provides them in Rust: structural_walk, structural_map, structural_visit, and structural_mutate. They work on every node through its registered fields and keep definitions and uses consistent. Each piece takes a Rust form:

Piece

Rust form

Nodes a callback receives

Selected by the type of the callback’s parameter

Walk callback result

WalkResult, or Result<WalkResult> to use ?

Traversal order

A required argument, WalkOrder::PreOrder or PostOrder

Result of a map

An Any, converted back with try_into

Definition region kind

A trailing DefRegionKind parameter

Errors

Result values

Structural equality

The global function ffi.StructuralEqual

Structural Walk#

A walk calls our callback on each node of the type it asks for. Let’s find the tensors that add_one writes:

fn written_tensors(func: Function) -> Result<Strings> {
    let mut writes = Vec::new();
    let on_store = |store: TensorStore| -> Result<WalkResult> {
        let dest: Var = store.dest.clone().try_cast()?;
        writes.push(dest.name.clone());
        Ok(WalkResult::Advance)
    };
    tvm_ffi::structural_walk(&func, on_store, WalkOrder::PostOrder)?;
    Ok(writes.into_iter().collect())
}

The output is:

['B']

on_store takes a TensorStore, so the walk calls it only on stores. Strings is our alias for Array<tvm_ffi::String>, which Python receives as a list of strings.

WalkResult::Advance continues into the node’s children, Skip does not descend into them, and Interrupt stops the whole walk.

The parameter list of a callback decides what it receives:

Callback

Receives

|store: TensorStore|

Each write, as a managed reference.

|store: &TensorStoreObj|

Each write, as a borrowed object, without changing its reference count.

|value: AnyView<'_>|

Every value the walk reaches.

|var: Var, kind: DefRegionKind|

Each variable, with its definition region kind.

(on_loop, on_store)

A tuple of callbacks, tried in order; only the first one whose type matches a node runs. Tuples hold up to 12 callbacks and can be nested.

WalkResult::interrupt_with stops the walk and returns a value, which lets a walk search for a node:

fn first_write(func: Function) -> Result<Option<TensorStore>> {
    let find_write = |store: TensorStore| WalkResult::interrupt_with(store);
    let found = tvm_ffi::structural_walk(&func, find_write, WalkOrder::PostOrder)?;
    match found {
        Some(interrupt) => Ok(Some(interrupt.value.try_into()?)),
        None => Ok(None),
    }
}

Python prints the dest.name of the store it receives. The output is:

B

structural_walk returns Ok(Some(interrupt)) when a callback stops it, with the payload in interrupt.value, and Ok(None) otherwise.

Definitions and Uses#

A callback receives the definition region kind when it takes a DefRegionKind as its last parameter. This walk lists the definition sites of add_one:

fn definitions(func: Function) -> Result<Strings> {
    let mut sites = Vec::new();
    let show_definition = |var: Var, kind: DefRegionKind| {
        if kind != DefRegionKind::None {
            sites.push(format!("{} {:?}", var.name.as_str(), kind));
        }
        WalkResult::Advance
    };
    tvm_ffi::structural_walk(&func, show_definition, WalkOrder::PreOrder)?;
    Ok(sites.iter().map(tvm_ffi::String::from).collect())
}

The output is:

A Pattern
n Pattern
B Pattern
n Pattern
i Simple

DefRegionKind::None marks a use site; Pattern and Simple are the two kinds of definition site. structural_map callbacks take the same trailing parameter.

Structural Map#

A map callback returns the replacement for its node. Let’s specialize add_one for n = 128, with n passed in from Python:

fn specialize(func: Function, n: Var, value: i64) -> Result<Function> {
    let specialize_var = |var: Var| -> Result<Expr> {
        if var.same_as(&n) {
            let ty: PrimType = var.ty.clone().try_cast()?;
            return make_const(ty, value);
        }
        Ok(var.into())
    };
    tvm_ffi::structural_map(func, specialize_var, WalkOrder::PostOrder)?.try_into()
}

The output is:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one(A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)):
    for i in range(128):
        B[i] = A[i] + T.float32(1.0)

make_const builds the constant 128 with the type of n, and every other variable is returned as it is. structural_map takes ownership of the managed reference we pass and returns the new root as an Any; Python still holds add_one, which is unchanged.

Structural Equality#

The crate has no Rust function for structural equality yet, so we call the global function ffi.StructuralEqual with its four arguments: the two values, map_free_vars, and skip_tensor_content. Let’s specialize add_one twice and compare the results:

fn compare_specialized(func: Function, n: Var) -> Result<Array<bool>> {
    let first = specialize(func.clone(), n.clone(), 128)?;
    let second = specialize(func, n, 128)?;
    let (map_free_vars, skip_tensor_content) = (false, false);
    let equal: bool = tvm_ffi::Function::get_global("ffi.StructuralEqual")?
        .call_tuple((&first, &second, map_free_vars, skip_tensor_content))?
        .try_into()?;
    Ok([first.same_as(&second), equal].into_iter().collect())
}

The output is:

[False, True]

The two functions are different objects, so same_as is false, but they are structurally equal.

Visitors and Mutators in Rust#

Stateful traversals use the same four structural utilities. State lives on a #[dispatch] struct, as in Dispatching on Node Types. #[dispatch(map)] and #[dispatch(visit)] group map_* and visit_* methods. #[dispatch(mutate, MutationView)] groups mutate_* methods and carries mutation permission through matching.

A visit or mutate callback takes control of its node’s children. Let’s record the loops around each write. Each handler receives the node and can take a StructuralVisitorView<'_, Self> for recursive traversal; the loop stack and collected writes remain on self:

#[derive(Default)]
struct LoopNest {
    loops: Vec<String>,
    writes: Vec<String>,
}

#[dispatch(visit)]
impl LoopNest {
    fn visit_for(
        &mut self,
        loop_: &ForObj,
        visitor: &mut StructuralVisitorView<'_, Self>,
    ) -> Result<Option<VisitInterrupt>> {
        self.loops.push(loop_.var.name.as_str().to_owned());
        let result = visitor.default_visit(self, loop_);
        self.loops.pop();
        result
    }

    fn visit_store(&mut self, store: &TensorStoreObj) -> Result<Option<VisitInterrupt>> {
        let loops = self.loops.join(", ");
        let write = format!("{} inside [{}]", script(&store.dest)?.as_str(), loops);
        self.writes.push(write);
        Ok(None)
    }
}

fn loops_around_writes(func: Function) -> Result<Strings> {
    let mut visitor = LoopNest::default();
    tvm_ffi::structural_visit(&func, &mut visitor)?;
    assert!(visitor.loops.is_empty());
    Ok(visitor.writes.iter().map(tvm_ffi::String::from).collect())
}

The output is:

['B inside [i]']

visit_for pushes the loop variable, lets default_visit(self, loop_) visit the loop’s children with their usual traversal, then pops the variable. It returns the traversal result, so an interrupt or error from a child reaches the caller. The view also provides visit(self, child) for a selected child and def_region_kind for the current node.

For mutation, a trailing &mut StructuralMutatorView<'_, Self> provides def_region_kind() and child-mutation methods. A MutationView<'_, For> carries the current loop and the engine’s permission to reuse it. Its typed fields can be borrowed directly; after that borrow ends, default_mutate consumes the handle to recurse with the same permission. The returned UnchangedOr<Any> becomes the common Any handler result without discarding an unchanged marker. Reuse still requires unique ownership, so an alias held by Python remains unchanged. Let’s reverse the loop by replacing each use of i with n - 1 - i. As before, .try_into()? extracts the built expression or returns its error:

#[derive(Default)]
struct ReverseLoops {
    reversed_index: HashMap<ObjectIdentity, Expr>,
}

#[dispatch(mutate, MutationView)]
impl ReverseLoops {
    fn mutate_for(
        &mut self,
        value: MutationView<'_, For>,
        mutator: &mut StructuralMutatorView<'_, Self>,
    ) -> Result<Any> {
        let index = (&value.end - 1 - &value.var).try_into()?;
        self.reversed_index.insert(ObjectIdentity::of(&value.var), index);
        mutator.default_mutate(self, value).map(Into::into)
    }

    fn mutate_var(
        &mut self,
        var: Var,
        mutator: &mut StructuralMutatorView<'_, Self>,
    ) -> Result<Any> {
        let is_use = mutator.def_region_kind() == DefRegionKind::None;
        Ok(match self.reversed_index.get(&ObjectIdentity::of(&var)) {
            Some(index) if is_use => index.clone().into(),
            _ => var.into(),
        })
    }
}

fn reverse_loops(func: Function) -> Result<Function> {
    tvm_ffi::structural_mutate(func, ReverseLoops::default())?.try_into()
}

The output is:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[n - 1 - i] = A[n - 1 - i] + T.float32(1.0)

mutate_for records the replacement, then default_mutate rewrites the loop’s children through the same struct, so mutate_var sees the replacement before any use of i.

Note

mutate_var replaces i only where def_region_kind() is None, so the loop’s own i stays a variable. ObjectIdentity keys the map by object identity, so only this loop’s i matches.

Note

TIRx also has traversals that track the facts valid at each node, such as loop ranges and branch conditions. In Rust they are StructuralVisitorWithConstraintContext and StructuralMutatorWithConstraintContext; constraints() returns the current facts.

With fields, typed callbacks, and the structural utilities, a Rust function works with the same IR that Python builds and prints, so an analysis or a transformation can move between the two languages one function at a time.

Quick Reference#

Construct

What it does

match_any!

Branches on the node type; the first matching arm runs

.try_cast(), downcast_ref

Downcasts to one type, as a managed reference in a Result or a borrowed object in an Option

.into()

Upcasts a managed reference to a parent type

Var::new, expression operators, .try_into()?

Creates a node; builds an expression with operators

structural_walk(&func, callbacks, WalkOrder::PreOrder)

Walks with typed callbacks that return WalkResult

structural_map(func, callbacks, order)

Rebuilds with replacements and returns an Any

A trailing DefRegionKind parameter

Receives the definition region kind

Function::get_global("ffi.StructuralEqual")

Compares by structure

#[dispatch](walk), (map), (visit), (mutate, MutationView)

Turns the methods of a struct into typed callbacks

StructuralVisitorView, StructuralMutatorView

Recurse with the same dispatcher and current definition region

static_init_block!, Function::register_global_typed

Registers Rust functions as global functions