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(handle)

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(handle)
class tvm.expr.And(handle)
class tvm.expr.BinaryOpExpr(handle)
class tvm.expr.Broadcast(handle)
class tvm.expr.Call(handle)
Extern = 0
ExternCPlusPlus = 1
Halide = 3
Intrinsic = 4
PureExtern = 2
PureIntrinsic = 5
class tvm.expr.Cast(handle)
class tvm.expr.CmpExpr(handle)
class tvm.expr.ConstExpr(handle)
class tvm.expr.Div(handle)
class tvm.expr.EQ(handle)
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 = <method-wrapper '__eq__' of type object at 0x8f8740>
class tvm.expr.Expr(handle)

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(handle)
class tvm.expr.GE(handle)
class tvm.expr.GT(handle)
class tvm.expr.IntImm(handle)
class tvm.expr.LE(handle)
class tvm.expr.LT(handle)
class tvm.expr.Let(handle)
class tvm.expr.Load(handle)
class tvm.expr.LogicalExpr(handle)
class tvm.expr.Max(handle)
class tvm.expr.Min(handle)
class tvm.expr.Mod(handle)
class tvm.expr.Mul(handle)
class tvm.expr.NE(handle)
class tvm.expr.Not(handle)
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 = <method-wrapper '__eq__' of type object at 0x8f8740>
class tvm.expr.Or(handle)
class tvm.expr.Ramp(handle)
class tvm.expr.Reduce(handle)
class tvm.expr.Select(handle)
class tvm.expr.Shuffle(handle)
class tvm.expr.StringImm(handle)
class tvm.expr.Sub(handle)
class tvm.expr.UIntImm(handle)
class tvm.expr.Var(handle)

Symbolic variable.

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(handle)
class tvm.stmt.AssertStmt(handle)
class tvm.stmt.AttrStmt(handle)
class tvm.stmt.Block(handle)
class tvm.stmt.Evaluate(handle)
class tvm.stmt.For(handle)
Parallel = 1
Serial = 0
Unrolled = 3
Vectorized = 2
class tvm.stmt.Free(handle)
class tvm.stmt.IfThenElse(handle)
class tvm.stmt.LetStmt(handle)
class tvm.stmt.Prefetch(handle)
class tvm.stmt.ProducerConsumer(handle)
class tvm.stmt.Provide(handle)
class tvm.stmt.Realize(handle)
class tvm.stmt.Stmt(handle)
class tvm.stmt.Store(handle)

tvm.ir_pass

tvm.ir_pass.Inline(*args) TVM PackedFunc Inline.
tvm.ir_pass.Simplify(*args) TVM PackedFunc Simplify.
tvm.ir_pass.ConvertSSA(*args) TVM PackedFunc ConvertSSA.
tvm.ir_pass.VerifySSA(*args) TVM PackedFunc VerifySSA.
tvm.ir_pass.CanonicalSimplify(*args) TVM PackedFunc CanonicalSimplify.
tvm.ir_pass.StorageFlatten(*args) TVM PackedFunc StorageFlatten.
tvm.ir_pass.VectorizeLoop(*args) TVM PackedFunc VectorizeLoop.
tvm.ir_pass.UnrollLoop(*args) TVM PackedFunc UnrollLoop.
tvm.ir_pass.ThreadSync(*args) TVM PackedFunc ThreadSync.
tvm.ir_pass.StorageRewrite(*args) TVM PackedFunc StorageRewrite.
tvm.ir_pass.MakeAPI(*args) TVM PackedFunc MakeAPI.
tvm.ir_pass.SplitHostDevice(*args) TVM PackedFunc SplitHostDevice.
tvm.ir_pass.InjectVirtualThread(*args) TVM PackedFunc InjectVirtualThread.
tvm.ir_pass.LoopPartition(*args) TVM PackedFunc LoopPartition.
tvm.ir_pass.RemoveNoOp(*args) TVM PackedFunc RemoveNoOp.
tvm.ir_pass.SplitPipeline(*args) TVM PackedFunc SplitPipeline.
tvm.ir_pass.LowerThreadAllreduce(*args) TVM PackedFunc LowerThreadAllreduce.
tvm.ir_pass.LowerIntrin(*args) TVM PackedFunc LowerIntrin.
tvm.ir_pass.LowerTVMBuiltin(*args) TVM PackedFunc LowerTVMBuiltin.
tvm.ir_pass.NarrowChannelAccess(*args) 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.

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

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