tvm_tirx.ir.transform

IR passes, pass contexts, sequential pipelines, and pass instrumentation.

A transformation in TIRx is any function that takes a module and returns a new one. We can write it as an ordinary Python function and call it directly. Passes are an optional, thin layer on top: a pass wraps a transformation with a name and an optimization level, and runs it under a shared PassContext. This lets us compose passes into a pipeline with Sequential, select which passes run through the context’s options, and observe each step with pass instruments. A pass only describes how a transformation is invoked; inside, different passes can work in completely different ways.

The example below wraps a function as a pass with function_pass(), which applies it to every function in the module. specialize assumes that the tensors have 64 elements. A built-in simplification pass then removes the if i < 100 check, which the specialized loop always satisfies. A PassPrintingInstrument prints the module before the first pass and after the last one:

import tvm_ffi
from tvm_tirx import ir, tirx
from tvm_tirx.ir.transform import PassContext, Sequential
from tvm_tirx.ir.transform.instrument import PassPrintingInstrument
from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.ir_module
class Module:
    @T.function
    def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
        for i in range(n):
            if i < 100:
                B[i] = A[i] + 1


def replace_n(var):
    return ir.const(64, var.ty) if var.same_as(n) else var


@tirx.transform.function_pass(opt_level=0)
def specialize(func, mod, ctx):
    return tvm_ffi.structural_map(func, (ir.Var, replace_n))


pipeline = Sequential([specialize, tirx.transform.StmtSimplify()])
printer = PassPrintingInstrument(["specialize"], ["tirx.StmtSimplify"])
with PassContext(opt_level=2, instruments=[printer]):
    result = pipeline(Module)

The output is:

Print IR before: specialize
from __future__ import annotations

# from tvm_tirx.script import tirx as T

n = T.dynamic("n", dtype="int32")
@T.ir_module
class Module:
    @T.function
    def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
        for i in range(n):
            if i < 100:
                B[i] = A[i] + T.float32(1.0)


Print IR after: tirx.StmtSimplify
from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.ir_module
class Module:
    @T.function
    def add_one(A: T.Tensor((64,), T.float32), B: T.Tensor((64,), T.float32)):
        for i in range(64):
            B[i] = A[i] + T.float32(1.0)

Module itself is unchanged, and result holds the transformed module.

Classes

ModulePass(pass_info)	A pass whose callback receives and returns a complete IRModule.
Pass(*args, **kwargs)	Base interface for callable IRModule transformations and analyses.
PassContext([opt_level, required_pass, ...])	Scope optimization levels, pass selection, instruments, and pass options.
PassInfo(opt_level, name)	Metadata of a pass: its name and optimization level.
PassInstrument()	A pass instrument implementation.
Sequential([passes, opt_level, name])	A pass that works on a sequence of pass objects.

Functions

PrintIR([header])	Log the header and module IR, returning the module unchanged.
module_pass([pass_func, opt_level, name])	Decorate a module pass.
pass_instrument([pi_cls])	Decorate a class as a pass instrument.

Namespaces

• tvm_tirx.ir.transform.instrument