Core TIRx IR

Core TIRx IR#

This page lists the core IR nodes and scalar operations in the top-level namespace. Layouts, execution scopes, visitors, compilation helpers, and tile-dispatch extensions are documented on their focused pages and excluded here so the same objects are not expanded twice.

For C++ construction, include tvm/tirx/expr.h for TensorVar, tensor loads, and tensor-region constructors. Include tvm/tirx/type.h for TensorType and TensorMapType. Tensor regions use the shared TensorRegion expression from tvm/ir/expr.h and TensorRegionType from tvm/ir/type.h.

Tensor-level IR nodes, layouts, operators, analysis, and compilation.

class tvm_tirx.tirx.TensorType(shape, dtype, scope=MemoryScope('global'), align=0, layout=None, *, data=None, loc: Location = UnknownLoc(), data_alignment=None)#

The shape, element type, and storage contract of a TIRx tensor expression.

Tensor variables expose the type’s properties and operations, such as tensor.shape and tensor.view(...). Strides and element offsets are described by layout.

Parameters:

shape (sequence[Expr] or Expr) – Logical dimension extents; a single extent denotes a one-dimensional tensor.
dtype (PrimType or str) – Primitive element type.
scope (MemoryScope or str, optional) – Storage scope, defaulting to global memory.
align (int, optional) – Data-pointer alignment in bytes; zero selects the runtime default.
layout (Layout or str, optional) – Explicit storage mapping. "default" creates a contiguous row-major layout; None leaves the layout unspecified.
data (Expr, optional) – Keyword-only pointer expression whose PtrType supplies the storage scope. The pointer itself is not stored in TensorType.
loc (Location, optional) – The source location of the type.
data_alignment (int, optional) – Keyword-only spelling of align; both must agree when explicitly set.

has_default_layout(analyzer=None) → bool#

Whether the layout matches implicit row-major storage with zero origin.

Equivalent split tiles are accepted when their extent covers this shape. An optional analyzer supplies existing bindings and constraints; otherwise symbolic analysis is initialized only when structural equality is insufficient.

effective_layout()#

Return the explicit layout or infer a compact global layout.

Inference uses the full tensor shape before a region is sliced and leaves this type unchanged. Other storage scopes and noncompact global tensors require an explicit layout.

is_scalar(expr, alloc_or_decl=True)#

Check if the tensor 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 tensor is a scalar, False otherwise.

ptr_to(expr, indices)#

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

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

view(expr, *args, **kwargs) → Var#

Creates a new view of the tensor. (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_ffi.dtype]).

Returns:: view – The corresponding view tensor.
Return type:: tvm_tirx.tirx.Var

permute(expr, *dims) → Var#

Permute the dimensions of the tensor.

Parameters:: dims (tuple of int) – The permutation of dimensions.
Returns:: permuted – The tensor with permuted dimensions.
Return type:: tvm_tirx.tirx.Var

rearrange(expr, 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) tensor collapses to a flat layout, a swizzled tensor keeps its swizzle, and a tmem tensor retains its backing address. It therefore does NOT flatten a swizzle atom — the same pattern on a swizzled SMEM tensor vs an unswizzled TMEM tensor 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=...).

tvm_tirx.tirx.tensor_data_ptr(tensor, *, ty=None, loc: Location = UnknownLoc())#

Project a tensor variable’s physical pointer.

The result type is inferred from its element type and storage scope. ty may supply an explicit result type; loc records the source location.

tvm_tirx.tirx.is_tensor_var(value) → bool#

Return whether value is an ordinary tvm_tirx.tirx.Var carrying TensorType.

Unlike isinstance(value, tvm_tirx.tirx.Var), this predicate distinguishes tensor variables from scalar or pointer variables.

class tvm_tirx.tirx.TensorMapType(loc: Location = UnknownLoc())#

The type of a tensor-map descriptor used by backend memory operations.

A pointer to this type represents descriptor storage, rather than a tensor’s element data. The descriptor’s shape and layout are defined by the operations that create it, not by fields on this type.

Parameters:: loc (tvm_tirx.ir.Location, optional) – The source location of the type.

tvm_tirx.tirx.convert(expr) → Expr#: Convert a scalar or sequence to primitive expressions.

class tvm_tirx.tirx.Cast(ty: str | PrimType, value, loc: Location = UnknownLoc())#

Convert a primitive value to another dtype while preserving its lane count.

Parameters:

ty (str | PrimType) – The target primitive type.
value (Expr) – The value to cast.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Add(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Add two primitive values, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Sub(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Subtract b from a, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Mul(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Multiply two primitive values, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Div(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Divide a by b; integer division truncates toward zero.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Mod(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the remainder of integer division truncated toward zero.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.FloorDiv(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Divide a by b, rounding the quotient down toward negative infinity.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.FloorMod(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute a - FloorDiv(a, b) * b, the remainder of floor division.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.LShift(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Shift the bits of a left by b positions.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.RShift(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Shift the bits of a right by b positions.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.BitwiseAnd(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the bitwise intersection of a and b.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.BitwiseOr(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the bitwise union of a and b.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.BitwiseXor(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the bitwise exclusive-or of a and b.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.BitwiseNot(a: Expr, loc: Location = UnknownLoc())#

Invert every bit of the input value.

Parameters:

a (Expr) – The input value.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Min(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Choose the smaller of two primitive values, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Max(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Choose the larger of two primitive values, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.EQ(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a equals b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.NE(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a differs from b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.LT(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a is less than b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.LE(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a is at most b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.GT(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a is greater than b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.GE(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compare whether a is at least b, elementwise for vectors.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.And(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the conjunction of two boolean values.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Or(a: Expr, b: Expr, loc: Location = UnknownLoc())#

Compute the disjunction of two boolean values.

Parameters:

a (Expr) – The left hand operand.
b (Expr) – The right hand operand.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Not(a: Expr, loc: Location = UnknownLoc())#

Negate a boolean value, elementwise for vectors.

Parameters:

a (Expr) – The input value
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Select(condition: Expr, true_value: Expr, false_value: Expr, loc: Location = UnknownLoc())#

Select between two values according to a boolean condition.

Note

Select may compute both true_value and false_value. Use tvm_tirx.ir.prim.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.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Let(var: Var, value: Expr, body: Expr, loc: Location = UnknownLoc())#

Bind var to value within an expression body and return the body result.

Parameters:

var (tvm_tirx.tirx.Var) – The variable in the binding.
value (Expr) – The value in to be bound.
body (Expr) – The body expression.
loc (Location, optional) – The location of this expression in the source code.

class tvm_tirx.tirx.Function(params, body: Stmt | Sequence[Stmt] | None, ret_type=None, attrs=None, loc: Location = UnknownLoc())#

A TIRx function definition or declaration.

The function’s ty is a FuncType derived from parameter types and ret_type. A present body is normalized to a SeqStmt.

Parameters:

params (list[tvm_tirx.ir.Var]) – Input parameters, in calling order.
body (Stmt or Sequence[Stmt] or None) – The function body, or None for a declaration without an implementation.
ret_type (tvm_tirx.ir.Type, optional) – The return type; defaults to MissingType until supplied or inferred.
attrs (tvm_tirx.ir.DictAttrs, optional) – Function attributes, empty by default.
loc (Location, optional) – The source location of the function.

with_body(new_body: Stmt | Sequence[Stmt] | None, loc: Location = UnknownLoc())#

Create a function with the same parameters, return type, and attributes but a new body.

Parameters:

new_body (Stmt | Sequence[Stmt] | None) – The new body, or None for a declaration.
loc (Location, optional) – The source location of the function.

Returns:

new_func – The created new function.

Return type:

Function

tvm_tirx.tirx.renew_def(func: Function) → Function#

Renew variable definitions and remap their uses in a function.

Definition regions follow the registered structural hooks. Free variables and function attributes are preserved, and the function signature is rebuilt.

Parameters:: func (Function) – The function to renew.
Returns:: The function with fresh variable definitions.
Return type:: Function

class tvm_tirx.tirx.IndexMap(initial_indices, final_indices, inverse_index_map)#

Map input index variables to expressions for output indices.

Use from_func() to derive a map from a Python callable, or supply the input variables and output expressions directly.

Parameters:

initial_indices (list[tvm_tirx.tirx.Var]) – Bound variables representing the input coordinates, in dimension order.
final_indices (list[Expr]) – Output coordinate expressions in terms of the input variables.
inverse_index_map (IndexMap or Callable or None) – A supplied inverse, or None to compute the inverse when requested. The caller is responsible for correctness of a supplied inverse.

static from_func(mapping_function: Callable, ndim: int | None = None, inverse_index_map: Callable | IndexMap | None = None, *, index_dtype: str = 'int64')#

Create an index map from a function

Parameters:

mapping_function (Callable) – The function to map from source indices to target indices. The function should accept tirx.Var parameters and return a either a tirx.Expr, or a list of tirx.Expr. Returning a tirx.Expr is equivalent to returning a list of length 1 containing that tirx.Expr.
ndim (Optional[int]) – The dimensionality of the tensor to which this transformation should be applied. If mapping_function uses variadic argument *args, ndim must be specified. If mapping_function does not use variadic arguments, ndim is optional.
inverse_index_map (Union[Callable, Optional[IndexMap]]) – The optional pre-defined inverse index map. When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. Otherwise, the inverse index map will be computed on the fly. It is the user’s responsibility to ensure the correctness of the pre-defined inverse index map.
index_dtype (str) – The default index dtype to use for input iters in the mapping function.

Returns:

index_map – Returns an IndexMap representing the mapping_function.

Return type:

IndexMap

is_equivalent_to(other_map: IndexMap, analyzer=None) → bool#

Return if the index maps are equivalent.

Parameters:

other_map (IndexMap) – The IndexMap to which the comparison should be made.
analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while comparing the mapped indices. When provided, its accumulated bindings and constraints are reused so that maps that are only equivalent under those bindings can be proven equal.

Returns:

is_equivalent – True if the two mappings represent the same transformation, otherwise False

Return type:

bool

map_indices(indices: list[Expr], analyzer=None) → list[Expr]#

Apply the index map to a set of indices

Parameters:

indices (List[Expr]) – The indices to be mapped
analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while simplifying mapped indices.

Returns:

result – The mapped indices

Return type:

List[Expr]

map_shape(shape: list[Expr], analyzer=None) → list[Expr]#

Apply the index map to a tensor shape

Parameters:

shape (List[Expr]) – The tensor shape to be mapped
analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while simplifying mapped shape expressions.

Returns:

result – The mapped shape

Return type:

List[Expr]

map_tensor(arr_src: Tensor) → Tensor#

Apply thie index map to transform the layout of the input Tensor

Parameters:: arr_src (runtime.Tensor) – The Tensor to be transformed
Returns:: arr_dst – The transformed Tensor
Return type:: runtime.Tensor

inverse(shape: list[Range | Expr], analyzer=None) → IndexMap#

Return the inverse of the map

Throws an error if the function is not bijective.

Parameters:

shape (List[Union[Range,Expr]]) – The region over which the inverse should be determined. Used for validating that the mapping is bijective over this range.
analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while deriving and validating the inverse.

Returns:

inverse – The inverse

Return type:

IndexMap

non_surjective_inverse(shape: list[Range | Expr], analyzer=None) → tuple[IndexMap, Expr]#

Return the inverse of the map

Can be applied to transformations that introduce padding.

Parameters:

shape (List[Union[Range,Expr]]) – The region over which the inverse should be determined. Used for determining the predicate.
analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while deriving the inverse and padding predicate.

Returns:

result – The inverse, and a predicate for which the inverse maps to a valid index in the input range.

Return type:

Tuple[IndexMap, Expr]

Examples

index_map = IndexMap.from_func(lambda i: [i//4, i%4])
inverse_map, predicate = index_map.non_surjective_inverse([14])
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"

tvm_tirx.tirx.call_packed_lowered(*args, loc: Location = UnknownLoc(), ty=None)#

Lowered version of call packed. The argument to a packed function can be an Expr or a tensor variable. The argument is the corresponding POD type when Expr is presented. When the argument is a tvm_tirx.tirx.Var carrying TensorType, 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 Var.) – Positional arguments.
loc (Location, optional) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

Core TIRx IR

Contents

Core TIRx IR#