tvm.ir#

Common data structures across all IR variants.

Classes:

Functions:

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.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(key)	Get a python int value of a key
get_int_tuple(key)	Get a python int tuple of a key
get_str(key)	Get a python int value of a key
keys()	Get list of names in the attribute.
list_field_info()	Get fields information

get_int(key)[源代码]#

Get a python int value of a key

参数:

key (str)

返回:

value

返回类型:

int

get_int_tuple(key)[源代码]#

Get a python int tuple of a key

参数:

key (str)

返回:

value

返回类型:

Tuple of int

get_str(key)[源代码]#

Get a python int value of a key

参数:

key (str)

返回:

value

返回类型:

int

keys()[源代码]#

Get list of names in the attribute.

返回:

keys -- List of keys

返回类型:

list of str

list_field_info()[源代码]#

Get fields information

返回:

infos -- List of field information

返回类型:

list of AttrFieldInfo

class tvm.ir.BaseExpr[源代码]#

Base class of all the expressions.

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

参数:

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

返回:

func -- A new copy of the function

返回类型:

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

返回:

func -- A new copy of the function

返回类型:

BaseFunc

without_attr(attr_key: str) → BaseFunc[源代码]#

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

参数:

attr_key (str) -- The attribute key to delete from the attrubte pairs.

返回:

func -- A new copy of the function

返回类型:

BaseFunc

class tvm.ir.CallingConv(*values)[源代码]#

Possible kinds of calling conventions.

class tvm.ir.DictAttrs[源代码]#

Dictionary attributes.

Methods:

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

get(key, default=None)[源代码]#

Get an element with a default value.

items()[源代码]#

Get items from the map.

keys()[源代码]#

Get list of names in the attribute.

返回:

keys -- List of keys

返回类型:

list of str

class tvm.ir.DummyGlobalInfo[源代码]#

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

参数:

name (str) -- The name of the 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.

参数:

• arg_types (List[tvm.ir.Type]) -- The argument types

• ret_type (tvm.ir.Type) -- The return type.

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.GlobalVar(name_hint: str, type_annot: Type | None = None)[源代码]#

A global variable in the IR.

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

参数:

name_hint (str) -- The name of the variable.

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.

参数:

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

Methods:

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

static from_expr(expr, functions=None)[源代码]#

Construct a module from a standalone expression.

参数:

• expr (RelaxExpr) -- The starting expression

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

返回:

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

返回类型:

Module

functions_items()[源代码]#

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

返回:

items -- The functions items.

返回类型:

List[Tuple[GlobalVar, Function]]

get_attr(attr_key)[源代码]#

Get the IRModule attribute.

参数:

attr_key (str) -- The attribute key.

返回:

attr_value -- Attribute value

返回类型:

Any

get_global_var(name)[源代码]#

Get a global variable in the function by name.

参数:

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

返回:

global_var -- The global variable mapped to name.

返回类型:

GlobalVar

抛出:

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

get_global_vars()[源代码]#

Collect all global vars defined in this module.

返回:

global_vars -- An array of global vars.

返回类型:

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.

参数:

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.

返回:

The updated module

返回类型:

IRModule

update(other)[源代码]#

Insert functions in another Module to current one.

参数:

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.

参数:

• 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

参数:

• name (str) -- The name for the global info.

• global_info (List[GlobalInfo]) -- The global info to be updated.

with_attr(attr_key, attr_value)[源代码]#

Copy the IRModule and add an attribute to it.

参数:

• attr_key (str) -- The attribute key.

• attr_value (Object) -- The new attribute value.

返回:

mod -- A new copy of the IRModule with the attribute

返回类型:

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

返回:

mod -- A new copy of the IRModule with the attribute

返回类型:

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

返回:

mod -- A new copy of the IRModule without the attribute

返回类型:

IRModule

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:

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.

参数:

• key (object) -- The attribute key.

• default (object) -- The default object.

返回:

value -- The result value.

返回类型:

object

items()[源代码]#

Get the items from the map

class tvm.ir.Node[源代码]#

Base class of all IR Nodes.

class tvm.ir.Op[源代码]#

Primitive operator in the IR.

Methods:

add_argument(name, type, description)	Add arguments information to the function.
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.
list_op_names()	List all the op names in the op registry.
reset_attr(attr_name)	Reset attribute about the operator.
set_attr(attr_name, value[, plevel])	Set attribute about the operator.
set_attrs_type_key(key)	Set the attribute type key of op.
set_num_inputs(n)	Set the support level of op.
set_support_level(level)	Set the support level of op.

add_argument(name, type, description)[源代码]#

Add arguments information to the function.

参数:

• name (str) -- The argument name.

• type (str) -- The argument type.

• description (str) -- The argument description.

static get(op_name)[源代码]#

Get the Op for a given name

参数:

op_name (str) -- The operator name

返回:

op -- The op of the corresponding name

返回类型:

Op

get_attr(attr_name)[源代码]#

Get additional attribute about the operator.

参数:

attr_name (str) -- The attribute name.

返回:

value -- The attribute value

返回类型:

object

has_attr(attr_name)[源代码]#

Check whether the operator has additional attribute.

参数:

attr_name (str) -- The attribute name.

返回:

value -- Whether the operator has additional attribute

返回类型:

bool

static list_op_names()[源代码]#

List all the op names in the op registry.

返回:

value -- The registered op names

返回类型:

List[str]

reset_attr(attr_name)[源代码]#

Reset attribute about the operator.

参数:

attr_name (str) -- The attribute name

set_attr(attr_name, value, plevel=10)[源代码]#

Set attribute about the operator.

参数:

• attr_name (str) -- The attribute name

• value (object) -- The attribute value

• plevel (int) -- The priority level

set_attrs_type_key(key)[源代码]#

Set the attribute type key of op.

参数:

key (str) -- The type key.

set_num_inputs(n)[源代码]#

Set the support level of op.

参数:

n (int) -- The input number.

set_support_level(level)[源代码]#

Set the support level of op.

参数:

level (int) -- The support level.

class tvm.ir.PointerType(element_type, storage_scope='')[源代码]#

PointerType used in the low-level TIR.

参数:

• 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.PrimExpr[源代码]#

Base class of all primitive expressions.

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

class tvm.ir.PrimType(dtype)[源代码]#

Primitive data type in the low level IR

参数:

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

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.

参数:

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

备注

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)

参数:

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

返回:

rng -- The constructed range.

返回类型:

Range

class tvm.ir.RelaxExpr[源代码]#

Base class of all non-primitive expressions.

Attributes:

checked_type	Get the checked type of tvm.relax.Expr.
struct_info	Get the struct info field

property checked_type#

Get the checked type of tvm.relax.Expr.

返回:

checked_type -- The checked type.

返回类型:

tvm.ir.Type

property struct_info: StructInfo | None#

Get the struct info field

返回:

struct_info -- The struct info if available.

返回类型:

tvm.relax.StructInfo

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.

参数:

spans (Array) -- The array of spans.

class tvm.ir.SourceName(name)[源代码]#

A identifier for a source location.

参数:

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.

参数:

• source (SourceName) -- The source name.

• lineno (int) -- The line number.

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

class tvm.ir.TupleType(fields)[源代码]#

The type of tuple values.

参数:

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

class tvm.ir.Type[源代码]#

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.

class tvm.ir.VDevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global')[源代码]#

tvm.ir.assert_structural_equal(lhs, rhs, map_free_vars=False)[源代码]#

Assert lhs and rhs are structurally equal to each other.

参数:

• 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.:

参见

structural_equal

tvm.ir.load_json(json_str) → Object[源代码]#

Load tvm object from json_str.

参数:

json_str (str) -- The json string

返回:

node -- The loaded tvm node.

返回类型:

Object

tvm.ir.make_node(type_key, **kwargs)[源代码]#

Make a new IR node by its type key and fields

参数:

• type_key (str) -- The type key of the node.

• **kwargs (dict) -- The fields of the node.

返回:

node -- The corresponding IR Node

返回类型:

Node

备注

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.

示例

The following code constructs a IntImm object

x = tvm.ir.make_node("IntImm", dtype="int32", value=10, span=None)
assert isinstance(x, tvm.tir.IntImm)
assert x.value == 10

tvm.ir.register_intrin_lowering(op_name, target, *, f=None, level=10)[源代码]#

Register Op lowering function

参数:

• 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

返回:

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

返回类型:

function

tvm.ir.register_op_attr(op_name, attr_key, value=None, level=10)[源代码]#

Register an operator property of an operator by name.

参数:

• 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

返回:

fregister -- Register function if value is not specified.

返回类型:

function

tvm.ir.save_json(node) → str[源代码]#

Save tvm object as json string.

参数:

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

返回:

json_str -- Saved json string.

返回类型:

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

A var-type node(e.g. tir::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.

参数:

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

返回:

result -- The comparison result.

返回类型:

bool

参见

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

参数:

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

返回:

result -- The hash result

返回类型:

int

参见

structrual_equal