Developer API

This page contains modules that are used by developers of TVM. Many of these APIs are PackedFunc registered in C++ backend.

tvm.node

Node is the base class of all TVM AST.

Normally user do not need to touch this api.

class tvm.node.NodeBase

NodeBase is the base class of all TVM language AST object.

same_as(other)

check object identity equality

tvm.node.Node

alias of tvm._ffi.node.NodeBase

tvm.register_node(type_key=None)

register node type

Parameters:type_key (str or cls) – The type key of the node

tvm.expr

Expression AST Node in TVM.

User do not need to deal with expression AST node directly. But they can be helpful for developer to do quick proptyping. While not displayed in the document and python file. Each expression node have subfields that can be visited from python side.

For example, you can use addexp.a to get the left operand of an Add node.

x = tvm.var("n")
y = x + 2
assert(isinstance(y, tvm.expr.Add))
assert(y.a == x)
class tvm.expr.Add(a, b)

Add node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.And(a, b)

And node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.BinaryOpExpr
class tvm.expr.Broadcast(value, lanes)

Broadcast node.

Parameters:
  • value (Expr) – The value of the expression.
  • lanes (int) – The lanes of the expression.
class tvm.expr.Call(dtype, name, args, call_type, func, value_index)

Call node.

Parameters:
  • dtype (str) – The return data type
  • name (str) – The name of the function
  • args (list of Expr) – The input arguments to the call
  • call_type (int) – The type of the call
  • func (Operation, optional) – Operation if call_type is Halide
  • value_index (int) – The output value index
Extern = 0
ExternCPlusPlus = 1
Halide = 3
Intrinsic = 4
PureExtern = 2
PureIntrinsic = 5
class tvm.expr.Cast(dtype, value)

Cast expression.

Parameters:
  • dtype (str) – The data type
  • value (Expr) – The value of the function.
class tvm.expr.CmpExpr
class tvm.expr.ConstExpr
class tvm.expr.Div(a, b)

Div node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.EQ(a, b)

EQ node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.EqualOp(a, b)

Deferred equal operator.

This is used to support sugar that a == b can either mean NodeBase.same_as or NodeBase.equal.

Parameters:
  • a (Expr) – Left operand.
  • b (Expr) – Right operand.
asnode()

Convert node.

same_as

Return self==value.

class tvm.expr.Expr

Base class of all tvm Expressions

class tvm.expr.ExprOp
astype(dtype)

Cast the expression to other type.

Parameters:dtype (str) – The type of new expression
Returns:expr – Expression with new type
Return type:Expr
equal(other)

Build an equal check expression with other expr.

Parameters:other (Expr) – The other expression
Returns:ret – The equality expression.
Return type:Expr
class tvm.expr.FloatImm(dtype, value)

Float constant.

Parameters:
  • dtype (str) – The data type
  • value (float) – The constant value.
class tvm.expr.GE(a, b)

GE node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.GT(a, b)

GT node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.IntImm(dtype, value)

Int constant.

Parameters:
  • dtype (str) – The data type
  • value (int) – The constant value.
class tvm.expr.LE(a, b)

LE node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.LT(a, b)

LT node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Let(var, value, body)

Let node.

Parameters:
  • var (Var) – The variable in the binding.
  • value (Expr) – The value in to be binded.
  • body (Expr) – The body expression.
class tvm.expr.Load(dtype, buffer_var, index, predicate)

Load node.

Parameters:
  • dtype (str) – The data type.
  • buffer_var (Var) – The buffer variable in the load expression.
  • index (Expr) – The index in the load.
  • predicate (Expr) – The load predicate.
class tvm.expr.LogicalExpr
class tvm.expr.Max(a, b)

Max node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Min(a, b)

Min node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Mod(a, b)

Mod node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Mul(a, b)

Mul node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.NE(a, b)

NE node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Not(a)

Not node.

Parameters:a (Expr) – The input value
class tvm.expr.NotEqualOp(a, b)

Deferred NE operator.

This is used to support sugar that a != b can either mean not NodeBase.same_as or make.NE.

Parameters:
  • a (Expr) – Left operand.
  • b (Expr) – Right operand.
asnode()

Convert node.

same_as

Return self==value.

class tvm.expr.Or(a, b)

Or node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.Ramp(base, stride, lanes)

Ramp node.

Parameters:
  • base (Expr) – The base expression.
  • stride (ramp stride) – The stride of the ramp.
  • lanes (int) – The lanes of the expression.
class tvm.expr.Reduce(combiner, src, rdom, condition, value_index)

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.
class tvm.expr.Select(condition, true_value, false_value)

Select node.

Parameters:
  • condition (Expr) – The condition expression.
  • true_value (Expr) – The value to take when condition is true.
  • false_value (Expr) – The value to take when condition is false.
class tvm.expr.Shuffle(vectors, indices)

Shuffle node.

Parameters:
  • vectors (Array of Expr) – The vectors
  • indices (Array of indices) – The indices
class tvm.expr.StringImm(value)

String constant.

Parameters:value (str) – The value of the function.
class tvm.expr.Sub(a, b)

Sub node.

Parameters:
  • a (Expr) – The left hand operand.
  • b (Expr) – The right hand operand.
class tvm.expr.UIntImm(dtype, value)

UInt constant.

Parameters:
  • dtype (str) – The data type
  • value (int) – The constant value.
class tvm.expr.Var(name, dtype)

Symbolic variable.

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

tvm.codegen

Code generation related functions.

tvm.codegen.build_module(lowered_func, target)

Build lowered_func into Module.

Parameters:
  • lowered_func (LoweredFunc) – The lowered function
  • target (str) – The target module type.
Returns:

module – The corressponding module.

Return type:

Module

tvm.stmt

Statement AST Node in TVM.

User do not need to deal with AST node directly. But they can be helpful for developer to do quick proptyping. While not displayed in the document and python file. Each statement node have subfields that can be visited from python side.

x = tvm.var("n")
a = tvm.var("array", tvm.handle)
st = tvm.make.Store(a, x + 1, 1)
assert isinstance(st, tvm.stmt.Store)
assert(st.buffer_var == a)
class tvm.stmt.Allocate(buffer_var, dtype, extents, condition, body)

Allocate node.

Parameters:
  • buffer_var (Var) – The buffer variable.
  • dtype (str) – The data type of the buffer.
  • extents (list of Expr) – The extents of the allocate
  • condition (Expr) – The condition.
  • body (Stmt) – The body statement.
class tvm.stmt.AssertStmt(condition, message, body)

AssertStmt node.

Parameters:
  • condition (Expr) – The assert condition.
  • message (Expr) – The error message.
  • body (Stmt) – The body statement.
class tvm.stmt.AttrStmt(node, attr_key, value, body)

AttrStmt node.

Parameters:
  • node (Node) – The node to annotate the attribute
  • attr_key (str) – Attribute type key.
  • value (Expr) – The value of the attribute
  • body (Stmt) – The body statement.
class tvm.stmt.Block(first, rest)

Block node.

Parameters:
  • first (Stmt) – The first statement.
  • rest (Stmt) – The following statement.
class tvm.stmt.Evaluate(value)

Evaluate node.

Parameters:value (Expr) – The expression to be evalued.
class tvm.stmt.For(loop_var, min_val, extent, for_type, device_api, body)

For node.

Parameters:
  • loop_var (Var) – The loop variable.
  • min_val (Expr) – The begining value.
  • extent (Expr) – The length of the loop.
  • for_type (int) – The for type.
  • device_api (int) – The device api type.
  • body (Stmt) – The body statement.
Parallel = 1
Serial = 0
Unrolled = 3
Vectorized = 2
class tvm.stmt.Free(buffer_var)

Free node.

Parameters:buffer_var (Var) – The buffer variable.
class tvm.stmt.IfThenElse(condition, then_case, else_case)

IfThenElse node.

Parameters:
  • condition (Expr) – The expression
  • then_case (Stmt) – The statement to execute if condition is true.
  • else_case (Stmt) – The statement to execute if condition is false.
class tvm.stmt.LetStmt(var, value, body)

LetStmt node.

Parameters:
  • var (Var) – The variable in the binding.
  • value (Expr) – The value in to be binded.
  • body (Stmt) – The body statement.
class tvm.stmt.Prefetch(func, value_index, dtype, bounds)

Prefetch node.

Parameters:
  • func (Operation) – The operation to create the function.
  • value_index (int) – The output value index
  • dtype (str) – The data type to be prefetched.
  • bounds (list of Range) – The bounds to be prefetched.
class tvm.stmt.ProducerConsumer(func, is_producer, body)

ProducerConsumer node.

Parameters:
  • func (Operation) – The Operation.
  • is_producer (bool) – Whether if the node is producer.
  • body (Stmt) – The body statement.
class tvm.stmt.Provide(func, value_index, value, args)

Provide node.

Parameters:
  • func (Operation) – The operation to create the function.
  • value_index (int) – The output value index
  • value (Expr) – The value to be stored.
  • args (list of Expr) – The index arguments of the Provide.
class tvm.stmt.Realize(func, value_index, dtype, bounds, condition, body)

Realize node.

Parameters:
  • func (Operation) – The operation to create the function.
  • value_index (int) – The output value index
  • dtype (str) – The data type of the operation.
  • bounds (list of range) – The bound of realize
  • condition (Expr) – The realize condition.
  • body (Stmt) – The realize body
class tvm.stmt.Stmt
class tvm.stmt.Store(buffer_var, value, index, predicate)

Store node.

Parameters:
  • buffer_var (Var) – The buffer Variable.
  • value (Expr) – The value we want to store.
  • index (Expr) – The index in the store expression.
  • predicate (Expr) – The store predicate.
tvm.stmt.stmt_list(stmt)

Make list of stmt from blocks.

Parameters:stmt (A block statement) –
Returns:stmt_list – The unpacked list of statements
Return type:list of Stmt
tvm.stmt.stmt_seq(*args)

Make sequence of statements

Parameters:args (list of Expr or Var) – List of statements to be combined as sequence.
Returns:stmt – The combined statement.
Return type:Stmt

tvm.ir_pass

tvm.ir_pass.Inline TVM PackedFunc Inline.
tvm.ir_pass.Simplify TVM PackedFunc Simplify.
tvm.ir_pass.ConvertSSA TVM PackedFunc ConvertSSA.
tvm.ir_pass.VerifySSA TVM PackedFunc VerifySSA.
tvm.ir_pass.CanonicalSimplify TVM PackedFunc CanonicalSimplify.
tvm.ir_pass.StorageFlatten TVM PackedFunc StorageFlatten.
tvm.ir_pass.VectorizeLoop TVM PackedFunc VectorizeLoop.
tvm.ir_pass.UnrollLoop TVM PackedFunc UnrollLoop.
tvm.ir_pass.ThreadSync TVM PackedFunc ThreadSync.
tvm.ir_pass.StorageRewrite TVM PackedFunc StorageRewrite.
tvm.ir_pass.MakeAPI TVM PackedFunc MakeAPI.
tvm.ir_pass.SplitHostDevice TVM PackedFunc SplitHostDevice.
tvm.ir_pass.InjectVirtualThread TVM PackedFunc InjectVirtualThread.
tvm.ir_pass.LoopPartition TVM PackedFunc LoopPartition.
tvm.ir_pass.RemoveNoOp TVM PackedFunc RemoveNoOp.
tvm.ir_pass.SplitPipeline TVM PackedFunc SplitPipeline.
tvm.ir_pass.LowerThreadAllreduce TVM PackedFunc LowerThreadAllreduce.
tvm.ir_pass.LowerIntrin TVM PackedFunc LowerIntrin.
tvm.ir_pass.LowerTVMBuiltin TVM PackedFunc LowerTVMBuiltin.
tvm.ir_pass.NarrowChannelAccess TVM PackedFunc NarrowChannelAccess.

Namespace of IR pass functions.

This namespace is used for developers. While you do not see any declarations. The functions are automatically exported from C++ side via PackedFunc.

Each api is a PackedFunc that can be called in a positional argument manner. You can read “include/tvm/pass.h” for the function signature of these functions.

tvm.ir_builder

Developer API of IR node builder make function.

class tvm.ir_builder.BufferVar(builder, buffer_var, content_type)

Buffer variable with content type, makes load store easily.

Do not create it directly, create use IRBuilder.

Examples

In the follow example, x is BufferVar. x[0] = ... directly emit a store to the IRBuilder, x[10] translates to Load.

# The following code generate IR for x[0] = x[
ib = tvm.ir_builder.create()
x = ib.pointer("float32")
x[0] = x[10] + 1
asnode()

Convert value to node

class tvm.ir_builder.IRBuilder

Auxiliary builder to build IR for testing and dev.

Examples

ib = tvm.ir_builder.create()
n = tvm.var("n")
A = ib.allocate("float32", n, name="A")
with ib.for_range(0, n, name="i") as i:
    with ib.if_scope((i % 2) == 0):
        A[i] = A[i] + 1
# The result stmt.
stmt = ib.get()
allocate(dtype, shape, name='buf', scope=None)

Create a allocate statement.

Parameters:
  • dtype (str) – The content data type.
  • shape (tuple of Expr) – The shape of array to be allocated.
  • name (str, optional) – The name of the buffer.
  • scope (str, optional) – The scope of the buffer.
Returns:

buffer – The buffer var representing the buffer.

Return type:

BufferVar

buffer_ptr(buf)

Create pointer variable corresponds to buffer ptr.

Parameters:buf (Buffer) – The buffer to be extracted.
Returns:ptr – The buffer var representing the buffer.
Return type:BufferVar
else_scope()

Create an else scope.

This can only be used right after an if scope.

Returns:else_scope – The result else scope.
Return type:WithScope

Examples

ib = tvm.ir_builder.create()
i = tvm.var("i")
x = ib.pointer("float32")
with ib.if_scope((i % 2) == 0):
    x[i] = x[i - 1] + 1
with ib.else_scope():
    x[i] = x[i - 1] + 2
emit(stmt)

Emit a statement to the end of current scope.

Parameters:stmt (Stmt or callable.) – The statement to be emitted or callable that build stmt given body.
for_range(begin, end, name='i', dtype='int32', for_type='serial')

Create a for iteration scope.

Parameters:
  • begin (Expr) – The min iteration scope.
  • end (Expr) – The end iteration scope
  • name (str, optional) – The name of iteration variable, if no input names, using typical index names i, j, k, then i_nidx
  • dtype (str, optional) – The data type of iteration variable.
  • for_type (str, optional) – The special tag on the for loop.
Returns:

loop_scope – The for scope, when enters returns loop_var

Return type:

With.Scope of Var

Examples

ib = tvm.ir_builder.create()
x = ib.pointer("float32")
with ib.for_range(1, 10, name="i") as i:
    x[i] = x[i - 1] + 1
get()

Return the builded IR.

Returns:stmt – The result statement.
Return type:Stmt
if_scope(cond)

Create an if scope.

Parameters:cond (Expr) – The condition.
Returns:if_scope – The result if scope.
Return type:WithScope

Examples

ib = tvm.ir_builder.create()
i = tvm.var("i")
x = ib.pointer("float32")
with ib.if_scope((i % 2) == 0):
    x[i] = x[i - 1] + 1
likely(expr)

Add likely tag for expression. :param expr: The expression. Usually a condition expression. :type expr: Expr

Returns:expr – The expression will likely tag.
Return type:Expr
new_scope()

Create new scope,

this is useful to set boundary of attr and allocate.

Returns:new_scope – The result new scope.
Return type:WithScope
pointer(content_type, name='ptr')

Create pointer variable with content type.

Parameters:
  • content_type (str) – The content data type.
  • name (str, optional) – The name of the pointer.
Returns:

ptr – The buffer var representing the buffer.

Return type:

BufferVar

scope_attr(node, attr_key, value)

Create an AttrStmt at current scope.

Parameters:
  • attr_key (str) – The key of the attribute type.
  • node (Node) – The attribute node to annottate on.
  • value (Expr) – Attribute value.

Examples

ib = tvm.ir_builder.create()
i = tvm.var("i")
x = ib.pointer("float32")
ib.scope_attr(x, "storage_scope", "global")
x[i] = x[i - 1] + 1
class tvm.ir_builder.WithScope(enter_value, exit_cb)

Auxiliary scope with

tvm.ir_builder.create()

Create a new IRBuilder

Returns:builder – The created IRBuilder
Return type:IRBuilder

tvm.make

namespace of IR node builder make function

This namespace is used for developers. While you do not see any declarations. The functions are automatically exported from C++ side via PackedFunc.

Each api is a PackedFunc that can be called in a positional argument manner. You can use make function to build the IR node.

tvm.make.node(type_key, **kwargs)

Make a new DSL node by its type key and fields

Parameters:
  • type_key (str) – The type key of the node.
  • **kwargs (dict) – The fields of the node.
Returns:

node – The corresponding DSL Node

Return type:

Node

Note

If the created node is instance of AttrsNode, then the creator function will also run bound checks and default value setup as supported by Attrs.

Example

The following code constructs a IntImm object

x = tvm.make.node("IntImm", dtype="int32", value=10)
assert isinstance(x, tvm.expr.IntImm)
assert x.value == 10
tvm.make.range_by_min_extent(min_value, extent)

Construct a Range by min and extent.

This constructs a range in [min_value, min_value + extent)

Parameters:
  • min_value (Expr) – The minimum value of the range.
  • extent (Expr) – The extent of the range.
Returns:

rng – The constructed range.

Return type:

Range

tvm.make.static_cast(dtype, expr)

Cast expr to dtype.

If expr is scalar and dtype is a corresponding vector type, a Broadcast is generated. Otherwise it is a Cast.

Parameters:
  • dtype (str) – The target data type.
  • expr (Expr) – The expression to be casted.
Returns:

casted – The casted expression.

Return type:

Expr