Control Flow

Note

All the code in this guide is in examples/programming_guides/control_flow.py.

This guide shows how to write conditions and loops in TIRx kernels. TIRx uses Python’s own control flow syntax: if, elif, else, for, while, break, and continue all mean what they mean in Python, and each thread executes them independently with its own values.

The examples use the following imports and a symbolic length n:

import torch

import tvm_tirx
from tvm_tirx.script import tirx as T

n = T.dynamic("n")

if Statements

An if statement runs its body only on the threads whose condition holds. The following kernel implements ReLU, which keeps positive values and replaces the rest with zero:

@T.function
def relu(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:
        if A[i] > 0.0:
            B[i] = A[i]
        else:
            B[i] = 0.0

The outer if i < n is the usual bounds check: the last block may have more threads than elements left, and the check keeps those threads from reading or writing past the end of the tensors. Inside it, if and else choose between two stores for each element, and elif can add more branches, just as in Python. A branch-free T.max(A[i], 0.0) would also compute ReLU; the if form shows how threads can take different paths.

for Loops

A for loop over range(...) runs its body once per iteration, in order, within each thread. T.unroll(...) iterates in the same way but asks the compiler to expand the loop into one copy of its body per iteration. The following kernel uses both to sum each row of an n x 64 matrix, one row per thread:

@T.function
def row_sum(A: T.Tensor((n, 64), T.float32), B: T.Tensor((n,), T.float32)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, 128), block=128))
    row = T.cuda.block_idx("x") * 128 + T.cuda.thread_idx("x")
    if row < n:
        total: T.float32 = 0.0
        for j in range(16):
            for k in T.unroll(4):
                total += A[row, j * 4 + k]
        B[row] = total

The outer range(16) loop stays a loop in the generated code, while the inner T.unroll(4) loop becomes four consecutive additions. Unrolling suits short loops with a fixed trip count, where it removes the loop overhead and lets the compiler schedule the statements together. total is a mutable local scalar, declared with a type annotation so that it keeps its value across iterations.

For nested loops over several dimensions, T.grid iterates over all index combinations, like nested range loops. The following kernel transposes a 64 x 64 matrix, with each thread moving a 4 x 4 tile:

@T.function
def transpose(A: T.Tensor((64, 64), T.float32), B: T.Tensor((64, 64), T.float32)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=16, block=16))
    row0 = T.cuda.block_idx("x") * 4
    col0 = T.cuda.thread_idx("x") * 4
    for r, c in T.grid(4, 4):
        B[col0 + c, row0 + r] = A[row0 + r, col0 + c]

for r, c in T.grid(4, 4) visits the 16 pairs (r, c) row by row, which is the same as for r in range(4) with for c in range(4) nested inside.

while Loops

A while loop repeats its body as long as its condition holds, which suits loops whose trip count depends on the data. The following kernel finds, for each query, the position where it would be inserted into a sorted tensor of 64 values, using binary search:

@T.function
def search_sorted(
    S: T.Tensor((64,), T.float32), Q: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.int32)
):
    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:
        lo: T.int32 = 0
        hi: T.int32 = 64
        while lo < hi:
            mid = (lo + hi) // 2
            if S[mid] < Q[i]:
                lo = mid + 1
            else:
                hi = mid
        B[i] = lo

Each thread narrows its own range [lo, hi) until it is empty, so different threads may run different numbers of iterations. lo and hi are mutable local scalars, so they keep their updates from one iteration to the next.

break and continue

break exits the innermost loop, and continue skips to its next iteration, just as in Python. Both work in for and while loops, and each thread decides on its own when to break or continue.

Early Thread Exit

T.gpu_thread_return() ends the current thread early. A common use is the bounds check, written as an early exit instead of wrapping the rest of the kernel in an if:

@T.function
def scale(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:
        T.gpu_thread_return()
    B[i] = A[i] * 2.0

Threads past the end of the tensors exit at T.gpu_thread_return(), and the remaining threads continue with the store. A plain Python return is different: it returns from the whole function and requires a value, so it cannot end a single GPU thread.