tvm.te

tvm.te#

Namespace for Tensor Expression Language

Functions:

`any`(*args[, span])	Create a new experssion of the union of all conditions in the arguments
`all`(*args[, span])	Create a new expression of the intersection of all conditions in the
`min_value`(dtype[, span])	minimum value of dtype
`max_value`(dtype[, span])	maximum value of dtype
`exp`(x, *[, ty, span])	Take exponential of input x.
`erf`(x, *[, ty, span])	Take gauss error function of the input x.
`tanh`(x, *[, ty, span])	Take hyperbolic tanh of input x.
`sigmoid`(x, *[, ty, span])	Quick function to get sigmoid
`log`(x, *[, ty, span])	Take log of input x.
`tan`(x, *[, ty, span])	Take tan of input x.
`cos`(x, *[, ty, span])	Take cos of input x.
`sin`(x, *[, ty, span])	Take sin of input x.
`sqrt`(x, *[, ty, span])	Take square root of input x.
`rsqrt`(x, *[, ty, span])	Take reciprocal of square root of input x.
`floor`(x[, span])	Take floor of float input x.
`ceil`(x[, span])	Take ceil of float input x.
`sinh`(x, *[, ty, span])	Take sinh of input x.
`cosh`(x, *[, ty, span])	Take cosh of input x.
`log2`(x, *[, ty, span])	Take log2 of input x.
`log10`(x, *[, ty, span])	Take log10 of input x.
`asin`(x, *[, ty, span])	Take asin of input x.
`asinh`(x, *[, ty, span])	Take asinh of input x.
`acos`(x, *[, ty, span])	Take acos of input x.
`acosh`(x, *[, ty, span])	Take acos of input x.
`atan`(x, *[, ty, span])	Take atan of input x.
`atanh`(x, *[, ty, span])	Take atanh of input x.
`trunc`(x[, span])	Get truncated value of the input.
`abs`(x[, span])	Get absolute value of the input element-wise.
`round`(x[, span])	Round elements of the array to the nearest integer.
`nearbyint`(x[, span])	Round elements of the array to the nearest integer.
`power`(x, y[, span])	x power y
`popcount`(x, *[, ty, span])	Count the number of set bits in input x.
`fmod`(x, y, *[, ty, span])	Return the remainder of x divided by y with the same sign as x.
`if_then_else`(cond, t, f[, span])	Conditional selection expression.
`isnan`(x[, span])	Check if input value is Nan.
`isfinite`(x[, span])	Check if input value is finite.
`isinf`(x[, span])	Check if input value is infinite.
`div`(a, b[, span])	Compute a / b as in C/C++ semantics.
`indexdiv`(a, b[, span])	Compute floor(a / b) where a and b are non-negative.
`indexmod`(a, b[, span])	Compute the remainder of indexdiv.
`truncdiv`(a, b[, span])	Compute the truncdiv of two expressions.
`truncmod`(a, b[, span])	Compute the truncmod of two expressions.
`floordiv`(a, b[, span])	Compute the floordiv of two expressions.
`floormod`(a, b[, span])	Compute the floormod of two expressions.
`logaddexp`(a, b[, span])	Compute the logaddexp of two expressions.
`comm_reducer`(fcombine, fidentity[, name])	Create a commutative reducer for reduction.
`min`(expr, axis[, where, init])	Create a min expression over axis.
`max`(expr, axis[, where, init])	Create a max expression over axis.
`sum`(expr, axis[, where, init])	Create a sum expression over axis.
`tag_scope`(tag)	The operator tag scope.
`placeholder`(shape[, dtype, name])	Construct an empty tensor object.
`compute`(shape, fcompute[, name, tag, attrs, ...])	Construct a new tensor by computing over the shape domain.
`scan`(init, update, state_placeholder[, ...])	Construct new tensors by scanning over axis.
`extern`(shape, inputs, fcompute[, name, ...])	Compute several tensors via an extern function.
`var`([name, dtype, span])	Create a new variable with specified name and dtype
`const`(value[, dtype, span])	Create a new constant with specified value and dtype
`thread_axis`([dom, tag, name, span])	Create a new IterVar to represent thread index.
`reduce_axis`(dom[, name, thread_tag, span])	Create a new IterVar for reduction.
`create_function`(ops[, index_dtype_override])	Create a TensorIR Function from tensor expression
`extern_function`(input_tensors, function, ...)	Compute tensors via a schedulable TIR Function

Classes:

`CommReducer`(lhs, rhs, result, identity_element)	Commutative reduce operator
`Reduce`(combiner, src, rdom, condition, ...)	Reduce node.
`TensorSlice`(tensor, indices)	Auxiliary data structure for enable slicing syntax from tensor.
`Tensor`(shape, dtype, op, value_index[, span, ty])	Tensor object, to construct, see function.Tensor
`PlaceholderOp`(name, tag, attrs, shape, dtype)	Placeholder operation.
`ComputeOp`(name, tag, attrs, axis, ...)	Scalar operation.
`ScanOp`(name, tag, attrs, scan_axis, init, ...)	Scan operation.
`ExternOp`(name, tag, attrs, inputs, ...)	External operation.

tvm.te.any(*args, span=None)#

Create a new experssion of the union of all conditions in the arguments

Parameters:

args (list) – List of symbolic boolean expressions
span (Optional[Span]) – The location of this operator in the source code.

Returns:

expr – Expression

Return type:

Expr

tvm.te.all(*args, span=None)#

Create a new expression of the intersection of all conditions in the: arguments

Parameters:

args (list) – List of symbolic boolean expressions
span (Optional[Span]) – The location of this operator in the source code.

Returns:

expr – Expression

Return type:

Expr

tvm.te.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.ir.Expr

tvm.te.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.ir.Expr

tvm.te.exp(x, *, ty=None, span=None)#

Take exponential of input x.

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

tvm.te.erf(x, *, ty=None, span=None)#

Take gauss error function of the input x.

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

tvm.te.tanh(x, *, ty=None, span=None)#

Take hyperbolic tanh of input x.

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

tvm.te.sigmoid(x, *, ty=None, span=None)#

Quick function to get sigmoid

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

tvm.te.log(x, *, ty=None, span=None)#

Take log of input x.

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

tvm.te.tan(x, *, ty=None, span=None)#

Take tan of input x.

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

tvm.te.cos(x, *, ty=None, span=None)#

Take cos of input x.

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

tvm.te.sin(x, *, ty=None, span=None)#

Take sin of input x.

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

tvm.te.sqrt(x, *, ty=None, span=None)#

Take square root of input x.

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

tvm.te.rsqrt(x, *, ty=None, span=None)#

Take reciprocal of square root of input x.

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

tvm.te.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.te.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.te.sinh(x, *, ty=None, span=None)#

Take sinh of input x.

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

tvm.te.cosh(x, *, ty=None, span=None)#

Take cosh of input x.

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

tvm.te.log2(x, *, ty=None, span=None)#

Take log2 of input x.

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

tvm.te.log10(x, *, ty=None, span=None)#

Take log10 of input x.

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

tvm.te.asin(x, *, ty=None, span=None)#

Take asin of input x.

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

tvm.te.asinh(x, *, ty=None, span=None)#

Take asinh of input x.

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

tvm.te.acos(x, *, ty=None, span=None)#

Take acos of input x.

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

tvm.te.acosh(x, *, ty=None, span=None)#

Take acos of input x.

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

tvm.te.atan(x, *, ty=None, span=None)#

Take atan of input x.

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

tvm.te.atanh(x, *, ty=None, span=None)#

Take atanh of input x.

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

tvm.te.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.te.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.te.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.te.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.te.power(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.te.popcount(x, *, ty=None, span=None)#

Count the number of set bits in input x.

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

tvm.te.fmod(x, y, *, ty=None, span=None)#

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.te.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.te.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.te.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.te.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.te.div(a, b, span=None)#

Compute a / b as in C/C++ semantics.

Parameters:

a (Expr) – The left hand operand, known to be non-negative.
b (Expr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

Note

When operands are integers, returns truncdiv(a, b, span).

tvm.te.indexdiv(a, b, span=None)#

Compute floor(a / b) where a and b are non-negative.

Parameters:

a (Expr) – The left hand operand, known to be non-negative.
b (Expr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

Note

Use this function to split non-negative indices. This function may take advantage of operands’ non-negativeness.

tvm.te.indexmod(a, b, span=None)#

Compute the remainder of indexdiv. a and b are non-negative.

Parameters:

a (Expr) – The left hand operand, known to be non-negative.
b (Expr) – The right hand operand, known to be non-negative.
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Expr

Note

Use this function to split non-negative indices. This function may take advantage of operands’ non-negativeness.

tvm.te.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.te.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.te.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.te.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.te.logaddexp(a, b, span=None)#

Compute the logaddexp 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

Alias of tvm.tirx.logaddexp()

tvm.te.comm_reducer(fcombine, fidentity, name='reduce')#

Create a commutative reducer for reduction.

Parameters:

fcombine (function(Expr -> Expr -> Expr)) – A binary function which takes two Expr as input to return a Expr.
fidentity (function(str -> Expr)) – A function which takes a type string as input to return a const Expr.

Returns:

reducer – A function which creates a reduce expression over axis. There are two ways to use it:

accept (expr, axis, where) to produce an Reduce Expr on specified axis;
simply use it with multiple Exprs.

Return type:

function

Example

n = te.var("n")
m = te.var("m")
mysum = te.comm_reducer(lambda x, y: x+y,
    lambda t: tvm.tirx.const(0, dtype=t), name="mysum")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), name="k")
B = te.compute((n,), lambda i: mysum(A[i, k], axis=k), name="B")

tvm.te.min(expr, axis, where=None, init=None, *args)#

Create a min expression over axis.

Parameters:

expr (Expr) – The source expression.
axis (IterVar) – The reduction IterVar axis
where (optional, Expr) – Filtering predicate of the reduction.

Returns:

value – The result value.

Return type:

Expr

Example

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this min reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
B = te.compute((m,), lambda i: te.min(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
min_res = te.min(m, n)

tvm.te.max(expr, axis, where=None, init=None, *args)#

Create a max expression over axis.

Parameters:

expr (Expr) – The source expression.
axis (IterVar) – The reduction IterVar axis
where (optional, Expr) – Filtering predicate of the reduction.

Returns:

value – The result value.

Return type:

Expr

Example

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this max reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
B = te.compute((m,), lambda i: te.max(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
max_res = te.max(m, n)

tvm.te.sum(expr, axis, where=None, init=None, *args)#

Create a sum expression over axis.

Parameters:

expr (Expr) – The source expression.
axis (IterVar) – The reduction IterVar axis
where (optional, Expr) – Filtering predicate of the reduction.

Returns:

value – The result value.

Return type:

Expr

Example

m = te.var("m")
n = te.var("n")
A = te.placeholder((m, n), name="A")
k = te.reduce_axis((0, n), name="k")

# there are two way to use this sum reducer:
# mode 1, accept (expr, axis, where) to produce an Reduce Expr
B = te.compute((m,), lambda i: te.sum(A[i, k], axis=k), name="B")

# mode 2, simply use it with multiple Exprs:
sum_res = te.sum(m, n)

class tvm.te.CommReducer(lhs: list[Var], rhs: list[Var], result: list[Expr], identity_element: list[Expr], span: Span | None = None)#

Commutative reduce operator

Parameters:

lhs (List[Var]) – The left arguments of the reducer.
rhs (List[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.

class tvm.te.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.te.TensorSlice(tensor, indices)#

Auxiliary data structure for enable slicing syntax from tensor.

Methods:

`asobject`()	Convert slice to object.
`expr_ty`()	Compile-time element type of the tensor.

Attributes:

dtype

Data content of the tensor.

asobject()#: Convert slice to object.

property dtype#: Data content of the tensor.

expr_ty()#: Compile-time element type of the tensor.

class tvm.te.Tensor(shape, dtype, op, value_index, span=<MISSING>, ty=<MISSING>)#

Tensor object, to construct, see function.Tensor

Attributes:

`ndim`	Dimension of the tensor.
`dtype`	Data content of the tensor.

Methods:

expr_ty()

Compile-time element type of the tensor.

property ndim#: Dimension of the tensor.

property dtype#: Data content of the tensor.

expr_ty()#: Compile-time element type of the tensor.

tvm.te.tag_scope(tag)#

The operator tag scope.

Parameters:: tag (str) – The tag name.
Returns:: tag_scope – The tag scope object, which can be used as decorator or context manger.
Return type:: TagScope

Example

n = te.var('n')
m = te.var('m')
l = te.var('l')
A = te.placeholder((n, l), name='A')
B = te.placeholder((m, l), name='B')
k = te.reduce_axis((0, l), name='k')

with tvm.te.tag_scope(tag='matmul'):
    C = te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k))

# or use tag_scope as decorator
@tvm.te.tag_scope(tag="conv")
def compute_relu(data):
    return te.compute(data.shape, lambda *i: tvm.tirx.Select(data(*i) < 0, 0.0, data(*i)))

tvm.te.placeholder(shape, dtype=None, name='placeholder')#

Construct an empty tensor object.

Parameters:

shape (Tuple of Expr) – The shape of the tensor
dtype (str, optional) – The data type of the tensor
name (str, optional) – The name hint of the tensor

Returns:

tensor – The created tensor

Return type:

Tensor

tvm.te.compute(shape, fcompute, name='compute', tag='', attrs=None, varargs_names=None)#

Construct a new tensor by computing over the shape domain.

The compute rule is result[axis] = fcompute(axis)

Parameters:

shape (Tuple of Expr) – The shape of the tensor
fcompute (lambda function of indices-> value) – Specifies the input source expression
name (str, optional) – The name hint of the tensor
tag (str, optional) – Additional tag information about the compute.
attrs (dict, optional) – The additional auxiliary attributes about the compute.
varargs_names (list, optional) – The names to use for each of the varargs. If not supplied, the varargs will be called i1, i2, …

Returns:

tensor – The created tensor

Return type:

Tensor

tvm.te.scan(init, update, state_placeholder, inputs=None, name='scan', tag='', attrs=None)#

Construct new tensors by scanning over axis.

Parameters:

init (Tensor or list of Tensor) – The initial condition of first init.shape[0] timestamps
update (Tensor or list of Tensor) – The update rule of the scan given by symbolic tensor.
state_placeholder (Tensor or list of Tensor) – The placeholder variables used by update.
inputs (Tensor or list of Tensor, optional) – The list of inputs to the scan. This is not required, but can be useful for the compiler to detect scan body faster.
name (str, optional) – The name hint of the tensor
tag (str, optional) – Additonal tag information about the compute.
attrs (dict, optional) – The additional auxiliary attributes about the compute.

Returns:

tensor – The created tensor or tuple of tensors contains multiple outputs.

Return type:

Tensor or list of Tensors

Example

# The following code is equivalent to numpy.cumsum
m = te.var("m")
n = te.var("n")
X = te.placeholder((m, n), name="X")
s_state = te.placeholder((m, n))
s_init = te.compute((1, n), lambda _, i: X[0, i])
s_update = te.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
res = tvm.te.scan(s_init, s_update, s_state, X)

tvm.te.extern(shape, inputs, fcompute, name='extern', dtype=None, in_buffers=None, out_buffers=None, tag='', attrs=None)#

Compute several tensors via an extern function.

Parameters:

shape (tuple or list of tuples.) – The shape of the outputs.
inputs (list of Tensor) – The inputs
fcompute (lambda function of inputs, outputs-> stmt) –
Specifies the IR statement or sequence of statements to do the computation. See the following note for function signature of fcompute
Note

Parameters
- ins (list of tvm.ir.Var) - Placeholder for each inputs
- outs (list of tvm.ir.Var) - Placeholder for each outputs
Returns
- stmt (tvm.tirx.Stmt or a sequence of statements) - The body that carries out array computation.
name (str, optional) – The name hint of the tensor
dtype (str or list of str, optional) – The data types of outputs, by default dtype will be same as inputs.
in_buffers (tvm.ir.Var or list of tvm.ir.Var, optional) – Input buffers.
out_buffers (tvm.ir.Var or list of tvm.ir.Var, optional) – Output buffers.

tag: str, optional: Additonal tag information about the compute.
attrs: dict, optional: The additional auxiliary attributes about the compute.

Returns:: tensor – The created tensor or tuple of tensors contains multiple outputs.
Return type:: Tensor or list of Tensors

Example

In the code below, C is generated by calling external PackedFunc tvm.contrib.cblas.matmul

A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern((n, m), [A, B],
               lambda ins, outs: tvm.tirx.call_packed(
                  "tvm.contrib.cblas.matmul",
                    ins[0], ins[1], outs[0], 0, 0), name="C")

tvm.te.var(name='tindex', dtype='int32', span=None)#

Create a new variable with specified name and dtype

Parameters:

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

Returns:

var – The result symbolic variable.

Return type:

tirx.Var

tvm.te.const(value, dtype='int32', span=None)#

Create a new constant with specified value and dtype

Parameters:

value (Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]) – The constant value.
dtype (str) – The data type
span (Optional[Span]) – The location of this variable in the source.

Returns:

const – The result constant expr.

Return type:

Expr

tvm.te.thread_axis(dom=None, tag='', name='', span=None)#

Create a new IterVar to represent thread index.

Parameters:

dom (Range or str) – The domain of iteration When str is passed, dom is set to None and str is used as tag
tag (str, optional) – The thread tag
name (str, optional) – The name of the var.
span (Optional[Span]) – The location of this variable in the source.

Returns:

axis – The thread itervar.

Return type:

IterVar

tvm.te.reduce_axis(dom, name='rv', thread_tag='', span=None)#

Create a new IterVar for reduction.

Parameters:

dom (Range) – The domain of iteration.
name (str) – The name of the variable.
thread_tag (Optional[str]) – The name of the thread_tag.
span (Optional[Span]) – The location of this variable in the source.

Returns:

axis – An iteration variable representing the value.

Return type:

IterVar

tvm.te.create_function(ops: list[Tensor | Var], index_dtype_override: str | None = None) → Function#

Create a TensorIR Function from tensor expression

Parameters:: ops (List[Union[_tensor.Tensor, tvm.tirx.Var]]) – The source expression.

Example

We define a matmul kernel using following code:

import tvm
from tvm import te
from tvm.te import create_function
import tvm.script

A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")
k = te.reduce_axis((0, 128), "k")
C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C")
func = create_function([A, B, C])
print(func.script())

If we want to use TensorIR schedule to do transformations on such kernel, we need to use create_function([A, B, C]) to create a schedulable Function. The generated function looks like:

@Ts.function
def tir_matmul(
    A: T.Tensor((128, 128)), B: T.Tensor((128, 128)), C: T.Tensor((128, 128))
) -> None:

    for i, j, k in T.grid(128, 128, 128):
        with Ts.sblock():
            vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
            with Ts.init():
                C[vi, vj] = 0.0
            C[vi, vj] += A[vi, vk] * B[vj, vk]

Returns:: func – The created function.
Return type:: tirx.Function

tvm.te.extern_function(input_tensors: list[Tensor], function: Function, **kwargs)#

Compute tensors via a schedulable TIR Function

Parameters:

input_tensors (list of Tensor) – Input tensors that map to the corresponding function input params.
function (Function) – The TIR Function

Returns:

tensor – The created tensor or tuple of tensors if it contains multiple outputs.

Return type:

Tensor or list of Tensors

Example

In the code below, a TVMScript defined TIR Function is inlined into a TE ExternOp. Applying te.create_function on this

A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")

@Ts.function
def before_split(A: T.Tensor((128, 128)), B: T.Tensor((128, 128))) -> None:

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

C = te.extern_function([A, B], func)

class tvm.te.PlaceholderOp(name, tag, attrs, shape, dtype)#: Placeholder operation.

class tvm.te.ComputeOp(name, tag, attrs, axis, reduce_axis, body)#: Scalar operation.

class tvm.te.ScanOp(name, tag, attrs, scan_axis, init, update, state_placeholder, inputs, spatial_axis_)#: Scan operation.

class tvm.te.ExternOp(name, tag, attrs, inputs, input_placeholders, output_placeholders, body)#: External operation.

tvm.te

Contents

tvm.te#