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.

Namespace for Tensor-level IR

tvm.tirx.Buffer#: alias of Var

class tvm.tirx.BufferAccessKind(value)#: Buffer access modes accepted by buffer_access_ptr().

class tvm.tirx.BufferType(dtype, storage_scope, shape, strides, elem_offset, data_alignment, offset_factor, layout, allocated_addr, span=<MISSING>)#: The structural type carried by an ordinary buffer variable.

tvm.tirx.buffer_data(buffer)#: Project the physical pointer associated with a buffer variable.

tvm.tirx.buffer_data_pointer_type(buffer)#: Return the pointer type produced by buffer_data().

tvm.tirx.is_buffer_var(value) → bool#

Return whether value is an ordinary tirx.Var carrying BufferType.

Use this predicate instead of isinstance(value, Buffer). Buffer is a source-compatibility alias for tvm.ir.Var and therefore does not discriminate buffer variables from scalar or pointer variables.

class tvm.tirx.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.CallEffectKind#: Possible kinds of tirx.Call effects.

class tvm.tirx.Let(var: Var, value: Expr, body: Expr, span: Span | None = None)#

Let node.

Parameters:

var (tirx.Var) – The variable in the binding.
value (Expr) – The value in to be bound.
body (Expr) – The body expression.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tirx.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.Stmt(span)#: Base class of all the statements.

class tvm.tirx.Bind(var: Var, value: Expr, span: Span | None = None)#

Bind node.

Bind a variable to a value in the enclosing scope. Bind has no body field. The bound variable is visible in all subsequent statements within the same enclosing scope (SeqStmt, ForNode.body, etc.).

Parameters:

var (tirx.Var) – The variable in the binding.
value (Expr) – The value to be bound.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.AssertStmt(kind: StringImm, condition: Expr, message_parts: list | None = None, span: Span | None = None)#

AssertStmt node.

Parameters:

kind (StringImm) – The error kind, e.g. “RuntimeError”, “TypeError”, “ValueError”.
condition (Expr) – The assert condition.
message_parts (list[StringImm]) – Error message fragments, concatenated at runtime when assertion fails.
span (Span | None) – The location of the stmt in the source code.

class tvm.tirx.ForKind(value)#: The kind of the for loop.

Note

ForKind can change the control flow semantics of the loop and need to be considered in all TIR passes.

class tvm.tirx.For(loop_var: Var, min: Expr, extent: Expr, kind: ForKind, body: Stmt, thread_binding: IterVar | None = None, annotations: Mapping[str, Object] | None = None, step: Expr | None = None, span: Span | None = None)#

For node.

Parameters:

loop_var (tirx.Var) – The loop variable.
min (Expr) – The beginning value.
extent (Expr) – The length of the loop.
kind (ForKind) – The type of the for.
body (Stmt) – The body statement.
thread_binding (Optional[tirx.IterVar]) – The thread this loop binds to. Only valid if kind is ThreadBinding
step (Expr) – The loop step. Default to none which represent one.
annotations (Optional[Mapping[str, Object]]) – Additional annotation hints.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.While(condition: Expr, body: Stmt, span: Span | None = None)#

While node.

Parameters:

condition (Expr) – The termination condition.
body (Stmt) – The body statement.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.Return(value: Expr, span: Span | None = None)#

Return node.

Parameters:

value (Expr) – The value to return.
span (Optional[Span]) – The location of this statement in the source code.

class tvm.tirx.Break(span: Span | None = None)#: Break node.

class tvm.tirx.Continue(span: Span | None = None)#: Continue node.

tvm.tirx.LetStmt#: alias of Bind

class tvm.tirx.BufferStore(buffer: Var, value: Expr, indices: list[Expr], span: Span | None = None)#

Buffer store node.

Parameters:

buffer (Buffer) – The buffer.
value (Expr) – The value we to be stored.
indices (List[Expr]) – The indices location to be stored.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.AllocBuffer(buffer: Var, *args, **kwargs)#

AllocBuffer node.

Allocates a buffer and declares it in scope.

Parameters:

buffer (Buffer) – The buffer being allocated and declared.
annotations (Optional[dict]) – Additional annotations about the allocation.
span (Optional[Span]) – The location of this AllocBuffer in the source code.

class tvm.tirx.AttrStmt(node: Any, attr_key: str, value: Expr, body: Stmt, span: Span | None = None)#

AttrStmt node.

Parameters:

node (Any) – The node to annotate the attribute
attr_key (str) – Attribute type key.
value (Expr) – The value of the attribute
body (Stmt) – The body statement.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.DeclBuffer(buffer: Var, *args, **kwargs)#

DeclBuffer node.

Parameters:

buffer (Buffer) – The buffer being declared.
data (Expr) – The physical data expression bound to the buffer view.
span (Optional[Span]) – The location of this DeclBuffer in the source code.

class tvm.tirx.SeqStmt(seq: list[Stmt], span: Span | None = None)#

Sequence of statements.

Parameters:

seq (List[Stmt]) – The statements
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.IfThenElse(condition: Expr, then_case: Stmt, else_case: Stmt | None, span: Span | None = None)#

IfThenElse node.

Parameters:

condition (Expr) – The expression
then_case (Stmt) – The statement to execute if condition is true.
else_case (Optional[Stmt]) – The statement to execute if condition is false.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tirx.Evaluate(value: Expr, span: Span | None = None)#

Evaluate node.

Parameters:

value (Expr) – The expression to be evaluated.
span (Optional[Span]) – The location of the stmt in the source code.

tvm.tirx.stmt_seq(*args: Expr | Stmt) → SeqStmt#

Make sequence of statements

Parameters:: *args (Union[Expr, Stmt]) – List of statements to be combined as sequence.
Returns:: stmt – The combined statement.
Return type:: Stmt

tvm.tirx.stmt_list(stmt: Stmt) → list[Stmt]#

Make list of stmt from blocks.

Parameters:: stmt (Stmt) – The input statement.
Returns:: stmt_list – The unpacked list of statements
Return type:: List[Stmt]

class tvm.tirx.BufferRegion(buffer: Var, region: list[Range])#

BufferRegion node.

Parameters:

buffer (Buffer) – The buffer of the buffer region
region (List[Range]) – The region array of the buffer region

class tvm.tirx.BufferRegionType#: The structural type of a BufferRegion expression.

class tvm.tirx.MatchBufferRegion(buffer: Var, source: BufferRegion)#

MatchBufferRegion node.

Parameters:

buffer (Buffer) – The target buffer
source (BufferRegion) – The region of source buffer

class tvm.tirx.SBlock(iter_vars: list[IterVar], reads: list[BufferRegion], writes: list[BufferRegion], name_hint: str, body: Stmt, init: Stmt | None = None, alloc_buffers: list[Var] | None = None, match_buffers: list[MatchBufferRegion] | None = None, annotations: Mapping[str, Object] | None = None, span: Span | None = None)#

SBlock node.

Parameters:

iter_vars (List[IterVar]) – The block Variable.
reads (List[BufferRegion]) – The read buffer regions of the block.
writes (List[BufferRegion]) – The write buffer regions of the block.
name_hint (str) – the name_hint of the block.
body (Stmt) – The body of the block.
init (Optional[Stmt]) – The init block of the reduction block
alloc_buffers (Optional[list[Buffer]]) – The buffer allocations
match_buffers (Optional[List[MatchBufferRegion]]) – The subregion buffer match
annotations (Optional[Mapping[str, Object]]) – Additional annotation hints.
span (Optional[Span]) – The location of this block in the source code.

class tvm.tirx.SBlockRealize(iter_values: list[Expr], predicate: Expr | bool, block: SBlock, span: Span | None = None)#

SBlockRealize node.

Parameters:

iter_values (List[Expr]) – The binding values of the block var.
predicate (Union[Expr, bool]) – The predicate of the block.
block (SBlock) – The block to realize
span (Optional[Span]) – The location of this block_realize in the source code.

class tvm.tirx.ScopeIdDefStmt(def_: ScopeIdDef, span: Span | None = None)#

ScopeIdDefStmt node.

Leaf statement that introduces scope-identifier vars (wg_id = Tx.warpgroup_id([N]), warp_id = Tx.warp_id_in_wg([4]), lane_id = Tx.lane_id([32]), …) at the kernel-body top level. The underlying ScopeIdDef carries the def vars, their extents, and the parent/child scope binding.

Note: the C++ field is named def (a Python keyword). Access it via getattr(stmt, "def") or stmt.__getattribute__("def") — the type-annotation alias here is purely for documentation.

Parameters:

def (ScopeIdDef) – The scope-id definition (def vars, extents, scope binding).
span (Optional[Span]) – The location of this statement in the source code.

class tvm.tirx.LambdaExpr(f_pred: Callable[[...], Expr])#

A reified Python lambda: bound variables and a body over them.

Used by tile primitive ops that take a per-element expression over the destination axes (e.g. tirx.tile.select).

apply(indices: list[Expr]) → Expr#: Substitute the bound variables with the given indices, returning the body.

TilePrimitiveCall node.

Parameters:

op (Op) – The operator.
args (List[Expr]) – The arguments.
workspace (Map[str, Buffer]) – The workspace.
config (Map[str, ObjectRef]) – The scheduler/config dictionary.
dispatch (Optional[str]) – The explicit variant name to dispatch to.
scope (ExecScope) – The cooperation scope of this call. Defaults to thread (an unscoped call).

replace(**changes: Any) → TilePrimitiveCall#

Return a copy of this call with selected fields replaced.

Every field that is not overridden in changes is preserved from self (including scope), so rebuilds never silently drop fields. The returned node is downcast to the registered subclass for op.

Parameters:: **changes (Any) – Field overrides; any of op, args, workspace, config, dispatch, scope.
Returns:: new_call – A new call with the requested fields replaced.
Return type:: TilePrimitiveCall

with_workspace(workspace: dict[str, Var]) → TilePrimitiveCall#: Return a copy with workspace replaced, preserving all other fields.

get_private_buffers(buffer_dict: dict[Any, tuple[Var, Stmt | None]], sctx: DispatchContext) → dict[str, Any]#

Create private (intermediate) buffers needed in this operator.

Parameters:

buffer_dict (Dict[Any, Tuple[Buffer, Optional[Stmt]]]) – A dictionary containing private buffers (and their init stmts) in other operators. Key can be anything to reference the buffer. This is used to reuse private buffers in other operators (like identity tensor etc.). If the buffer is not found in the buffer_dict, it will be created and added to the buffer_dict. If the buffer is found in the buffer_dict but smaller than required, it will be enlarged and updated.
sctx (DispatchContext) – The dispatch context. This is used to get the target and reuse op dispatch implementations.
Returns
-------
private_buffer_refs (Dict[str, Any]) – The references to private buffers created in this operator. Key will be the name to add into workspace. private buffer can be accessed by buffer_dict[private_buffer_refs[name]]

class tvm.tirx.PrimFunc(params, body, ret_type=None, attrs=None, span=None)#

A function declaration expression.

Parameters:

params (List[Union[tvm.tirx.Var, tvm.tirx.Buffer]]) – List of input parameters to the function.
body (tvm.tirx.Stmt) – The body of the function.
ret_type (tvm.ir.Type) – The return type annotation of the function.
attrs (Optional[tvm.Attrs]) – Attributes of the function, can be None
span (Optional[Span]) – The location of this itervar in the source code.

with_body(new_body, span=None)#

Create a new PrimFunc with the same set signatures but a new body.

Parameters:

new_body (Stmt) – The new body.
span (Optional[Span]) – The location of this itervar in the source code.

Returns:

new_func – The created new function.

Return type:

PrimFunc

specialize(param_map: Mapping[Var, Expr | Var])#

Specialize parameters of PrimFunc

Parameters:: param_map (Mapping[tirx.Var, Union[Expr, Buffer]]) – The mapping from function params to the instance

Examples

We can define a Meta TIR function with symbolic shape:

@T.prim_func(s_tir=True)
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
    A = T.match_buffer(a, (m, n), "float32")
    B = T.match_buffer(b, (m, n), "float32")

    for i, j in T.grid(m, n):
        with T.sblock():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

Then we can make it specialized with given shapes or buffers.

a, _, m, n = mem_copy.params
func = mem_copy.specialize({a: tirx.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

The specialized function:

@T.prim_func(s_tir=True)
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    B = T.match_buffer(b, (16, 16), "float32")

    for i, j in T.grid(16, 16):
        with T.sblock():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

Returns:: func – The new function with parameter specialized
Return type:: PrimFunc

class tvm.tirx.TensorIntrin(desc, impl)#

A tensor intrinsic.

Parameters:

desc (PrimFunc) – The function to describe the computation.
impl (PrimFunc) – The function of the implementation for the execution.

static register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False)#

Parameters:

name (str) – The name of the TensorIntrin to register.
desc (PrimFunc) – The function to describe the computation.
impl (PrimFunc) – The function of the implementation for the execution.
override (bool) – Whether override existing intrinsic.

static get(name: str, allow_missing: bool = False) → TensorIntrin | None#

Look up a tensor intrinsic by its name.

Parameters:

name (str) – The name of the TensorIntrin to look up.
allow_missing (bool) – Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
exist. (doesn't)

Returns:

result – The TensorIntrin with the specified name, or None if not found.

Return type:

Optional[TensorIntrin]

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

A mapping from multi-dimensional indices to another set of multi-dimensional indices

Parameters:

initial_indices (List[tirx.Var]) – Variables representing the indices prior to remapping.
final_indices (List[Expr]) – Expressions defining the indices after remapping.
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.

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 buffer 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.arith.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.arith.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 buffer shape

Parameters:

shape (List[Expr]) – The buffer shape to be mapped
analyzer (Optional[tvm.arith.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.arith.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.arith.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.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:

Core TIRx IR

Contents

Core TIRx IR#