Metaprogramming

Note

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

This guide shows how to use Python to generate, reuse, and specialize TIRx kernels. TIRx functions are written in Python, so ordinary Python code runs before compilation and decides what each kernel contains: which dtypes and sizes it uses, which helper code it includes, and which values are fixed.

The examples use the following imports:

from __future__ import annotations

import torch

import tvm_tirx
from tvm_tirx.script import tirx as T

Parameterizing Kernels with Captured Values

A Python function can build kernels. The inner @T.function can read any value from the enclosing Python scope, such as a shape variable, a dtype, a constant, or a launch size. Each call to the outer function produces a new kernel with those values filled in:

def make_add(dtype, value, block=256):
    n = T.dynamic("n")

    @T.function
    def add(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype)):
        T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, block), block=block))
        i = T.cuda.block_idx("x") * block + T.cuda.thread_idx("x")
        if i < n:
            B[i] = A[i] + value

    return add


add_two_f16 = make_add(T.float16, 2, block=128)
add_two_f16_compiled = tvm_tirx.compile(add_two_f16, target="cuda")
a = torch.arange(10, dtype=torch.float16, device="cuda")
b = torch.empty_like(a)
add_two_f16_compiled(a, b)
print(b)

The parser reads n, dtype, value, and block from the enclosing scope when it parses add. In add_two_f16, the tensors have element type T.float16, the constant becomes T.float16(2.0), and each block has 128 threads. n stays a symbolic dimension, so the kernel still accepts tensors of any length. The script prints:

tensor([ 2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.], device='cuda:0',
       dtype=torch.float16)

Calling make_add(T.float32, 1) instead produces a separate float32 kernel from the same definition. The parser needs the function’s source code, so define kernel generators in a Python file rather than in an interactive session that has no source to read.

Reusable Code Fragments

Kernels often repeat the same few statements. A function marked with @T.inline holds such a fragment, and a kernel calls it like a regular function. In the following kernel, each thread handles two neighboring elements and calls add_one_if_valid once for each:

@T.inline
def add_one_if_valid(A, B, i, n):
    if i < n:
        B[i] = A[i] + 1


n = T.dynamic("n")


@T.function
def add_one_pairs(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=T.ceildiv(n, 512), block=256))
    i = (T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x")) * 2
    add_one_if_valid(A, B, i, n)
    add_one_if_valid(A, B, i + 1, n)

During parsing, TVMScript expands each add_one_if_valid(...) call into the helper’s statements, with the arguments substituted. The body of add_one_pairs is therefore equivalent to:

i = (T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x")) * 2
if i < n:
    B[i] = A[i] + 1
if i + 1 < n:
    B[i + 1] = A[i + 1] + 1

No function call remains in the compiled kernel, and other kernels can call add_one_if_valid in the same way.

Compile-Time Specialization

A symbolic dimension such as n varies at runtime. When a size should instead be fixed before compilation, @T.jit lets us declare it as a T.constexpr parameter and set it later through .specialize(...):

@T.jit
def add_one_fixed(
    A: T.Tensor((N,), T.float32),
    B: T.Tensor((N,), T.float32),
    *,
    N: T.constexpr,
):
    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_256 = add_one_fixed.specialize(N=256)

Unlike @T.function, which parses its definition immediately, @T.jit waits until .specialize supplies the constexpr values. add_one_256 is an ordinary TIRx function with 256 fixed in its parameter types, launch configuration, and bounds check, and it compiles with tvm_tirx.compile like any other function. The annotations refer to N before it has a value, which is why the file starts with from __future__ import annotations.