IR Overview#

One primary design goal of TIRx is to give developers and agents easy access to every stage of compilation, not only its input and output. By providing a consistent way to inspect and modify programs at each stage, TIRx makes it easier to understand and customize the compiler.

TIRx represents a program as a tree of intermediate representation (IR) nodes, such as functions, statements, and values. The tree is similar to the abstract syntax tree (AST) of Python code, but in a more structured, typed form. A compiler pipeline analyzes and transforms this tree step by step to produce target code.

IR objects are immutable: an analysis reads the tree, and a transformation builds a new tree instead of modifying the existing one, sharing the parts that did not change.

The tvm_tirx.script interfaces print the IR as TVMScript, a Python-based notation for tensor programs. For example, this TVMScript definition describes an add-one computation on a GPU:

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)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, 256), block=256))
    i = T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x")
    if i < n:
        B[i] = A[i] + 1

add_one adds one to every element of a tensor A of length n and writes the result to B. Each GPU thread computes its index i and handles one element. At a high level, TIRx IR contains the following elements:

Types (ir.Type) describe the values that expressions produce, such as the scalar type T.float32 and the tensor type T.Tensor((n,), T.float32).
Expressions (ir.Expr) compute values, such as A[i] + 1 and the call T.cuda.thread_idx("x") to a registered operator (ir.Op).
Statements (ir.Stmt) describe what the program does and in what order, such as the if branch and the write B[i] = A[i] + 1.
Functions (ir.BaseFunc) wrap a body of statements with parameters, such as add_one.
Modules (ir.IRModule) group named functions into a complete program.

Exploring the IR Tree#

To get a visual sense of the IR structure, we can use ir.format_tree to display an IR object and its fields as a text tree. Let’s start with the GPU function’s body, follow it down to the addition in B[i] = A[i] + 1, and print that expression:

from tvm_tirx import ir

region = add_one.body.seq[0]          # The device region.
branch = region.body.seq[1]          # The if statement after the index binding.
store = branch.then_case.seq[0]      # The write to B[i].
addition = store.value
print(ir.format_tree(addition))

The output is:

Add
├── a: TensorLoad(source=A, indices=[i])
└── b: FloatImm(value=1.0)

The root is the addition. Its a field holds the tensor read, and its b field holds the constant. Names such as A and i refer to variables already defined in the function. The display uses short class names: Add, for example, is ir.prim.Add.

The ir.format_tree function provides options to control how the tree is displayed. For example, we can show expression types by expanding the ty field for ir.Expr objects:

print(ir.format_tree(addition, expand=[(ir.Expr, "ty")]))

The output is:

Add
├── ty: PrimType(dtype=float32)
├── a: TensorLoad(source=A, indices=[i])
│   └── ty: PrimType(dtype=float32)
└── b: FloatImm(value=1.0)
    └── ty: PrimType(dtype=float32)

Here, each PrimType(dtype=float32) tells us that the expression it belongs to produces a float32 value. Field selectors can be names, such as "a", or (class, field) pairs that target a particular class. We use a class selector here to override the default rule that hides Expr.ty.

We can also make the display more compact by folding selected fields onto the same line as their parent. For example, folding a and b displays this addition on a single line:

print(ir.format_tree(addition, fold=["a", "b"]))

The output is:

Add(a=TensorLoad(source=A, indices=[i]), b=FloatImm(value=1.0))

This guide uses format_tree to inspect the objects in each example. Its API reference describes the remaining options for expanding, folding, and hiding fields.

Class Hierarchy#

The following diagram shows the class hierarchy of the IR nodes:

ffi.Object                   TVM-FFI base object
├── ir.Expr                  Expressions
│   ├── ir.Constant          Literal constants
│   │   ├── ir.IntImm        Integer constants
│   │   ├── ir.FloatImm      Floating-point constants
│   │   ├── ir.StringImm     String constants
│   │   ├── ir.DataTypeImm   Dtype constants
│   │   └── ir.GenericConst  Constants with arbitrary payloads
│   ├── ir.Var               Local variables
│   ├── ir.GlobalVar         Global function references
│   ├── ir.Tuple             Tuple construction
│   ├── ir.TupleGetItem      Tuple field access
│   ├── ir.Call              Calls
│   ├── ir.Op                Registered operators
│   ├── ir.TensorLoad        Tensor element reads
│   ├── ir.TensorRegion      Tensor regions
│   ├── ir.BaseFunc          Function definitions
│   │   └── tirx.Function    Tensor-program functions
│   ├── ir.StagingExpr       Expressions removed during compilation
│   │   └── ir.LambdaExpr    Construction-time lambdas
│   ├── ir.OpaqueExpr        Opaque construction-time expressions
│   ├── prim.Cast            Primitive dtype conversion
│   ├── prim.Add             Primitive addition, a + b
│   ├── prim.Sub             Primitive subtraction, a - b
│   ├── prim.Mul             Primitive multiplication, a * b
│   ├── prim.Div             Primitive division in C semantics
│   ├── prim.Mod             Primitive remainder in C semantics
│   ├── prim.FloorDiv        Primitive floor division, a // b
│   ├── prim.FloorMod        Primitive floor remainder, a % b
│   ├── prim.Min             Primitive minimum, min(a, b)
│   ├── prim.Max             Primitive maximum, max(a, b)
│   ├── prim.LShift          Primitive left shift, a << b
│   ├── prim.RShift          Primitive right shift, a >> b
│   ├── prim.BitwiseAnd      Primitive bitwise AND, a & b
│   ├── prim.BitwiseOr       Primitive bitwise OR, a | b
│   ├── prim.BitwiseXor      Primitive bitwise XOR, a ^ b
│   ├── prim.BitwiseNot      Primitive bitwise NOT, ~a
│   ├── prim.EQ              Primitive equal, a == b
│   ├── prim.NE              Primitive not equal, a != b
│   ├── prim.LT              Primitive less than, a < b
│   ├── prim.LE              Primitive less or equal, a <= b
│   ├── prim.GT              Primitive greater than, a > b
│   ├── prim.GE              Primitive greater or equal, a >= b
│   ├── prim.And             Primitive logical AND
│   ├── prim.Or              Primitive logical OR
│   ├── prim.Not             Primitive logical NOT
│   ├── prim.Select          Primitive conditional value selection
│   ├── prim.Let             Primitive expression-local binding
│   ├── prim.Ramp            Primitive vector arithmetic progression
│   ├── prim.Broadcast       Primitive scalar broadcast to vector
│   └── prim.Shuffle         Primitive vector lane selection
├── ir.Type                  Types
│   ├── ir.MissingType       Unresolved type information
│   ├── ir.NoneType          The type of None
│   ├── ir.AnyType           Any value
│   ├── ir.OpaqueType        Opaque construction-time types
│   ├── ir.PrimType          Scalar and vector dtypes
│   ├── ir.StringType        String values
│   ├── ir.PtrType           Pointer types
│   ├── ir.TupleType         Tuple field types
│   ├── ir.FuncType          Function signatures
│   ├── ir.TensorRegionType  Tensor region values
│   ├── tirx.TensorType      Tensor types
│   └── tirx.TensorMapType   Tensor map types
├── ir.Stmt                  Statements
│   ├── ir.SeqStmt           Statement sequences
│   ├── ir.Bind              Variable definitions
│   ├── ir.RegionStmt        Operator-defined lexical regions
│   ├── ir.AssertStmt        Program assertions
│   ├── ir.TensorStore       Tensor element writes
│   ├── ir.Evaluate          Evaluation with an unused result
│   ├── ir.If                Conditional branches
│   ├── ir.For               Counted loops
│   ├── ir.While             Conditional loops
│   ├── ir.Return            Function returns
│   ├── ir.Break             Loop exits
│   └── ir.Continue          Next loop iterations
├── ir.Location              Source locations
│   ├── ir.UnknownLoc        Unavailable source information
│   ├── ir.SourceLoc         Source ranges
│   └── ir.CallSiteLoc       Callee and caller locations
└── ir.IRModule              Global function definitions

These IR nodes are designed to be generic, so that new IR variants can be built on top of them, and their namespaces show their roles:

Namespace	Holds
`ir`	The common nodes shared by all variants.
`ir.prim`	Primitive expressions (shown as `prim`), which compute scalar or vector values of a primitive type (`ir.PrimType`), such as `float32` or `int32`. They concisely represent dynamic indices and element computations, such as `A[i] + 1` in `add_one`.
`tirx`	The types and operations specific to tensor programs, built on `ir`, such as `tirx.Function` and `tirx.TensorType`.

TIRx is designed for extension: downstream projects build on the IR the same way tirx builds on ir, registering new operators, defining new types, and adding their own TVMScript namespaces from Python, C++, or Rust, while reusing TIRx as a foundation.

Note

IR nodes can also carry source locations (ir.Location). For example, an expression’s loc field records where it came from in the source, so a diagnostic can point back to the original code.

How This Guide Is Organized#

The rest of this guide has two parts: the first explains what the IR contains, and the second covers the tools we use to inspect and transform it.

Understanding the IR

Types and Expressions explains how values, computations, and tensor shapes are represented.
Statements covers how a body binds variables, writes tensors, branches, and loops.
Operators, Calls, and Regions covers operators, the main extension point, and how calls and regions apply them.
Functions and Modules shows how functions are defined and organized in modules.
Printing and Parsing covers utilities for parsing and printing TVMScript.

Working with IR

Structural Traversal, Rewriting, and Comparison introduces structural walks, maps, visitors, mutators, and equality for traversing, rewriting, and comparing IR.
Symbolic Integer Analysis covers symbolic analysis and simplification of integer indices and bounds, which compiler analyses and transformations can use.
Working with IR in C++ and Working with IR in Rust show how to access and transform the same IR from C++ and Rust.

IR Overview

Contents

IR Overview#

Exploring the IR Tree#

Class Hierarchy#

How This Guide Is Organized#