tvm.te#

Namespace for Tensor Expression Language

Classes:

Functions:

class tvm.te.ComputeOp[源代码]#

Scalar operation.

class tvm.te.ExternOp[源代码]#

External operation.

class tvm.te.PlaceholderOp[源代码]#

Placeholder operation.

class tvm.te.ScanOp[源代码]#

Scan operation.

Attributes:

scan_axis

Represent the scan axis, only defined when it is a ScanOp

property scan_axis#

Represent the scan axis, only defined when it is a ScanOp

class tvm.te.Tensor[源代码]#

Tensor object, to construct, see function.Tensor

Attributes:

axis	Axis of the tensor.
ndim	Dimension of the tensor.
op	The corressponding Operation.
shape	The output shape of the tensor.
value_index	The output value index the tensor corresponds to.

property axis#

Axis of the tensor.

property ndim#

Dimension of the tensor.

property op#

The corressponding Operation.

property shape#

The output shape of the tensor.

property value_index#

The output value index the tensor corresponds to.

class tvm.te.TensorSlice(tensor, indices)[源代码]#

Auxiliary data structure for enable slicing syntax from tensor.

Methods:

asobject()

Convert slice to object.

Attributes:

dtype

Data content of the tensor.

asobject()[源代码]#

Convert slice to object.

property dtype#

Data content of the tensor.

tvm.te.abs(x, span=None)[源代码]#

Get absolute value of the input element-wise.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.acos(x)[源代码]#

Take acos of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.acosh(x)[源代码]#

Take acos of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.add(lhs, rhs, span=None)[源代码]#

Generic add operator.

参数:

• lhs (object) -- The left operand.

• rhs (object) -- The right operand.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

op -- The result Expr of add operaton.

返回类型:

tvm.Expr

tvm.te.all(*args, span=None)[源代码]#

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

arguments

参数:

• args (list) -- List of symbolic boolean expressions

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

expr -- Expression

返回类型:

Expr

tvm.te.any(*args, span=None)[源代码]#

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

参数:

• args (list) -- List of symbolic boolean expressions

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

expr -- Expression

返回类型:

Expr

tvm.te.asin(x)[源代码]#

Take asin of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.asinh(x)[源代码]#

Take asinh of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.atan(x)[源代码]#

Take atan of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.atanh(x)[源代码]#

Take atanh of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.ceil(x, span=None)[源代码]#

Take ceil of float input x.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.comm_reducer(fcombine, fidentity, name='reduce')[源代码]#

Create a commutative reducer for reduction.

参数:

• 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.

返回:

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

1. accept (expr, axis, where) to produce an Reduce Expr on specified axis;

2. simply use it with multiple Exprs.

返回类型:

function

示例

n = te.var("n")
m = te.var("m")
mysum = te.comm_reducer(lambda x, y: x+y,
    lambda t: tvm.tir.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.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)

参数:

• 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, ...

返回:

tensor -- The created tensor

返回类型:

Tensor

tvm.te.const(value, dtype='int32', span=None)[源代码]#

Create a new constant with specified value and dtype

参数:

• value (Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]) -- The constant value.

• dtype (str) -- The data type

• span (Optional[Span]) -- The location of this variable in the source.

返回:

const -- The result constant expr.

返回类型:

PrimExpr

tvm.te.cos(x)[源代码]#

Take cos of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.cosh(x)[源代码]#

Take cosh of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.create_prim_func(ops: List[Tensor | Var], index_dtype_override: str | None = None) → PrimFunc[源代码]#

Create a TensorIR PrimFunc from tensor expression

参数:

ops (List[Union[_tensor.Tensor, tvm.tir.Var]]) -- The source expression.

示例

We define a matmul kernel using following code:

import tvm
from tvm import te
from tvm.te import create_prim_func
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_prim_func([A, B, C])
print(func.script())

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

@T.prim_func
def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.match_buffer(b, (128, 128))
    C = T.match_buffer(c, (128, 128))

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

返回:

func -- The created function.

返回类型:

tir.PrimFunc

tvm.te.div(a, b, span=None)[源代码]#

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

参数:

• a (PrimExpr) -- The left hand operand, known to be non-negative.

• b (PrimExpr) -- The right hand operand, known to be non-negative.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

备注

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

tvm.te.erf(x)[源代码]#

Take gauss error function of the input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.exp(x)[源代码]#

Take exponential of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

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.

参数:

• 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 to do the computation. See the following note for function signature of fcompute

  备注

  Parameters

  • ins (list of tvm.tir.Buffer) - Placeholder for each inputs

  • outs (list of tvm.tir.Buffer) - Placeholder for each outputs

  Returns

  • stmt (tvm.tir.Stmt) - The statement 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.tir.Buffer or list of tvm.tir.Buffer, optional) -- Input buffers.

• out_buffers (tvm.tir.Buffer or list of tvm.tir.Buffer, optional) -- Output buffers.

tag: str, optional

Additonal tag information about the compute.

attrs: dict, optional

The additional auxiliary attributes about the compute.

返回:

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

返回类型:

Tensor or list of Tensors

示例

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.tir.call_packed(
                  "tvm.contrib.cblas.matmul",
                    ins[0], ins[1], outs[0], 0, 0), name="C")

tvm.te.extern_primfunc(input_tensors: List[Tensor], primfunc: PrimFunc, **kwargs)[源代码]#

Compute tensors via a schedulable TIR PrimFunc

参数:

• input_tensors (list of Tensor) -- Input tensors that map to the corresponding primfunc input params.

• primfunc (PrimFunc) -- The TIR PrimFunc

返回:

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

返回类型:

Tensor or list of Tensors

示例

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

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

@T.prim_func
def before_split(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (128, 128))
    B = T.match_buffer(b, (128, 128))
    for i, j in T.grid(128, 128):
        with T.block("B"):
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj] * 2.0

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

tvm.te.floor(x: PrimExprWithOp, span=None)[源代码]#

Take floor of float input x.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.floordiv(a, b, span=None)[源代码]#

Compute the floordiv of two expressions.

参数:

• a (PrimExpr) -- The left hand operand

• b (PrimExpr) -- The right hand operand

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

tvm.te.floormod(a, b, span=None)[源代码]#

Compute the floormod of two expressions.

参数:

• a (PrimExpr) -- The left hand operand

• b (PrimExpr) -- The right hand operand

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

tvm.te.fmod(x, y)[源代码]#

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

参数:

• x (PrimExpr) -- Input argument.

• y (PrimExpr) -- Input argument.

返回:

z -- The result.

返回类型:

PrimExpr

tvm.te.if_then_else(cond, t, f, span=None)[源代码]#

Conditional selection expression.

参数:

• cond (PrimExpr) -- The condition

• t (PrimExpr) -- The result expression if cond is true.

• f (PrimExpr) -- The result expression if cond is false.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

result -- The result of conditional expression.

返回类型:

Node

备注

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.indexdiv(a, b, span=None)[源代码]#

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

参数:

• a (PrimExpr) -- The left hand operand, known to be non-negative.

• b (PrimExpr) -- The right hand operand, known to be non-negative.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

备注

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.

参数:

• a (PrimExpr) -- The left hand operand, known to be non-negative.

• b (PrimExpr) -- The right hand operand, known to be non-negative.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

备注

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

tvm.te.isfinite(x, span=None)[源代码]#

Check if input value is finite.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.isinf(x, span=None)[源代码]#

Check if input value is infinite.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.isnan(x, span=None)[源代码]#

Check if input value is Nan.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.log(x)[源代码]#

Take log of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.log10(x)[源代码]#

Take log10 of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.log2(x)[源代码]#

Take log2 of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

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

Create a max expression over axis.

参数:

• expr (PrimExpr) -- The source expression.

• axis (IterVar) -- The reduction IterVar axis

• where (optional, Expr) -- Filtering predicate of the reduction.

返回:

value -- The result value.

返回类型:

PrimExpr

示例

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
# tvm.max represents tvm.te.max or tvm.tir.max.
B = te.compute((m,), lambda i: tvm.max(A[i, k], axis=k), name="B")

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

tvm.te.max_value(dtype: str, span: Span | None = None) → Any[源代码]#

maximum value of dtype

参数:

• dtype (str) -- The data type.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

value -- The maximum value of dtype.

返回类型:

tvm.Expr

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

Create a min expression over axis.

参数:

• expr (PrimExpr) -- The source expression.

• axis (IterVar) -- The reduction IterVar axis

• where (optional, Expr) -- Filtering predicate of the reduction.

返回:

value -- The result value.

返回类型:

PrimExpr

示例

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
# tvm.min represents tvm.te.min or tvm.tir.min.
B = te.compute((m,), lambda i: tvm.min(A[i, k], axis=k), name="B")

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

tvm.te.min_value(dtype, span=None)[源代码]#

minimum value of dtype

参数:

• dtype (str) -- The data type.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

value -- The minimum value of dtype.

返回类型:

tvm.Expr

tvm.te.multiply(lhs, rhs, span=None)[源代码]#

Generic multiply operator.

参数:

• lhs (object) -- The left operand.

• rhs (object) -- The right operand.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

op -- The result Expr of multiply operaton.

返回类型:

tvm.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

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.placeholder(shape, dtype=None, name='placeholder')[源代码]#

Construct an empty tensor object.

参数:

• 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

返回:

tensor -- The created tensor

返回类型:

Tensor

tvm.te.popcount(x)[源代码]#

Count the number of set bits in input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.power(x, y, span=None)[源代码]#

x power y

参数:

• x (PrimExpr) -- Input argument.

• y (PrimExpr) -- The exponent

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

z -- The result.

返回类型:

PrimExpr

tvm.te.reduce_axis(dom, name='rv', thread_tag='', span=None)[源代码]#

Create a new IterVar for reduction.

参数:

• 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.

返回:

axis -- An iteration variable representing the value.

返回类型:

IterVar

tvm.te.round(x, span=None)[源代码]#

Round elements of the array to the nearest integer.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.rsqrt(x)[源代码]#

Take reciprocal of square root of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.scan(init, update, state_placeholder, inputs=None, name='scan', tag='', attrs=None)[源代码]#

Construct new tensors by scanning over axis.

参数:

• 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.

返回:

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

返回类型:

Tensor or list of Tensors

示例

# 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.sigmoid(x)[源代码]#

Quick function to get sigmoid

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.sin(x)[源代码]#

Take sin of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.sinh(x)[源代码]#

Take sinh of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.size_var(name='size', dtype='int32', span=None)[源代码]#

Create a new variable represents a tensor shape size, which is non-negative.

参数:

• name (str) -- The name

• dtype (str) -- The data type

• span (Optional[Span]) -- The location of this variable in the source.

返回:

var -- The result symbolic shape variable.

返回类型:

SizeVar

tvm.te.sqrt(x)[源代码]#

Take square root of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.subtract(lhs, rhs, span=None)[源代码]#

Generic subtract operator.

参数:

• lhs (object) -- The left operand.

• rhs (object) -- The right operand.

• span (Optional[Span]) -- The location of this operator in the source.

返回:

op -- The result Expr of subtract operaton.

返回类型:

tvm.Expr

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

Create a sum expression over axis.

参数:

• expr (PrimExpr) -- The source expression.

• axis (IterVar) -- The reduction IterVar axis

• where (optional, Expr) -- Filtering predicate of the reduction.

返回:

value -- The result value.

返回类型:

PrimExpr

示例

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
# tvm.sum represents tvm.te.sum or tvm.tir.sum.
B = te.compute((m,), lambda i: tvm.sum(A[i, k], axis=k), name="B")

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

tvm.te.tag_scope(tag)[源代码]#

The operator tag scope.

参数:

tag (str) -- The tag name.

返回:

tag_scope -- The tag scope object, which can be used as decorator or context manger.

返回类型:

TagScope

示例

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.tir.Select(data(*i) < 0, 0.0, data(*i)))

tvm.te.tan(x)[源代码]#

Take tan of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.tanh(x)[源代码]#

Take hyperbolic tanh of input x.

参数:

x (PrimExpr) -- Input argument.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.thread_axis(dom=None, tag='', name='', span=None)[源代码]#

Create a new IterVar to represent thread index.

参数:

• 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.

返回:

axis -- The thread itervar.

返回类型:

IterVar

tvm.te.trace(args, trace_action='tvm.default_trace_action')[源代码]#

Trace tensor data at the runtime.

The trace function allows to trace specific tensor at the runtime. The tracing value should come as last argument. The trace action should be specified, by default tvm.default_trace_action is used.

参数:

• args (list of Expr or Buffers.) -- Positional arguments.

• trace_action (str.) -- The name of the trace action.

返回:

call -- The call expression.

返回类型:

PrimExpr

参见

tvm.tir.call_packed

Creates packed function.

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.

参数:

• x (PrimExpr) -- Input argument.

• span (Optional[Span]) -- The location of this operator in the source code.

返回:

y -- The result.

返回类型:

PrimExpr

tvm.te.truncdiv(a, b, span=None)[源代码]#

Compute the truncdiv of two expressions.

参数:

• a (PrimExpr) -- The left hand operand

• b (PrimExpr) -- The right hand operand

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

备注

This is the default integer division behavior in C.

tvm.te.truncmod(a, b, span=None)[源代码]#

Compute the truncmod of two expressions.

参数:

• a (PrimExpr) -- The left hand operand

• b (PrimExpr) -- The right hand operand

• span (Optional[Span]) -- The location of this operator in the source.

返回:

res -- The result expression.

返回类型:

PrimExpr

备注

This is the default integer division behavior in C.

tvm.te.var(name='tindex', dtype='int32', span=None)[源代码]#

Create a new variable with specified name and dtype

参数:

• name (str) -- The name

• dtype (str) -- The data type

• span (Optional[Span]) -- The location of this variable in the source.

返回:

var -- The result symbolic variable.

返回类型:

tir.Var