tvm.ir#
Common data structures across all IR variants.
Classes:
|
Array container of TVM. |
|
Attribute node, which is mainly use for defining attributes of relay operators. |
|
Base class of all the expressions. |
|
Base class of all functions. |
|
Possible kinds of calling conventions. |
|
This object contains a list of ConstantPoolInfo objects to be used as read-only memory in the compilation |
|
ConstantPoolInfo object holds information related to RO memory pools where the statically sized allocate nodes are pooled into. |
|
Relay ADT constructor. |
Dictionary attributes. |
|
|
Environment function. |
|
Function type. |
Base node for all global info that can appear in the IR |
|
|
A global type variable that is used for defining new types or type aliases. |
|
A global variable in the IR. |
|
IRModule that holds functions and type definitions. |
|
Incomplete type during type inference. |
|
Map container of TVM. |
|
Base class of all IR Nodes. |
|
Primitive operator in the IR. |
|
PointerType used in the low-level TIR. |
|
PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into. |
|
PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into. |
|
Base class of all primitive expressions. |
|
Primitive data type in the low level IR |
|
Represent a range in TVM. |
Base class of all non-primitive expressions. |
|
|
Reference Type in relay. |
|
A sequence of source spans |
|
A identifier for a source location. |
|
Specifies a location in a source program. |
|
The quantized type of a tensor, with scale, zero point, and datatype |
|
A concrete TensorType in Relay. |
|
Affine types of a node with multiple outputs |
|
The type of tuple values. |
|
The base class of all types. |
|
Type function application. |
Abstract class representing a type constraint. |
|
|
Stores the definition for an Algebraic Data Type (ADT) in Relay. |
|
Possible kinds of TypeVars. |
|
User defined type relation, it is an input-output relation on types. |
|
Type parameter in functions. |
|
|
|
This object contains a list of WorkspacePoolInfo objects to be used as workspace memory in the compilation |
|
WorkspacePoolInfo object holds information related to RW memory pools where the statically sized allocate nodes will pooled into. |
Functions:
|
Assert lhs and rhs are structurally equal to each other. |
|
Load tvm object from json_str. |
|
Make a new IR node by its type key and fields |
|
Register Op lowering function |
|
Register an operator property of an operator by name. |
|
Save tvm object as json string. |
|
Check structural equality of lhs and rhs. |
|
Compute structural hash of node |
- 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 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:
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.
Get fields information
- class tvm.ir.BaseFunc[源代码]#
Base class of all functions.
Attributes:
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.
- class tvm.ir.CallingConv(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]#
Possible kinds of calling conventions.
- class tvm.ir.ConstantMemoryPools(pools: List[ConstantPoolInfo])[源代码]#
This object contains a list of ConstantPoolInfo objects to be used as read-only memory in the compilation
- 参数:
pools (List[ConstantPoolInfo]) -- The list of ConstantPoolInfo objects to be used with the compilation
- 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.
- 参数:
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.Constructor(name_hint, inputs, belong_to)[源代码]#
Relay ADT constructor.
- 参数:
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.EnvFunc[源代码]#
Environment function.
This is a global function object that can be serialized by its name.
Methods:
get
(name)Get a static env function
- 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
- 参数:
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.GlobalInfo[源代码]#
Base node for all global info that can appear in the IR
Methods:
same_as
(other)Overload with structural equality.
- class tvm.ir.GlobalTypeVar(name_hint, kind=TypeKind.AdtHandle)[源代码]#
A global type variable that is used for defining new types or type aliases.
- 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.
Methods:
astext
([show_meta_data, annotate])Get the text format of the expression.
- astext(show_meta_data: bool = True, annotate: Callable[[Object], str] | None = None) str [源代码]#
Get the text format of the expression.
- 参数:
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.
- 返回:
text -- The text format of the expression.
- 返回类型:
备注
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.IRModule(functions=None, type_definitions=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:
astext
([show_meta_data, annotate])Get the text format of the expression.
from_expr
(expr[, functions, type_defs])Construct a module from a standalone expression.
Get items in self.functions.items() in alphabetical order.
get_attr
(attr_key)Get the IRModule attribute.
get_constructor
(tag)Look up an ADT constructor by tag.
get_global_type_var
(name)Get a global type variable in the function by name.
Collect all global type vars defined in this module.
get_global_var
(name)Get a global variable in the function by name.
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.
- astext(show_meta_data=True, annotate=None)[源代码]#
Get the text format of the expression.
- 参数:
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.
- 返回:
text -- The text format of the expression.
- 返回类型:
备注
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.
- static from_expr(expr, functions=None, type_defs=None)[源代码]#
Construct a module from a standalone expression.
- 参数:
- 返回:
mod -- A module containing the passed definitions, where expr is set as the entry point (wrapped in a function if necessary)
- 返回类型:
Module
- get_constructor(tag)[源代码]#
Look up an ADT constructor by tag.
- 参数:
tag (int) -- The tag for a constructor.
- 返回:
constructor -- The constructor associated with the given tag,
- 返回类型:
- 抛出:
tvm.error.TVMError if the corresponding constructor cannot be found. --
- get_global_type_var(name)[源代码]#
Get a global type variable in the function by name.
- 参数:
name (str) -- The name of the global type variable.
- 返回:
global_type_var -- The global variable mapped to
name
.- 返回类型:
- 抛出:
tvm.error.TVMError if we cannot find corresponding global type var. --
- get_global_type_vars()[源代码]#
Collect all global type vars defined in this module.
- 返回:
global_type_vars -- An array of global type vars.
- 返回类型:
- 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.
- 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.relay.Function) -- 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.
- class tvm.ir.IncompleteType(kind=TypeKind.Type)[源代码]#
Incomplete type during type inference.
- kindOptional[TypeKind]
The kind of the incomplete type.
- 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:
- class tvm.ir.Op[源代码]#
Primitive operator in the IR.
Methods:
add_argument
(name, type, description)Add arguments information to the function.
add_type_rel
(rel_name[, type_rel_func])Attach the type function corresponding to the return type.
astext
([show_meta_data, annotate])Get the text format of the expression.
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 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 the support level of op.
set_support_level
(level)Set the support level of op.
- add_type_rel(rel_name, type_rel_func=None)[源代码]#
Attach the type function corresponding to the return type.
- 参数:
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++:
When type_rel_func is not None
OpAddTypeRel on C++ side will adjust type_rel_func with TypeReporter to calling convention of relay type system.
type_rel_func returns output argument's type, return None means can't infer output's type.
only support single output operators for now, the last argument is output tensor.
- when type_rel_func is None, will call predefined type_rel_funcs in relay
according to
tvm.relay.type_relation.
+ rel_name.
- astext(show_meta_data=True, annotate=None)[源代码]#
Get the text format of the expression.
- 参数:
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.
- 返回:
text -- The text format of the expression.
- 返回类型:
备注
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.
- 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
- 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.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.PoolInfoProperties(size_hint_bytes: int | None = -1, clock_frequency_hz: int | None = -1, read_bandwidth_bytes_per_cycle: int | None = -1, write_bandwidth_bytes_per_cycle: int | None = -1, read_latency_cycles: int | None = 0, write_latency_cycles: int | None = 0, target_burst_bytes=None)[源代码]#
PoolInfo object holds information related to memory pools where the statically sized allocate nodes will pooled into.
- 参数:
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.
- 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.
- 参数:
备注
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.
- class tvm.ir.RelayExpr[源代码]#
Base class of all non-primitive expressions.
Attributes:
Get the checked type of tvm.relay.Expr.
Get the struct info field
- property checked_type#
Get the checked type of tvm.relay.Expr.
- 返回:
checked_type -- The checked type.
- 返回类型:
tvm.relay.Type
- property struct_info: StructInfo | None#
Get the struct info field
- 返回:
struct_info -- The struct info if available.
- 返回类型:
- 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.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
- 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).
- 参数:
shape (List[tvm.ir.PrimExpr]) -- The shape of the Tensor
dtype (Optional[str]) -- The content data type.
Attributes:
Get shape of the type as concrete tuple of int.
- class tvm.ir.TupleAffineType(types)[源代码]#
Affine types of a node with multiple outputs
- 参数:
types (List[TensorAffineType]) -- The shape of the Tensor
- 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 Relay types by referential equality.
- class tvm.ir.TypeCall(func, args)[源代码]#
Type function application.
- 参数:
func (tvm.ir.Type) -- The function.
args (List[tvm.ir.Type]) -- The arguments.
- 返回:
type_call -- The type function application.
- 返回类型:
- 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.
- 参数:
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.TypeKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]#
Possible kinds of TypeVars.
- 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.
- 参数:
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
- 返回:
type_relation -- The type relation.
- 返回类型:
- 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.
- class tvm.ir.WorkspaceMemoryPools(pools: List[WorkspacePoolInfo])[源代码]#
This object contains a list of WorkspacePoolInfo objects to be used as workspace memory in the compilation
- 参数:
pools (List[WorkspacePoolInfo]) -- The list of ConstantPoolInfo objects to be used with the compilation
- 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.
- 参数:
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.
- 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.:
- 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
- 参数:
- 返回:
node -- The corresponding IR 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
- tvm.ir.register_op_attr(op_name, attr_key, value=None, level=10)[源代码]#
Register an operator property of an operator by name.
- 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.
- 返回类型:
- 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.
- 参数:
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.
- 返回类型:
参见
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
- 返回类型:
参见
structrual_equal