tvm.ir

Common data structures across all IR variants.

Classes:

SourceName(name)

A identifier for a source location.

Span(source_name, line, end_line, column, ...)

Specifies a location in a source program.

Node

Base class of all IR Nodes, implements astext function.

EnvFunc

Environment function.

Type

The base class of all types.

TypeKind(value)

Possible kinds of TypeVars.

PrimType(dtype)

Primitive data type in the low level IR

PointerType(element_type[, storage_scope])

PointerType used in the low-level TIR.

TypeVar(name_hint[, kind])

Type parameter in functions.

GlobalTypeVar(name_hint[, kind])

A global type variable that is used for defining new types or type aliases.

TupleType(fields)

The type of tuple values.

TypeConstraint

Abstract class representing a type constraint.

FuncType(arg_types, ret_type[, type_params, ...])

Function type.

IncompleteType([kind])

Incomplete type during type inference.

RelayRefType(value)

Reference Type in relay.

TensorType(shape[, dtype])

A concrete TensorType in Relay.

TensorAffineType(scale, zero_point, dtype[, ...])

The quantized type of a tensor, with scale, zero point, and datatype

TupleAffineType(types)

Affine types of a node with multiple outputs

TypeCall(func, args)

Type function application.

TypeRelation(func, args, num_inputs, attrs)

User defined type relation, it is an input-output relation on types.

BaseExpr

Base class of all the expressions.

PrimExpr

Base class of all primitive expressions.

RelayExpr

Base class of all non-primitive expressions.

GlobalVar(name_hint[, type_annot])

A global variable in the IR.

Range(begin[, end, span])

Represent a range in TVM.

Op()

Primitive operator in the IR.

CallingConv(value)

Possible kinds of calling conventions.

BaseFunc

Base class of all functions.

Constructor(name_hint, inputs, belong_to)

Relay ADT constructor.

TypeData(header, type_vars, constructors)

Stores the definition for an Algebraic Data Type (ADT) in Relay.

IRModule([functions, type_definitions])

IRModule that holds functions and type definitions.

Attrs

Attribute node, which is mainly use for defining attributes of relay operators.

DictAttrs

Dictionary attributes.

Array

Array container of TVM.

Map

Map container of TVM.

PoolInfo()

PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into.

WorkspacePoolInfo(pool_name, targets[, ...])

WorkspacePoolInfo object holds information related to RW memory pools where the statically sized allocate nodes will pooled into.

ConstantPoolInfo(pool_name, targets[, ...])

ConstantPoolInfo object holds information related to RO memory pools where the statically sized allocate nodes are pooled into.

WorkspaceMemoryPools(pools)

This object contains a list of WorkspacePoolInfo objects to be used as workspace memory in the compilation

ConstantMemoryPools(pools)

This object contains a list of ConstantPoolInfo objects to be used as read-only memory in the compilation

PoolInfoProperties([size_hint_bytes, ...])

PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into.

Functions:

load_json(json_str)

Load tvm object from json_str.

save_json(node)

Save tvm object as json string.

structural_equal(lhs, rhs[, map_free_vars])

Check structural equality of lhs and rhs.

assert_structural_equal(lhs, rhs[, ...])

Assert lhs and rhs are structurally equal to each other.

structural_hash(node[, map_free_vars])

Compute structural hash of node

register_op_attr(op_name, attr_key[, value, ...])

Register an operator property of an operator by name.

register_intrin_lowering(op_name, target, *)

Register Op lowering function

make_node(type_key, **kwargs)

Make a new IR node by its type key and fields
class tvm.ir.SourceName(name)

A identifier for a source location.

Parameters

name (str) – The name of the source.

class tvm.ir.Span(source_name, line, end_line, column, end_column)

Specifies a location in a source program.

Parameters

  • source (SourceName) – The source name.

  • lineno (int) – The line number.

  • col_offset (int) – The column offset of the location.

class tvm.ir.Node

Base class of all IR Nodes, implements astext function.

Methods:

astext([show_meta_data, annotate])

Get the text format of the expression.
astext(show_meta_data=True, annotate=None)

Get the text format of the expression.

Parameters

  • show_meta_data (bool) – Whether to include meta data section in the text if there is meta data.

  • annotate (Optional[Object->str]) – Optionally annotate function to provide additional information in the comment block.

Returns

text – The text format of the expression.

Return type

str

Notes

The meta data section is necessary to fully parse the text format. However, it can contain dumps that are big (e.g constant weights), so it can be helpful to skip printing the meta data section.

class tvm.ir.EnvFunc

Environment function.

This is a global function object that can be serialized by its name.

Methods:

get(name)

Get a static env function
static get(name)

Get a static env function

Parameters

name (str) – The name of the function.

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

Object

tvm.ir.save_json(node)

Save tvm object as json string.

Parameters

node (Object) – A TVM object to be saved.

Returns

json_str – Saved json string.

Return type

str

tvm.ir.structural_equal(lhs, rhs, map_free_vars=False)

Check structural equality of lhs and rhs.

The structural equality is recursively defined in the DAG of IRNodes. There are two kinds of nodes:

  • Graph node: a graph node in lhs can only be mapped as equal to one and only one graph node in rhs.

  • Normal node: equality is recursively defined without the restriction of graph nodes.

Vars(tir::Var, TypeVar) and non-constant relay expression nodes are graph nodes. For example, it means that %1 = %x + %y; %1 + %1 is not structurally equal to %1 = %x + %y; %2 = %x + %y; %1 + %2 in relay.

A var-type node(e.g. tir::Var, TypeVar) can be mapped as equal to another var with the same type if one of the following condition holds:

  • They appear in a same definition point(e.g. function argument).

  • They points to the same VarNode via the same_as relation.

  • They appear in a same usage point, and map_free_vars is set to be True.

The rules for var are used to remap variables occurs in function arguments and let-bindings.

Parameters

  • lhs (Object) – The left operand.

  • rhs (Object) – The left operand.

  • map_free_vars (bool) – Whether or not shall we map free vars that does not bound to any definitions as equal to each other.

Returns

result – The comparison result.

Return type

bool

See also

structural_hash, assert_strucural_equal

tvm.ir.assert_structural_equal(lhs, rhs, map_free_vars=False)

Assert lhs and rhs are structurally equal to each other.

Parameters

  • lhs (Object) – The left operand.

  • rhs (Object) – The left operand.

  • map_free_vars (bool) – Whether or not shall we map free vars that does not bound to any definitions as equal to each other.

:raises ValueError : if assertion does not hold.:

See also

structural_equal

tvm.ir.structural_hash(node, map_free_vars=False)

Compute structural hash of node

The structural hash value is recursively defined in the DAG of IRNodes. There are two kinds of nodes:

  • Normal node: the hash value is defined by its content and type only.

  • Graph node: each graph node will be assigned a unique index ordered by the first occurence during the visit. The hash value of a graph node is combined from the hash values of its contents and the index.

structural_hash is made to be concistent with structural_equal. If two nodes are structurally equal to each other, then their structural hash (with the same map_free_vars option) should be equal to each other as well.

If the structural hash of two nodes equals to each other, then it is highly likely(except for rare hash value collison cases) that the two nodes are structurally equal to each other.

Parameters

  • node (Object) – The input to be hashed.

  • map_free_vars (bool) – If map_free_vars is set to true, we will hash free variables by the order of their occurences. Otherwise, we will hash by their in-memory pointer address.

Returns

result – The hash result

Return type

int

See also

structrual_equal

class tvm.ir.Type

The base class of all types.

Methods:

same_as(other)

Compares two Relay types by referential equality.
same_as(other)

Compares two Relay types by referential equality.

class tvm.ir.TypeKind(value)

Possible kinds of TypeVars.

class tvm.ir.PrimType(dtype)

Primitive data type in the low level IR

Parameters

dtype (str) – The runtime data type relates to the primtype.

class tvm.ir.PointerType(element_type, storage_scope='')

PointerType used in the low-level TIR.

Parameters

  • element_type (tvm.ir.Type) – The type of pointer’s element.

  • storage_scope (str) – The storage scope into which the pointer addresses.

class tvm.ir.TypeVar(name_hint, kind=TypeKind.Type)

Type parameter in functions.

A type variable represents a type placeholder which will be filled in later on. This allows the user to write functions which are generic over types.

Parameters

  • name_hint (str) – The name of the type variable. This name only acts as a hint, and is not used for equality.

  • kind (Optional[TypeKind]) – The kind of the type parameter.

class tvm.ir.GlobalTypeVar(name_hint, kind=TypeKind.AdtHandle)

A global type variable that is used for defining new types or type aliases.

Parameters

  • name_hint (str) – The name of the type variable. This name only acts as a hint, and is not used for equality.

  • kind (Optional[TypeKind]) – The kind of the type parameter.

class tvm.ir.TupleType(fields)

The type of tuple values.

Parameters

fields (List[Type]) – The fields in the tuple

class tvm.ir.TypeConstraint

Abstract class representing a type constraint.

class tvm.ir.FuncType(arg_types, ret_type, type_params=None, type_constraints=None)

Function type.

A function type consists of a list of type parameters to enable the definition of generic functions, a set of type constraints which we omit for the time being, a sequence of argument types, and a return type.

We can informally write them as: forall (type_params), (arg_types) -> ret_type where type_constraints

Parameters

  • arg_types (List[tvm.relay.Type]) – The argument types

  • ret_type (tvm.relay.Type) – The return type.

  • type_params (Optional[List[tvm.relay.TypeVar]]) – The type parameters

  • type_constraints (Optional[List[tvm.relay.TypeConstraint]]) – The type constraints.

class tvm.ir.IncompleteType(kind=TypeKind.Type)

Incomplete type during type inference.

kindOptional[TypeKind]

The kind of the incomplete type.

class tvm.ir.RelayRefType(value)

Reference Type in relay.

Parameters

value (Type) – The value type.

class tvm.ir.TensorType(shape, dtype='float32')

A concrete TensorType in Relay.

This is the type assigned to tensors with a known dtype and shape. For example, a tensor of float32 and (5, 5).

Parameters

  • shape (List[tvm.ir.PrimExpr]) – The shape of the Tensor

  • dtype (Optional[str]) – The content data type.

Attributes:

concrete_shape

Get shape of the type as concrete tuple of int.
property concrete_shape

Get shape of the type as concrete tuple of int.

Returns

shape – The concrete shape of the Type.

Return type

List[int]

:raises TypeError : If the shape is symbolic:

class tvm.ir.TensorAffineType(scale, zero_point, dtype, axis=- 1)

The quantized type of a tensor, with scale, zero point, and datatype

The real space value is calculated as x = x_q * scale + zero_point

Parameters

  • scale (Expr) – The scale

  • zero_point (Expr) – The zero_point

  • dtype (str) – The content data type.

  • axis (int) – The axis for per-channel quantization.

class tvm.ir.TupleAffineType(types)

Affine types of a node with multiple outputs

Parameters

types (List[TensorAffineType]) – The shape of the Tensor

class tvm.ir.TypeCall(func, args)

Type function application.

Parameters
Returns

type_call – The type function application.

Return type

TypeCall

class tvm.ir.TypeRelation(func, args, num_inputs, attrs)

User defined type relation, it is an input-output relation on types.

TypeRelation is more generalized than TypeCall as it allows inference

of both inputs and outputs.

Parameters

  • func (EnvFunc) – User defined relation function.

  • args ([tvm.ir.Type]) – List of types to the func.

  • num_inputs (int) – Number of input arguments in args, this act as a hint for type inference.

  • attrs (Attrs) – The attribute attached to the relation information

Returns

type_relation – The type relation.

Return type

tvm.ir.TypeRelation

class tvm.ir.BaseExpr

Base class of all the expressions.

class tvm.ir.PrimExpr

Base class of all primitive expressions.

PrimExpr is used in the low-level code optimizations and integer analysis.

class tvm.ir.RelayExpr

Base class of all non-primitive expressions.

Attributes:

checked_type

Get the checked type of tvm.relay.Expr.
property checked_type

Get the checked type of tvm.relay.Expr.

Returns

checked_type – The checked type.

Return type

tvm.relay.Type

class tvm.ir.GlobalVar(name_hint, type_annot=None)

A global variable in the IR.

GlobalVar is used to refer to the global functions stored in the IRModule.

Parameters

name_hint (str) – The name of the variable.

class tvm.ir.Range(begin, end=None, span=None)

Represent a range in TVM.

You do not need to create a Range explicitly. Python lists and tuples will be converted automatically to a Range in API functions.

Parameters

  • begin (PrimExpr) – The begin value of the range when end is None. Otherwise it is the length of the range.

  • end (Optional[PrimExpr]) – The end value of the range.

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

Note

The constructor creates the range [begin, end) if the end argument is not None. Otherwise, it creates [0, begin).

Methods:

from_min_extent(min_value, extent[, span])

Construct a Range by min and extent.
static from_min_extent(min_value, extent, span=None)

Construct a Range by min and extent.

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

Parameters

  • min_value (PrimExpr) – The minimum value of the range.

  • extent (PrimExpr) – The extent of the range.

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

Returns

rng – The constructed range.

Return type

Range

class tvm.ir.Op

Primitive operator in the IR.

Methods:

get(op_name)

Get the Op for a given name

get_attr(attr_name)

Get additional attribute about the operator.

has_attr(attr_name)

Check whether the operator has additional attribute.

set_attr(attr_name, value[, plevel])

Set attribute about the operator.

reset_attr(attr_name)

Reset attribute about the operator.

add_type_rel(rel_name[, type_rel_func])

Attach the type function corresponding to the return type.

add_argument(name, type, description)

Add arguments information to the function.

set_support_level(level)

Set the support level of op.

set_num_inputs(n)

Set the support level of op.

set_attrs_type_key(key)

Set the attribute type key of op.

list_op_names()

List all the op names in the op registry.
static get(op_name)

Get the Op for a given name

Parameters

op_name (str) – The operator name

Returns

op – The op of the corresponding name

Return type

Op

get_attr(attr_name)

Get additional attribute about the operator.

Parameters

attr_name (str) – The attribute name.

Returns

value – The attribute value

Return type

object

has_attr(attr_name)

Check whether the operator has additional attribute.

Parameters

attr_name (str) – The attribute name.

Returns

value – Whether the operator has additional attribute

Return type

bool

set_attr(attr_name, value, plevel=10)

Set attribute about the operator.

Parameters

  • attr_name (str) – The attribute name

  • value (object) – The attribute value

  • plevel (int) – The priority level

reset_attr(attr_name)

Reset attribute about the operator.

Parameters

attr_name (str) – The attribute name

add_type_rel(rel_name, type_rel_func=None)

Attach the type function corresponding to the return type.

Parameters

  • rel_name (str) – The type relation name to register.

  • type_rel_func (Optional[function (args: List[Type], attrs: Attrs) -> Type]) –

    The backing relation function which can solve an arbitrary relation on variables. Differences with type_rel_func in C++:

    1. When type_rel_func is not None

      1. OpAddTypeRel on C++ side will adjust type_rel_func with TypeReporter to calling convention of relay type system.

      2. type_rel_func returns output argument’s type, return None means can’t infer output’s type.

      3. only support single output operators for now, the last argument is output tensor.

    2. when type_rel_func is None, will call predefined type_rel_funcs in relay

      according to tvm.relay.type_relation. + rel_name.

add_argument(name, type, description)

Add arguments information to the function.

Parameters

  • name (str) – The argument name.

  • type (str) – The argument type.

  • description (str) – The argument description.

set_support_level(level)

Set the support level of op.

Parameters

level (int) – The support level.

set_num_inputs(n)

Set the support level of op.

Parameters

n (int) – The input number.

set_attrs_type_key(key)

Set the attribute type key of op.

Parameters

key (str) – The type key.

static list_op_names()

List all the op names in the op registry.

Returns

value – The registered op names

Return type

List[str]

tvm.ir.register_op_attr(op_name, attr_key, value=None, level=10)

Register an operator property of an operator by name.

Parameters

  • op_name (str) – The name of operator

  • attr_key (str) – The attribute name.

  • value (object, optional) – The value to set

  • level (int, optional) – The priority level

Returns

fregister – Register function if value is not specified.

Return type

function

tvm.ir.register_intrin_lowering(op_name, target, *, f=None, level=10)

Register Op lowering function

Parameters

  • op_name (str) – The op name

  • target (str) – The target string for given intrinsic lowering function

  • f (function, optional) – The function to be registered.

  • level (int) – The priority level

Returns

fregister – Register op lowering function if f is not specified.

Return type

function

class tvm.ir.CallingConv(value)

Possible kinds of calling conventions.

class tvm.ir.BaseFunc

Base class of all functions.

Attributes:

attrs

Return the attrs member of the function.

Methods:

with_attr(attr_key_or_dict[, attr_value])

Create a new copy of the function and update the attribute.
property attrs

Return the attrs member of the function.

with_attr(attr_key_or_dict, attr_value=None)

Create a new copy of the function and update the attribute.

Parameters

  • attr_key_or_dict (Union[str, dict]) – The attribute key to use or a dict containing multiple key value pairs.

  • attr_value (Object) – The new attribute value.

Returns

func – A new copy of the function

Return type

Function

class tvm.ir.Constructor(name_hint, inputs, belong_to)

Relay ADT constructor.

Parameters

  • name_hint (str) – Name of constructor (only a hint).

  • inputs (List[Type]) – Input types.

  • belong_to (GlobalTypeVar) – Denotes which ADT the constructor belongs to.

class tvm.ir.TypeData(header, type_vars, constructors)

Stores the definition for an Algebraic Data Type (ADT) in Relay.

Note that ADT definitions are treated as type-level functions because the type parameters need to be given for an instance of the ADT. Thus, any global type var that is an ADT header needs to be wrapped in a type call that passes in the type params.

Parameters

  • header (GlobalTypeVar) – The name of the ADT. ADTs with the same constructors but different names are treated as different types.

  • type_vars (List[TypeVar]) – Type variables that appear in constructors.

  • constructors (List[Constructor]) – The constructors for the ADT.

class tvm.ir.IRModule(functions=None, type_definitions=None)

IRModule that holds functions and type definitions.

IRModule is the basic unit for all IR transformations across the stack.

Parameters

functions (Optional[dict]) – Map of global var to BaseFunc

Methods:

update(other)

Insert functions in another Module to current one.

update_func(var, func)

Update the function corresponding to a global variable in the module.

get_global_var(name)

Get a global variable in the function by name.

get_global_vars()

Collect all global vars defined in this module.

get_global_type_vars()

Collect all global type vars defined in this module.

get_global_type_var(name)

Get a global type variable in the function by name.

get_constructor(tag)

Look up an ADT constructor by tag.

from_expr(expr[, functions, type_defs])

Construct a module from a standalone expression.

script([tir_prefix, show_meta])

Print IRModule into TVMScript

get_attr(attr_key)

Get the IRModule attribute.

with_attr(attr_key, attr_value)

Copy the IRModule and add an attribute to it.
update(other)

Insert functions in another Module to current one.

Parameters

other (IRModule) – The module to merge into the current Module.

update_func(var, func)

Update the function corresponding to a global variable in the module.

Parameters

  • var (GlobalVar) – The global variable.

  • func (tvm.relay.Function) – The function to be inserted.

get_global_var(name)

Get a global variable in the function by name.

Parameters

name (str) – The name of the global variable.

Returns

global_var – The global variable mapped to name.

Return type

GlobalVar

Raises

tvm.error.TVMError if we cannot find corresponding global var. –

get_global_vars()

Collect all global vars defined in this module.

Returns

global_vars – An array of global vars.

Return type

Array[GlobalVar]

get_global_type_vars()

Collect all global type vars defined in this module.

Returns

global_type_vars – An array of global type vars.

Return type

Array[GlobalTypeVar]

get_global_type_var(name)

Get a global type variable in the function by name.

Parameters

name (str) – The name of the global type variable.

Returns

global_type_var – The global variable mapped to name.

Return type

GlobalTypeVar

Raises

tvm.error.TVMError if we cannot find corresponding global type var. –

get_constructor(tag)

Look up an ADT constructor by tag.

Parameters

tag (int) – The tag for a constructor.

Returns

constructor – The constructor associated with the given tag,

Return type

Constructor

Raises

tvm.error.TVMError if the corresponding constructor cannot be found. –

static from_expr(expr, functions=None, type_defs=None)

Construct a module from a standalone expression.

Parameters

  • expr (RelayExpr) – The starting expression

  • global_funcs (Optional[dict]) – Map of global vars to function definitions

  • type_defs (Optional[dict]) – Map of global type vars to type definitions

Returns

mod – A module containing the passed definitions, where expr is set as the entry point (wrapped in a function if necessary)

Return type

Module

script(tir_prefix: str = 'T', show_meta: bool = False) → str

Print IRModule into TVMScript

Parameters

  • tir_prefix (str) – The tir namespace prefix

  • show_meta (bool) – Whether to show meta information

Returns

script – The TVM Script of the IRModule

Return type

str

get_attr(attr_key)

Get the IRModule attribute.

Parameters

attr_key (str) – The attribute key.

Returns

attr_value – Attribute value

Return type

Any

with_attr(attr_key, attr_value)

Copy the IRModule and add an attribute to it.

Parameters

  • attr_key (str) – The attribute key.

  • attr_value (Object) – The new attribute value.

Returns

mod – A new copy of the IRModule with the attribute

Return type

IRModule

class tvm.ir.Attrs

Attribute node, which is mainly use for defining attributes of relay operators.

Used by function registered in python side, such as compute, schedule and alter_layout. Attrs is passed as the first argument to these functions.

Methods:

list_field_info()

Get fields information

keys()

Get list of names in the attribute.

get_int_tuple(key)

Get a python int tuple of a key

get_int(key)

Get a python int value of a key

get_str(key)

Get a python int value of a key
list_field_info()

Get fields information

Returns

infos – List of field information

Return type

list of AttrFieldInfo

keys()

Get list of names in the attribute.

Returns

keys – List of keys

Return type

list of str

get_int_tuple(key)

Get a python int tuple of a key

Parameters

key (str) –

Returns

value

Return type

Tuple of int

get_int(key)

Get a python int value of a key

Parameters

key (str) –

Returns

value

Return type

int

get_str(key)

Get a python int value of a key

Parameters

key (str) –

Returns

value

Return type

int

class tvm.ir.DictAttrs

Dictionary attributes.

Methods:

keys()

Get list of names in the attribute.

items()

Get items from the map.
keys()

Get list of names in the attribute.

Returns

keys – List of keys

Return type

list of str

items()

Get items from the map.

tvm.ir.make_node(type_key, **kwargs)

Make a new IR 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 IR 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.ir.make_node("IntImm", dtype="int32", value=10)
assert isinstance(x, tvm.tir.IntImm)
assert x.value == 10
class tvm.ir.Array

Array container of TVM.

You do not need to create Array explicitly. Normally python list and tuple will be converted automatically to Array during tvm function call. You may get Array in return values of TVM function call.

class tvm.ir.Map

Map container of TVM.

You do not need to create Map explicitly. Normally python dict will be converted automatically to Map during tvm function call. You can use convert to create a dict[Object-> Object] into a Map

Methods:

items()

Get the items from the map

get(key[, default])

Get an element with a default value.
items()

Get the items from the map

get(key, default=None)

Get an element with a default value.

Parameters

  • key (object) – The attribute key.

  • default (object) – The default object.

Returns

value – The result value.

Return type

object

class tvm.ir.PoolInfo

PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into. This is a base class for WorkspacePoolInfo and ConstantPoolInfo.

class tvm.ir.WorkspacePoolInfo(pool_name: str, targets, pool_info_properties=None)

WorkspacePoolInfo object holds information related to RW memory pools where the statically sized allocate nodes will pooled into.

Parameters

  • pool_name (str) – The name of the memory pool

  • targets (list[Target]) – A list of targets which could access the pool

  • pool_info_properties (PoolInfoProperties) – The properties of the pool.

class tvm.ir.ConstantPoolInfo(pool_name: str, targets, constant_info_arr=None, pool_info_properties=None)

ConstantPoolInfo object holds information related to RO memory pools where the statically sized allocate nodes are pooled into.

Parameters

  • pool_name (str) – The name of the memory pool

  • targets (list[Target]) – describes which targets could access the pool

  • pool_info_properties (PoolInfoProperties) – The properties of the pool.

class tvm.ir.WorkspaceMemoryPools(pools: List[tvm.ir.memory_pools.PoolInfo])

This object contains a list of WorkspacePoolInfo objects to be used as workspace memory in the compilation

Parameters

pools (List[WorkspacePoolInfo]) – The list of ConstantPoolInfo objects to be used with the compilation

class tvm.ir.ConstantMemoryPools(pools: List[tvm.ir.memory_pools.ConstantPoolInfo])

This object contains a list of ConstantPoolInfo objects to be used as read-only memory in the compilation

Parameters

pools (List[ConstantPoolInfo]) – The list of ConstantPoolInfo objects to be used with the compilation

class tvm.ir.PoolInfoProperties(size_hint_bytes: Optional[int] = - 1, clock_frequency_hz: Optional[int] = - 1, read_bandwidth_bytes_per_cycle: Optional[int] = - 1, write_bandwidth_bytes_per_cycle: Optional[int] = - 1, read_latency_cycles: Optional[int] = 0, write_latency_cycles: Optional[int] = 0, target_burst_bytes=None)

PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into.

Parameters

  • size_hint_bytes (Optional[int]) – The expected size hint to be used by the allocator. The default value would be -1 which means the pool is not size restricted.

  • clock_frequency_hz (Optional[int]) – The clock frequency that the memory pool runs at in Hz. If not specified/known, this will default to -1 indicating it hasn’t been defined.

  • read_bandwidth_bytes_per_cycle (Optional[int]) – The read bandwidth of the memory pool in bytes/cycle. If not specified/known, this will default to -1 indicating it hasn’t been defined.

  • write_bandwidth_bytes_per_cycle (Optional[int]) – The write bandwidth of the memory pool in bytes/cycle. If not specified/known, this will default to -1 indicating it hasn’t been defined.

  • read_latency_cycles (Optional[int]) – The read latency of the memory pool in cycles. If not specified/known, this will default to 0.

  • write_latency_cycles (Optional[int]) – The write latency of the memory pool in cycles. If not specified/known, this will default to 0.

  • target_burst_bytes (Optional[Union[Dict[Target, int], None]]) – The burst length of the memory pool in bytes per target. If not specified/known for a given target, a burst length of 1 byte will be assumed.

tvm.instrument

Common pass instrumentation across IR variants.

Classes:

PassInstrument()

A pass instrument implementation.

PassTimingInstrument()

A wrapper to create a passes time instrument that implemented in C++

Functions:

pass_instrument([pi_cls])

Decorate a pass instrument.
class tvm.instrument.PassInstrument

A pass instrument implementation.

To use, a user class can either subclass from PassInstrument directly, or can apply the pass_instrument() wrapper. In either case, the enter_pass_ctx, exit_pass_ctx, should_run, run_before_pass, and run_after_pass methods can be defined to adjust the instrument’s behavior. See the no-op implementations in this class definition for more information on each.

Methods:

enter_pass_ctx()

Called when entering the instrumented context.

exit_pass_ctx()

Called when exiting the instrumented context.

should_run(mod, info)

Determine whether to run the pass or not.

run_before_pass(mod, info)

Instrument before the pass runs.

run_after_pass(mod, info)

Instrument after the pass runs.
enter_pass_ctx()

Called when entering the instrumented context.

Returns

Return type

None

exit_pass_ctx()

Called when exiting the instrumented context.

Returns

Return type

None

should_run(mod, info)

Determine whether to run the pass or not.

Called once for each pass that is run while the instrumented context is active.

Parameters
Returns

should_run – True to run the pass, or False to skip the pass.

Return type

bool

run_before_pass(mod, info)

Instrument before the pass runs.

Called once for each pass that is run while the instrumented context is active.

Parameters
Returns

Return type

None

run_after_pass(mod, info)

Instrument after the pass runs.

Called once for each pass that is run while the instrumented context is active.

Parameters
Returns

Return type

None

class tvm.instrument.PassTimingInstrument

A wrapper to create a passes time instrument that implemented in C++

Methods:

render()

Retrieve rendered time profile result :returns: string -- The rendered string result of time profiles :rtype: string
static render()

Retrieve rendered time profile result :returns: string – The rendered string result of time profiles :rtype: string

Examples

timing_inst = PassTimingInstrument()
with tvm.transform.PassContext(instruments=[timing_inst]):
    relay_mod = relay.transform.InferType()(relay_mod)
    relay_mod = relay.transform.FoldScaleAxis()(relay_mod)
    # before exiting the context, get profile results.
    profiles = timing_inst.render()
tvm.instrument.pass_instrument(pi_cls=None)

Decorate a pass instrument.

Parameters

pi_class (class) – Instrument class. See example below.

Examples

@tvm.instrument.pass_instrument
class SkipPass:
    def __init__(self, skip_pass_name):
        self.skip_pass_name = skip_pass_name

    # Uncomment to customize
    # def enter_pass_ctx(self):
    #    pass

    # Uncomment to customize
    # def exit_pass_ctx(self):
    #    pass

    # If pass name contains keyword, skip it by return False. (return True: not skip)
    def should_run(self, mod, pass_info)
        if self.skip_pass_name in pass_info.name:
            return False
        return True

    # Uncomment to customize
    # def run_before_pass(self, mod, pass_info):
    #    pass

    # Uncomment to customize
    # def run_after_pass(self, mod, pass_info):
    #    pass

skip_annotate = SkipPass("AnnotateSpans")
with tvm.transform.PassContext(instruments=[skip_annotate]):
    tvm.relay.build(mod, "llvm")

tvm.transform

Common pass infrastructure across IR variants.

Classes:

ModulePass

A pass that works on tvm.IRModule.

Pass

The base class of all passes.

PassContext([opt_level, required_pass, ...])

The basis where a Relay optimization/analysis runs on.

PassInfo(opt_level, name[, required])

The class contains the meta data required by a pass.

Sequential([passes, opt_level, name, required])

A pass that works on a sequence of pass objects.

Functions:

PrintIR([header, show_meta_data])

A special trace pass that prints the header and IR.

module_pass([pass_func, opt_level, name, ...])

Decorate a module pass.
class tvm.transform.ModulePass

A pass that works on tvm.IRModule. Users don’t need to interact with this class directly. Instead, a module pass should be created through module_pass, because the design of the module_pass API is flexible enough to handle the creation of a module pass in different manners. In addition, all members of a module pass can be accessed from the base class. The same rule applies to FunctionPass as well.

class tvm.transform.Pass

The base class of all passes. All methods here are just simple wrappers that are implemented in the backend. They are defined for users to conveniently interact with the base class.

Attributes:

info

Get the pass meta.
property info

Get the pass meta.

class tvm.transform.PassContext(opt_level=2, required_pass=None, disabled_pass=None, instruments=None, config=None)

The basis where a Relay optimization/analysis runs on. Each pass context contains a number of auxiliary information that is used to help an optimization pass. Such information includes the error reporter to record the errors of during the optimization, etc.

opt_levelOptional[int]

The optimization level of this pass.

required_passOptional[Union[List[str], Set[str], Tuple[str]]]

The list of passes that are required by a certain pass.

disabled_passOptional[Union[List[str], Set[str], Tuple[str]]]

The list of passes that are disabled.

instrumentsOptional[Sequence[PassInstrument]]

The list of pass instrument implementations.

configOptional[Dict[str, Object]]

Additional configurations for specific passes.

Methods:

override_instruments(instruments)

Override instruments within this PassContext.

current()

Return the current pass context.

list_configs()

List all registered PassContext configuration names and metadata.
override_instruments(instruments)

Override instruments within this PassContext.

If there are existing instruments, their exit_pass_ctx callbacks are called. Then switching to new instruments and calling new enter_pass_ctx callbacks.

instrumentsSequence[PassInstrument]

The list of pass instrument implementations.

static current()

Return the current pass context.

static list_configs()

List all registered PassContext configuration names and metadata.

Returns

configs

Return type

Dict[str, Dict[str, str]]

class tvm.transform.PassInfo(opt_level, name, required=None)

The class contains the meta data required by a pass. It is the container of information needed by running an optimization or analysis. This class can be extended by adding new members when more meta data is needed.

Parameters

  • opt_level (int) – The optimization level of this pass.

  • name (str) – The pass name.

  • required (List[str]) – The list of passes that are required by a certain pass.

tvm.transform.PrintIR(header='', show_meta_data=False)

A special trace pass that prints the header and IR.

Parameters

  • header (str) – The header to be displayed along with the dump.

  • show_meta_data (bool) – A boolean flag to indicate if meta data should be printed.

Returns

Return type

The pass

class tvm.transform.Sequential(passes=None, opt_level=0, name='sequential', required=None)

A pass that works on a sequence of pass objects. Multiple passes can be executed sequentially using this class.

Note that users can also provide a series of passes that they don’t want to apply when running a sequential pass. Pass dependency will be resolved in the backend as well.

Parameters

  • passes (Optional[List[Pass]]) – A sequence of passes candidate for optimization.

  • opt_level (Optional[int]) – The optimization level of this sequential pass. The opt_level of a default sequential pass is set to 0. Note that some of the passes within the Sequantial may still not be executed if their opt_level is higher than the provided opt_level.

  • name (Optional[str]) – The name of the sequential pass.

  • required (Optional[List[str]]) – The list of passes that the sequential pass is dependent on.

tvm.transform.module_pass(pass_func=None, opt_level=None, name=None, required=None)

Decorate a module pass.

This function returns a callback when pass_func is provided. Otherwise, it serves a decorator function.

pass_func can also be a class type with a method transform_module. This function will create a decorated ModulePass using transform_module as the pass function.

Parameters

  • pass_func (Optional[Callable[(Module, PassContext) ->Module]]) – The transformation function or class.

  • opt_level (int) – The optimization level of this module pass.

  • name (Optional[str]) – The name of the module pass. The name could be empty. In this case, the name of the optimization function will be used as the pass name.

  • required (Optional[List[str]]) – The list of passes that the module pass is dependent on.

Returns

create_module_pass – A decorator will be returned if pass_func is not provided, otherwise return the decorated result. The returned decorator has two behaviors depending on the input: A new ModulePass will be returned when we decorate a pass function. A new ModulePass class will be returned when we decorate a class type.

Return type

Union[Callable, ModulePass]

Examples

The following code block decorates a module pass class.

@relay.transform.module_pass
class CustomPipeline:
    def __init__(self, enable_fold):
        self.enable_fold = enable_fold
        self.cse = relay.transform.EliminateCommonSubexpr()
        self.const_fold = relay.transform.FoldConstant()

    def transform_module(self, mod, ctx):
        mod = self.cse(mod, ctx)
        if self.enable_fold:
            mod = self.const_fold(mod, ctx)
        return mod

# create an instance of customized pipeline
pipeline = CustomPipeline(enable_fold=False)
assert isinstance(pipeline, transform.ModulePass)
# run the pipeline.
output_module = pipeline(input_module)

The following code creates a module pass by decorating a user defined transform function.

@relay.transform.module_pass(opt_level=2)
def transform(mod, ctx):
    tp = relay.TensorType((10,), "float32")
    x = relay.var("x", tp)
    gv = relay.GlobalVar("var")
    func = relay.Function([x], relay.abs(x))
    new_mod = tvm.IRModule({gv: func})
    new_mod.update(mod)
    return new_mod

module_pass = transform
assert isinstance(module_pass, transform.ModulePass)
assert module_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = module_pass(m)
# Now a function abs should be added to the module m.