Symbolic Integer Analysis

Tensor programs compute indices from symbolic values such as a tensor length n, whose actual value is known only at runtime. Many analyses and transformations still need to reason about these indices at compile time: to simplify an index calculation, or to prove that an access stays in bounds so that a runtime check can be removed. TIRx provides a set of lightweight symbolic integer analysis utilities to help us with this reasoning and simplification:

• Simplification: an Analyzer rewrites an integer expression into a simpler equivalent form.

• Proofs: the analyzer checks whether a condition, such as i // 4 < n, always holds.

• Facts: a ConstraintContext records what is known at a point in the program, such as a loop’s range or a branch condition, and the analyzer uses those facts in its simplifications and proofs.

To see these utilities at work, consider a variant of add_one that works on a two-dimensional tensor. A single loop visits all n * 4 elements, and the indices i // 4 and i % 4 select the row and column of each one:

from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.function
def add_one_2d(A: T.Tensor((n, 4), T.float32), B: T.Tensor((n, 4), T.float32)):
    for i in range(n * 4):
        B[i // 4, i % 4] = A[i // 4, i % 4] + 1

Before generating code, a compiler typically flattens each tensor to one dimension. The function keeps its interface, while the body accesses one-dimensional views of the same storage, declared with T.decl_tensor. Element (r, c) of an (n, 4) tensor sits at position r * 4 + c in such a view, so the program becomes:

@T.function
def add_one_1d(A: T.Tensor((n, 4), T.float32), B: T.Tensor((n, 4), T.float32)):
    A_flat = T.decl_tensor(A.data, ty_args=[T.Tensor((n * 4,), T.float32)])
    B_flat = T.decl_tensor(B.data, ty_args=[T.Tensor((n * 4,), T.float32)])
    for i in range(n * 4):
        B_flat[i // 4 * 4 + i % 4] = A_flat[i // 4 * 4 + i % 4] + 1

The flattened index computes a division and a remainder for every element, yet it always equals i. Recognizing this, and many similar facts about index expressions, is what symbolic integer analysis is for.

Note

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

Simplify and Prove

An Analyzer answers questions about integer expressions. Let’s start with the index of the flattened program. We read it from the store in add_one_1d and ask the analyzer to simplify it:

from tvm_tirx import ir, sym

analyzer = sym.Analyzer()
loop_1d = add_one_1d.body.seq[-1]  # follows the two decl_tensor statements
index = loop_1d.body.seq[0].indices[0]  # i // 4 * 4 + i % 4
print(analyzer.simplify(index))

The output is:

i

The analyzer recognizes that taking the row i // 4, multiplying it by the row length, and adding the column i % 4 gives back i. This holds for any i, so the analyzer needs no further information, and a pass can replace the index with i and drop the division and the remainder.

Not every question can be answered so directly. In add_one_2d, the row index i // 4 must stay below n for the access to be valid. can_prove checks whether a condition always holds:

loop = add_one_2d.body.seq[0]  # for i in range(n * 4)
i = loop.var
print(analyzer.can_prove(i // 4 < n))

The output is:

False

Here the analyzer knows nothing about i. As far as it can tell, i could be any integer. The bound depends on the loop, and the analyzer needs that information as a fact.

Note

A False result means the analyzer could not establish the condition, not that the condition is false.

Constraint Context

A ConstraintContext holds the facts known at a point in the program. Inside the loop, i lies in [0, n * 4). We record that range and synchronize the analyzer with the context:

loop_context = sym.ConstraintContext()
loop_context.append(sym.RangeConstraint(i, ir.Range(loop.begin, loop.end)))
analyzer.sync_constraint_context(loop_context)
print(analyzer.can_prove(i // 4 < n))

The output is:

False

Still not proved. When n is zero or less, the range [0, n * 4) is empty and the loop never runs. Inside the loop body, though, we know the loop is running, so its extent is positive: n * 4 > 0. With that fact, the proof succeeds:

loop_context.append(sym.PredicateConstraint(loop.end > 0))
analyzer.sync_constraint_context(loop_context)
print(analyzer.can_prove(i // 4 < n))

The output is:

True

Facts also come from branches. Inside if i % 4 < 2, the condition holds, so it becomes a fact. We add it to a snapshot of the loop context, so the original context stays unchanged:

branch_context = loop_context.snapshot()
branch_context.append(sym.PredicateConstraint(i % 4 < 2))
analyzer.sync_constraint_context(branch_context)
print(analyzer.can_prove(i % 4 < 3))
analyzer.sync_constraint_context(loop_context)
print(analyzer.can_prove(i % 4 < 3))

The output is:

True
False

Inside the branch, the analyzer proves i % 4 < 3. Synchronizing with loop_context again models leaving the branch, and the fact no longer applies.

Note

An analyzer uses only the facts of the context it was last synchronized with, which lets the same analyzer serve different points in a program.

Structural Walk and Map with Constraint Context

Building contexts by hand works for one query, but a pass over a whole function would have to rebuild the facts at every node. The structural utilities in tvm_tirx.tirx collect the facts for us: each works like its TVM-FFI counterpart and also gives our callback the facts valid at each node.

Utility	Works like	Facts reach the callback as
structural_walk_with_constraint_context	structural_walk	Callback is f(value, constraints)
structural_map_with_constraint_context	structural_map	Callback is f(value, constraints)
structural_visit_with_constraint_context	structural_visit	visitor.constraints
structural_mutate_with_constraint_context	structural_mutate	mutator.constraints

Here constraints is a ConstraintContext, which our callback reads but must not add facts to.

Compilers often insert bounds checks to stay safe, such as the if i // 4 < n below. Here the loop already guarantees the condition, so the check only costs time on every iteration:

@T.function
def add_one_guarded(A: T.Tensor((n, 4), T.float32), B: T.Tensor((n, 4), T.float32)):
    for i in range(n * 4):
        if i // 4 < n:
            B[i // 4, i % 4] = A[i // 4, i % 4] + 1

A walk can check whether each condition is provable from the facts at that point:

from tvm_tirx import tirx

guard_analyzer = sym.Analyzer()
proven = []


def check_guard(value, constraints):
    if isinstance(value, ir.If):
        guard_analyzer.sync_constraint_context(constraints)
        proven.append(guard_analyzer.can_prove(value.condition))


tirx.structural_walk_with_constraint_context(add_one_guarded, check_guard, order="pre")
print(proven)

The output is:

[True]

The traversal collected the loop range and the fact that the loop runs, n * 4 > 0, the same facts we added by hand. The callback sees the facts that hold at the if itself, so the branch’s own condition is not a fact there.

A map with the same callback structure removes the check by returning the branch body in place of the if:

def remove_redundant_guard(value, constraints):
    if isinstance(value, ir.If):
        guard_analyzer.sync_constraint_context(constraints)
        if guard_analyzer.can_prove(value.condition):
            return value.then_case
    return value


simplified = tirx.structural_map_with_constraint_context(
    add_one_guarded, remove_redundant_guard, order="pre"
)
print(simplified.script())

The output is:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one_guarded[n](A: T.Tensor((n, 4), T.float32), B: T.Tensor((n, 4), T.float32)):
    for i in range(n * 4):
        B[i // 4, i % 4] = A[i // 4, i % 4] + T.float32(1.0)

A transformation like remove_redundant_guard is an ordinary function, and we can call it directly. When it helps, we can also wrap it as a pass with function_pass() and compose it with other passes, as described in the pass API.

Beyond simplification and proofs, tvm_tirx.sym also provides bounds, divisibility, iteration-map, and region analyses, documented in the tvm_tirx.sym reference.