Core TVMScript#

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

from tvm.script import tirx as Tx

Tx.alloc_buffer(...)

Tile primitives and backend-specific namespaces are documented separately in Tile Primitive 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#

The entry point of TVM parser for tirx.

tvm.tirx.script.parser.entry.prim_func(func: Callable | None = None, private: bool = False, check_well_formed=True, s_tir: bool = False, persistent: bool = False) → PrimFunc | Callable

The parsing method for tirx prim func, by using @prim_func as decorator.

Parameters:

func (Callable) – The function to be parsed as prim func. (Listed as optional to allow the decorator to be used without arguments, like @prim_func, or with an argument, @prim_func(private=True))
private (bool, optional) – Whether the function should be treated as private. A private function has no global symbol attribute; if the function is not private, it will have a global symbol matching the function name.

Returns:

res – The parsed tirx prim func.

Return type:

Union[PrimFunc, Callable]

tvm.tirx.script.parser.entry.inline(*args, definition_depth: int | None = None, defining_var_table=None) → Callable

Decorator for inline function definitions with Python LEGB scoping.

@T.inline follows Python’s lexical scoping with late binding: - At definition time, record which scopes are visible. - At call time, read current values from those scopes.

Example:

import tvm
from tvm.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.prim_func(s_tir=True)
def use(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
    capture(A, B)               # Produces B[()] = A[128]

class tvm.tirx.script.parser.entry.TIRJit(func: Callable, check_well_formed: bool = True, is_stir: bool = False, persistent: bool = False, private: bool = False)

Top-level kernel decorator with compile-time .specialize() params.

Parses the function body lazily: parsing is deferred until .specialize() supplies concrete values for the params annotated as T.constexpr. The return type of .specialize() is a tvm.tirx.PrimFunc, identical in type to what @T.prim_func produces today.

Constexpr params are removed from the resulting PrimFunc’s parameter list; their values are baked into the IR (e.g. into T.Buffer((M, K), ...) shape annotations and into the body).

specialize(**specialization_kwargs) → PrimFunc

Build a PrimFunc by binding constexprs and absent optional params.

Parameters:: **specialization_kwargs – One value per T.constexpr-annotated parameter. A T.Optional parameter may additionally be supplied as None to remove it from the resulting PrimFunc ABI. Omitting an optional parameter keeps it as a normal runtime parameter.
Returns:: A concrete TIRx PrimFunc, identical in type to the output of @T.prim_func.
Return type:: PrimFunc

tvm.tirx.script.parser.entry.jit(func: Callable | None = None, private: bool = False, check_well_formed: bool = True, is_stir: bool = False, persistent: bool = False) → TIRJit | Callable

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

Use @T.jit (instead of @T.prim_func) when the kernel takes compile-time parameters annotated with T.constexpr or runtime parameters that may be removed with T.Optional. The resulting object exposes .specialize(**specialization_kwargs), which returns a tvm.tirx.PrimFunc.

Example:

from tvm.script import tirx as T

@T.jit
def add(
    A: T.Buffer((N,), "float32"),
    B: T.Buffer((N,), "float32"),
    *,
    N: T.constexpr,
):
    ...

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

@T.jit
def guarded(optional: T.Optional(T.handle), out: T.handle):
    if optional is not None:
        ...

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

class tvm.tirx.script.parser.entry.TIRMacro(*args, **kwargs)

Specialization of the ScriptMacro class for TIR.

Apache-compatible hygienic macro. Distinct from TIRInline (which uses Python LEGB late binding) so upstream code that relies on capture-at-definition-time semantics keeps working.

call_count

Counter for the number of times this macro has been invoked. Used to generate unique block name suffixes.

Type:: int

parse_macro(parser: Parser) → None

The main macro parsing function. Different scripts may have different ways to parse a macro, and to return a value to the evaluator.

Parameters:

parser (Parser) – The parser with the appropriate frame already created and populated depending macro’s hygiene settings,

Returns:

The return value depends on the specifics of the particular script. It can be
”None” or any other value or any type.

tvm.tirx.script.parser.entry.macro(*args, hygienic: bool = True) → Callable

Decorator for macro definitions with hygienic capture.

Parameters:: hygienic (bool) – Specifies whether the macro is hygienic or not. A hygienic macro resolves symbols at definition time; a non-hygienic macro at use time. Defaults to True.

class tvm.tirx.script.parser.entry.BufferProxy: Buffer proxy class for constructing tirx buffer.

class tvm.tirx.script.parser.entry.PtrProxy: Ptr proxy class for constructing tirx pointer.

Core IR builder#

IRBuilder for TIR

The buffer declaration function.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The type of the buffer prior to flattening.
dtype (str) – The data type in the content of the buffer.
data (tirx.Var) – The pointer to the head of the data.
strides (List[Expr]) – The strides of each dimension.
elem_offset (Expr) – The offset in terms of number of dtype elements (including lanes).
scope (str) – The optional storage scope of buffer data pointer.
align (int) – The alignment requirement of data pointer in bytes.
offset_factor (int) – The factor of elem_offset field.
buffer_name (str) – The name of the buffer.

Returns:

res – The declared buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.prim_func(is_private: bool = False, s_tir: bool = False, persistent: bool = False, *, private: bool | None = None) → PrimFuncFrame

The primitive function statement.

Parameters:

is_private (bool) – Whether the PrimFunc is annotated as private.
s_tir (bool) – Whether this PrimFunc uses s_tir (apache-derived TIR) semantics: parser fills layout=None on buffers, ScriptComplete wraps body in a root SBlock. Default (False) selects tirx semantics: parser fills DefaultLayout(shape) and no root-block wrapping.
persistent (bool) – Whether this is a persistent kernel.
private (bool) – Alias for is_private (used in decorator syntax).

Returns:

res – The PrimFuncFrame.

Return type:

frame.PrimFuncFrame

tvm.tirx.script.builder.ir.arg(name: str, obj: Var) → Var

The PrimFunc arguments adding function.

Parameters:

name (str) – The name of the argument.
var (Union[tirx.Var, Buffer]) – The argument of tirx.Var or Buffer.

Returns:

res – The argument.

Return type:

Union[tirx.Var, Buffer]

tvm.tirx.script.builder.ir.func_name(name: str) → None

The PrimFunc naming statement.

Parameters:: name (str) – The name of the PrimFunc.

tvm.tirx.script.builder.ir.func_attr(attrs: dict[str, Any]) → None

The PrimFunc annotation statement.

Parameters:: attrs (Dict[str, Any]) – The annotations of the PrimFunc.

tvm.tirx.script.builder.ir.func_ret(ret_type: Type | None) → Type

The PrimFunc return type statement.

Parameters:: ret_type (Type) – The return type of the PrimFunc.
Returns:: res – The return type.
Return type:: Type

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

tvm.tirx.script.builder.ir.sblock(name: str = '', no_realize: bool = False, exec_scope: str = '') → SBlockFrame

The sblock declaration statement.

Parameters:

name (str) – The name of the sblock.
no_realize (bool) – The flag whether to construct SBlockRealize or SBlock.
exec_scope (str) – The execution scope of the block.

Returns:

res – The SBlockFrame.

Return type:

frame.SBlockFrame

tvm.tirx.script.builder.ir.block_name_suffix_context(block_suffix: str)

Context manager to set block name suffix during macro expansion.

Parameters:: block_suffix (str) – The suffix to append to block names (e.g., “_1”, “_2”).
Yields:: None

tvm.tirx.script.builder.ir.init() → BlockInitFrame

The block initialization statement.

Returns:: res – The BlockInitFrame.
Return type:: frame.BlockInitFrame

tvm.tirx.script.builder.ir.where(predicate: Expr | int) → None

The block predicate statement.

Parameters:: predicate (Union[Expr, Literal[0, 1]]) – The predicate condition.

tvm.tirx.script.builder.ir.reads(*buffer_slices: list[BufferRegion | TensorLoad]) → None

The block buffer region reading statement.

Parameters:: buffer_slices (List[Union[BufferRegion, TensorLoad]]) – The array of buffer regions to read.

tvm.tirx.script.builder.ir.writes(*buffer_slices: list[BufferRegion | TensorLoad]) → None

The block buffer region writing statement.

Parameters:: buffer_slices (List[Union[BufferRegion, TensorLoad]]) – The array of buffer regions to write.

tvm.tirx.script.builder.ir.sblock_attr(attrs: dict[str, Any]) → None

The block annotation statement (for non-tirx SBlock usage).

Parameters:: attrs (Dict[str, Any]) – The annotation of the block.

Statement-level buffer allocation (creates an AllocBuffer IR node).

Emits an AllocBuffer statement and returns the Buffer directly:

buf = T.alloc_buffer((128, 128))

For SBlock-level buffer allocation (added to SBlock.alloc_buffers), use T.sblock_alloc_buffer() instead.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[tirx.Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

Buffer

SBlock-level buffer allocation function.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The type of the buffer prior to flattening.
dtype (str) – The data type in the content of the buffer.
data (tirx.Var) – The pointer to the head of the data.
strides (List[Expr]) – The strides of each dimension.
elem_offset (Expr) – The offset in terms of number of dtype elements (including lanes).
scope (str) – The optional storage scope of buffer data pointer.
align (int) – The alignment requirement of data pointer in bytes.
offset_factor (int) – The factor of elem_offset field.
layout (Optional[Union[str, Layout]]) – The layout of the buffer.
allocated_addr (Optional[Union[int, Tuple[int]]]) – The address of the allocated buffer. Might be multi-dimensional. There can be pooled storage scopes on some devices. For example, the Trainium device has a pooled storage scope for the SRAN buffers. (“trn.sbuf”) CUDA has a pooled storage scope for the shared memory (“shared.dyn”)

Returns:

res – The allocated buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.wg_reg_tile(elem_per_thread: int, dtype: str = 'float32') → Var

Warpgroup-wide (128, elem_per_thread) register tile in local scope.

Sugar for the recurring pattern:

T.alloc_buffer(
    (128, elem_per_thread), dtype,
    layout=wg_local_layout(elem_per_thread),
    scope="local",
)

Used to stage a tcgen05 load: each of the 128 threads in a warpgroup owns one row of elem_per_thread contiguous elements.

class tvm.tirx.script.builder.ir.axis

The axis class

static spatial(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The spatial block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

static reduce(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The reduced block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

static scan(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The scanning block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

static opaque(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The opaque block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

static remap(kinds: str, bindings: list[Expr], dtype: str = 'int32') → list[Var] | Var

The block axis remapping function.

Parameters:

kinds (str) – The types of the iteration variables.
bindings (List[Expr]) – The binding values of the iteration variables.
dtype (str) – The data types of the iteration variables.

Returns:

res – The iteration variables.

Return type:

tirx.Var

static S(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The spatial block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

static R(dom: Range | list[Expr] | tuple[Expr], binding: Expr, dtype: str = 'int32') → Var

The reduced block axis defining function.

Parameters:

dom (Union[Range, List[Expr], Tuple[Expr]]) – The domain of the iteration variable.
binding (Expr) – The binding value of the iteration variable.
dtype (str) – The data type of the iteration variable.

Returns:

res – The iteration variable.

Return type:

tirx.Var

The serial For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.
step (Expr) – The optional step value of iteration.
unroll (bool or int, optional) – If True, adds {"pragma_unroll": True} annotation, which asks CUDA codegen to emit #pragma unroll while preserving the loop as a C++ for. If False, adds {"disable_unroll": True} annotation. If a positive integer, emits #pragma unroll N. Boolean values are handled separately from integers, so False keeps disabling unrolling.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds. Bounds that do not already have this dtype are converted (literals are retyped, other expressions get a Cast). Note T.thread_binding does not support this; its loop var is always int32.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.parallel(start: Expr, stop: Expr | None = None, *, annotations: dict[str, Any] | None = None, step: Expr | None = None, dtype: str | None = None) → ForFrame

The parallel For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.
step (Expr) – The optional step value of iteration.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.vectorized(start: Expr, stop: Expr | None = None, *, annotations: dict[str, Any] | None = None, step: Expr | None = None, dtype: str | None = None) → ForFrame

The vectorized For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.
step (Expr) – The optional step value of iteration.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.unroll(start: Expr, stop: Expr | None = None, *, annotations: dict[str, Any] | None = None, step: Expr | None = None, dtype: str | None = None) → ForFrame

The unrolled For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.
step (Expr) – The optional step value of iteration.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.thread_binding(start: Expr, stop: Expr | None = None, thread: str | None = None, *, annotations: dict[str, Any] | None = None) → ForFrame

The thread-binding For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
thread (str) – The thread for loop variable to bind.
annotations (Dict[str, Any]) – The optional annotations of the For statement.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.grid(*extents: tuple[Expr | tuple[Expr, Expr]], dtype: str | None = None) → ForFrame

The grid For statement.

Parameters:

extents (Tuple[Union[Expr, Tuple[Expr, Expr]]]) – If a single Expr is provided, it is used as the extent of the iteration. If a tuple of two Expr is provided, the first is the start of the iteration, and the second is the extent of the iteration.
dtype (str, optional) – The dtype of every loop variable, either "int32" or "uint32". When omitted each loop variable takes the dtype of its own extent.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.tirx.script.builder.ir.Assert(condition: Expr, message, error_kind: str = 'RuntimeError') → AssertFrame

Create an assertion statement.

Parameters:

condition (Expr) – The Expr to test.
message (str or list[str]) – The error message when the assertion fails. Can be a single string or a list of string parts (fragments stored separately in the IR for binary size reduction through string reuse).
error_kind (str) – The error kind (e.g. “RuntimeError”, “TypeError”, “ValueError”).

Returns:

res – The result AssertFrame.

Return type:

frame.AssertFrame

tvm.tirx.script.builder.ir.attr(node_or_dict: Any, attr_key: str | None = None, value: Expr | str | None = None) → AttrFrame | _FrameScope

Create an attribute node, or multiple attribute nodes from a dict.

Usage 1 — single attr:

with T.attr(node, key, value):
    ...

Usage 2 — dict sugar (node defaults to 0):

with T.attr({"key1": value1, "key2": value2}):
    ...

Parameters:

node_or_dict (Any) – If a dict, each key-value pair becomes an AttrStmt with node=0. Otherwise the node to annotate.
attr_key (str, optional) – Attribute type key (required when node_or_dict is not a dict).
value (Union[Expr, str], optional) – The attribute value (required when node_or_dict is not a dict).

Returns:

res – A single AttrFrame, or a _FrameScope wrapping multiple AttrFrames.

Return type:

Union[frame.AttrFrame, _FrameScope]

tvm.tirx.script.builder.ir.hint(message: str = '', **attrs) → HintFrame

Universal directive primitive for the sketch language.

Parameters:

message (str) – Free-form directive string that the agent interprets.
**attrs – Optional structured key-value attributes for known patterns.

Returns:

res – Usable as context manager (with T.hint(“msg”):) or bare statement (T.hint(“msg”)).

Return type:

frame.HintFrame

tvm.tirx.script.builder.ir.While(condition: Expr) → WhileFrame

Create a while node.

Parameters:: condition (Expr) – The termination condition of the loop.
Returns:: res – The result WhileFrame.
Return type:: frame.WhileFrame

tvm.tirx.script.builder.ir.Return(value: Expr) → None: Create a return node.

tvm.tirx.script.builder.ir.Break() → None: Create a break node.

tvm.tirx.script.builder.ir.Continue() → None: Create a continue node.

tvm.tirx.script.builder.ir.If(condition: Expr) → IfFrame

Create an if node.

Parameters:: condition (Expr) – The condition of if statement, executes the true branch if the condition is true, otherwise jump into the false branch.
Returns:: res – The result IfFrame.
Return type:: frame.IfFrame

tvm.tirx.script.builder.ir.Then() → ThenFrame

Create a then.

Returns:: res – The result ThenFrame.
Return type:: frame.ThenFrame

tvm.tirx.script.builder.ir.Else() → ElseFrame

Create an else.

Returns:: res – The result ElseFrame.
Return type:: frame.ElseFrame

tvm.tirx.script.builder.ir.decl_buffer(shape, dtype='float32', data=None, strides=None, elem_offset=None, byte_offset=None, scope='global', align=0, offset_factor=0, layout='default', allocated_addr=None) → Var

Create a buffer declaration node.

When data is provided, creates a DeclBuffer (alias to existing data). When data is None, creates an AllocBuffer (new allocation).

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The type of the buffer prior to flattening.
dtype (str) – The data type in the content of the buffer.
data (tirx.Var) – The pointer to the head of the data.
strides (List[Expr]) – The strides of each dimension.
elem_offset (Expr) – The offset in terms of number of dtype elements (including lanes).
byte_offset (Expr) – The offset in terms of number of bytes.
scope (str) – The optional storage scope of buffer data pointer.
align (int) – The alignment requirement of data pointer in bytes.
offset_factor (int) – The factor of elem_offset field.
layout (Layout) – The layout of the buffer.

Returns:

res – The declared buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.launch_thread(thread: IterVar | str, extent: Expr) → LaunchThreadFrame

Launch a thread.

Parameters:

thread (Union[IterVar, str]) – The iteration variable.
extent (Expr) – The extent of environment thread.

Returns:

res – The result LaunchThreadFrame.

Return type:

frame.LaunchThreadFrame

Examples

from tvm.script.ir_builder import tirx as T brow = T.env_thread(“blockIdx.y”) T.launch_thread(brow, 1)

tvm.tirx.script.builder.ir.env_thread(thread_tag: str, dtype: str = 'int32') → IterVar

Bind a var to thread env

Parameters:

thread_tag (str) – The thread type tag.
dtype (str) – The data type of the thread env.

Returns:

res – The result iteration variable gets bound to the thread env.

Return type:

IterVar

tvm.tirx.script.builder.ir.buffer_store(buffer: Var, value: Expr, indices: list[Expr | slice]) → None

Buffer store node.

Parameters:

buffer (Buffer) – The buffer.
value (Expr) – The value to be stored.
indices (List[Union[Expr, slice]]) – The indices location to be stored.

tvm.tirx.script.builder.ir.evaluate(value: Expr) → None

Evaluate the input expression.

Parameters:: value (Expr) – The input expression to evaluate.

tvm.tirx.script.builder.ir.boolean(expr: Expr | None = None) → Expr

Construct a new tirx.Var with type boolean or cast expression to type boolean.

Parameters:: expr (Expr) – The expression to be cast.
Returns:: res – The new tirx.Var with type boolean or casted expression with type boolean.
Return type:: Expr

tvm.tirx.script.builder.ir.handle(dtype: str | None = None, storage_scope: str = '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.script.builder.ir.void(expr: Expr | None = None) → Expr

Construct a new tirx.Var with type void or cast expression to type void.

Parameters:: expr (Expr) – The expression to be cast.
Returns:: res – The new tirx.Var with type void or casted expression with type void.
Return type:: Expr

tvm.tirx.script.builder.ir.var(dtype: str, name: str = '') → Var

Construct a new tirx.Var.

Parameters:

dtype (str) – The dtype of the Var.
name (str) – The name of the Var.

Returns:

res – The result tirx.Var.

Return type:

tirx.Var

tvm.tirx.script.builder.ir.ptr(dtype: str, storage_scope: str = '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:

tirx.Var

tvm.tirx.script.builder.ir.min(a: Expr, b: Expr) → Expr

Compute the minimum value of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand

Returns:

res – The result expression.

Return type:

Expr

tvm.tirx.script.builder.ir.max(a: Expr, b: Expr) → Expr

Compute the maximum value of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand

Returns:

res – The result expression.

Return type:

Expr

tvm.tirx.script.builder.ir.iter_var(v: Var | str, dom: Range, iter_type: str, thread_tag: str) → IterVar

The iteration variable.

Parameters:

var (Union[tirx.Var, str]) – The internal variable that is used for iteration.
dom (Range) – The domain of the iteration.
iter_type (str) – The iteration type.
thread_tag (str) – The thread type tag.

Returns:

res – The iteration variable.

Return type:

IterVar

tvm.tirx.script.builder.ir.comm_reducer(combiner: Callable, identity: list[Expr]) → CommReducer

Create a CommReducer from lambda inputs/outputs and the identities

Parameters:

combiner (Callable) – A binary function which takes two Expr as input to return a Expr.
identity (List[Expr]) – A list of types of output Expr.

Returns:

res – The CommReducer.

Return type:

CommReducer

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

tvm.tirx.script.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.script.builder.ir.buffer_var(dtype: str, storage_scope: str = '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:

tirx.Var

tvm.tirx.script.builder.ir.abs(x, span=None)

Get absolute value of the input element-wise.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.fabs(x, span=None)

Get absolute value of the input element-wise.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.acos(x)

Take acos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.acosh(x)

Take acos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.address_of(obj: Var | TensorLoad, span: Span | None = None) → Expr

Returns the address of a buffer element or addressable variable.

Parameters:

obj (Union[Buffer, TensorLoad, tirx.Var]) – The buffer, buffer load, or addressable variable.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.asin(x)

Take asin of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.asinh(x)

Take asinh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.atan(x)

Take atan of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.atan2(x1, x2)

Take arctan2(x1, x2).

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.atanh(x)

Take atanh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.bitwise_and(x, y, span=None)

Take bitwise and of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.bitwise_not(x, span=None)

Take bitwise not of input value

Parameters:

x (Expr) – Input operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.bitwise_or(x, y, span=None)

Take bitwise or of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.bitwise_xor(x, y, span=None)

Take bitwise xor of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.ceil(x, span=None)

Take ceil of float input x.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.clz(x)

Count leading zero bits of an integer x.

Parameters:: x (Expr) – Input 32 or 64 bit integer. The result is undefined if the input is 0.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.copysign(x1, x2)

Change the sign of x1 to that of x2, element-wise.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.cos(x)

Take cos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.cosh(x)

Take cosh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.erf(x)

Take gauss error function of the input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.exp(x)

Take exponential of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.exp2(x)

Calculate 2**x

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.exp10(x)

Calculate 10**x

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.floor(x: ExprWithOp, span=None)

Take floor of float input x.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.ceildiv(lhs, rhs, span=None)

Generic ceildiv operator.

Parameters:

lhs (object) – The left operand.
rhs (object) – The right operand.
span (Optional[Span]) – The location of this operator in the source.

Returns:

op – The result Expr of ceildiv operaton.

Return type:

tvm.Expr

tvm.tirx.script.builder.ir.floordiv(a, b, span=None)

Compute the floordiv of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

tvm.tirx.script.builder.ir.floormod(a, b, span=None)

Compute the floormod of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

tvm.tirx.script.builder.ir.fmod(x, y)

Return the remainder of x divided by y with the same sign as x.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.fma(x, y, z)

Take fused multiply-add of input x, y, z.

Parameters:

x (Expr) – First input argument.
y (Expr) – Second input argument.
z (Expr) – Third input argument.

Returns:

out – The result of x * y + z.

Return type:

Expr

tvm.tirx.script.builder.ir.filter(var, pred, *, span=None)

Thread-set filter escape hatch.

Use this wrapper only when the predicate is not in the canonical thread-filter grammar (see src/tirx/analysis/filter_canonical.h). Canonical predicates – pure conjunctions of scopeid_var <op> const comparisons plus bare T.cuda.elect_sync() calls – are recognized by the lowering pass directly from if cond:, so the wrapper is redundant for them.

When wrapped: var (a ScopeIdDef-declared scope identifier) tells the compiler which active-set axis to collapse to a singleton when the opaque predicate evaluates true; pred is preserved verbatim and evaluated at runtime.

The legacy three-argument range form filter(var, lo, hi) has been removed – write lo <= var and var < hi (or var == lo when hi == lo + 1) at the call site instead.

tvm.tirx.script.builder.ir.selector(var, pred, span=None)

Analysis-only active-thread selector.

selector(var, pred) denotes the unique value of var in the current active domain for which pred is true. It is intended for compiler metadata and should not survive to executable codegen.

tvm.tirx.script.builder.ir.hypot(x1, x2)

Equivalent to sqrt(x1**2 + x2**2), element-wise.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.if_then_else(cond, t, f, span=None)

Conditional selection expression.

Parameters:

cond (Expr) – The condition
t (Expr) – The result expression if cond is true.
f (Expr) – The result expression if cond is false.
span (Optional[Span]) – The location of this operator in the source.

Returns:

result – The result of conditional expression.

Return type:

Node

Note

Unlike Select, if_then_else will not execute the branch that does not satisfy the condition. You can use it to guard against out of bound access. Unlike Select, if_then_else cannot be vectorized if some lanes in the vector have different conditions.

tvm.tirx.script.builder.ir.infinity(dtype: str, span: Span | None = None) → Any

infinity value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The infinity value of dtype.

Return type:

tvm.Expr

tvm.tirx.script.builder.ir.isfinite(x, span=None)

Check if input value is finite.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.isinf(x, span=None)

Check if input value is infinite.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.isnan(x, span=None)

Check if input value is Nan.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.isnullptr(x, span=None)

Check if input value is nullptr.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.ldexp(x1, x2)

Returns x1 * (2 ** x2).

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.likely(cond, span=None)

Mark condition as likely.

Parameters:

cond (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The marked expression.

Return type:

Expr

tvm.tirx.script.builder.ir.log(x)

Take log of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.log1p(x)

Take log(x + 1) with respect to input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.log2(x)

Take log2 of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.log10(x)

Take log10 of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.lookup_param(param_name, span=None)

Returns the param by name

Parameters:

param_name (str) – The name of param.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.max_value(dtype: str, span: Span | None = None) → Any

maximum value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The maximum value of dtype.

Return type:

tvm.Expr

tvm.tirx.script.builder.ir.min_value(dtype, span=None)

minimum value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The minimum value of dtype.

Return type:

tvm.Expr

tvm.tirx.script.builder.ir.nearbyint(x, span=None)

Round elements of the array to the nearest integer. This intrinsic uses llvm.nearbyint instead of llvm.round which is faster but will results different from te.round. Notably nearbyint rounds according to the rounding mode, whereas te.round (llvm.round) ignores that. For differences between the two see: https://en.cppreference.com/w/cpp/numeric/math/round https://en.cppreference.com/w/cpp/numeric/math/nearbyint

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.nextafter(x1, x2)

Return the next floating-point value after x1 towards x2.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.popcount(x)

Count the number of set bits in input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.pow(x, y, span=None)

x power y

Parameters:

x (Expr) – Input argument.
y (Expr) – The exponent
span (Optional[Span]) – The location of this operator in the source code.

Returns:

z – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.q_multiply_shift(x, y, q, s)

Execute a multiplication between two Q-numbers x and y followed by a right shift s. The mathematical expression is:

out = round(x*y*2^-s)

More about Q-numbers here: https://en.wikipedia.org/wiki/Q_(number_format) The rounding rule is to the nearest value, rounding half up (i.e., round(x.1) = x and round (x.5) = x+1)

Parameters:

x (Expr) – First Q-number
y (Expr) – Second Q-number
q (Expr) – Number of fractional bits in x and y. Needs to be > 0
s (Expr) – Integer shift

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.q_multiply_shift_per_axis(x: Expr, y: Expr, ls: Expr, rs: Expr, q: IntImm, is_lshift_required: IntImm, is_rshift_required: IntImm)

Execute a multiplication between two Q-numbers x and y

Parameters:

x (Expr) – First Q-number.
y (Expr) – Second Q-number.
ls (Expr) – Integer left shift.
rs (Expr) – Integer right shift.
q (IntImm) – Number of fractional bits in x and y. Needs to be > 0.
is_lshift_required (IntImm) – Whether we need to do left shift or not.
is_rshift_required (IntImm) – Whether we need to do right shift or not.

Returns:

z – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.continue_loop(span=None)

Create a tir intrinsic call to represent continue expression

Parameters:: span (Optional[Span]) – The location of this operator in the source code.
Returns:: ret – The continue expression
Return type:: Expr

tvm.tirx.script.builder.ir.break_loop(span=None)

Create a tir intrinsic call to represent break expression

Parameters:: span (Optional[Span]) – The location of this operator in the source code.
Returns:: ret – The break expression
Return type:: Expr

tvm.tirx.script.builder.ir.reinterpret(dtype, value, span: Span | None = None) → Expr

Reinterpret a value as an exact primitive or pointer type.

Parameters:

dtype (str or tvm.ir.Type) – The data type.
value (Expr) – The input value.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The reinterpret cast value of dtype.

Return type:

tvm.Expr

tvm.tirx.script.builder.ir.round(x, span=None)

Round elements of the array to the nearest integer.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.rsqrt(x)

Take reciprocal of square root of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.shift_left(x, y, span=None)

Return the result of x left shifted by y bits.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.shift_right(x, y, span=None)

Return the result of x right shifted by y bits.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.sigmoid(x)

Quick function to get sigmoid

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.sin(x)

Take sin of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.sinh(x)

Take sinh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.sqrt(x)

Take square root of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.tan(x)

Take tan of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.tanh(x)

Take hyperbolic tanh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.tirx.script.builder.ir.thread_return()

TVM intrinsic to call thread_return()

Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.trunc(x, span=None)

Get truncated value of the input.

The truncated value of the scalar x is the nearest integer i which is closer to zero than x is.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

Expr

tvm.tirx.script.builder.ir.truncdiv(a, b, span=None)

Compute the truncdiv of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

Note

This is the default integer division behavior in C.

tvm.tirx.script.builder.ir.truncmod(a, b, span=None)

Compute the truncmod of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

Note

This is the default integer division behavior in C.

tvm.tirx.script.builder.ir.tvm_access_ptr(ptype, data, offset, extent, rw_mask)

Get head access address with memory access pattern info

Parameters:

ptype (Expr, PrimType, or str) – The data type of pointer. If a PrimType or str, it is wrapped via type_annotation() so that the lowering rule (which reads args[0].dtype() for the cast type) sees the intended dtype instead of void from a raw StringImm.
data (DType*) – The data of pointer.
offset (int) – The offset of pointer.
extent (int) – The extent of pointer.
rw_mask (int) – The read write mask.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.ptr_byte_offset(data, byte_offset, dtype)

Cast data + byte_offset to dtype*.

byte_offset is always in bytes. Use this when the source CUDA shape needs an explicitly typed local pointer derived from a byte-addressed base.

tvm.tirx.script.builder.ir.tvm_throw_last_error()

Throw TVMGetLastError()

Returns:: ret – The return expression
Return type:: Expr

tvm.tirx.script.builder.ir.print_buffer(buffer_var, dtype, is_string, is_scalar, dim_num, *shape): Print out buffer memory during runtime.

tvm.tirx.script.builder.ir.tvm_stack_alloca(dtype_str, num)

Return new on stack dtype[num]

Parameters:

dtype_str (str) – The data type of array.
num (int) – The size of array.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_stack_make_shape(*args)

Allocate a shape tuple on stack, return the handle

Parameters:: args (int) – The tuple shape.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset)

Allocate a Tensor(DLTensor) on stack, return the handle

Parameters:

data (Expr) – The data of array.
shape (Expr) – The shape of array.
strides (Expr) – The strides of array.
ndim (Expr) – The dimensions of array.
arr_dtype (Expr) – The data type of array.
elem_offse (Expr) – The element offset of array.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.call_packed(*args, span=None)

Build expression by call an external packed function.

The argument to packed function can be Expr or Buffer. The argument is the corresponding POD type when Expr is presented.

When the argument is Buffer, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is Tensor.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.call_cpacked(*args, span=None)

Build expression by call an external packed function.

Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.call_packed_lowered(*args, span=None)

Lowered version of call packed. The argument to packed function can be Expr or Buffer. The argument is the corresponding POD type when Expr is presented. When the argument is Buffer, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is Tensor.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.call_cpacked_lowered(*args, span=None)

Lowered version of call c-packed. Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.call_extern(dtype, func_name, *args, span=None)

Build expression by calling a extern function.

Parameters:

dtype (str) – The data type of the result.
func_name (str) – The extern function name.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.call_intrin(dtype: str | Type, func_name, *args, attrs=None, span=None)

Build expression by calling an intrinsic function.

Intrinsics can be overloaded with multiple data types via the intrinsic translation rule.

Parameters:

dtype (str or tvm.ir.Type) – The data type of the result.
func_name (str) – The intrinsic function name.
args (list) – Positional arguments.
attrs (Optional[tvm.ir.Attrs or Dict[str, Object]]) – Additional attributes for the call.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.call_llvm_intrin(dtype, name, *args, span=None)

Build expression by calling a llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.call_llvm_pure_intrin(dtype, name, *args, span=None)

Build expression by calling a pure llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.call_pure_extern(dtype, func_name, *args, span=None)

Build expression by calling a pure extern function.

Parameters:

dtype (str) – The data type of the result.
func_name (str) – The extern function name.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_tuple(*value)

Create a tuple structure in value field of AttrStmt

Parameters:: value (Expr) – The value in tuple.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.handle_add_byte_offset(handle, offset)

Add offset to handle

Parameters:

handle (Expr) – The handle.
offset (int) – The offset.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_struct_set(arr, index, field, value)

Set value in struct field in array

Parameters:

arr (StructType*) – The array of struct.
index (int) – The index of struct.
field (int) – The field of struct.
value (Expr) – The value to be set in field.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_struct_get(arr, index, field, dtype)

Get struct field value in array

Parameters:

dtype (str) – The date type of the result.
arr (StructType*) – The array of struct.
index (int) – The index of struct.
field (int) – The field of struct.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_thread_invariant(cond)

Mark condition as thread invariant.

Parameters:: cond (Expr) – The condition.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.tvm_thread_allreduce(*freduce_args)

Perform allreduce inside threadblock.

Parameters:: freduce_args (Expr) – The args.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout)

TVM intrinsic for tensor core load operators

Parameters:

fragment (tirx.Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_mma_sync(fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c)

TVM intrinsic for tensor core mma_sync operators

Parameters:

fragment_d (tirx.Var) – The wmma fragment_d.
index_d (Expr) – The fragment_d index.
fragment_a (tirx.Var) – The wmma fragment_a.
index_a (Expr) – The fragment_a index.
fragment_b (tirx.Var) – The wmma fragment_b.
index_b (Expr) – The fragment_b index.
fragment_c (tirx.Var) – The wmma fragment_c.
index_c (Expr) – The fragment_c index.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_bmma_sync(fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c)

TVM intrinsic for tensor core bmma_sync operators

Parameters:

fragment_d (tirx.Var) – The bwmma fragment_d.
index_d (Expr) – The fragment_d index.
fragment_a (tirx.Var) – The bwmma fragment_a.
index_a (Expr) – The fragment_a index.
fragment_b (tirx.Var) – The bwmma fragment_b.
index_b (Expr) – The fragment_b index.
fragment_c (tirx.Var) – The bwmma fragment_c.
index_c (Expr) – The fragment_c index.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_fill_fragment(fragment, m, n, k, index, value)

TVM intrinsic for tensor core fill_fragment operators

Parameters:

fragment (tirx.Var) – The wmma fragment
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
value (Expr) – The value to be filled in fragment.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_store_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout)

TVM intrinsic for tensor core store operators

Parameters:

fragment (tirx.Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1)

Perform synchronization in specified scope.

Parameters:

storage_scope (str) – The storage scope to perform synchronization.
is_load (bool) – Whether to perform load synchronization. (for global sync only)
num_blocks (int) – The number of blocks to synchronize. (for global sync only)

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_kernel_replace_point(): Mark where a transform should replace generated kernel initialization.

tvm.tirx.script.builder.ir.tvm_global_barrier_kinit()

Initialize the global barrier.

Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.tvm_warp_shuffle(mask, value, warp_id, width, warp_size)

Exchange value between threads inside a warp.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
warp_id (Expr) – The source lane index to fetch value.
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_warp_shuffle_up(mask, value, offset, width, warp_size)

Copy value from a lane with lower (by offset) index relative to caller.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
offset (Expr) – The difference between source lane index and destination lane index: offset = dst_lane_idx - src_lane_idx
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_warp_shuffle_down(mask, value, offset, width, warp_size)

Copy value from a lane with higher (by offset) index relative to caller.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
offset (Expr) – The difference between source lane index and destination lane index: offset = src_lane_idx - dst_lane_idx
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_warp_shuffle_xor(mask, value, lane_mask, width, warp_size)

Copy value from a lane with index computed by src_lane_idx ^ lane_mask.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
lane_mask (Expr) – The mask to compute source lane index:
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.tvm_warp_activemask()

Return a 32-bit mask indicates currently active threads in a calling warp.

Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.vectorlow(dtype, vec)

Get the low level half of the vector

Parameters:

dtype (str) – The data type of the result.
vec (list) – The input vector.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.vectorhigh(dtype, vec)

Get the high level half of the vector

Parameters:

dtype (str) – The data type of the result.
vec (list) – The input vector.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.vectorcombine(dtype, vec1, vec2)

Concat two vectors

Parameters:

vec1 (list) – The input vector.
vec2 (list) – The input vector.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.dp4a(vec1, vec2, acc=0)

Dot product of two int8x4 vectors and add an optional accumulator

Parameters:

vec1 (int8x4) – The input vector.
vec2 (int8x4) – The input vector.
acc (int32) – The accumulator.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.assume(cond=None)

Provide a true statement that can be used for simplifications

Parameters:: cond (Expr) – The constraint condition.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.undef()

Returns an initialized but arbitrary value

Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.tvm_call_packed(*args, span=None)

Build expression by call an external packed function.

The argument to packed function can be Expr or Buffer. The argument is the corresponding POD type when Expr is presented.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.tvm_call_cpacked(*args, span=None)

Build expression by call an external packed function.

Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.tvm_call_packed_lowered(*args, span=None)

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.tvm_call_cpacked_lowered(*args, span=None)

Lowered version of call c-packed. Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Buffer.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Expr

See also

te.extern: Create tensor with extern function call.

tvm.tirx.script.builder.ir.TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint)

Backend function to allocate temporal workspace

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
nbytes (int) – The size of the space requested.
dtype_code_hint (int) – The type code of the array elements. Only used in certain backends such as OpenGL.
dtype_bits_hint (int) – The type bits of the array elements. Only used in certain backends such as OpenGL.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.TVMBackendFreeWorkspace(device_type, device_id, ptr)

Backend function to free temporal workspace.

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
ptr (tirx.Var) – The result allocated space pointer.

Returns:

call – The call expression.

Return type:

Expr

tvm.tirx.script.builder.ir.start_profile_intrinsic(id)

Start profile intrinsic. :param id: The intrinsic id. :type id: int

Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.end_profile_intrinsic(id)

End profile intrinsic. :param id: The intrinsic id. :type id: int

Returns:: call – The call expression.
Return type:: Expr

class tvm.tirx.script.builder.ir.meta_var(value: Any)

A value used only for TVMScript parser-time metaprogramming.

Assignments unwrap this object without emitting an IR binding. The shared wrapper is exposed as I.meta_var; dialect namespaces may provide compatibility aliases to the same implementation. For Relax, this is the explicit opt-out from default primitive binding emission.

Parameters:: value (Any) – The parser-time value.

tvm.tirx.script.builder.ir.llvm_lookup_intrinsic_id(name)

Lookup LLVM intrinsic id by name.

Parameters:: name (str) – The name of the intrinsic.
Returns:: intrin_id – The intrinsic id.
Return type:: int

tvm.tirx.script.builder.ir.type_annotation(dtype)

Create a type annotation expression

Parameters:: dtype (Expr) – The data type.
Returns:: call – The call expression.
Return type:: Expr

tvm.tirx.script.builder.ir.broadcast: alias of Broadcast

tvm.tirx.script.builder.ir.ramp: alias of Ramp

tvm.tirx.script.builder.ir.cast(value, dtype, span=None): Cast an expression to the requested data type.

A canonical local variable in the IR.

Parameters:

name (str) – The name of the variable.
ty (Optional[Type or str]) – The exact type of the variable. A string denotes a primitive dtype.
span (Optional[Span]) – Span that points to the original source code.

access_ptr(access_mask, ptr_type='handle', content_lanes=1, offset=0, extent=None)

Get an access pointer to the head of buffer.

This is the recommended method to get buffer data ptress when interacting with external functions.

Parameters:

access_mask (int) – The access pattern MASK. Indicate whether the access will read or write to the data content.
ptr_type (str or tvm.ir.Type, optional) – The data type of the result pointer. Do not specify unless we want to cast pointer to specific type.
content_lanes (int, optional) – The number of lanes for the data type. This value is greater than one for vector types.
offset (Expr, optional) – The offset of pointer. We can use it to offset by the number of elements from the address of ptr.
extent (Expr, optional) – The extent of pointer.

Examples

# Get access ptr for read
buffer.access_ptr("r")
# Get access ptr for read/write with bitmask
buffer.access_ptr(BufferAccessKind.READ | BufferAccessKind.WRITE)
# Get access ptr for read/write with str flag
buffer.access_ptr("rw")
# Get access ptr for read with offset
buffer.access_ptr("r", offset = 100)
# Get access ptr for read with extent
buffer.access_ptr("r", extent = 100)

property byte_offset: Get the byte offset of the buffer.

byte_offset_of(indices, inner=True)

Get the byte offset of the buffer at the given indices. Note that indices subject to buffer’s layout mapping.

Parameters:

indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.
inner (bool, optional) – If False, the offset is relative to the original buffer. Default is True.

Returns:

offset – The byte offset of the buffer at the given indices.

Return type:

Expr

chunk(spec) → ChunkIndexer

Split dims into equal contiguous chunks and pick a chunk per dim — rank-preserving. Index the result with [picks].

spec is a per-dim tuple (length = rank). Each entry is None (leave the dim) or a positive int n (split that dim, extent E with E % n == 0, into n equal chunks of E // n). Then chunk(spec)[picks] takes one entry per dim: a chunked dim’s pick is the chunk index (int / Expr) and narrows that dim to the chunk’s [c*E//n : (c+1)*E//n) range — the dim is kept at E // n, no dimension is added; an unchunked dim’s pick is a normal index (: / int / slice). The result is the same BufferRegion as the hand-written slice — one line instead of the c*k : (c+1)*k arithmetic:

X[.., c * k : (c + 1) * k, ..]        # before (k = E // n)
X.chunk((None, .., n, ..))[.., c, ..]  # after (k inferred)

elem_offset_of(indices, inner=True)

Get the element offset of the buffer at the given indices. Note that indices subject to buffer’s layout mapping.

Parameters:

indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.
inner (bool, optional) – If False, the offset is relative to the original buffer. Default is True.

Returns:

offset – The element offset of the buffer at the given indices.

Return type:

Expr

get_flattened_buffer()

Generate a Buffer that is a flattened version of this buffer.

Returns:: flattened – The corresponding flat buffer.
Return type:: Buffer

is_scalar(alloc_or_decl=True)

Check if the buffer is a scalar.

Parameters:: alloc_or_decl (bool, optional) – Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar, False for decl_scalar.
Returns:: bool
Return type:: True if the buffer is a scalar, False otherwise.

local(*shape, layout=None) → Var

Create a thread-local view of this buffer.

By default, both the inferred and explicit-shape forms address the raw physical storage span. local()[k] is the k-th physical storage element, including any gaps or layout offset, while local(d0, d1, ...) is a row-major reshape of that same span. Pass layout= to request a mediated view explicitly. This is an escape hatch whose shape is interpreted by the supplied layout. When that shape is explicit, the parent buffer does not need a layout.

When called with no shape arguments, auto-infers a 1D shape from the span of the parent layout’s non-thread component (i.e. self.layout.storage().span()). The explicit-layout= form instead infers the parent layout’s storage().size() for compatibility. Either inference requires the parent buffer to have a layout.

Parameters:

shape (tuple of Expr) – The shape of the local view for indexing. Without layout=, its product must equal the per-thread physical storage span. With an explicit layout, the shape is not constrained by the raw span. If omitted, a matching 1D shape is computed automatically.
layout (optional) – Override layout. If None, the default (identity) layout is used.

Returns:

local – The corresponding local buffer.

Return type:

DeclBufferFrame

offset_of(indices)

Determine the offset of the provided indices in the flattened buffer.

Parameters:: indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.
Returns:: flattened_indices – The offset indices of the element in the flattened buffer.
Return type:: List[Expr]

permute(*dims) → Var

Permute the dimensions of the buffer.

Parameters:: dims (tuple of int) – The permutation of dimensions.
Returns:: permuted – The buffer with permuted dimensions.
Return type:: DeclBufferFrame

ptr_to(indices)

Get the pointer to the buffer at the given indices (logical indices).

Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.

rearrange(pattern: str = <object object>, /, **sizes) → Var

einops-style relayout in one line: buf.rearrange("b (2 r) -> 2 b r").

A pure reshape+permute+reshape over the SAME physical bytes, spelled as an einops pattern. Lowers to view (split lhs groups) → permute (reorder to rhs atom order) → view (merge rhs groups), so it inherits whatever the underlying axis machinery does: a plain (unswizzled) buffer collapses to a flat layout, a swizzled buffer keeps its swizzle, and a tmem buffer carries allocated_addr through. It therefore does NOT flatten a swizzle atom — the same pattern on a swizzled SMEM buffer vs an unswizzled TMEM buffer legitimately yields different physical layouts (that is the point: rearrange acts on the operand, not a string).

pattern is "lhs -> rhs"; each side is space-separated axis names, with (a b) grouping a product axis. Every lhs group’s product must equal that input dim; at most one axis per group may be unknown (inferred from the dim), the rest supplied via **sizes. Cannot express a replica (R[...]), a stride-fiction/padded view, or a reshape crossing a swizzle-atom boundary — keep those as explicit view(layout=...).

scope(): Return the storage scope associated with this buffer. :returns: scope – The storage scope associated with this buffer. :rtype: str

property sub: SubIndexer

buf.sub[2, 4:8, ::4].

Unlike plain buf[...] (BufferLoad for scalar indices, extent-1 BufferRegion dims for tile-primitive operands), sub follows numpy basic-indexing semantics as a view constructor: an integer index removes the dim (select), a:b narrows it, and a::s takes every s-th element (requires the extent divisible by s and a < s). Trailing dims are kept whole.

Type:: Numpy-style view indexer

tile(*specs) → TileIndexer

Chunk a dim: split it into factors, pick a chunk, keep the rest.

Rank-preserving — the picked dim’s remaining factors merge back into that one dim, and every other dim is untouched, so N dims in gives N dims out. Chunk multiple dims by chaining (dims never shift): buf.tile(0, (nx, -1))[cx, :].tile(1, (-1, ny))[:, cy].

tirx.Call as tile(dim, factors) for one dim, or pass several (dim, factors) specs as sugar for a chain. factors is the tuple the dim splits into (row-major, like unflatten(); one -1 inferred). The indexer takes one entry per factor: an int / Expr picks it (fixing the chunk, dropping the axis, folding its offset) and : keeps it. At least one factor per dim must be picked — a pure keep-everything split is unflatten(), not a chunk:

# 64 rows split into (stripe, warp, row) = (-1, WARPS, 4); this
# warp's 16 interleaved rows (stripe x row merged):
buf.tile(1, (-1, WARPS, 4))[:, warp, :]

tile(d, (n, -1))[c, :]   # contiguous block c
tile(d, (-1, n))[:, c]   # round-robin chunk c

A picked index may be a dynamic Expr (e.g. a warp id); picking several factors of one dim is allowed.

view(*args, **kwargs) → Var

Creates a new view of the buffer. (used by parser)

Supported signatures are view(*shape, layout=None), where shape can contain -1 to indicate that the dimension size is auto-inferred, and view(dtype: Union[str, tvm.DataType]).

Returns:: view – The corresponding view buffer.
Return type:: DeclBufferFrame

vload(begin, dtype=None)

Generate an Expr that loads dtype from begin index.

Parameters:

begin (Array of Expr) – The beginning index in unit of Buffer.dtype
dtype (str) – The data type to be loaded, can be vector type which have lanes that is multiple of Buffer.dtype

Returns:

load – The corresponding load expression.

Return type:

Expr

vstore(begin, value)

Generate a Stmt that store value into begin index.

Parameters:

begin (Array of Expr) – The beginning index in unit of Buffer.dtype
value (Expr) – The value to be stored.

Returns:

store – The corresponding store stmt.

Return type:

Stmt

with_allocated_addr(allocated_addr): Return a new buffer with the allocated address.

with_dtype(dtype): Return a new buffer with the dtype.

class tvm.tirx.script.builder.ir.Reduce(combiner: CommReducer, src: list[Expr], rdom: list[IterVar], condition: Expr, value_index: int, init: list[Expr] | None = None, span: Span | None = None)

Reduce node.

Parameters:

combiner (CommReducer) – The combiner.
src (list of Expr) – The source expression.
rdom (list of IterVar) – The iteration domain
condition (Expr) – The reduce condition.
value_index (int) – The value index.
init (list of Expr) – The initial value for output. This can be an int, float, or TE tensor-load Call.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.FloatImm(dtype: str | PrimType, value: float, span: Span | None = None)

Float constant.

Parameters:

dtype (str) – The data type
value (float) – The constant value.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.IntImm(dtype: str | PrimType, value: int, span: Span | None = None)

Int constant.

Parameters:

dtype (str) – The data type
value (int) – The constant value.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.StringImm(value: str, span: Span | None = None)

String constant.

Parameters:

value (str) – The value of the function.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Cast(dtype: str | PrimType, value, span: Span | None = None)

Cast expression.

Parameters:

dtype (str) – The data type
value (Expr) – The value of the function.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Add(a: Expr, b: Expr, span: Span | None = None)

Add node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Sub(a: Expr, b: Expr, span: Span | None = None)

Sub node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Mul(a: Expr, b: Expr, span: Span | None = None)

Mul node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Div(a: Expr, b: Expr, span: Span | None = None)

Div node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Mod(a: Expr, b: Expr, span: Span | None = None)

Mod node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.FloorDiv(a: Expr, b: Expr, span: Span | None = None)

FloorDiv node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.FloorMod(a: Expr, b: Expr, span: Span | None = None)

FloorMod node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Min(a: Expr, b: Expr, span: Span | None = None)

Min node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Max(a: Expr, b: Expr, span: Span | None = None)

Max node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.EQ(a: Expr, b: Expr, span: Span | None = None)

EQ node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.NE(a: Expr, b: Expr, span: Span | None = None)

NE node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.LT(a: Expr, b: Expr, span: Span | None = None)

LT node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.LE(a: Expr, b: Expr, span: Span | None = None)

LE node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.GT(a: Expr, b: Expr, span: Span | None = None)

GT node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.GE(a: Expr, b: Expr, span: Span | None = None)

GE node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.And(a: Expr, b: Expr, span: Span | None = None)

And node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Or(a: Expr, b: Expr, span: Span | None = None)

Or node.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Not(a: Expr, span: Span | None = None)

Not node.

Parameters:

a (Expr) – The input value
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Select(condition: Expr, true_value: Expr, false_value: Expr, span: Span | None = None)

Select node.

Note

Select may compute both true_value and false_value. Use tvm.tirx.if_then_else instead if you want to get a conditional expression that only evaluates the correct branch.

Parameters:

condition (Expr) – The condition expression.
true_value (Expr) – The value to take when condition is true.
false_value (Expr) – The value to take when condition is false.
span (Optional[Span]) – The location of this expression in the source code.

tvm.tirx.script.builder.ir.BufferLoad(buffer: Var, indices: list[Expr], span: Span | None = None) → TensorLoad

Construct a validated buffer load.

Parameters:

buffer (Buffer) – The buffer to be loaded.
indices (List[Expr]) – The buffer indices to load values from.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Ramp(base: Expr, stride: Expr, lanes: Expr, span: Span | None = None)

Ramp node.

Parameters:

base (Expr) – The base expression.
stride (Expr) – The stride of the ramp.
lanes (Expr) – The lanes of the expression.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Broadcast(value: Expr, lanes: Expr, span: Span | None = None)

Broadcast node.

Parameters:

value (Expr) – The value of the expression.
lanes (Expr) – The lanes of the expression.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Shuffle(vectors: list[Expr], indices: list[Expr], span: Span | None = None)

Shuffle node.

Parameters:

vectors (List[Expr]) – The vectors
indices (List[Expr]) – The indices
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.script.builder.ir.Call(op: Expr | str, args: list[Expr] | tuple[Expr, ...], attrs: Attrs | dict | None = None, ty_args: list[tvm.ir.Type] | tuple[tvm.ir.Type, ...] | None = None, span: Span | None = None, ret_ty: Type | str | None = None): Core function call node.

class tvm.tirx.script.builder.ir.CallEffectKind: Possible kinds of tirx.Call effects.

tvm.tirx.script.builder.ir.Bind(value: Expr, type_annotation: Type | None = None, *, var: Var | None = None) → Var

Create a Bind (variable binding).

Emits a flat Bind statement to the current frame and returns the bound variable.

Parameters:

value (Expr) – The value to be bound.
type_annotation (Optional[Type] = None) – The type annotation of the binding. Usually it is used for fine-grained var typing, particularly, PointerType.
var (Optional[tirx.Var] = None) – The variable to bind. If not specified, a new variable will be created.

Returns:

var – The bound variable.

Return type:

tirx.Var

tvm.tirx.script.builder.ir.bind(value: Expr, type_annotation: Type | None = None, *, var: Var | None = None) → Var

Create a Bind (variable binding).

Emits a flat Bind statement to the current frame and returns the bound variable.

Parameters:

value (Expr) – The value to be bound.
type_annotation (Optional[Type] = None) – The type annotation of the binding. Usually it is used for fine-grained var typing, particularly, PointerType.
var (Optional[tirx.Var] = None) – The variable to bind. If not specified, a new variable will be created.

Returns:

var – The bound variable.

Return type:

tirx.Var

class tvm.tirx.script.builder.ir.LocalVectorAnnotation(dtype: str, shape: tuple)

Marker for local vector/tensor allocation via type annotation subscript.

Created when a DtypeConstructor is subscripted, e.g. T.float32[N] or T.float32[M, N]. The parser’s visit_ann_assign recognises this object and lowers it to T.alloc_local(shape=..., dtype=...).

class tvm.tirx.script.builder.ir.DtypeConstructor(ffi_name: str, dtype_str: str)

Callable + subscriptable dtype object.

Replaces the plain functions previously returned by func_gen.

T.float32() — same FFI call as before (returns tirx.Var).
T.float32[N] — returns LocalVectorAnnotation("float32", (N,)).
T.float32[M, N] — returns LocalVectorAnnotation("float32", (M, N)).
x: T.float32 — parser calls this object, gets a tirx.Var.

tvm.tirx.script.builder.ir.Let(expr: Expr, where: dict[Var, Expr]) → Expr: Create a Let expression binding

class tvm.tirx.script.builder.ir.IterVar(dom: Range, var: Var | str, iter_type: int, thread_tag: str = '', span: Span | None = None)

Represent iteration variable.

IterVar represents axis iterations in the computation.

Parameters:

dom (Range) – The domain of the iteration.
var (Union[tirx.Var, str]) – The internal variable that is used for iteration.
iter_type (int) – The iteration type.
thread_tag (str) – The thread type tag.
span (Optional[Span]) – The location of this expression in the source code.

See also

te.thread_axis: Create thread axis IterVar.
te.reduce_axis: Create reduce axis IterVar.

expr_ty() → PrimType: Compile-time type of the iteration variable.

class tvm.tirx.script.builder.ir.CommReducer(lhs: list[Var], rhs: list[Var], result: list[Expr], identity_element: list[Expr], span: Span | None = None)

Commutative reduce operator

Parameters:

lhs (List[tirx.Var]) – The left arguments of the reducer.
rhs (List[tirx.Var]) – The right arguments of the reducer.
result (List[Expr]) – The reduction results.
identity_element (List[Expr]) – The identity elements.
span (Optional[Span]) – The location of this expression in the source code.

tvm.tirx.script.builder.ir.Range(begin: Expr, end: Expr) → Range

Create a Range object.

Parameters:

begin (Expr) – The begin value of the range.
end (Optional[Expr]) – The end value of the range.

tvm.tirx.script.builder.ir.vscale(): Get the target’s vscale value. It will be lowered to llvm.vscale intrinsic (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic) :returns: call – tirx.Call to the vscale intrinsic :rtype: Expr

tvm.tirx.script.builder.ir.get_active_lane_mask(dtype, base, limit)

Calculate a predicate mask given an upper bound (limit) and a current value (base).

It will be lowered to the llvm.get.active.lane.mask intrinsic. (https://llvm.org/docs/LangRef.html#llvm-get-active-lane-mask-intrinsics)

Parameters:

dtype (str) – The data type of the result.
base (Expr) – An expression reprsenting the base.
limit (Expr) – An expression representing the limit.

tvm.tirx.script.builder.ir.masked_load(dtype, buffer, *indices_and_mask)

Load vector lanes selected by a predicate mask.

Parameters:

dtype (str) – The vector data type to load.
buffer (Buffer) – The buffer to load.
indices_and_mask (Expr) – The buffer indices followed by a boolean lane mask. The mask must match the lane count and scalability of the loaded vector.

Returns:

call – A tirx.masked_load call with result type dtype.

Return type:

Expr

tvm.tirx.script.builder.ir.masked_store(buffer, value, *indices_and_mask)

Store vector lanes selected by a predicate mask.

Parameters:

buffer (Buffer) – The buffer to update.
value (Expr) – The vector value to store.
indices_and_mask (Expr) – The buffer indices followed by a boolean lane mask. The mask must match the lane count and scalability of value.

Returns:

call – A void-typed tirx.masked_store call.

Return type:

Expr

tvm.tirx.script.builder.ir.call_kernel(kernel, launch_args: list[int | Expr | list[int | Expr]], *args: list[Any], **kwargs: dict[str, Any])

tirx.Call an external kernel.

Parameters:

kernel (Any) – The external kernel to call.
launch_args (List[Union[int, tirx.Expr, List[Union[int, tirx.Expr]]]]) – The launch arguments. A list of integers for grid size, block size, and shared memory size. The actual requirements depend on the kernel.
args (List[tirx.Expr]) – The arguments to pass to the kernel.
kwargs (Dict[str, Any]) – Additional keyword arguments to pass to the kernel or compilation.

tvm.tirx.script.builder.ir.ignore_loop_partition(predicate) → Expr

Annotate a predicate not be considered as target condition of loop partition.

Parameters:: predicate (Expr) – The annotated predicate expression.

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

A memory layout that swizzles a tile layout.

per_element / swizzle_len / atom_len / swizzle_inner carry the swizzle (formerly the standalone SwizzleLayout); tile_layout is the tiled memory map the swizzle is applied to. A bare swizzle is a ComposeLayout over a trivial identity tile.

class tvm.tirx.script.builder.ir.ExecScope(name: str)

An execution scope, identified by one of {cluster, cta, warpgroup, warp, thread}. The ctor FATALs on any other name.

property name: str: Human-readable name of this scope (derived from kind).

class tvm.tirx.script.builder.ir.Iter(extent: Expr, stride: Expr, axis: Axis | str): A memory layout that tiles data across devices.

class tvm.tirx.script.builder.ir.Layout

verify_well_formed() → bool

Verify if the layout is well-formed.

Returns:: True if the layout is well-formed, False otherwise
Return type:: bool

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]

apply_to_shape(coord: list[Expr], input_shape: list[Expr]) → list[Expr]

Compute the per-shard value that each shard would take if coord were interpreted against input_shape.

Tries self.group(input_shape) first. On success, each group owns exactly one input_shape entry, so coord[d] can be split within that group’s shard extents (bounds stay local to one input dim — simpler analyzer simplification, no cross-dim complications).

Falls back to FlattenCoord(coord, input_shape) + SplitCoord on self’s raw shard shape when the group call fails (e.g. when input_shape does not align with the layout’s factor boundaries).

Returns a list of length len(self.shard); each entry is the value that shard would iterate.

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]

direct_sum(left: TileLayout, left_shape: list[Expr], right_shape: list[Expr]) → TileLayout | ComposeLayout

Direct-sum on the tiling domain (unscaled composition): A + B.

This layout is treated as the right addend B grouped by right_shape. The left layout is treated as A grouped by left_shape. The resulting layout is evaluated over the interleaved domain S_A ⊗ S_B, without span scaling (unlike tiling).

is_tile_inner(tile_layout: TileLayout | ComposeLayout, tiled_shape: list[Expr], inner_shape: list[Expr]) → TileLayout | None

Check if a layout is the inner layout of a tiled layout.

Parameters:

tile_layout (Union[TileLayout, ComposeLayout]) – The tiled layout to check
tiled_shape (List[Expr]) – The shape of the tiled layout
inner_shape (List[Expr]) – The shape of the inner layout

Returns:

The outer layout if it is the inner layout of the tiled layout, None otherwise

Return type:

Optional[TileLayout]

is_tile_outer(tile_layout: TileLayout | ComposeLayout, tiled_shape: list[Expr], outer_shape: list[Expr]) → Layout | None

Check if a layout is the outer layout of a tiled layout.

Parameters:

tile_layout (Union[TileLayout, ComposeLayout]) – The tiled layout to check
tiled_shape (List[Expr]) – The shape of the tiled layout
outer_shape (List[Expr]) – The shape of the outer layout

Returns:

The inner layout if it is the outer layout of the tiled layout, None otherwise

Return type:

Optional[Layout]

is_direct_sum_right(sum_layout: TileLayout | ComposeLayout, interleaved_shape: list[Expr], right_shape: list[Expr]) → TileLayout | None

Check if this layout is the right addend B in a direct-sum A + B.

Returns the left addend A if recognized, otherwise None.

is_direct_sum_left(sum_layout: TileLayout | ComposeLayout, interleaved_shape: list[Expr], left_shape: list[Expr]) → Layout | None

Check if this layout is the left addend A in a direct-sum A + B.

Returns the right addend B if recognized, otherwise None.

slice(shape: list[Expr], region: list[tuple[Expr, Expr]]) → Layout | None

Slice the layout with a given shape and region.

Parameters:

shape (List[Expr]) – The shape of the layout
region (List[Tuple[Expr, Expr], tvm.ir.Range]) – The region to slice, each element is (begin, end)

Returns:

The sliced layout, or None if slicing is not possible

Return type:

Optional[Layout]

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_swizzle() → bool: Check if the layout is a bare swizzle (ComposeLayout over a trivial tile).

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

is_trainium() → bool: Check if the layout is trainium layout.

unpack(num: int) → 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
Returns:: The unpacked layout
Return type:: Layout

broadcast(num: int, position: int = -1, axis: 'Axis' | str = 'm') → Layout

Insert a stride-0 broadcast dim of extent num at position.

position follows Python list-insert semantics (negative indices count from the end; -1 appends after the last shard dim). The new dim has stride 0 — accessing along it doesn’t move the byte offset, so the same physical element is “seen” num times.

Useful for layouts where a consumer reads the same SMEM datum multiple times (e.g. sf_reuse over MMA-K steps).

pack(num: int) → 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
Returns:: The packed layout
Return type:: Layout

class tvm.tirx.script.builder.ir.ScopeIdDef(def_ids: list[Var], extents: list[Expr] | None, parent: str, cur: str, preferred_extents: list[Expr] | None = None)

Definition of scope identifiers with their extents and parent-child relationships.

The constructor accepts parent and cur as scope-name strings; they are converted by the FFI into the closed ScopeBinding enum and stored on the scope field (an int value of that enum).

extents=None defers the extent: the value is inferred from sibling ScopeIdDef relationships at LowerTIRx entry via the verifier’s closure. Deferred form requires def_ids to contain exactly one Var.

tvm.tirx.script.builder.ir.TensorMap() → Var

Create a TIRx var that represents a CUDA tensor-map descriptor.

The host/runtime ABI passes a handle to descriptor storage. CUDA kernel codegen lowers this type to const __grid_constant__ CUtensorMap when it appears as a kernel parameter.

class tvm.tirx.script.builder.ir.TileLayout(spec: _LayoutSpec)

A memory layout that tiles data across devices.

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_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]]

group_many(shapes: Sequence[Sequence[Expr]]) → tuple[TileLayout, list[list[int]]]

Group the layout by the minimal common refinement of several shapes.

Repeated cumulative product boundaries are retained, so an extent-one dimension in any input shape becomes a real unit iterator in the refined layout. This operation only splits existing shard iterators; it does not canonicalize or reorder the layout.

Parameters:: shapes (Sequence[Sequence[Expr]]) – Logical shapes with provably equal total products.
Returns:: The commonly refined layout and one separator list per input shape.
Return type:: Tuple[TileLayout, List[List[int]]]

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

classmethod trainium(annotation: str, shape: tuple[Expr], is_psum: bool = False) → TileLayout: Create a TileLayout from an annotation string and a shape.

to_psum() → TileLayout: Convert the layout to a psum 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.script.builder.ir.add_to_parent(stmt: Stmt) → None: Add a statement to the parent frame.

tvm.tirx.script.builder.ir.alloc_cast_frag(src, dtype)

Allocate a register frag holding src value-cast to dtype.

Inherits src’s logical shape and its (lane, register) layout — only the element dtype changes — so Tx.cast(dst, src) is a per-thread element-wise cast with no cross-lane movement. .permute(...) the result to the axis order a downstream consumer (e.g. stmatrix via Tx.copy(dispatch="ldstmatrix")) expects.

Parameters:

src (Buffer) – Source register frag (e.g. from alloc_tcgen05_ldst_frag).
dtype (str) – Destination element dtype.

Returns:

Fresh local frag, src.shape shaped, src.layout, dtype-cast.

Return type:

Buffer

Statement-level buffer allocation (creates an AllocBuffer IR node).

Emits an AllocBuffer statement and returns the Buffer directly:

buf = T.alloc_buffer((128, 128))

For SBlock-level buffer allocation (added to SBlock.alloc_buffers), use T.sblock_alloc_buffer() instead.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[tirx.Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.alloc_scalar(dtype: str = 'float32', scope: str = 'global') → TensorLoad: Allocate a zero-dimensional buffer (scalar).

Statement-level buffer allocation (creates an AllocBuffer IR node).

Emits an AllocBuffer statement and returns the Buffer directly:

buf = T.alloc_buffer((128, 128))

For SBlock-level buffer allocation (added to SBlock.alloc_buffers), use T.sblock_alloc_buffer() instead.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[tirx.Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.cluster_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a kernel→cluster scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.cta_id(extents: list[Expr | int] | None = None, preferred=None, dtype: str = 'int32') → Var | list[Var]

Define a kernel→cta scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.cta_id_in_cluster(extents: list[Expr | int] | None = None, preferred=None, dtype: str = 'int32') → Var | list[Var]

Define a cluster→cta scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) → TensorLoad: Declare a zero-dimensional buffer (scalar) from a pointer.

tvm.tirx.script.builder.ir.device_entry() → None

Mark the device-region entry within the enclosing PrimFunc body.

Flat marker (no with). Subsequent statements in the function body accumulate into an AttrStmt("tirx.device_entry", True, body=...); the wrapping is closed by the PrimFunc frame at function end.

Anything written before this marker is host code (e.g. T.match_buffer); anything after is device code.

Example:

@T.prim_func
def kernel(...):
    A = T.match_buffer(...)
    T.device_entry()           # device region starts here
    bx = T.cta_id([SM_COUNT])  # standalone scope-id def
    ...

tvm.tirx.script.builder.ir.lane_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a warp→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.local_scalar(dtype: str = 'float32') → TensorLoad: Allocate a zero-dimensional buffer in local memory.

tvm.tirx.script.builder.ir.meta_class(cls)

Decorator for utility classes used inside @T.prim_func.

Instances of decorated classes are treated as parser meta values.

tvm.tirx.script.builder.ir.register_script_namespace(name: str, namespace: object) → object: Register a TVMScript namespace on the TIRx builder facade.

class tvm.tirx.script.builder.ir.scalar_wrapper(scalar: TensorLoad): Internal wrapper to allow IRBuilder auto-naming on scalar assignment.

tvm.tirx.script.builder.ir.shared_scalar(dtype: str = 'float32') → TensorLoad: Allocate a zero-dimensional buffer in shared memory.

Statement-level buffer allocation (creates an AllocBuffer IR node).

Emits an AllocBuffer statement and returns the Buffer directly:

buf = T.alloc_buffer((128, 128))

For SBlock-level buffer allocation (added to SBlock.alloc_buffers), use T.sblock_alloc_buffer() instead.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[tirx.Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.thread_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a cta→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.thread_id_in_wg(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a warpgroup→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

Statement-level buffer allocation (creates an AllocBuffer IR node).

Emits an AllocBuffer statement and returns the Buffer directly:

buf = T.alloc_buffer((128, 128))

For SBlock-level buffer allocation (added to SBlock.alloc_buffers), use T.sblock_alloc_buffer() instead.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[tirx.Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

Buffer

tvm.tirx.script.builder.ir.warp_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a cta→warp scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.warp_id_in_wg(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a warpgroup→warp scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.warpgroup_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | list[Var]

Define a cta→warpgroup scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.tirx.script.builder.ir.match_buffer(param, shape=None, dtype='float32', data=None, strides=None, elem_offset=None, scope='global', align=-1, offset_factor=0, layout='default')#: Bind a function parameter or an existing buffer region to a TIRx buffer. shape is required for a function parameter and is inferred from a BufferRegion when omitted. layout accepts a layout object, a registered layout string, or None.

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

tvm.tirx.script.builder.ir.alloc_tcgen05_ldst_frag(instr_shape, tensor_shape, dtype)#

Allocate a local register fragment whose layout matches a tcgen05.{ld,st} atom. instr_shape accepts "32x32b", "16x64b", "16x128b", or "16x256b". For example, a two-CTA Layout-B accumulator and its readback fragment can be allocated as:

C = tmem_pool.alloc_tcgen05_mma_D(
    (64, 128), "float32", M=128, cta_group=2)
frag = Tx.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32")
Tx.tile.wg.copy_async(frag[:, :], C[:, :])

Core TVMScript

Contents

Core TVMScript#

Parser entry points#

Core IR builder#