Core TVMScript#

TIRx kernels use tvm_tirx.script.tirx for the parser and core IR builders:

from tvm_tirx.script import tirx as Tx

Tx.alloc_tensor(...)

Tile ops and backend-specific namespaces are documented separately in Tensor Instruction Authoring API, CUDA Authoring and Support APIs, and Direct PTX Instructions. For the relationship between these authoring layers and TIRx IR, see The Programming Model.

Parser entry points#

Public canonical TVMScript dialect namespace.

tvm_tirx.script.tirx.function(function: LambdaType | None = None, **options: Any) → Any

Parse a Python function into a function of the selected IR language variant.

Parameters:

function (Callable, optional) – The function to be parsed. May be omitted to use the decorator with keyword options, such as @T.function(private=True).
private (bool, optional) – Whether the function should be treated as private. A private function has no global symbol attribute; a public function has a global symbol matching its name. Defaults to False.
check_well_formed (bool, optional) – Whether to check that the constructed function is well formed. Defaults to True. Module construction does not validate the completed module; call the language analysis API explicitly.
persistent (bool, optional) – For T.function, mark the resulting function as a persistent kernel. Defaults to False. See tvm_tirx.tirx.script.ir_builder.function().
pure (bool, optional) – For R.function, declare whether the function is pure, meaning that it has no observable side effects. Defaults to True. See tvm.relax.script.ir_builder.function_().
**options – Keyword options are forwarded to the selected language variant’s function_() hook, except check_well_formed, which controls parser validation. Options supported by only one language variant are not shared between T.function and R.function.

Returns:

result – The parsed function, or a decorator when function is omitted. Class members retain their Python functions until the enclosing module is constructed.

Return type:

Function or relax.Function or Callable

tvm_tirx.script.tirx.inline(function: LambdaType | None = None, **options: Any) → Callable[[...], Any]

Decorate a helper that constructs IR in its caller’s active frames.

Parameters:

function (Callable, optional) – The helper function. May be omitted to supply keyword options.
hygienic (bool, optional) – Whether the helper resolves symbols in its definition environment instead of its calling environment. Defaults to True. T.macro and R.macro capture values at definition time; T.inline refreshes captured closure cells when called.
**options – Keyword configuration for this helper decorator. hygienic controls name lookup as described above; other options are accepted but are not consumed or forwarded to builder hooks.

Returns:

result – The construction helper, or a decorator when function is omitted.

Return type:

Callable

Notes

T.inline follows Python lexical scoping with late binding of captured closure cells. Its return statements produce Python values, as do those of R.macro. T.macro emits returns in the active primitive function.

Examples

An inline helper can read values from its enclosing scope:

import tvm_tirx
from tvm_tirx.script import tirx as T
x_value = 128

@T.inline
def capture(A, B):
    B[()] = A[x_value]  # x_value resolved from enclosing scope

@T.function
def use(A: T.Tensor((1024,), "int32"), B: T.Tensor((), "int32")) -> None:
    capture(A, B)       # Produces B[()] = A[128]

tvm_tirx.script.tirx.jit(func: LambdaType | None = None, *, private: bool = False, check_well_formed: bool = True, persistent: bool = False) → JITFunction | Callable[[LambdaType], JITFunction]

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:

Deferred kernel, or its definition-site decorator when func is None.

Return type:

JITFunction or Callable[[types.FunctionType], JITFunction]

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.

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)

tvm_tirx.script.tirx.macro(function: LambdaType | None = None, **options: Any) → Callable[[...], Any]

Decorate a helper that constructs IR in its caller’s active frames.

Parameters:

function (Callable, optional) – The helper function. May be omitted to supply keyword options.
hygienic (bool, optional) – Whether the helper resolves symbols in its definition environment instead of its calling environment. Defaults to True. T.macro and R.macro capture values at definition time; T.inline refreshes captured closure cells when called.
**options – Keyword configuration for this helper decorator. hygienic controls name lookup as described above; other options are accepted but are not consumed or forwarded to builder hooks.

Returns:

result – The construction helper, or a decorator when function is omitted.

Return type:

Callable

Notes

Examples

An inline helper can read values from its enclosing scope:

import tvm_tirx
from tvm_tirx.script import tirx as T
x_value = 128

@T.inline
def capture(A, B):
    B[()] = A[x_value]  # x_value resolved from enclosing scope

@T.function
def use(A: T.Tensor((1024,), "int32"), B: T.Tensor((), "int32")) -> None:
    capture(A, B)       # Produces B[()] = A[128]

Core IR builder#

Common statement frames and construction helpers live in tvm_tirx.script.ir_builder.frame and tvm_tirx.script.ir_builder.stmt. The TIRx namespace re-exports them alongside its function, tensor/layout, and execution scope builders. Native shared statement builders live under include/tvm/script/ir_builder and src/script/ir_builder; dialect-specific thread-placement and tensor alias policies remain in the TIRx/S-TIR builders.

Concrete TIRx types, tensors, allocations and construction metadata.

class tvm_tirx.tirx.script.ir_builder.ir.Axis(*args: Any, **kwargs: Any)

A named physical coordinate used by a layout.

Use Axis.get(name) or a registered attribute such as Axis.m. Axis metadata identifies memory or thread coordinates and any owning backend; value @ axis attaches a stride or offset to that coordinate.

Examples

from tvm_tirx.script import tirx as T

memory = T.Axis.get("m")  # Also available as T.Axis.m.
layout = T.TileLayout(T.S[8 : 2 @ memory])
assert int(layout.apply(3)["m"]) == 6

property name: str: Return the canonical axis name.

classmethod get(name: str) → Axis

Get or create the axis singleton named name.

Unknown names are registered without thread or memory attributes.

is_thread() → bool: Check if the axis is a thread axis.

is_memory() → bool: Check if the axis is a memory axis.

property backend: str | None: Backend owning this axis, or None for a generic axis.

get_scope() → ThreadScope | None: Get the scope of the axis.

get_subscope() → ThreadScope | None: Get the subscope of the axis.

class tvm_tirx.tirx.script.ir_builder.ir.ComposeLayout(per_element: int, swizzle_len: int, atom_len: int, tile_layout: TileLayout, swizzle_inner: bool = True)

Apply an XOR swizzle to the memory coordinate of a tile layout.

All three bit counts must be nonnegative.

Parameters:

per_element (int) – Number of low address bits left unchanged by the swizzle.
swizzle_len (int) – Number of bits in each XOR group; zero disables swizzling.
atom_len (int) – Separation between the bit groups; must be at least swizzle_len.
tile_layout (TileLayout) – The underlying mapping to memory coordinates.
swizzle_inner (bool, optional) – If true, XOR the higher group into the lower group; otherwise XOR the lower group into the higher group. Defaults to true.

Examples

from tvm_tirx.script import tirx as T

tile = T.TileLayout(T.S[8, 8])
layout = T.ComposeLayout(0, 2, 3, tile)
# Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
assert int(layout.apply(1, 2)["m"]) == 11

class tvm_tirx.tirx.script.ir_builder.ir.FloatImm(ty: str | dtype | PrimType, value: float, loc: Location = UnknownLoc())

Float constant.

Parameters:

ty (str, tvm_ffi.dtype or PrimType) – The data type
value (float) – The constant value.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.script.ir_builder.ir.IntImm(ty: str | dtype | PrimType, value: int, loc: Location = UnknownLoc())

Int constant.

Parameters:

ty (str, tvm_ffi.dtype or PrimType) – The data type
value (int) – The constant value.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.script.ir_builder.ir.Iter(extent: Expr, stride: Expr, axis: Axis | str)

One extent-and-stride factor contributing to a layout axis.

Parameters:

extent (Expr) – Number of positions in this factor.
stride (Expr) – Increment on the output axis for each position.
axis (Axis or str) – Output coordinate receiving the contribution, resolved by name for strings.

tvm_tirx.tirx.script.ir_builder.ir.Lambda(parameter_types, function, *, ret_type=None)

Build a shared staging lambda from explicit types and a Python callable.

Scalar constructors such as T.float32 may be used as parameter types. An optional return annotation checks the body type without inserting casts.

class tvm_tirx.tirx.script.ir_builder.ir.Layout(*args: Any, **kwargs: Any)

Base class for mappings from logical indices to named physical axes.

Concrete layouts describe memory offsets and thread ownership. apply() returns coordinates keyed by axis name; size() describes the input size and span() the covered output span.

check_well_formed() → None

Check whether the layout is well formed.

Raises:: ValueError, RuntimeError – If the layout is not well formed.

size(axis_name: str | None = None)

Get the size of the layout.

Parameters:: axis_name (Optional[str]) – The name of the axis to get the size of. If not provided, the default input size will be returned.

span(axis_name: str | None = None)

Get the span of the layout.

Parameters:: axis_name (Optional[str]) – The name of the axis to get the span of. If not provided, the default span will be returned.

apply(*coord: list[Expr], shape: list[Expr] | None = None) → dict[str, Expr]

Apply the layout on the input coordinate and get the mapped output.

Input cases: - coord is a single element -> will be treated as a 1D coordinate - coord is a list of elements -> will be treated as a multi-dimensional coordinate - shape is provided -> turn the coord with shape into a 1D coordinate - shape is not provided -> use the default shape

Returns:: The mapped output (axis name -> value on the axis)
Return type:: Dict[str, Expr]

canonicalize() → Layout

Canonicalize the layout by simplifying and fusing iterators where possible.

Returns:: The canonicalized layout
Return type:: Layout

tile(outer: TileLayout, outer_shape: list[Expr], inner_shape: list[Expr]) → TileLayout | ComposeLayout

Tile the current layout with an outer layout.

Parameters:

outer (TileLayout) – The outer layout to tile with
outer_shape (List[Expr]) – The shape of the outer layout
inner_shape (List[Expr]) – The shape of the inner layout

Returns:

The resulting tiled layout

Return type:

Union[TileLayout, ComposeLayout]

tile_to(to_shape: list[Expr], current_shape: list[Expr]) → Layout

Tile the current layout to the given shape.

Parameters:

to_shape (List[Expr]) – The shape to tile to
current_shape (List[Expr]) – The current shape of the layout

is_trivial() → bool: Check if the layout is trivial.

unpack(num: int, axis: str = 'm') → Layout

Unpack the layout, where a single element in the layout is unpacked into num contiguous elements.

Parameters:

num (int) – The number of elements to unpack into
axis (str) – Storage coordinate whose element size changes. Ownership coordinates are preserved. Tensor-memory dtype views use TCol.

Returns:

The unpacked layout

Return type:

Layout

pack(num: int, axis: str = 'm') → Layout

Pack the layout, where num contiguous elements in the layout are packed into a single element.

Parameters:

num (int) – The number of elements to pack into
axis (str) – Storage coordinate whose element size changes. Ownership coordinates are preserved. Tensor-memory dtype views use TCol.

Returns:

The packed layout

Return type:

Layout

class tvm_tirx.tirx.script.ir_builder.ir.MemoryScope(*args: Any, **kwargs: Any)

A registered memory-space identity.

MemoryScope(name) looks up an existing scope and raises ValueError for an unknown name. Registration alone does not enable backend support.

Examples

custom = MemoryScope.register("my_accelerator.scratch")
assert MemoryScope("my_accelerator.scratch") == custom

property name: str: Return the canonical registered spelling.

classmethod register(name: str) → MemoryScope: Register a new canonical name, raising ValueError for a duplicate.

tvm_tirx.tirx.script.ir_builder.ir.Ptr: alias of PtrType

class tvm_tirx.tirx.script.ir_builder.ir.Range(begin: Expr, end: Expr | None = None, loc: Location = UnknownLoc())

A half-open interval stored as begin and extent.

Use from_min_extent() when the extent is already known.

Parameters:

begin (Expr) – The lower bound when end is provided; otherwise the extent of [0, begin).
end (Expr, optional) – The exclusive upper bound. With two arguments, extent is end - begin.
loc (Location, optional) – The source location of the range.

static from_min_extent(min_value: Expr, extent: Expr, loc: Location = UnknownLoc()) → Range

Construct a Range by min and extent.

This constructs a range in [min_value, min_value + extent)

Parameters:

min_value (Expr) – The minimum value of the range.
extent (Expr) – The extent of the range.
loc (Location, optional) – The location of this node in the source code.

Returns:

rng – The constructed range.

Return type:

Range

static get_extent_from_begin_end(begin: Expr, end: Expr) → Expr: Extract the width from endpoints, reusing a shallow addition operand.

tvm_tirx.tirx.script.ir_builder.ir.Tensor: alias of TensorType

class tvm_tirx.tirx.script.ir_builder.ir.TensorLoad(source: Expr, indices: list[Expr], loc: Location = UnknownLoc())

Load an element or vector from an indexed expression source.

The source type’s registered inference hook validates the inputs and derives the primitive result type.

Parameters:

source (Expr) – The tensor or other expression supporting indexed loads.
indices (list[Expr]) – Indices interpreted by the source type, in dimension order.
loc (Location, optional) – The source location of the load.

tvm_tirx.tirx.script.ir_builder.ir.TensorMap: alias of TensorMapType

class tvm_tirx.tirx.script.ir_builder.ir.TensorRegion(source: Expr, region: list[Range], ty: Type, loc: Location = UnknownLoc())

A region of an arbitrary tensor expression.

Parameters:

source (Expr) – The source expression.
region (list[Range]) – The ranges describing the region.
ty (tvm_tirx.ir.Type) – The result type, including any dialect-specific subscript semantics.
loc (Location, optional) – The location of the expression in the source code.

class tvm_tirx.tirx.script.ir_builder.ir.ThreadScope(*args: Any, **kwargs: Any)

A registered thread scope, queried by its qualified backend.scope name.

Obtain a tag with ThreadScope.get("cuda.warp"), for example; unknown names raise an error. backend identifies its owner and rank orders the hierarchy, with smaller ranks denoting wider scopes. The registration index does not determine the hierarchy.

classmethod get(name: str) → ThreadScope: Look up a registered tag without creating unknown scopes.

property name: str: Canonical qualified name, such as cuda.warp.

property backend: str: Backend owning this tag.

property local_name: str: Backend-local name, such as warp.

property rank: int: Explicit hierarchy rank; smaller ranks denote wider scopes.

class tvm_tirx.tirx.script.ir_builder.ir.TileLayout(spec: _LayoutSpec)

Map logical indices to memory and ownership coordinates using strided factors.

shard holds the factors for distinct logical elements; replica holds factors for copies of those elements. offset gives an origin per axis. Use from_iters() to construct these fields directly.

Parameters:: spec (layout specification) – Compose S[shape:stride] shard factors, optional R[shape:stride] replica factors, and offsets with +. Integer strides use the memory axis m; stride @ Axis.name selects another axis. Omitted strides use contiguous row-major memory order.

Examples

from tvm_tirx.script import tirx as T

layout = T.TileLayout(T.S[4, 8])  # Contiguous 4-by-8 tile.
assert int(layout.apply(1, 2)["m"]) == 10

static from_iters(shard: Sequence[Iter] = (), replica: Sequence[Iter] = (), offset: dict[Axis | str, Expr] | None = None) → TileLayout: Construct a TileLayout from pre-built Iter objects.

is_contiguous(analyzer=None) → bool

Whether memory-only tile strides are contiguous, regardless of origin.

Replicas and non-memory axes are rejected. An optional analyzer supplies existing constraints; symbolic proof is otherwise initialized lazily.

is_trivial() → bool: Check if the layout is trivial.

group(shape: list[Expr]) → tuple[Layout, list[int]]

Group the current layout by the given shape.

Parameters:: shape (List[Expr]) – The shape to group by
Returns:: The grouped layout and the separators
Return type:: Tuple[Layout, List[int]]

get_scope() → tuple[ThreadScope, ThreadScope] | None: Get the scope pair of the layout.

permute_dims(perm: list[int]) → TileLayout: Permute the dimensions of the layout.

permute_by_groups(seps: list[int], perm: list[int]) → TileLayout

Permute groups of shard iters defined by seps.

seps follows the convention of group()’s second return value: seps[0] == 0 and group i covers shard indices [seps[i], seps[i + 1]). The number of groups is len(seps) - 1.

Parameters:

seps (list[int]) – Group boundary positions in the shard list.
perm (list[int]) – Permutation of range(len(seps) - 1) selecting the new group order.

tvm_tirx.tirx.script.ir_builder.ir.Tuple(*fields: Type) → Type: Construct a tuple type for a TIRx function or binding annotation.

class tvm_tirx.tirx.script.ir_builder.ir.Var(name: str | None = None, ty: Type | str | None = None, loc: Location = UnknownLoc(), *, name_hint: str | None = None)

A canonical local variable in the IR.

Parameters:

name (str, optional) – The display name. Variables with the same name remain distinct bindings.
ty (Optional[Type or str]) – The variable type, or MissingType when omitted. A string denotes a primitive dtype; "handle" denotes a pointer to void.
loc (Location, optional) – Location that points to the original source code.
name_hint (str, optional) – Keyword-only alternative to name; supply exactly one name spelling.

tvm_tirx.tirx.script.ir_builder.ir.alloc_local(shape, dtype='float32', data=None, *, scope='local', align=-1, layout=<MISSING>, allocated_addr=None, annotations=None) → AlreadyEmitted[Var]: Emit a tensor allocation and return its bound variable in a receipt.

tvm_tirx.tirx.script.ir_builder.ir.alloc_scalar(dtype: str = 'float32', scope: str | MemoryScope = MemoryScope('global'), *, annotations: dict[str, Any] | None = None) → AlreadyEmitted[TensorLoad]: Allocate a zero-dimensional tensor (scalar), with optional allocation annotations.

tvm_tirx.tirx.script.ir_builder.ir.alloc_shared(shape, dtype='float32', data=None, *, scope='shared', align=-1, layout=<MISSING>, allocated_addr=None, annotations=None) → AlreadyEmitted[Var]: Emit a tensor allocation and return its bound variable in a receipt.

tvm_tirx.tirx.script.ir_builder.ir.alloc_tensor(placement=None, *, ty_args, attrs=None, annotations: dict[str, Any] | None = None, loc: Location = UnknownLoc(), ty=None) → Call | AlreadyEmitted[Var]

Allocate the tensor described by one explicit Tensor type argument.

Optional placement is an ordinary tuple operand. Allocation annotations belong to attrs; annotations accepts a Python dictionary whose scalar values are normalized to typed IR constants. Under an active builder, emit one binding and return its tensor variable in an AlreadyEmitted receipt. Parsed assignments consume this receipt automatically; programmatic code uses T.bind_(receipt) or receipt.value to obtain the variable without emitting another binding. Outside a builder return the allocation Call. Construct ir.Call("tirx.alloc_tensor", ...) directly when a pure expression is needed inside a builder.

Tensor allocation conveniences such as T.alloc_local and T.alloc_shared use the same receipt convention. Scalar allocation and declaration helpers instead return AlreadyEmitted[TensorLoad] in an active builder, consumed through the same binding protocol.

Examples

Parsed TVMScript uses an ordinary assignment. The allocation creates one binding, and subsequent tensor accesses use its variable.

from tvm_tirx import ir
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder

@T.function(private=True)
def example():
    tensor = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="local")])
    tensor[0] = 1

tensor_type = T.Tensor((4,), "int32", scope="local")
with IRBuilder() as builder:
    raw = ir.Call("tirx.alloc_tensor", [], ty_args=[tensor_type])
    receipt = T.alloc_tensor(ty_args=[tensor_type])
    tensor = T.bind_(receipt, name="tensor")
    assert tensor.same_as(receipt.value)
binding = builder.get()
assert isinstance(binding, ir.Bind)
assert binding.var.same_as(tensor)
assert not binding.value.same_as(raw)

tvm_tirx.tirx.script.ir_builder.ir.assign(dst: TensorLoad | Var, value: Any) → AlreadyEmitted[Stmt]

Store a value into scalar storage and return its emission receipt.

Parameters:

dst (TensorLoad | tvm_tirx.tirx.Var) – A tensor load identifying the destination tensor and indices, or a tensor variable with constant shape (1,). Immutable scalar variables and whole multi-element tensors are not assignment destinations.
value (Any) – The scalar value to store, converted by the ordinary tensor-store API.

Returns:

result – Receipt for the exact emitted store. Parser consumption does not emit it again. Active builder source locations are preserved.

Return type:

AlreadyEmitted[Stmt]

tvm_tirx.tirx.script.ir_builder.ir.decl_scalar(dtype, data, scope, layout=None) → AlreadyEmitted[TensorLoad]: Declare a zero-dimensional tensor (scalar) from a pointer.

tvm_tirx.tirx.script.ir_builder.ir.decl_tensor(data, *, ty_args, attrs=None, loc: Location = UnknownLoc(), ty=None) → Call: Declare a pointer-backed tensor with one explicit Tensor type argument.

tvm_tirx.tirx.script.ir_builder.ir.handle(dtype: str | None = None, storage_scope: str | MemoryScope = MemoryScope('global')) → Var

Create a TIR var that represents a pointer.

Parameters:

dtype (str | None) – The data type of the pointer. If omitted, construct an opaque handle.
storage_scope (str) – The storage scope of the pointer.

Returns:

res – The new tirx.Var with type handle or casted expression with type handle.

Return type:

Expr

tvm_tirx.tirx.script.ir_builder.ir.index_map(mapping: Callable, *, inverse_index_map: Callable | None = None, index_dtype: str = 'int64') → IndexMap: Create a TIR Index mapping

tvm_tirx.tirx.script.ir_builder.ir.local_scalar(dtype: str = 'float32', init: Any | None = None, *, name: str | None = None, annotations: dict[str, Any] | None = None) → AlreadyEmitted[TensorLoad]

Allocate scalar local storage, optionally naming and initializing it.

Parameters:

dtype (str) – The scalar data type.
init (Any) – Initial value. None leaves the storage uninitialized; otherwise one tensor store is emitted immediately after the allocation.
name (str | None) – Optional name for the allocated tensor variable, set in place.
annotations (dict[str, Any] | None) – Annotations passed through to the allocation.

Returns:

scalar – A receipt for the load from index zero of the one-element local tensor.

Return type:

AlreadyEmitted[TensorLoad]

tvm_tirx.tirx.script.ir_builder.ir.meta_class(cls)

Decorator for utility classes used inside @T.function.

Instances of decorated classes are treated as parser meta values.

tvm_tirx.tirx.script.ir_builder.ir.meta_var(value: T) → T

Return a Python metadata value without binding, naming or relocating it.

Parameters:

value (T) – Any host or IR object, including an unpackable sequence.

Returns:

T – The exact input object. No frame is required, no IR is emitted, and existing names and source locations are retained.
.. code:: python – # Source and generated Python (the value retains its identity) value = I.meta_var(existing_value) a, b = I.meta_var((left, right))

tvm_tirx.tirx.script.ir_builder.ir.ptr(dtype: str, storage_scope: str | MemoryScope = MemoryScope('global')) → Var

The pointer declaration function.

Parameters:

dtype (str) – The data type of the pointer.
storage_scope (str) – The storage scope of the pointer.

Returns:

res – The pointer.

Return type:

tvm_tirx.tirx.Var

tvm_tirx.tirx.script.ir_builder.ir.shared_scalar(dtype: str = 'float32', *, annotations: dict[str, Any] | None = None) → AlreadyEmitted[TensorLoad]: Allocate a zero-dimensional tensor in shared memory.

tvm_tirx.tirx.script.ir_builder.ir.smem(shape, dtype='float32', data=None, *, scope='shared', align=-1, layout=<MISSING>, allocated_addr=None, annotations=None) → AlreadyEmitted[Var]: Emit a tensor allocation and return its bound variable in a receipt.

tvm_tirx.tirx.script.ir_builder.ir.target(target_config: dict | str, host: dict | str | Target | None = None) → Target

Create a target

Parameters:

target_config (Union[Dict, str]) – The target configuration.
host (Optional[Union[Dict, str, Target]]) – The target configuration.

Returns:

res – The target.

Return type:

Target

tvm_tirx.tirx.script.ir_builder.ir.tmem(shape, dtype='float32', data=None, *, scope='tmem', align=-1, layout=<MISSING>, allocated_addr=None, annotations=None) → AlreadyEmitted[Var]: Emit a tensor allocation and return its bound variable in a receipt.

tvm_tirx.tirx.script.ir_builder.ir.Func: alias of FuncType

tvm_tirx.tirx.script.ir_builder.ir.String: alias of StringType

tvm_tirx.tirx.script.ir_builder.ir.Prim: alias of PrimType

tvm_tirx.tirx.script.ir_builder.ir.Opaque: alias of OpaqueType

class tvm_tirx.tirx.script.ir_builder.ir.NoneType: The concrete type whose only value is None.

tvm_tirx.tirx.script.ir_builder.ir.Missing: alias of MissingType

class tvm_tirx.tirx.script.ir_builder.ir.LetAnnotation(type_spec=None)#: Marker used by Tx.let and Tx.let[dtype] annotations to construct an explicit LetStmt.

Tracing ordinary Python#

The native builders also support frontends that execute ordinary Python to construct IR. IRBuilder.line_trace_scope() records file-backed Python source lines. Script-facing void calls automatically emit into an active IRBuilder, with or without line tracing:

from tvm_tirx.script import tirx as Tx
from tvm_tirx.script.ir_builder import IRBuilder
from tvm_tirx.tirx.script import ir_builder as I

with IRBuilder() as builder:
    with I.function_():
        I.func_name_("synchronize")
        with builder.line_trace_scope():
            with Tx.device_entry(
                launch=Tx.cuda.LaunchConfig(grid=1, block=32)
            ):
                Tx.ptx["bar.sync"](0)
function = builder.get()

Auto emission applies to script-facing calls through Tx, including PTX and backend namespaces. A void call returns an AlreadyEmitted receipt; passing that receipt to evaluate or emit_ does not emit it again. Calls producing values remain expressions. Raw IR constructors such as tvm_tirx.ir.Call, tvm_tirx.tirx.op and raw backend APIs remain pure even inside an IRBuilder. Line tracing only supplies source locations; it neither enables nor disables emission. The trace scope restores its previous trace and location state on exit, including exceptional exits, and applies only to its owning builder. Explicit with_loc and at_ locations compose with the traced lines.

Line tracing excludes TVM and standard-library implementation files. Frontends can supply additional directories using line_trace_scope(exclude_paths=(... ,)). User helpers retain their own source line and the caller’s location through CallSiteLoc. These locations cover whole lines; the TVMScript parser records expression ranges. Python loops and branches execute during construction; use the native statement frames to construct runtime control flow.

Tensor parameters and shared dimensions#

Traced frontends can bind native types directly with I.arg_. A shared symbolic dimension imposes the same size constraint on multiple tensor arguments:

n = Tx.Var("n", Tx.int64)
# Inside an I.function_() frame:
x = I.arg_("x", Tx.Tensor((n,), Tx.float32))
y = I.arg_("y", Tx.Tensor((n,), Tx.float32))

The compiled host entry accepts tensor arguments, including PyTorch tensors, and checks their dtype, rank and shared dimensions. Use independent symbolic dimensions for unrelated sizes. If the same dimension is also an explicit scalar parameter, I.arg_("n", n) binds that existing Var. x.data and x.ptr_to(indices) provide pointers inside the function. Tx.Ptr retains its raw-pointer ABI and does not imply tensor argument conversion.

Explicit register fragments#

A register fragment is an ordinary local tensor with an explicit layout. For a two-CTA Layout-B accumulator, define the TMEM and register layouts outside the kernel:

from tvm_tirx.backend.cuda.axis import tid_in_wg
from tvm_tirx.backend.cuda.layout import tmem_mma_operand_layout
from tvm_tirx.tirx.layout import TileLayout, S

layout_b = tmem_mma_operand_layout(
    "D", (64, 128), "float32", M=128, cta_group=2)
register_layout = TileLayout(S[(64, 2, 64) : (1 @ tid_in_wg, 64 @ tid_in_wg, 1)])

Inside the kernel, declare the accumulator over its allocated TMEM address, allocate the register fragment, and prepare matching instruction views:

C = Tx.cuda.decl_tmem(tmem_address, ty_args=[
    Tx.Tensor((64, 128), "float32", scope="tmem", layout=layout_b)
])
frag = Tx.alloc_tensor(ty_args=[
    Tx.Tensor((64, 128), "float32", scope="local", layout=register_layout)
])
C_words = C.rearrange("(wh r) (wl c) -> (wl wh) r c",
                      wh=2, r=32, wl=2)
frag_words = frag.rearrange("(wh r) (wl c) -> (wl wh) r c",
                            wh=2, r=32, wl=2)
Tx.cuda.tile.apply(
    Tx.ptx["tcgen05.ld.sync.aligned.32x32b.x64.b32"],
    frag_words, C_words, scope="warpgroup")
Tx.ptx.tcgen05.wait__ld.sync.aligned()

tmem_address must name a live allocation with at least 64 columns. The readback requires a full active warpgroup. Complete the producing MMA’s commit/mbarrier wait and the required tcgen05.fence::after_thread_sync before the load.

Layout B stores logical C[i,j] in TMEM lane i + 64*(j//64) and column j%64. The register layout assigns that same element to thread i + 64*(j//64) and local slot j%64. The views above both have shape (4,32,64): warp=2*wl+wh chooses the physical 32-lane slab, r selects its row, and c selects its 64 columns. Each warp issues one .32x32b.x64 instruction and each lane receives 64 registers. No temporary copy is needed. wait::ld makes the loaded registers available to subsequent instructions; it does not wait for an earlier MMA.

See tcgen05.ld / tcgen05.st for other native shapes, pack/unpack, and a complete store/load example.

Core TVMScript

Contents

Core TVMScript#

Parser entry points#

Core IR builder#

Tracing ordinary Python#

Tensor parameters and shared dimensions#

Explicit register fragments#