Structural Traversal, Rewriting, and Comparison#
Most compiler analyses and transformations come down to two patterns: traversing a program to collect information, or rewriting part of it. TIRx provides first-class support for both through its foundation, TVM-FFI.
Every TIRx IR node registers structural hooks that tell these utilities how to traverse and rewrite it, including which fields define variables. The hooks are compiled, so traversal runs at native speed. Objects without hooks fall back to their reflected fields, and new IR types can register the same hooks to get the same support. On top of these hooks, TVM-FFI provides:
Walk: call our function on every node of a given type, to collect information or search for a match.
Map: replace matching nodes and get back a new program, rebuilt around the replacements.
Equality and Hashing: compare and hash programs by structure rather than object identity, to check results in tests or to deduplicate equal programs.
Visit and Mutate: let our function decide when and whether to process a node’s children, for advanced use cases that need fine-grained control over the traversal order.
We start with walk and map, which cover most everyday tasks, then compare programs, look at how variables are defined and used, and finish with visit and mutate.
Throughout this chapter, we use the add_one function from Quick Start:
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
Note
The runnable code in this chapter is under examples/ir_guides in the repository.
Structural Walk#
Let’s start by inspecting a program. Before we traverse add_one, it
helps to see what we will be traversing.
ir.format_tree prints the tree of IR
nodes that add_one is made of:
import tvm_ffi
from tvm_tirx import ir
print(ir.format_tree(add_one))
The output is:
Function
├── params[0]: A = Var(ty=TensorType(shape=[n], dtype=float32))
├── params[1]: B = Var(ty=TensorType(shape=[n], dtype=float32))
└── body: SeqStmt
└── seq[0]: For(begin=IntImm(value=0), end=n)
├── var: i = Var(ty=int32)
└── body: SeqStmt
└── seq[0]: TensorStore(dest=B, indices=[i])
└── value: Add
├── a: TensorLoad(source=A, indices=[i])
└── b: FloatImm(value=1.0)
Most analyses care about only a few kinds of node in this tree. To find the
tensors that add_one writes, for example, we only need the
TensorStore nodes.
tvm_ffi.structural_walk visits every
node under a root and calls our function on the nodes that match a type, so
we pair the store type with a callback:
writes = []
def on_store(store):
writes.append(store.dest.name)
tvm_ffi.structural_walk(add_one, (ir.TensorStore, on_store))
print(writes)
The output is:
['B']
The walker descends through the whole tree on its own, from the function
down to the loop body. Our callback runs only once, on the TensorStore
node for B[i] = A[i] + 1.
Matching Nodes to Callbacks#
Here we matched a single type, but an analysis often needs several kinds of node at once, or every node. The callback argument accepts these forms:
Argument |
Calls |
|---|---|
|
|
|
|
|
|
|
|
Matching follows the class hierarchy, so a type also matches its
subclasses: (ir.Stmt, f) receives every statement. When several pairs
match a node, they are tried in order and only the first one runs. A
callback that returns None lets the walk continue; the next sections
show what else it can return.
Traversal Order#
Our first walk did not depend on the order in which nodes are reached,
but many analyses do. An analysis that tracks state, for example, may need
to see a loop before the statements inside it. By default, the walker calls
a callback after visiting the node’s children (post-order), and
order="pre" calls it before them. Let’s record every statement both
ways:
for order in ("pre", "post"):
trace = []
tvm_ffi.structural_walk(
add_one, (ir.Stmt, lambda stmt: trace.append(type(stmt).__name__)), order=order
)
print(order, trace)
The output is:
pre ['SeqStmt', 'For', 'SeqStmt', 'TensorStore']
post ['TensorStore', 'SeqStmt', 'For', 'SeqStmt']
The function body and the loop body are each a SeqStmt, a sequence of
statements. Pre-order reaches the loop before the write inside it;
post-order reaches the write first.
Skipping and Stopping#
So far, every walk has covered the whole tree. Often we can do less: an
analysis may not care what is inside a loop, and a search can end as soon
as it finds a match. A callback steers the walk through its return value.
In pre-order, returning WalkResult.SKIP tells the walker not to descend
into the node’s children:
seen = []
def skip_loops(stmt):
seen.append(type(stmt).__name__)
if isinstance(stmt, ir.For):
return tvm_ffi.WalkResult.SKIP
return None
tvm_ffi.structural_walk(add_one, (ir.Stmt, skip_loops), order="pre")
print(seen)
The output is:
['SeqStmt', 'For']
The walk records the loop but never enters it, so the loop body and the write are not reached.
For a search, a callback returns VisitInterrupt(payload) instead. This
stops the whole walk, and structural_walk returns the interrupt, so the
payload carries the result back to us. Here we look for the first write:
def find_write(store):
return tvm_ffi.VisitInterrupt(store)
found = tvm_ffi.structural_walk(add_one, (ir.TensorStore, find_write))
print(found.value.dest.name)
The output is:
B
The walk ends at the first write, and found.value holds it. A walk that
finishes without an interrupt returns None, so the result also tells us
whether anything was found.
In summary, a walk callback steers the traversal with its return value:
Return value |
Effect on the walk |
|---|---|
|
Continue. |
|
In pre-order, do not descend into the node’s children. In post-order, the children are already visited, so it has no effect. |
|
Stop the whole walk; |
Structural Map#
Walks let us inspect a program; a transformation also needs to rewrite it.
tvm_ffi.structural_map takes the same
callbacks as a walk, but each one returns a replacement for its node. The
mapper rebuilds the parents of everything that changed and returns the new
root.
Suppose add_one is always called on tensors of 128 elements. Knowing
n ahead of time, we can specialize the function by replacing n with
a constant:
def specialize(var):
if var.same_as(n):
return ir.const(128, var.ty)
return var
add_one_128 = tvm_ffi.structural_map(add_one, (ir.Var, specialize))
print(add_one_128.script())
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)
Every occurrence of n, in both parameter shapes and the loop bound,
becomes 128. Because IR is immutable, the map returns a new function
and leaves add_one unchanged.
Note
A map callback must always return a node, even when it has nothing to
change: returning its input keeps that node as it is, which is what
specialize does for every other variable. Returning None,
WalkResult.SKIP, or a VisitInterrupt raises a TypeError.
Immutability also keeps rewrites cheap: a map does not copy the whole program. It rebuilds only the nodes on the path from each replacement up to the root, and the new program shares every unchanged part with the original. The original program remains valid, so we can keep it to inspect or compare.
Structural Equality and Hashing#
To test a transformation, we compare its result with a program we write down ourselves. The two are different objects even when they describe the same computation, so we compare them by structure, not by identity. Let’s build the same function twice with a small helper:
def make_add(value):
n = T.dynamic("n")
@T.function
def add(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
for i in range(n):
B[i] = A[i] + value
return add
first, second = make_add(1), make_add(1)
print(first.same_as(second))
print(tvm_ffi.structural_equal(first, second))
print(tvm_ffi.structural_hash(first) == tvm_ffi.structural_hash(second))
The output is:
False
True
True
The two functions are different objects, yet
tvm_ffi.structural_equal finds them
equal, and tvm_ffi.structural_hash
gives them the same hash. Each call to make_add also creates its own
n, A, B, and i. The comparison matches each variable defined
in one program with its counterpart in the other, so these different
variable objects do not make the programs unequal.
Note
A variable that is not defined anywhere in the compared values is
free and has no definition to match. By default, a free variable is
equal only to itself. map_free_vars=True pairs free variables up as
well:
x, y = ir.Var("x", "int32"), ir.Var("y", "int32")
print(tvm_ffi.structural_equal(x + 1, y + 1))
print(tvm_ffi.structural_equal(x + 1, y + 1, map_free_vars=True))
The output is:
False
True
x + 1 and y + 1 differ by default; with the flag, x pairs
with y.
In a test, we usually want more than a yes-or-no answer.
ir.assert_structural_equal
raises an error that shows where the programs differ. Let’s compare
first with make_add(2), its + 2 variant:
ir.assert_structural_equal(first, make_add(2))
The error reads, in part:
StructuralEqual check failed, caused by lhs at <root>.body.seq[0].body.seq[0].value.b.value:
...
B[i] = A[i] + T.float32(1.0)
^^^^^^^^^^^^^^
...
and rhs at <root>.body.seq[0].body.seq[0].value.b.value:
...
B[i] = A[i] + T.float32(2.0)
^^^^^^^^^^^^^^
The error names the access path of the first difference, the chain of fields from the function down to the constant, and underlines that constant in both programs.
Equality and hashing are also useful inside the compiler itself, whenever we
need to detect that two functions or subtrees are the same program.
tvm_ffi.StructuralKey wraps a value so
that a Python set or dictionary compares it by structure, which lets us
deduplicate equal programs:
programs = [make_add(1), make_add(1), make_add(2)]
unique = {tvm_ffi.StructuralKey(program) for program in programs}
print(len(unique))
The output is:
2
The two make_add(1) functions collapse into one entry, and
make_add(2) keeps its own.
Definitions and Uses#
A variable (ir.Var) appears in a program in
two roles: a definition site introduces it, and a use site reads
it. In add_one, the parameters define A, B, and n, and the
loop defines i; range(n) and B[i] = A[i] + 1 use them. Each
variable has one definition site and any number of use sites.
Every site refers to the same ir.Var object, so the object alone cannot tell us its role; its position can. Renaming, finding unused variables, and comparing programs depend on that position.
TIRx IR nodes use their structural hooks to mark which of their fields are definition sites. As the utilities traverse the program, they report the role at each position as a definition region kind:
Kind |
Meaning |
In |
|---|---|---|
|
A use site. |
|
|
Defines exactly the variable at the site. |
The loop variable |
|
Defines every variable in the site by matching, including those in its type. |
The parameters |
Note
A DEF_PATTERN site is a pattern matched against the values passed
in. The parameter A: T.Tensor((n,), T.float32) defines both A
and its length n, so n needs no separate declaration. A later
appearance of n, as in B’s shape, adds a runtime check: calling
add_one with an A of length 100 binds n to 100, so a B of
length 50 raises an error.
The utilities keep definitions and uses consistent for us. When
structural_map changes the type of a
variable at its definition site, it creates a new variable with the new
type and repoints every use site to it. Replacing n with 128, for
example, changes A’s type to T.Tensor((128,), T.float32), so the
result has a new A, and A[i] in the body refers to it.
To receive the kind in our own function, we register it through
with_def_region_kind instead of callbacks, and the function is
called as fn(var, kind). The kind gives our function additional context
about each site, which analyses and rewrites can use to treat definitions
and uses differently. The following walk prints each definition site in
add_one:
def show_definition(var, kind):
if kind != tvm_ffi.DefRegionKind.NONE:
print(var.name, kind.name)
tvm_ffi.structural_walk(add_one, with_def_region_kind=(ir.Var, show_definition), order="pre")
The output is:
A DEF_PATTERN
n DEF_PATTERN
B DEF_PATTERN
n DEF_PATTERN
i DEF_SIMPLE
n appears twice, once for each shape: the first defines it, and the
second is the runtime check. Every other site is a use site, reported as
NONE.
structural_map accepts the same form. A rename leaves each variable’s
type unchanged, so we create the new variables ourselves: a new variable at
each definition site, returned as the same object at every use site:
new_vars = {}
def rename(var, kind):
if kind != tvm_ffi.DefRegionKind.NONE and var not in new_vars:
new_vars[var] = ir.Var(var.name + "1", var.ty)
return new_vars.get(var, var)
renamed = tvm_ffi.structural_map(add_one, with_def_region_kind=(ir.Var, rename), order="pre")
print(renamed.script())
The output is:
from __future__ import annotations
# from tvm_tirx.script import tirx as T
@T.function
def add_one[n1](A1: T.Tensor((n1,), T.float32), B1: T.Tensor((n1,), T.float32)):
for i1 in range(n1):
B1[i1] = A1[i1] + T.float32(1.0)
Pre-order traversal reaches each definition site before its use sites, so
the new variable exists by the time a use needs it. The
var not in new_vars check keeps the first replacement for n, so both
shapes share n1.
Advanced Control: Visit and Mutate#
Walk and map recurse for us, and that covers most tasks. Some analyses, however, need to know where they are in the program, not just which node they are looking at. To report the loops around each write, for example, a callback must know when the traversal enters a loop and when it leaves again.
tvm_ffi.structural_visit and
tvm_ffi.structural_mutate hand that
control to us. A matching callback also receives the visitor or mutator and
decides when, or whether, to process the node’s children; if it never does,
the children are skipped. Let’s record the loops around each write:
loops = []
writes = []
def visit_for(loop, visitor):
loops.append(loop.var.name)
result = visitor.default_visit(loop)
loops.pop()
return result
def visit_store(store, visitor):
writes.append((store.dest.name, list(loops)))
tvm_ffi.structural_visit(add_one, [(ir.For, visit_for), (ir.TensorStore, visit_store)])
print(writes)
The output is:
[('B', ['i'])]
visit_for pushes the loop variable, lets
default_visit visit the
loop’s children, then pops the variable, so loops always names the loops
around the current node. default_visit processes a node’s children with
their usual traversal, without calling the node’s own callback again;
default_mutate does
the same for a mutator.
Note
Because the callback now drives the recursion, it also takes on what the
walker used to do for us. Return the result of default_visit, as
visit_for does, so that a VisitInterrupt from a child reaches the
caller.
Mutate brings the same control to rewriting. A mutate callback returns the
new node, as in a map, but it can also set up state before its children are
rewritten. Let’s reverse the loop by replacing each use of i with
n - 1 - i:
reversed_index = {}
def reverse_loop(loop, mutator):
reversed_index[loop.var] = loop.end - 1 - loop.var
return mutator.default_mutate(loop)
def reverse_uses(var, mutator):
is_use = mutator.def_region_kind() == tvm_ffi.DefRegionKind.NONE
if is_use and var in reversed_index:
return reversed_index[var]
return var
add_one_reversed = tvm_ffi.structural_mutate(
add_one, [(ir.For, reverse_loop), (ir.Var, reverse_uses)]
)
print(add_one_reversed.script())
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)
When the mutate reaches the loop, reverse_loop records the replacement
n - 1 - i for i in the reversed_index dictionary, then lets
default_mutate rewrite the loop’s children, so the replacement is ready
before any use of i is reached. The loop’s own i is a definition and
must stay a variable, so reverse_uses replaces i only where
def_region_kind() is
NONE, that is, at uses.
These utilities cover the common ways to traverse, rewrite, and compare IR, but they are a convenience, not a requirement. IR objects are ordinary objects whose fields we can read directly, so we can still write our own recursive visitors and mutators by hand when a task calls for it, and mix them with the structural utilities as needed.