tvm.ir

Common data structures across all IR variants.

Classes:

`Attrs`	Attribute node, which is mainly use for defining attributes of operators.
`DictAttrs`(__dict__)	Dictionary attributes.
`EnvFunc`(name, func)	Environment function.
`Node`	Base class of all IR Nodes.
`SourceName`(name)	A identifier for a source location.
`Span`(source_name, line, end_line, column, ...)	Specifies a location in a source program.
`SequentialSpan`(spans)	A sequence of source spans
`Array`(input_list)	Array container that represents a sequence of values in the FFI.
`Map`(input_dict)	Map container.
`BaseExpr`([span])	Base class of all the expressions.
`GlobalVar`(name_hint)	A global variable in the IR.
`PrimExpr`(dtype[, span])	Base class of all primitive expressions.
`Range`(begin[, end, span])	Represent a range in TVM.
`RelaxExpr`(struct_info_[, span])	Base class of all non-primitive expressions.
`BaseFunc`(struct_info_, attrs[, span])	Base class of all functions.
`CallingConv`(value)	Possible kinds of calling conventions.
`GlobalInfo`	Base node for all global info that can appear in the IR
`DummyGlobalInfo`()
`VDevice`([target, vdevice_id, memory_scope])
`IRModule`([functions, attrs, global_infos])	IRModule that holds functions and type definitions.
`Op`()	Primitive operator in the IR.
`FuncType`(arg_types, ret_type)	Function type.
`PointerType`(element_type[, storage_scope])	PointerType used in the low-level TIR.
`PrimType`(dtype)	Primitive data type in the low level IR
`TupleType`(fields)	The type of tuple values.
`Type`([span])	The base class of all types.

Functions:

`make_node`(type_key, **kwargs)	Make a new IR node by its type key and fields
`assert_structural_equal`(lhs, rhs[, ...])	Assert lhs and rhs are structurally equal to each other.
`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.
`structural_hash`(node[, map_free_vars])	Compute structural hash of node
`register_intrin_lowering`(op_name, target, *)	Register Op lowering function
`register_op_attr`(op_name, attr_key[, value, ...])	Register an operator property of an operator by name.

class tvm.ir.Attrs

Attribute node, which is mainly use for defining attributes of 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:

`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

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

Dictionary attributes.

Methods:

`keys`()	Get list of names in the attribute.
`get`(key[, default])	Get an element with a default value.
`items`()	Get items from the map.

keys()

Get list of names in the attribute.

Returns:: keys – List of keys
Return type:: list of str

get(key, default=None): Get an element with a default value.

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("ir.IntImm", dtype="int32", value=10, span=None)
assert isinstance(x, tvm.tirx.IntImm)
assert x.value == 10

class tvm.ir.EnvFunc(name, func)

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.

class tvm.ir.Node: Base class of all IR Nodes.

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.SequentialSpan(spans)

A sequence of source spans

This span is specific for an expression, which is from multiple expressions after an IR transform.

Parameters:: spans (Array) – The array of spans.

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.load_json(json_str) → Object

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) → str

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(tirx::Var, relax::Var) are graph nodes.

A var-type node(e.g. tirx::Var) 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 free variables (i.e. variables without a definition site) should be mapped as equal to each other.

Returns:

result – The comparison result.

Return type:

bool

See also

structural_hash, assert_strucural_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 occurrence 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 occurrences. 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.Array(input_list: Iterable[T])

Array container that represents a sequence of values in the FFI.

tvm_ffi.convert() will map python list/tuple to this class.

Parameters:: input_list – The list of values to be stored in the array.

Examples

import tvm_ffi

a = tvm_ffi.Array([1, 2, 3])
assert tuple(a) == (1, 2, 3)

Notes

For structural equality and hashing, use structural_equal and structural_hash APIs.

See also

tvm_ffi.convert()

class tvm.ir.Map(input_dict: Mapping[K, V])

Map container.

tvm_ffi.convert() will map python dict to this class.

Parameters:: input_dict – The dictionary of values to be stored in the map.

Examples

import tvm_ffi

amap = tvm_ffi.Map({"a": 1, "b": 2})
assert len(amap) == 2
assert amap["a"] == 1
assert amap["b"] == 2

Notes

For structural equality and hashing, use structural_equal and structural_hash APIs.

See also

tvm_ffi.convert()

Methods:

`keys`()	Return a dynamic view of the map's keys.
`values`()	Return a dynamic view of the map's values.
`items`()	Get the items from the map.
`get`()	Get an element with a default value.

keys() → KeysView[K]: Return a dynamic view of the map’s keys.

values() → ValuesView[V]: Return a dynamic view of the map’s values.

items() → ItemsView[K, V]: Get the items from the map.

get(key: K) → V | None

get(key: K, default: V | _DefaultT) → V | _DefaultT

Get an element with a default value.

Parameters:

key – The attribute key.
default – The default object.

Returns:

The result value.

Return type:

value

class tvm.ir.BaseExpr(span=<object object>): Base class of all the expressions.

class tvm.ir.GlobalVar(name_hint: str)

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.PrimExpr(dtype, span=<object object>)

Base class of all primitive expressions.

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

class tvm.ir.Range(begin: PrimExpr, end: PrimExpr | None = None, span: Span | None = 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 node 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: PrimExpr, extent: PrimExpr, span: Span | None = None) → Range

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 node in the source code.

Returns:

rng – The constructed range.

Return type:

Range

class tvm.ir.RelaxExpr(struct_info_, span=<object object>)

Base class of all non-primitive expressions.

Attributes:

struct_info

Get the struct info field

property struct_info: StructInfo | None

Get the struct info field

Returns:: struct_info – The struct info if available.
Return type:: tvm.relax.StructInfo

class tvm.ir.BaseFunc(struct_info_, attrs, span=<object object>)

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.
`with_attrs`(attr_map)	Copy the IRModule and add the given attribute map to it.
`without_attr`(attr_key)	Create a new copy of the function with an attribute without provided key.

property attrs: Return the attrs member of the function.

with_attr(attr_key_or_dict, attr_value=None) → BaseFunc

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:

BaseFunc

with_attrs(attr_map: DictAttrs | dict[str, Object]) → BaseFunc

Copy the IRModule and add the given attribute map to it. :param attr_map: The attribute map :type attr_map: Union[DictAttrs, Dict[str, Object]]

Returns:: func – A new copy of the function
Return type:: BaseFunc

without_attr(attr_key: str) → BaseFunc

Create a new copy of the function with an attribute without provided key.

Parameters:: attr_key (str) – The attribute key to delete from the attrubte pairs.
Returns:: func – A new copy of the function
Return type:: BaseFunc

class tvm.ir.CallingConv(value): Possible kinds of calling conventions.

class tvm.ir.GlobalInfo

Base node for all global info that can appear in the IR

Methods:

same_as(other)

Overload with structural equality.

same_as(other): Overload with structural equality.

class tvm.ir.DummyGlobalInfo

class tvm.ir.VDevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global')

class tvm.ir.IRModule(functions=None, attrs=None, global_infos=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:

`functions_items`()	Get items in self.functions.items() in alphabetical order.
`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.
`update_global_info`(name, global_info)	Update global info 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.
`replace_global_vars`(replacements)	Replace GlobalVar instances within the module
`from_expr`(expr[, functions])	Construct a module from a standalone expression.
`get_attr`(attr_key)	Get the IRModule attribute.
`with_attr`(attr_key, attr_value)	Copy the IRModule and add an attribute to it.
`without_attr`(attr_key)	Copy the IRModule and remove an attribute key and its associated value.
`with_attrs`(attr_map)	Copy the IRModule and add the given attribute map to it.

functions_items()

Get items in self.functions.items() in alphabetical order.

Returns:: items – The functions items.
Return type:: List[Tuple[GlobalVar, Function]]

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.ir.BaseFunc) – The function to be inserted.

update_global_info(name, global_info)

Update global info in the module

Parameters:

name (str) – The name for the global info.
global_info (List[GlobalInfo]) – The global info to be updated.

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]

replace_global_vars(replacements: dict[str | GlobalVar, str | GlobalVar]) → IRModule

Replace GlobalVar instances within the module

Replace GlobalVars within the IRModule. Since the IRModule may contain internal references to a GlobalVar, either in TIR or in Relax, this method should be used whenever replacing or renaming a GlobalVar.

Parameters:: replacements (Dict[Union[str, _expr.GlobalVar], Union[str, _expr.GlobalVar]]) – A dictionary where each key is a GlobalVar to be replaced, and the corresponding value is the GlobalVar with which to replace it.
Returns:: The updated module
Return type:: IRModule

static from_expr(expr, functions=None)

Construct a module from a standalone expression.

Parameters:

expr (RelaxExpr) – The starting expression
global_funcs (Optional[dict]) – Map of global vars to function 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

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

without_attr(attr_key: str) → IRModule

Copy the IRModule and remove an attribute key and its associated value. :param attr_key: The attribute key. :type attr_key: str

Returns:: mod – A new copy of the IRModule without the attribute
Return type:: IRModule

with_attrs(attr_map: DictAttrs | dict[str, Object]) → IRModule

Copy the IRModule and add the given attribute map to it. :param attr_map: The attribute map :type attr_map: Union[DictAttrs, Dict[str, Object]]

Returns:: mod – A new copy of the IRModule with the attribute
Return type:: IRModule

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_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_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_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

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

class tvm.ir.FuncType(arg_types, ret_type)

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.

Parameters:

arg_types (List[tvm.ir.Type]) – The argument types
ret_type (tvm.ir.Type) – The return type.

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.PrimType(dtype)

Primitive data type in the low level IR

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

class tvm.ir.TupleType(fields)

The type of tuple values.

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

class tvm.ir.Type(span=<object object>)

The base class of all types.

Methods:

same_as(other)

Compares two TVM types by referential equality.

same_as(other): Compares two TVM types by referential equality.

tvm.ir.diagnostics

The diagnostic interface to TVM, used for reporting and rendering diagnostic information by the compiler. This module exposes three key abstractions: a Diagnostic, the DiagnosticContext, and the DiagnosticRenderer.

tvm.ir.diagnostics.get_global_func(name: str, allow_missing: bool = False) → Function | None

Get a global function by name.

Parameters:

name – The name of the global function
allow_missing – Whether allow missing function or raise an error.

Returns:

The function to be returned, None if function is missing.

Return type:

func

Examples

import tvm_ffi

@tvm_ffi.register_global_func("demo.echo")
def echo(x):
    return x


f = tvm_ffi.get_global_func("demo.echo")
assert f(123) == 123

See also

tvm_ffi.register_global_func()

tvm.ir.diagnostics.register_global_func(func_name: str | Callable[[...], Any], f: Callable[[...], Any] | None = None, override: bool = False) → Any

Register global function.

Parameters:

func_name – The function name
f – The function to be registered.
override – Whether override existing entry.

Returns:

Register function if f is not specified.

Return type:

fregister

Examples

import tvm_ffi

# we can use decorator to register a function
@tvm_ffi.register_global_func("mytest.echo")
def echo(x):
    return x


# After registering, we can get the function by its name
f = tvm_ffi.get_global_func("mytest.echo")
assert f(1) == 1

# we can also directly register a function
tvm_ffi.register_global_func("mytest.add_one", lambda x: x + 1)
f = tvm_ffi.get_global_func("mytest.add_one")
assert f(1) == 2

See also

tvm_ffi.get_global_func(), tvm_ffi.remove_global_func()

class tvm.ir.diagnostics.Object

Base class of all TVM FFI objects.

This is the root Python type for objects backed by the TVM FFI runtime. Each instance references a handle to a C++ runtime object. Python subclasses typically correspond to C++ runtime types and are registered via tvm_ffi.register_object().

Notes

Equality of two Object instances uses underlying handle identity unless an overridden implementation is provided on the concrete type. Use same_as() to check whether two references point to the same underlying object.
Subclasses that omit __slots__ get __slots__ = () injected automatically by the metaclass. To allow a per-instance __dict__, declare __slots__ = ("__dict__",) explicitly in the class body.
Most users interact with subclasses (e.g. Tensor, Function) rather than Object directly.

Examples

Constructing objects is typically performed by Python wrappers that call into registered constructors on the FFI side.

import tvm_ffi.testing

# Acquire a testing object constructed through FFI
obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=12)
assert isinstance(obj, tvm_ffi.Object)
assert obj.same_as(obj)

Subclasses can declare explicit slots when needed.

@tvm_ffi.register_object("my.MyObject")
class MyObject(tvm_ffi.Object):
    __slots__ = ()

Subclasses that need a per-instance __dict__ (e.g. for attribute caching) can opt in explicitly.

@tvm_ffi.register_object("my.MyDynObject")
class MyDynObject(tvm_ffi.Object):
    __slots__ = ("__dict__",)

same_as(other: object) → bool

Return True if both references point to the same object.

This checks identity of the underlying FFI handle rather than performing a structural, value-based comparison.

Parameters:: other (object) – The object to compare against.
Return type:: bool

Examples

import tvm_ffi.testing

x = tvm_ffi.testing.create_object("testing.TestObjectBase")
y = x
z = tvm_ffi.testing.create_object("testing.TestObjectBase")
assert x.same_as(y)
assert not x.same_as(z)

tvm.ir.diagnostics.get_renderer()

Get the diagnostic renderer.

Returns:: renderer
Return type:: DiagnosticRenderer

class tvm.ir.diagnostics.DiagnosticLevel(value): The diagnostic level, see diagnostic.h for more details.

class tvm.ir.diagnostics.Diagnostic(level, span, message): A single diagnostic object from TVM.

class tvm.ir.diagnostics.DiagnosticRenderer(render_func)

A diagnostic renderer, which given a diagnostic context produces a “rendered” form of the diagnostics for either human or computer consumption.

render(ctx)

Render the provided context.

Params

ctx: DiagnosticContext: The diagnostic context to render.

class tvm.ir.diagnostics.DiagnosticContext(module, renderer)

A diagnostic context which records active errors and contains a renderer.

emit(diagnostic): Emit a diagnostic.

render(): Render the current context using its renderer member.