Statements#
This chapter introduces statements, the IR nodes that describe actions: binding a value to a variable, writing to a tensor, repeating work in a loop, or choosing between branches. TIRx provides a generic set of statements that appear in most high-level kernel programs. Statements run in order, and they contain expressions: a loop holds its bounds, and a store holds the value it writes.
A function’s body is a tree of statements. This add-one function
loops over its input and writes each result to B:
from tvm_tirx import ir
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
print(ir.format_tree(add_one.body))
The output is:
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)
The body is a sequence holding one For loop. The loop’s body is
another sequence holding one TensorStore, which writes the
expression A[i] + 1 to B[i].
ir.Stmt is the base class of all statements. The following table
lists the statement classes, and the sections below introduce them in
turn, ending with regions, the main place for extensions:
Class |
Represents |
Fields |
|---|---|---|
|
Statements that run in order |
|
|
A variable definition |
|
|
A write to a tensor element |
|
|
An expression run for its effect |
|
|
A loop over an integer range |
|
|
A loop that runs while a condition holds |
|
|
A choice between two sequences |
|
|
Skip to the next iteration of the innermost loop |
None |
|
Leave the innermost loop |
None |
|
A run-time check that raises an error |
|
|
A function return |
|
|
An operator-defined statement with a nested body |
|
Sequences, Bindings, and Stores#
ir.SeqStmt holds statements that run one after another, in its
seq field. A function body, a loop body, and each branch of an
if are all sequences.
ir.Bind defines a variable: its var field holds the new
variable, and its value field holds the expression that computes
it. ir.TensorStore writes value to the element of dest
selected by indices. In this add-one loop, the sum is bound to a
name before it is written:
from tvm_tirx import ir
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):
value = A[i] + 1
B[i] = value
loop = add_one.body[0]
print(ir.format_tree(loop.body))
add_one.body[0] is shorthand for add_one.body.seq[0], the first
statement in the sequence. The output is:
SeqStmt
├── seq[0]: Bind
│ ├── var: value = Var(ty=float32)
│ └── value: Add
│ ├── a: TensorLoad(source=A, indices=[i])
│ └── b: FloatImm(value=1.0)
└── seq[1]: TensorStore(dest=B, indices=[i])
└── value: value
The Bind defines value, and the TensorStore that follows it
uses it. A variable defined by a statement in a sequence is visible to
the statements after it in the same sequence, including any nested
inside them, but not outside that sequence. Each variable is defined
only once in a function.
ir.Evaluate runs an expression and discards its result. Its most
common use is to hold a Call made only for its effect.
T.cuda.cta_sync() is such a call: it makes each thread in a GPU
block wait until all threads in the block reach it:
@T.function
def sync_block(A: T.Tensor((128,), T.float32)):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
T.cuda.cta_sync()
region = sync_block.body[0]
print(ir.format_tree(region.body[0], hide=["attrs", "ty_args"]))
The output is:
Evaluate
└── value: Call(op=tirx.cuda.cta_sync)
└── args: []
The barrier produces no value; Evaluate turns the call into a
statement.
Loops#
ir.For repeats its body once for each value of its loop
variable var, from begin up to but not including end, in
steps of step. The loop defines var, which is visible only
inside body:
from tvm_tirx import ir
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
print(ir.format_tree(add_one.body[0]))
The output is:
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)
The tree omits step and kind because they have their default
values: a step of one, and iterations that run in order. Other values
of kind (ir.ForKind) ask for the iterations to be run in
parallel, vectorized, or unrolled.
ir.While repeats its body as long as its condition is
true. The condition is checked before each iteration. This loop halves
a value until it is at most one:
@T.function
def halve_until_small(X: T.Tensor((1,), T.float32)):
while X[0] > 1.0:
X[0] = X[0] * 0.5
print(ir.format_tree(halve_until_small.body[0]))
The output is:
While
├── condition: GT
│ ├── a: TensorLoad(source=X, indices=[IntImm(value=0)])
│ └── b: FloatImm(value=1.0)
└── body: SeqStmt
└── seq[0]: TensorStore(dest=X, indices=[IntImm(value=0)])
└── value: Mul
├── a: TensorLoad(source=X, indices=[IntImm(value=0)])
└── b: FloatImm(value=0.5)
Branches#
ir.If runs then_case when its condition is true, and
else_case otherwise. Both cases are sequences, and else_case
is None when the branch has no else. This loop applies ReLU,
keeping positive values and replacing the rest with zero:
from tvm_tirx import ir
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
def relu(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
for i in range(n):
if A[i] > 0:
B[i] = A[i]
else:
B[i] = 0.0
loop = relu.body[0]
print(ir.format_tree(loop.body[0]))
The output is:
If
├── condition: GT
│ ├── a: TensorLoad(source=A, indices=[i])
│ └── b: FloatImm(value=0.0)
├── then_case: SeqStmt
│ └── seq[0]: TensorStore(dest=B, indices=[i])
│ └── value: TensorLoad(source=A, indices=[i])
└── else_case: SeqStmt
└── seq[0]: TensorStore(dest=B, indices=[i])
└── value: FloatImm(value=0.0)
The condition is a boolean expression, here the comparison
A[i] > 0. For each element, only one of the two stores runs.
Loop Exits#
ir.Continue skips to the next iteration of the innermost enclosing
loop, and ir.Break leaves that loop. This loop records the index of
the first positive element of A:
from tvm_tirx import ir
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
def first_positive(A: T.Tensor((n,), T.int32), B: T.Tensor((1,), T.int32)):
B[0] = -1
for i in range(n):
if A[i] <= 0:
continue
B[0] = i
break
loop = first_positive.body[1]
print(ir.format_tree(loop.body))
The output is:
SeqStmt
├── seq[0]: If(else_case=None)
│ ├── condition: LE
│ │ ├── a: TensorLoad(source=A, indices=[i])
│ │ └── b: IntImm(value=0)
│ └── then_case: SeqStmt
│ └── seq[0]: Continue()
├── seq[1]: TensorStore(dest=B, indices=[IntImm(value=0)])
│ └── value: i
└── seq[2]: Break()
Continue and Break have no fields.
Assertions and Returns#
ir.AssertStmt checks a condition at run time and raises an
error when it is false. ir.Return ends the function and gives its
value to the caller. This function checks that its input is not
empty, then returns the first element plus one:
from tvm_tirx import ir
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
def first_plus_one(A: T.Tensor((n,), T.float32)) -> T.float32:
assert n > 0, "the input must contain at least one element"
return A[0] + 1
print(ir.format_tree(first_plus_one.body))
The output is:
SeqStmt
├── seq[0]: AssertStmt
│ ├── condition: GT
│ │ ├── a: n
│ │ └── b: IntImm(value=0)
│ ├── error_kind: StringImm(value='RuntimeError')
│ └── message_parts[0]: StringImm(value='the input must contain at least one element')
└── seq[1]: Return
└── value: Add
├── a: TensorLoad(source=A, indices=[IntImm(value=0)])
└── b: FloatImm(value=1.0)
The assertion’s error_kind names the error it raises, and its
message_parts hold the message as a list of string fragments. The
fragments are joined only when the assertion fails, through the TVM FFI
C API. Many assertions can therefore share common fragments, such as
an argument name or a shared phrase, which keeps error messages
compact in the compiled code. Calling the function with an empty
tensor fails with this message before the tensor read. Otherwise, the
function returns A[0] + 1.
Regions#
Regions are the main place for extensions. ir.RegionStmt is a
statement whose meaning is defined by an operator rather than by a
dedicated node class. A library registers a region operator to add a
new kind of statement, and the IR represents every such statement in
the same form:
opis the operator that defines the region’s meaning.argsandattrsare the operator’s operands and attributes, evaluated in the enclosing scope.bodyis the sequence of statements inside the region.body_paramsare variables the region defines on entry, visible only insidebody.result_varsare variables the region defines for the statements that follow it.
This GPU version of add-one places its computation in a device region, where each thread handles one element:
from tvm_tirx import ir
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
def gpu_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
region = gpu_add_one.body[0]
print(ir.format_tree(
region, hide=["attrs", "args", "body_params", "result_vars", "value"]))
The output is:
RegionStmt(op=tirx.device_entry)
└── body: SeqStmt
├── seq[0]: Bind
│ └── var: i = Var(ty=int32)
└── seq[1]: If(else_case=None)
├── condition: LT
│ ├── a: i
│ └── b: n
└── then_case: SeqStmt
└── seq[0]: TensorStore(dest=B, indices=[i])
The tirx.device_entry operator marks its body as code that runs on
the GPU, and its operands carry the launch configuration. The view
hides those operands and the value expressions to focus on the
statements inside the region.
Note
Regions and calls are the two ways operators extend TIRx: a call applies an operator as an expression, and a region applies one as a statement.