Printing and Parsing

Printing and Parsing#

Every TIRx program, at any stage of compilation, can be written in TVMScript, an embedded Python syntax. The printer turns IR into TVMScript, and the parser turns TVMScript back into IR. This lets one syntax serve three roles:

  • DSL frontend: developers and agents write tensor programs directly in TVMScript.

  • IR constructor for tests: a test writes its input and expected output as TVMScript, then compares them structurally with the result of a pass.

  • Diagnostic printer: when compilation fails, the printer shows the program as it stands at that stage, in the same syntax the user wrote.

Note

The runnable code in this chapter is under examples/ir_guides in the repository.

How the Parser Works#

We value predictability over syntactic sugar. Each syntax form has one fixed meaning:

  • an annotation is a type;

  • an assignment binds a variable;

  • a subscript store writes a tensor;

  • a call is an ordinary Python call.

Types such as T.Tensor and operators such as T.cuda.thread_idx are plain functions that return IR, so the parser has no rules specific to them.

The parser is a thin layer over an IR builder, the Python API that constructs IR. The builder works with frames: entering a function, loop, or region opens a frame, and the statements built inside it become its body. Parsing is transpilation: the parser reads the Python syntax tree of a TVMScript definition, rewrites each syntax form into the matching builder call, then compiles and runs the result as an ordinary Python function in the definition’s original scope.

For example, this definition:

@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

becomes this builder program:

with _I0.IRBuilder() as _builder0:
    with T.function_() as _fn0:

        def _declare0():
            T.func_name_('add_one')
            A = T.arg_('A', T.Tensor((n,), T.float32))
            B = T.arg_('B', T.Tensor((n,), T.float32))
        _declare0()

        def _build0():
            A, B = _fn0.params
            with T.for_(T.range_(n), names='i') as i:
                T.setitem_(value=A[i] + 1, target=B, key=i)
        _build0()
_result0 = _fn0.function
_result0.__name__ = 'add_one'
T.check_well_formed_(_result0)

Concretely, the parser rewrites each syntax form into builder calls. The table also lists forms that this example does not use:

TVMScript

Builder program

@T.function definition

A T.function_() frame, named by T.func_name_('add_one')

Parameter A: T.Tensor((n,), T.float32)

T.arg_('A', T.Tensor((n,), T.float32))

for i in range(n):

A T.for_(T.range_(n), names='i') frame

x = value

x = T.bind_(value, name='x')

if cond: and else:

A T.if_ frame with T.then_ and T.else_

B[i] = value

T.setitem_(value=value, target=B, key=i)

Comparison i < n

T.lt_(i, n)

Arithmetic A[i] + 1

Unchanged Python

The signature is declared before the body is built, and the finished function is checked for well-formedness.

Arithmetic such as A[i] + 1 stays as written: A is the IR variable returned by arg_, so Python’s [] and + build the IR expression directly. Because the program runs as normal Python, n and other names come from the surrounding scope. The real output also passes source positions to each call, which is how parsed IR gets its loc.

Because the builder program runs as ordinary Python, the parser supports metaprogramming out of the box. A TVMScript function can read any Python value from its enclosing scope, and a Python function can build a new TVMScript function on each call:

def make_add(value, dtype):
    @T.function
    def add(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype)):
        for i in range(n):
            B[i] = A[i] + value

    return add


add_two = make_add(2, T.int32)
add_three = make_add(3, T.float32)

add_two adds 2 to int32 tensors, and add_three adds T.float32(3.0) to float32 tensors. The parser needs no template language for this: value and dtype are evaluated like any other Python name when the builder program runs.

The builder is also a public API, so we can construct the same function directly, without the parser:

from tvm_tirx.script.ir_builder import IRBuilder

with IRBuilder() as builder:
    with T.function_():
        T.func_name_("add_one")
        A = T.arg_("A", T.Tensor((n,), T.float32))
        B = T.arg_("B", T.Tensor((n,), T.float32))
        with T.serial(n) as i:
            T.assign(B[i], A[i] + 1)
built = builder.get()

ir.assert_structural_equal(built, add_one)

builder.get() returns a tirx.Function that is structurally equal to add_one. Hand-written builder code uses the user-facing API, such as T.serial and T.assign, rather than the protocol hooks the parser emits.

Writing builder code directly is sometimes helpful because it stays unambiguously in the Python domain. A plain Python for or if runs while the IR is built, and IR control flow comes only from explicit frames such as T.serial and T.if_ with T.then_. Developers and agents therefore know exactly what each line executes.

Note

To record source locations, wrap the construction in builder.line_trace_scope(), and each statement then records its source line in loc.

Region Syntax#

Regions are built with frames in the same way. A region operator such as T.device_entry opens a frame, and the statements inside its with block become the region’s body:

with 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

As a shorthand, a region call written as a plain statement opens the same frame for the rest of its enclosing block, so T.device_entry(...) followed by the same statements at the same indentation produces the same ir.RegionStmt.

Extending the Syntax#

Because the parser only rewrites syntax into builder calls, extending TVMScript rarely touches the parser. Operators and types registered by an extension package are plain functions that return IR, so they work in TVMScript as soon as the package is imported. For example, an operator registered under an extension namespace my_ext is written my_ext.cmul(a, b), like any built-in operator.

As an advanced topic, a new language variant can also be added if necessary. The builder hooks that the parser emits for each syntax form, such as bind_, setitem_, and for_, form the parser protocol, defined in tvm_tirx.script.ir_builder.parser_protocol. A variant implements these hooks on its own namespace and registers it with the parser, and the same parser then builds its IR.

Diagnostics#

When compilation fails, the most useful view is usually the program as it stands at that stage. Calling script() on the module at that point shows it in the same syntax used to write it.

The printer can also point at specific objects. obj_to_underline marks objects in the output, and obj_to_annotate attaches a message to them:

loop = add_one.body.seq[0]
addition = loop.body.seq[0].value
print(
    add_one.script(
        obj_to_underline=[addition],
        obj_to_annotate={addition: "result dtype: float32"},
    )
)

The output ends with:

@T.function
def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    for i in range(n):
        B[i] = A[i] + T.float32(1.0)  # result dtype: float32
               ^^^^^^^^^^^^^^^^^^^^^

show() accepts the same options and displays the result with syntax highlighting.