tvm_tirx.tirx.script.jit

Contents

tvm_tirx.tirx.script.jit#

tvm_tirx.tirx.script.jit(func=None, *, private=False, check_well_formed=True, persistent=False)#

Decorator: capture the kernel and defer parsing until .specialize().

Use @T.jit (instead of @T.function) when the kernel takes compile-time parameters annotated with T.constexpr or optional runtime parameters that may specialize to ir.NoneType. The resulting object exposes .specialize(**const_args), which returns a tvm_tirx.tirx.Function.

Parameters:

  • func (types.FunctionType or None, optional) – Function supplied by a definition-site @T.jit application. None, the default, returns a decorator for @T.jit(**options).

  • private (bool, optional) – Omit the function’s public global symbol when True. Default is False.

  • check_well_formed (bool, optional) – Validate each constructed specialization. Default is True; this option controls the parser separately from the function builder kwargs.

  • persistent (bool, optional) – Mark constructed specializations as persistent kernels. Default is False.

Returns:

JITFunction or Callable[[types.FunctionType], JITFunction] – Deferred kernel, or its definition-site decorator when func is None.

Raises:

  • TypeError – If the decorated value is not a Python function.

  • SyntaxError – If application occurs after definition or uses an unsupported bare or preconfigured callable alias instead of a qualified namespace decorator.

  • OSError – If the function’s source cannot be recovered.

Return type:

JITFunction | Callable[[LambdaType], JITFunction]

Examples

Specialize compile-time dimensions before compiling the kernel:

from __future__ import annotations

from tvm_tirx.script import tirx as T

@T.jit
def add(
    A: T.Tensor((N,), "float32"),
    B: T.Tensor((N,), "float32"),
    *,
    N: T.constexpr,
):
    for i in T.serial(N):
        B[i] = A[i] + 1.0

kernel = add.specialize(N=1024)  # returns a Function

@T.jit
def guarded(
    optional: T.Tensor((1,), "int32") | None,
    output: T.Tensor((1,), "int32"),
):
    if T.constexpr(optional is not None):
        output[0] = optional[0]
    else:
        output[0] = 0

present = guarded.specialize()
absent = guarded.specialize(optional=None)
Contents