tvm

The user facing API for computation declaration.

tvm.load_json(json_str) Load tvm object from json_str.
tvm.save_json(node) Load tvm object as json string.
tvm.var([name, dtype]) Create a new variable with specified name and dtype
tvm.const(value[, dtype]) construct a constant
tvm.convert(value) Convert value to TVM node or function.
tvm.placeholder(shape[, dtype, name]) Construct an empty tensor object.
tvm.compute(shape, fcompute[, name, tag, attrs]) Construct a new tensor by computing over the shape domain.
tvm.scan(init, update, state_placeholder[, …]) Construct new tensors by scanning over axis.
tvm.extern(shape, inputs, fcompute[, name, …]) Compute several tensor via extern function.
tvm.decl_buffer(shape[, dtype, name, data, …]) Decleare a new symbolic buffer.
tvm.reduce_axis(dom[, name]) Create a new IterVar for reduction.
tvm.select(cond, t, f) Construct a select branch.
tvm.thread_axis([dom, tag, name]) Create a new IterVar to represent thread index.
tvm.comm_reducer(fcombine, fidentity[, name]) Create a commutative reducer for reduction.
tvm.sum(expr, axis[, where]) Create a sum expression over axis.
tvm.min(expr, axis[, where]) Create a min expression over axis.
tvm.max(expr, axis[, where]) Create a max expression over axis.
tvm.tag_scope(tag) The operator tag scope.
tvm.load_json(json_str)

Load tvm object from json_str.

Parameters:json_str (str) – The json string
Returns:node – The loaded tvm node.
Return type:Node
tvm.save_json(node)

Load tvm object as json string.

Parameters:node (Node) – A TVM Node object to be saved.
Returns:json_str – Saved json string.
Return type:str
tvm.var(name='tindex', dtype='int32')

Create a new variable with specified name and dtype

Parameters:
  • name (str) – The name
  • dtype (int) – The data type
Returns:

var – The result symbolic variable.

Return type:

Var

tvm.const(value, dtype=None)

construct a constant

tvm.convert(value)

Convert value to TVM node or function.

Parameters:value (python value) –
Returns:tvm_val – Converted value in TVM
Return type:Node or Function
tvm.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.compute(shape, fcompute, name='compute', tag='', attrs=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) – Additonal tag information about the compute.
  • attrs (dict, optional) – The additional auxiliary attributes about the compute.
Returns:

tensor – The created tensor

Return type:

Tensor

tvm.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 it it contains multiple outputs.

Return type:

Tensor or list of Tensors

Example

# The following code is equivalent to numpy.cumsum
m = tvm.var("m")
n = tvm.var("n")
X = tvm.placeholder((m, n), name="X")
s_state = tvm.placeholder((m, n))
s_init = tvm.compute((1, n), lambda _, i: X[0, i])
s_update = tvm.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
res = tvm.scan(s_init, s_update, s_state, X)
tvm.extern(shape, inputs, fcompute, name='extern', dtype=None, in_buffers=None, out_buffers=None, tag='', attrs=None)

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

    Note

    Parameters

    • ins (list of Buffer) - Placeholder for each inputs
    • outs (list of Buffer) - Placeholder for each outputs

    Returns

    • stmt (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 (Buffer or list of Buffer, optional) – Input buffers.
  • out_buffers (Buffer or list of Buffers, 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 it it 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 = tvm.placeholder((n, l), name='A')
B = tvm.placeholder((l, m), name='B')
C = tvm.extern((n, m), [A, B],
               lambda ins, outs: tvm.call_packed(
                  "tvm.contrib.cblas.matmul",
                    ins[0], ins[1], outs[0], 0, 0), name="C")
tvm.decl_buffer(shape, dtype=None, name='buffer', data=None, strides=None, elem_offset=None, scope='', data_alignment=-1, offset_factor=0)

Decleare 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 (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 Var for elem_offset if elem_offset is not None.
Returns:

buffer – The created buffer

Return type:

Buffer

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.

tvm.reduce_axis(dom, name='rv')

Create a new IterVar for reduction.

Parameters:
  • dom (Range) – The domain of iteration.
  • name (str) – The name of the variable.
Returns:

axis – An iteration variable representing the value.

Return type:

IterVar

tvm.select(cond, t, f)

Construct a select branch.

Parameters:
  • cond (Expr) – The condition
  • t (Expr) – The result expression if cond is true.
  • f (Expr) – The result expression if cond is false.
Returns:

node – The tvm.expr.Select node

Return type:

Node

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

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.
Returns:

axis – The thread itervar.

Return type:

IterVar

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

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

Return type:

function

Example

n = tvm.var('n')
m = tvm.var('m')
mysum = tvm.comm_reducer(lambda x, y: x+y,
    lambda t: tvm.const(0, dtype=t), name="mysum")
A = tvm.placeholder((n, m), name='A')
k = tvm.reduce_axis((0, m), name='k')
B = tvm.compute((n,), lambda i: mysum(A[i, k], axis=k), name='B')
tvm.sum(expr, axis, where=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 = tvm.var("m")
n = tvm.var("n")
A = tvm.placeholder((m, n), name="A")
k = tvm.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 = tvm.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.min(expr, axis, where=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 = tvm.var("m")
n = tvm.var("n")
A = tvm.placeholder((m, n), name="A")
k = tvm.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 = tvm.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.max(expr, axis, where=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 = tvm.var("m")
n = tvm.var("n")
A = tvm.placeholder((m, n), name="A")
k = tvm.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 = tvm.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.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 = tvm.var('n')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((n, l), name='A')
B = tvm.placeholder((m, l), name='B')
k = tvm.reduce_axis((0, l), name='k')

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

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