tvm.tir¶

Namespace for Tensor-level IR

class tvm.tir.Buffer¶

Symbolic data buffer in TVM.

Buffer provide a way to represent data layout specialization of data structure in TVM.

Do not construct directly, use decl_buffer() instead. See the documentation of decl_buffer() for more details.

See also

decl_buffer: Declare a buffer

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, 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(Buffer.READ | Buffer.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)

vload(begin, dtype=None, predicate=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
predicate (Optional[PrimExpr]) – A vector mask of boolean values indicating which lanes of a vector are to be loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

Returns

load – The corresponding load expression.

Return type

Expr

vstore(begin, value, predicate=None)¶

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.
predicate (Optional[PrimExpr]) – A vector mask of boolean values indicating which lanes of a vector are to be stored. The number lanes of the mask must be equal to the number of lanes in value.

Returns

store – The corresponding store stmt.

Return type

Stmt

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

get_flattened_buffer()¶

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

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

offset_of(indices)¶

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

Parameters: indices (Union[PrimExpr, List[PrimExpr]]) – 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[PrimExpr]

tvm.tir.decl_buffer(shape, dtype=None, name='buffer', data=None, strides=None, elem_offset=None, scope='', data_alignment=- 1, offset_factor=0, buffer_type='', axis_separators=None, span=None)¶

Declare a new symbolic buffer.

Normally buffer is created automatically during lower and build. This is only needed if user want to specify their own buffer layout.

See the note below for detailed discussion on usage of buffer.

Parameters

shape (tuple of Expr) – The shape of the buffer.
dtype (str, optional) – The data type of the buffer.
name (str, optional) – The name of the buffer.
data (tir.Var, optional) – The data pointer in the buffer.
strides (array of Expr) – The stride of the buffer.
elem_offset (Expr, optional) – The beginning offset of the array to data. In terms of number of elements of dtype.
scope (str, optional) – The storage scope of the buffer, if not global. If scope equals empty string, it means it is global memory.
data_alignment (int, optional) – The alignment of data pointer in bytes. If -1 is passed, the alignment will be set to TVM’s internal default.
offset_factor (int, optional) – The factor of elem_offset field, when set, elem_offset is required to be multiple of offset_factor. If 0 is pssed, the alignment will be set to 1. if non-zero is passed, we will created a tir.Var for elem_offset if elem_offset is not None.
buffer_type (str, optional, {"", "auto_broadcast"}) – auto_broadcast buffer allows one to implement broadcast computation without considering whether dimension size equals to one. TVM maps buffer[i][j][k] -> buffer[i][0][k] if dimension j’s shape equals 1.
axis_separators (list of int, optional) – If passed, a list of separators between groups of axes, each of which is flattened to an output axis. For flat memory spaces, should either be None, or an empty list.
span (Optional[Span]) – The location of the decl_buffer creation in the source.

Returns

buffer – The created buffer

Return type

tvm.tir.Buffer

Example

Here’s an example of how broadcast buffer can be used to define a symbolic broadcast operation,

m0, m1, m2 = te.var("m0"), te.var("m1"), te.var("m2")
n0, n1, n2 = te.var("n0"), te.var("n1"), te.var("n2")
o0, o1, o2 = te.var("o0"), te.var("o1"), te.var("o2")
A = te.placeholder((m0, m1, m2), name='A')
B = te.placeholder((n0, n1, n2), name='B')
C = te.compute((o0, o1, o2), lambda i, j, k: A[i, j, k] + B[i, j, k], name='C')
Ab = tvm.tir.decl_buffer(A.shape, A.dtype, name="Ab", buffer_type="auto_broadcast")
Bb = tvm.tir.decl_buffer(B.shape, B.dtype, name="Bb", buffer_type="auto_broadcast")
s = te.create_schedule(C.op)
fadd = tvm.build(s, [A, B, C], target='llvm', name='bcast_add', binds={A:Ab, B:Bb})
dev = tvm.cpu(0)
a = tvm.nd.array(np.random.uniform(size=(2, 4, 3)).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=(2, 1, 3)).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros((2, 4, 3), dtype=C.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

Note

Buffer data structure reflects the DLTensor structure in dlpack. While DLTensor data structure is very general, it is usually helpful to create function that only handles specific case of data structure and make compiled function benefit from it.

If user pass strides and elem_offset is passed as None when constructing the function, then the function will be specialized for the DLTensor that is compact and aligned. If user pass a fully generic symbolic array to the strides, then the resulting function becomes fully generic.

class tvm.tir.DataProducer¶

class tvm.tir.Layout¶

Layout is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. For example, NCHW16c can describe a 5-D tensor of [batch_size, channel, height, width, channel_block]. Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel).

See also

layout: Declare a layout

index_of(axis)¶

Get the index of an axis

Parameters: axis (str) – The axis name, need to be [a-z,A-Z]
Returns: index – The index of the axis, -1 if not found.
Return type: int

factor_of(axis)¶

Get the factor size of the subordinate axis.

Parameters: axis (str) – The axis name, need to be [a-z,A-Z]
Returns: factor – the size of the subordinate-axis of axis (if axis is a primal-axis), or the size of axis itself (if axis is a subordinate-axis). Return -1 if axis is not in the layout.
Return type: int

class tvm.tir.BijectiveLayout¶

Bijective mapping for two layouts (src-layout and dst-layout). It provides shape and index conversion between each other.

Do not construct directly, use bijective_layout instead. See the documentation of bijective_layout for more details.

Parameters

src_layout (str or Layout) – source layout.
dst_layout (str or Layout) – destination layout.

See also

bijective_layout: Declare a layout

forward_index(index)¶

Given the indices of the src-layout, infer the dst index.

Parameters: index (Array of Expr) – The indices in src-layout.
Returns: dst_index – The inferred indices in dst-layout.
Return type: Array of Expr

backward_index(index)¶

Given the indices of the dst-layout, infer the src index.

Parameters: index (Array of Expr) – The indices in dst-layout.
Returns: src_index – The inferred indices in src-layout.
Return type: Array of Expr

forward_shape(shape)¶

Given the shape of the src-layout, infer the dst shape.

Parameters: shape (Array of Expr) – The shape in src-layout.
Returns: dst_shape – The inferred shape in dst-layout.
Return type: Array of Expr

backward_shape(shape)¶

Given the shape of the dst-layout, infer the src shape.

Parameters: shape (Array of Expr) – The shape in dst-layout.
Returns: src_shape – The inferred shape in src-layout.
Return type: Array of Expr

tvm.tir.bijective_layout(src_layout: Union[str, tvm.tir.data_layout.Layout], dst_layout: Union[str, tvm.tir.data_layout.Layout]) → tvm.tir.data_layout.BijectiveLayout¶

Create a bijective layout mapping.

Parameters

src_layout (str or Layout) – source layout.
dst_layout (str or Layout) – destination layout.

Returns

bijective_layout – The created bijective layout

Return type

BijectiveLayout

tvm.tir.layout(layout_str: str, dtype: str = 'int32') → tvm.tir.data_layout.Layout¶

Create a layout node from a string.

Parameters

layout_str (str) – A layout representation is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. For example, NCHW16c can describe a 5-D tensor of [batch_size, channel, height, width, channel_block]. Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel).
dtype (str) – The dtype of generated axes vars in the returned layout. It is required to be integer type.

Returns

layout – The created layout

Return type

Layout

class tvm.tir.Var(name: str, dtype: Union[str, tvm.ir.type.Type], span: Optional[tvm.ir.base.Span] = None)¶

Symbolic variable.

Parameters

name (str) – The name
dtype (Union[str, ir.Type]) – The data type
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tir.SizeVar(name: str, dtype: Union[str, tvm.ir.type.Type], span: Optional[tvm.ir.base.Span] = None)¶

Symbolic variable to represent a tensor index size: which is greater or equal to zero.

Parameters

name (str) – The name
dtype (Union[str, ir.Type]) – The data type
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tir.Reduce(combiner: tvm.tir.expr.CommReducer, src: List[tvm.ir.expr.PrimExpr], rdom: List[tvm.tir.expr.IterVar], condition: tvm.ir.expr.PrimExpr, value_index: int, init: Optional[List[tvm.ir.expr.PrimExpr]] = None, span: Optional[tvm.ir.base.Span] = None)¶

Reduce node.

Parameters

combiner (CommReducer) – The combiner.
src (list of Expr) – The source expression.
rdom (list of IterVar) – The iteration domain
condition (PrimExpr) – 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 ProducerLoad
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tir.FloatImm(dtype: str, value: float, span: Optional[tvm.ir.base.Span] = 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.tir.IntImm(dtype: str, value: int, span: Optional[tvm.ir.base.Span] = 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.tir.StringImm(value: str, span: Optional[tvm.ir.base.Span] = 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.tir.Cast(dtype, value, span: Optional[tvm.ir.base.Span] = None)¶

Cast expression.

Parameters

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

class tvm.tir.Add(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Add node.

Parameters

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

class tvm.tir.Sub(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Sub node.

Parameters

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

class tvm.tir.Mul(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Mul node.

Parameters

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

class tvm.tir.Div(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Div node.

Parameters

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

class tvm.tir.Mod(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Mod node.

Parameters

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

class tvm.tir.FloorDiv(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

FloorDiv node.

Parameters

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

class tvm.tir.FloorMod(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

FloorMod node.

Parameters

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

class tvm.tir.Min(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Min node.

Parameters

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

class tvm.tir.Max(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Max node.

Parameters

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

class tvm.tir.EQ(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

EQ node.

Parameters

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

class tvm.tir.NE(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

NE node.

Parameters

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

class tvm.tir.LT(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

LT node.

Parameters

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

class tvm.tir.LE(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

LE node.

Parameters

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

class tvm.tir.GT(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

GT node.

Parameters

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

class tvm.tir.GE(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

GE node.

Parameters

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

class tvm.tir.And(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

And node.

Parameters

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

class tvm.tir.Or(a: tvm.ir.expr.PrimExpr, b: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Or node.

Parameters

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

class tvm.tir.Not(a: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Not node.

Parameters

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

class tvm.tir.Select(condition: tvm.ir.expr.PrimExpr, true_value: tvm.ir.expr.PrimExpr, false_value: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Select node.

Note

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

Parameters

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

class tvm.tir.BufferLoad(buffer: tvm.tir.buffer.Buffer, indices: List[tvm.ir.expr.PrimExpr], predicate: Optional[tvm.ir.expr.PrimExpr] = None, span: Optional[tvm.ir.base.Span] = None)¶

Buffer load node.

Parameters

buffer (Buffer) – The buffer to be loaded.
indices (List[PrimExpr]) – The buffer indices to load values from.
span (Optional[Span]) – The location of this expression in the source code.
predicate (Optional[PrimExpr]) – A vector mask of boolean values indicating which lanes of a vector are to be loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

class tvm.tir.ProducerLoad(producer: tvm.tir.buffer.DataProducer, indices: List[tvm.ir.expr.PrimExpr], span: Optional[tvm.ir.base.Span] = None)¶

Producer load node.

Parameters

producer (DataProducer) – The buffer to be loaded.
indices (List[PrimExpr]) – The buffer indices.
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tir.Ramp(base: tvm.ir.expr.PrimExpr, stride: tvm.ir.expr.PrimExpr, lanes: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Ramp node.

Parameters

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

class tvm.tir.Broadcast(value: tvm.ir.expr.PrimExpr, lanes: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Broadcast node.

Parameters

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

class tvm.tir.Shuffle(vectors: List[tvm.ir.expr.PrimExpr], indices: List[tvm.ir.expr.PrimExpr], span: Optional[tvm.ir.base.Span] = None)¶

Shuffle node.

Parameters

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

class tvm.tir.Call(dtype: str, op: Union[tvm.ir.op.Op, str], args: List[tvm.ir.expr.PrimExpr], span: Optional[tvm.ir.base.Span] = None)¶

tir.Call node.

Parameters

dtype (str) – The return data type
op (Union[Op, str]) – The function to be called, or the name to the global tvm.Op
args (list of Expr) – The input arguments to the call
span (Optional[Span]) – The location of this expression in the source code.

class tvm.tir.CallEffectKind¶: Possible kinds of tir.Call effects.

class tvm.tir.Let(var: tvm.tir.expr.Var, value: tvm.ir.expr.PrimExpr, body: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Let node.

Parameters

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

class tvm.tir.IterVar(dom: tvm.ir.expr.Range, var: Union[tvm.tir.expr.Var, str], iter_type: int, thread_tag: str = '', span: Optional[tvm.ir.base.Span] = None)¶

Represent iteration variable.

IterVar represents axis iterations in the computation.

Parameters

dom (Range) – The domain of the iteration.
var (Union[tir.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.

class tvm.tir.CommReducer(lhs: List[tvm.tir.expr.Var], rhs: List[tvm.tir.expr.Var], result: List[tvm.ir.expr.PrimExpr], identity_element: List[tvm.ir.expr.PrimExpr], span: Optional[tvm.ir.base.Span] = None)¶

Commutative reduce operator

Parameters

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

class tvm.tir.Any(span: Optional[tvm.ir.base.Span] = None)¶

Any node.

spanOptional[Span]: The location of this expression in the source code.

class tvm.tir.Stmt¶: Base class of all the statements.

class tvm.tir.LetStmt(var: tvm.tir.expr.Var, value: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

LetStmt node.

Parameters

var (tir.Var) – The variable in the binding.
value (PrimExpr) – The value in to be bound.
body (Stmt) – The body statement.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.AssertStmt(condition: tvm.ir.expr.PrimExpr, message: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

AssertStmt node.

Parameters

condition (PrimExpr) – The assert condition.
message (PrimExpr) – The error message.
body (tvm.tir.Stmt) – The body statement.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.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.tir.For(loop_var: tvm.tir.expr.Var, min: tvm.ir.expr.PrimExpr, extent: tvm.ir.expr.PrimExpr, kind: tvm.tir.stmt.ForKind, body: tvm.tir.stmt.Stmt, thread_binding: Optional[tvm.tir.expr.IterVar] = None, annotations: Optional[Mapping[str, tvm.runtime.object.Object]] = None, span: Optional[tvm.ir.base.Span] = None)¶

For node.

Parameters

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

class tvm.tir.While(condition: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

While node.

Parameters

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

class tvm.tir.BufferStore(buffer: tvm.tir.buffer.Buffer, value: tvm.ir.expr.PrimExpr, indices: List[tvm.ir.expr.PrimExpr], predicate: Optional[tvm.ir.expr.PrimExpr] = None, span: Optional[tvm.ir.base.Span] = None)¶

Buffer store node.

Parameters

buffer (Buffer) – The buffer.
value (PrimExpr) – The value we to be stored.
indices (List[PrimExpr]) – The indices location to be stored.
predicate (Optional[PrimExpr]) – A vector mask of boolean values indicating which lanes of a vector are to be stored. The number lanes of the mask must be equal to the number of lanes in value.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.BufferRealize(buffer: tvm.tir.buffer.Buffer, bounds: List[tvm.ir.expr.Range], condition: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

Buffer realize node.

Parameters

buffer (Buffer) – The buffer.
bounds (List[Range]) – The value we to be stored.
condition (PrimExpr) – The realize condition.
body (Stmt) – The body of the statement.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.ProducerStore(producer: tvm.tir.buffer.DataProducer, value: tvm.ir.expr.PrimExpr, indices: List[tvm.ir.expr.PrimExpr], span: Optional[tvm.ir.base.Span] = None)¶

ProducerStore node.

Parameters

producer (DataProducer) – The data producer.
value (PrimExpr) – The value to be stored.
indices (list of Expr) – The index arguments of the store.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.Allocate(buffer_var: tvm.tir.expr.Var, dtype: str, extents: List[tvm.ir.expr.PrimExpr], condition: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, annotations: Optional[Mapping[str, tvm.runtime.object.Object]] = None, span: Optional[tvm.ir.base.Span] = None)¶

Allocate node.

Parameters

buffer_var (tir.Var) – The buffer variable.
dtype (str) – The data type of the buffer.
extents (list of Expr) – The extents of the allocate
condition (PrimExpr) – The condition.
body (Stmt) – The body statement.
annotations (Optional[Mapping[str, Object]]) – Additional annotation hints
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.AllocateConst(buffer_var: tvm.tir.expr.Var, dtype: str, extents: List[tvm.ir.expr.PrimExpr], data_or_idx: Union[tvm.runtime.ndarray.NDArray, int], body: tvm.tir.stmt.Stmt, annotations: Optional[Mapping[str, tvm.runtime.object.Object]] = None, span: Optional[tvm.ir.base.Span] = None)¶

Allocate constant node.

Parameters

buffer_var (tir.Var) – The buffer variable.
dtype (str) – The data type of the buffer.
extents (list of Expr) – The extents of the allocate
data_or_idx (Union[NDArray, int]) – If an NDArray, this is the const data associated with the constant. If an integer, this is the index into the “constants” attribute of the IRModule that contains the AllocateConst.
body (Stmt) – The body statement.
annotations (Optional[Mapping[str, Object]]) – Additional annotations about the allocation.
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.AttrStmt(node: tvm.runtime.object.Object, attr_key: str, value: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

AttrStmt node.

Parameters

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

class tvm.tir.DeclBuffer(buffer: tvm.tir.buffer.Buffer, body: tvm.tir.stmt.Stmt, span: Optional[tvm.ir.base.Span] = None)¶

DeclBuffer node.

Parameters

buffer (Buffer) – The buffer being declared.
body (Stmt) – The body statement to be executed.
span (Optional[Span]) – The location of this DeclBuffer in the source code.

class tvm.tir.ProducerRealize(producer: tvm.tir.buffer.DataProducer, bounds: List[tvm.ir.expr.Range], condition: tvm.ir.expr.PrimExpr, body: tvm.tir.stmt.Stmt, storage_scope: str = '', span: Optional[tvm.ir.base.Span] = None)¶

ProducerRealize node.

Parameters

producer (DataProducer) – The data producer.
bounds (List[Range]) – The bound of realize
condition (PrimExpr) – The realize condition.
body (Stmt) – The realize body
storage_scope (str) – The storage scope associated with this realization
span (Optional[Span]) – The location of the stmt in the source code.

class tvm.tir.SeqStmt(seq: List[tvm.tir.stmt.Stmt], span: Optional[tvm.ir.base.Span] = None)¶

Sequence of statements.

Parameters

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

class tvm.tir.IfThenElse(condition: tvm.ir.expr.PrimExpr, then_case: tvm.tir.stmt.Stmt, else_case: Optional[tvm.tir.stmt.Stmt], span: Optional[tvm.ir.base.Span] = None)¶

IfThenElse node.

Parameters

condition (PrimExpr) – 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.tir.Evaluate(value: tvm.ir.expr.PrimExpr, span: Optional[tvm.ir.base.Span] = None)¶

Evaluate node.

Parameters

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

class tvm.tir.Prefetch(buffer: tvm.tir.buffer.Buffer, bounds: List[tvm.ir.expr.Range], span: Optional[tvm.ir.base.Span] = None)¶

Prefetch node.

Parameters

buffer (Buffer) – The buffer to be prefetched.
bounds (List[Range]) – The bounds to be prefetched.
span (Optional[Span]) – The location of the stmt in the source code.

tvm.tir.stmt_seq(*args: Union[tvm.ir.expr.PrimExpr, tvm.tir.stmt.Stmt]) → tvm.tir.stmt.SeqStmt¶

Make sequence of statements

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

tvm.tir.stmt_list(stmt: tvm.tir.stmt.Stmt) → List[tvm.tir.stmt.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.tir.BufferRegion(buffer: tvm.tir.buffer.Buffer, region: List[tvm.ir.expr.Range])¶

BufferRegion node.

Parameters

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

class tvm.tir.MatchBufferRegion(buffer: tvm.tir.buffer.Buffer, source: tvm.tir.stmt.BufferRegion)¶

MatchBufferRegion node.

Parameters

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

class tvm.tir.Block(iter_vars: List[tvm.tir.expr.IterVar], reads: List[tvm.tir.stmt.BufferRegion], writes: List[tvm.tir.stmt.BufferRegion], name_hint: str, body: tvm.tir.stmt.Stmt, init: Optional[tvm.tir.stmt.Stmt] = None, alloc_buffers: Optional[List[tvm.tir.buffer.Buffer]] = None, match_buffers: Optional[List[tvm.tir.stmt.MatchBufferRegion]] = None, annotations: Optional[Mapping[str, tvm.runtime.object.Object]] = None, span: Optional[tvm.ir.base.Span] = None)¶

Block 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.tir.BlockRealize(iter_values: List[tvm.ir.expr.PrimExpr], predicate: Union[tvm.ir.expr.PrimExpr, bool], block: tvm.tir.stmt.Block, span: Optional[tvm.ir.base.Span] = None)¶

BlockRealize node.

Parameters

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

class tvm.tir.PrimFunc(params, body, ret_type=None, buffer_map=None, attrs=None, span=None)¶

A function declaration expression.

Parameters

params (List[Union[tvm.tir.Var, tvm.tir.Buffer]]) – List of input parameters to the function.
body (tvm.tir.Stmt) – The body of the function.
ret_type (tvm.ir.Type) – The return type annotation of the function.
buffer_map (Map[tvm.tir.Var, tvm.tir.Buffer]) – The buffer binding map.
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[tvm.tir.expr.Var, Union[tvm.ir.expr.PrimExpr, tvm.tir.buffer.Buffer]])¶

Specialize parameters of PrimFunc

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

Examples

We can define a Meta TIR function with symbolic shape:

@T.prim_func
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.block():
            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: tir.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

The specialized function:

@T.prim_func
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.block():
            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.tir.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: tvm.tir.function.PrimFunc, impl: tvm.tir.function.PrimFunc, override: bool = False)¶

Register a tensor intrinsic with its name.

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) → Optional[tvm.tir.function.TensorIntrin]¶

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.tir.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[tir.Var]) – Variables representing the indices prior to remapping.
final_indices (List[PrimExpr]) – 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: Optional[int] = None, inverse_index_map: Optional[Union[Callable, tvm.tir.function.IndexMap]] = 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 tir.Var parameters and return a either a tir.PrimExpr, or a list of tir.PrimExpr. Returning a tir.PrimExpr is equivalent to returning a list of length 1 containing that tir.PrimExpr.
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.

Returns

index_map – Returns an IndexMap representing the mapping_function.

Return type

IndexMap

static from_func_with_separators(mapping_function: Callable, ndim: Optional[int] = None, inverse_index_map: Optional[Union[Callable, tvm.tir.function.IndexMap]] = 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 tir.Var parameters and return either a tir.PrimExpr or a list. Each element of the returned list should be either a tir.PrimExpr or the object IndexMap.AXIS_SEPARATOR. Returning a tir.PrimExpr is equivalent to returning a list of length 1 containing that tir.PrimExpr.
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

ret – Returns a tuple whose first element is an IndexMap representing the mapping_function, and whose second index is a list of indices at which IndexMap.AXIS_SEPARATOR occurred.

Return type

Tuple[IndexMap, List[int]]

is_equivalent_to(other_map: tvm.tir.function.IndexMap) → bool¶

Return if the index maps are equivalent.

Parameters: other_map (IndexMap) – The IndexMap to which the comparison should be made.
Returns: is_equivalent – True if the two mappings represent the same transformation, otherwise False
Return type: bool

map_indices(indices: List[tvm.ir.expr.PrimExpr]) → List[tvm.ir.expr.PrimExpr]¶

Apply the index map to a set of indices

Parameters: indices (List[PrimExpr]) – The indices to be mapped
Returns: result – The mapped indices
Return type: List[PrimExpr]

map_shape(shape: List[tvm.ir.expr.PrimExpr]) → List[tvm.ir.expr.PrimExpr]¶

Apply the index map to a buffer shape

Parameters: shape (List[PrimExpr]) – The buffer shape to be mapped
Returns: result – The mapped shape
Return type: List[PrimExpr]

map_ndarray(arr_src: tvm.runtime.ndarray.NDArray) → tvm.runtime.ndarray.NDArray¶

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

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

inverse(shape: List[Union[tvm.ir.expr.Range, tvm.ir.expr.PrimExpr]]) → tvm.tir.function.IndexMap¶

Return the inverse of the map

Throws an error if the function is not bijective.

Parameters: shape (List[Union[Range,PrimExpr]]) – The region over which the inverse should be determined. Used for validating that the mapping is bijective over this range.
Returns: inverse – The inverse
Return type: IndexMap

non_surjective_inverse(shape: List[Union[tvm.ir.expr.Range, tvm.ir.expr.PrimExpr]]) → Tuple[tvm.tir.function.IndexMap, tvm.ir.expr.PrimExpr]¶

Return the inverse of the map

Can be applied to transformations that introduce padding.

Parameters: shape (List[Union[Range,PrimExpr]]) – The region over which the inverse should be determined. Used for determining the 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, PrimExpr]

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.tir.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 recieve an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is NDArray.

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