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.

handle
class tvm.expr.And(a, b)

And node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.BinaryOpExpr
handle
class tvm.expr.Broadcast(value, lanes)

Broadcast node.

Parameters
  • value (Expr) – The value of the expression.

  • lanes (int) – The lanes of the expression.

handle
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
handle
class tvm.expr.Cast(dtype, value)

Cast expression.

Parameters
  • dtype (str) – The data type

  • value (Expr) – The value of the function.

handle
class tvm.expr.CmpExpr
handle
class tvm.expr.ConstExpr
handle
class tvm.expr.Div(a, b)

Div node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.EQ(a, b)

EQ node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

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

Return self==value.

class tvm.expr.Expr

Base class of all tvm Expressions

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

handle
class tvm.expr.FloorDiv(a, b)

FloorDiv node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.FloorMod(a, b)

FloorMod node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.GE(a, b)

GE node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.GT(a, b)

GT node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.IntImm(dtype, value)

Int constant.

Parameters
  • dtype (str) – The data type

  • value (int) – The constant value.

handle
class tvm.expr.LE(a, b)

LE node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.LT(a, b)

LT node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

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

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

handle
class tvm.expr.LogicalExpr
handle
class tvm.expr.Max(a, b)

Max node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.Min(a, b)

Min node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.Mod(a, b)

Mod node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.Mul(a, b)

Mul node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.NE(a, b)

NE node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.Not(a)

Not node.

Parameters

a (Expr) – The input value

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

Return self==value.

class tvm.expr.Or(a, b)

Or node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

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

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

handle
class tvm.expr.Select(condition, true_value, false_value)

Select node.

Note

Select may compute both true_value and false_value. Use tvm.if_then_else instead if you want to get a conditional expression that only evaluates the correct branch.

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.

handle
class tvm.expr.Shuffle(vectors, indices)

Shuffle node.

Parameters
  • vectors (Array of Expr) – The vectors

  • indices (Array of indices) – The indices

handle
class tvm.expr.StringImm(value)

String constant.

Parameters

value (str) – The value of the function.

handle
class tvm.expr.Sub(a, b)

Sub node.

Parameters
  • a (Expr) – The left hand operand.

  • b (Expr) – The right hand operand.

handle
class tvm.expr.UIntImm(dtype, value)

UInt constant.

Parameters
  • dtype (str) – The data type

  • value (int) – The constant value.

handle
class tvm.expr.Var(name, dtype)

Symbolic variable.

Parameters
  • name (str) – The name

  • dtype (int) – The data type

handle

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.

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

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

handle
class tvm.stmt.Block(first, rest)

Block node.

Parameters
  • first (Stmt) – The first statement.

  • rest (Stmt) – The following statement.

handle
class tvm.stmt.Evaluate(value)

Evaluate node.

Parameters

value (Expr) – The expression to be evalued.

handle
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
handle
class tvm.stmt.Free(buffer_var)

Free node.

Parameters

buffer_var (Var) – The buffer variable.

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

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

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

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

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

handle
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

handle
class tvm.stmt.Stmt
handle
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.

handle
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.SkipVectorize

TVM PackedFunc SkipVectorize.

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/ir_pass.h” for the function signature and “src/api/api_pass.cc” for the PackedFunc’s body 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