tvm.ir

Contents

tvm.ir#

Common data structures across all IR variants.

Functions:

const(value[, dtype, loc])

construct a constant

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.

is_prim_expr(value)

Return whether an expression has a primitive result type.

is_prim_var(value)

Return whether a value is an ordinary variable with a primitive type.

reinfer_type(call)

Derive a Call's result type from its current inputs without changing the Call.

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

Register an operator property of an operator by name.

stmt_seq(*args)

Make sequence of statements

stmt_list(stmt)

Make list of stmt from blocks.

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.

Scriptable()

Marker for display methods installed by the TVMScript owner.

SourceName(name)

An identifier for a source location.

Location()

Non-null source-location metadata attached to IR nodes.

SourceLoc(source_name, start_line, ...)

A source range with unchanged frontend coordinate units and endpoints.

UnknownLoc()

The canonical immutable location used when source information is unavailable.

CallSiteLoc(callee, caller)

A callee's location together with the location of its caller.

AnyType([loc])

The top type, which admits any value.

FuncType(arg_types, ret_type)

Function type.

MissingType()

Type information that has not been supplied or computed.

OpaqueType()

Type marker for opaque values that must be removed from finished IR.

PointerType(element_type[, storage_scope])

PointerType used in the low-level TIR.

PrimType(dtype)

Primitive data type in the low level IR

StringType()

Semantic string type, independent of the function's calling convention.

TensorRegionType()

The type of a multi-dimensional tvm.ir.TensorRegion.

TupleType(fields[, loc])

The type of tuple values.

Type([loc])

The base class of all types.

Call(op, args[, attrs, ty_args, loc, ty])

Core function call node.

Constant([loc, ty])

Base class of literal constants.

GenericConst(value, ty[, loc])

A literal payload with an explicit expression type.

DataTypeImm(value[, loc])

A data type literal whose expression type is AnyType.

StringImm(value[, loc])

A string literal with StringType.

ExprOperand()

Python operator surface for anything that denotes an expression.

ExprWithOp([loc, ty])

Common type-directed operator behavior for core expressions.

GlobalVar(name_hint)

A global variable in the IR.

LambdaExpr(parameter_types, function, *[, ...])

A typed staging expression representing a lambda computation.

StagingExpr([loc, ty])

A traversable expression eliminated before executable IR.

OpaqueExpr([loc, ty])

Base class for opaque values that must be removed from finished IR.

Range(begin[, end, loc])

Represent a range in TVM.

TensorLoad(*args, **kwargs)

An indexed load from an expression source.

TensorRegion(source, region, ty[, loc])

A region of an arbitrary tensor expression.

Tuple(fields[, loc])

Tuple expression that groups several fields together.

TupleGetItem(tuple_value, index[, loc])

Get the index-th item from a tuple.

Var([name, ty, loc, name_hint])

A canonical local variable in the IR.

BaseFunc(attrs[, loc, ty])

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

IRModule([functions, attrs, global_infos])

IRModule that holds functions and type definitions.

Op()

Primitive operator in the IR.

Stmt(loc)

Base class of all the statements.

SeqStmt(seq[, loc])

Sequence of statements.

Bind(var, value[, loc])

Bind node.

Evaluate(value[, loc])

Evaluate node.

Return(value[, loc])

Return node.

If(condition, then_case, else_case[, loc])

If node.

ForKind(value)

The kind of the for loop.

For(loop_var, min, extent, kind, body[, ...])

For node.

While(condition, body[, loc])

While node.

Break([loc])

Break node.

Continue([loc])

Continue node.

AssertStmt(kind, condition[, message_parts, loc])

AssertStmt node.

RegionStmt(op, args, body_params, attrs, body)

An operation with enclosing-scope operands and one lexical body.

TensorStore(dest, indices, value[, loc])

Store a primitive value into a tensor destination.

Array(input_list)

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

Map(input_dict)

Map container.
tvm.ir.const(value, dtype=None, loc: Location = UnknownLoc())#

construct a constant

Parameters:

  • value (number) – The content of the constant number.

  • dtype (str or None, optional) – The data type.

  • loc (Location, optional) – The location of the constant value in the source.

Returns:

const_val – The result expression.

Return type:

tvm.Expr

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, loc=tvm.ir.UnknownLoc())
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.Scriptable#

Marker for display methods installed by the TVMScript owner.

Methods:

script(*[, name, show_meta, ir_prefix, ...])

Print TVM IR into TVMScript text format

show([style, black_format, name, show_meta, ...])

A sugar for print highlighted TVM script.
script(*, name: str | None = None, show_meta: bool = False, ir_prefix: str = 'I', module_alias: str = 'cls', int_dtype: str = 'int32', float_dtype: str = 'void', verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, show_all_ty: bool = True, extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None) → str#

Print TVM IR into TVMScript text format

Parameters:

  • name (Optional[str] = None) – The name of the object

  • show_meta (bool = False) – Whether to print the meta data of the object

  • ir_prefix (str = "I") – The prefix of AST nodes from tvm.ir

  • module_alias (str = "cls") – Retained for compatibility. Cross-function calls use the module name directly.

  • int_dtype (str = "int32") – The default data type of integer

  • float_dtype (str = "void") – The default data type of float

  • verbose_expr (bool = False) – Whether to print the detailed definition of each variable in the expression

  • indent_spaces (int = 4) – The number of spaces for indentation

  • print_line_numbers (bool = False) – Whether to print line numbers

  • num_context_lines (int = -1) – The number of lines of context to print before and after the line to underline.

  • show_all_ty (bool = True) – If True (default), annotate all variable bindings with the struct info of that variable. If False, only add annotations where required for unambiguous round-trip of Relax -> TVMScript -> Relax.

  • extra_config (Optional[dict] = None) – Dialect-specific configuration passed through to PrinterConfig.extra_config. Keys are conventionally namespaced as “<dialect>.<knob>”, e.g. {"tirx.prefix": "T"}.

  • path_to_underline (Optional[List[AccessPath]] = None) – Object path to be underlined

  • path_to_annotate (Optional[Dict[AccessPath, str]] = None) – Object path to be annotated

  • obj_to_underline (Optional[List[Object]] = None) – Object to be underlined

  • obj_to_annotate (Optional[Dict[Object, str]] = None) – Object to be annotated

Returns:

script – The TVM Script of the given TVM IR

Return type:

str

show(style: str | None = None, black_format: bool | None = None, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = 'I', module_alias: str = 'cls', int_dtype: str = 'int32', float_dtype: str = 'void', verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, show_all_ty: bool = True, extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None) → None#

A sugar for print highlighted TVM script.

Parameters:

  • style (str, optional) – Pygmentize printing style, auto-detected if None. See tvm.script.printer.highlight.cprint for more details.

  • black_format (Optional[bool]) –

    If true, use the formatter Black to format the TVMScript. If false, do not apply the auto-formatter.

    If None (default), determine the behavior based on the environment variable “TVM_BLACK_FORMAT”. If this environment variable is unset, set to the empty string, or set to the integer zero, black auto-formatting will be disabled. If the environment variable is set to a non-zero integer, black auto-formatting will be enabled.

    Note that the “TVM_BLACK_FORMAT” environment variable only applies to the .show() method, and not the underlying .script() method. The .show() method is intended for human-readable output based on individual user preferences, while the .script() method is intended to provided a consistent output regardless of environment.

  • name (Optional[str] = None) – The name of the object

  • show_meta (bool = False) – Whether to print the meta data of the object

  • ir_prefix (str = "I") – The prefix of AST nodes from tvm.ir

  • module_alias (str = "cls") – Retained for compatibility. Cross-function calls use the module name directly.

  • int_dtype (str = "int32") – The default data type of integer

  • float_dtype (str = "void") – The default data type of float

  • verbose_expr (bool = False) – Whether to print the detailed definition of each variable in the expression

  • indent_spaces (int = 4) – The number of spaces for indentation

  • print_line_numbers (bool = False) – Whether to print line numbers

  • num_context_lines (int = -1) – The number of lines of context to print before and after the line to underline.

  • show_all_ty (bool = True) – If True (default), annotate all variable bindings with the struct info of that variable. If False, only add annotations where required for unambiguous round-trip of Relax -> TVMScript -> Relax.

  • extra_config (Optional[dict] = None) – Dialect-specific configuration passed through to PrinterConfig.extra_config.

  • path_to_underline (Optional[List[AccessPath]] = None) – Object path to be underlined

  • path_to_annotate (Optional[Dict[AccessPath, str]] = None) – Object path to be annotated

  • obj_to_underline (Optional[List[Object]] = None) – Object to be underlined

  • obj_to_annotate (Optional[Dict[Object, str]] = None) – Object to be annotated

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

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

class tvm.ir.SourceName(name)#

An identifier for a source location.

Parameters:

name (str) – The name of the source.

class tvm.ir.Location#

Non-null source-location metadata attached to IR nodes.

IR constructors default to the shared UNKNOWN_LOC.

class tvm.ir.SourceLoc(source_name, start_line, start_column, end_line, end_column)#

A source range with unchanged frontend coordinate units and endpoints.

Parameters:

  • source_name (SourceName) – The name of the source.

  • start_line (int) – The starting line number.

  • start_column (int) – The starting column offset.

  • end_line (int) – The ending line number.

  • end_column (int) – The ending column offset.

class tvm.ir.UnknownLoc#

The canonical immutable location used when source information is unavailable.

class tvm.ir.CallSiteLoc(callee: Location, caller: Location)#

A callee’s location together with the location of its caller.

class tvm.ir.AnyType(loc: Location = UnknownLoc())#

The top type, which admits any value.

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:
class tvm.ir.MissingType#

Type information that has not been supplied or computed.

Unlike AnyType or Void, this is not a concrete type and must be resolved before a boundary that requires fully typed IR.

class tvm.ir.OpaqueType#

Type marker for opaque values that must be removed from finished IR.

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.

Methods:

matches_code(*codes)

Return whether this type has any of the given DLPack dtype codes.

matches_element_type(code, bits)

Return whether this type has the given scalar element code and bits.

is_scalar()

Return whether this type has exactly one fixed lane.
matches_code(*codes) → bool#

Return whether this type has any of the given DLPack dtype codes.

matches_element_type(code, bits: int) → bool#

Return whether this type has the given scalar element code and bits.

is_scalar() → bool#

Return whether this type has exactly one fixed lane.

class tvm.ir.StringType#

Semantic string type, independent of the function’s calling convention.

class tvm.ir.TensorRegionType#

The type of a multi-dimensional tvm.ir.TensorRegion.

class tvm.ir.TupleType(fields, loc: Location = UnknownLoc())#

The type of tuple values.

Parameters:

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

class tvm.ir.Type(loc=<MISSING>)#

The base class of all types.

__expr_methods__ names methods exposed on expressions of this type. Each method receives the type instance and then the expression operand: expr.method(...) is equivalent to expr.ty.method(expr, ...). __expr_properties__ maps read-only expression property names to getters taking (type_instance, expression). Ordinary expression attributes take precedence, then declared properties, then methods. Subclasses inherit these declarations or replace them explicitly.

Methods:

missing()

Construct a MissingType for missing type information.

Missing()

Construct a MissingType for missing type information.

same_as(other)

Compares two TVM types by referential equality.

is_base_of(derived)

Check if this Relax type is a base of another Relax type.
static missing()#

Construct a MissingType for missing type information.

static Missing()#

Construct a MissingType for missing type information.

same_as(other)#

Compares two TVM types by referential equality.

is_base_of(derived: Type) → bool#

Check if this Relax type is a base of another Relax type.

class tvm.ir.Call(op: Expr | str, args: list[Expr] | tuple[Expr, ...], attrs: Attrs | dict | None = None, ty_args: list[tvm.ir.Type] | tuple[tvm.ir.Type, ...] | None = None, loc: Location = UnknownLoc(), ty: Type | str | None = None)#

Core function call node.

Omitted or None ty uses available result inference, or a missing type when no deduction is available. Explicit types, including Type.missing(), are preserved exactly. Inference errors propagate. Construction permits provisional IR; validate() checks the operator contract explicitly after inputs are ready.

Methods:

validate()

Check the registered operator contract without changing this Call.
validate() → None#

Check the registered operator contract without changing this Call.

class tvm.ir.Constant(loc=<MISSING>, ty=<MISSING>)#

Base class of literal constants.

class tvm.ir.GenericConst(value, ty: Type, loc: Location = UnknownLoc())#

A literal payload with an explicit expression type.

class tvm.ir.DataTypeImm(value: str | dtype, loc: Location = UnknownLoc())#

A data type literal whose expression type is AnyType.

Parameters:

  • value (str or tvm.DataType) – The represented data type.

  • loc (Location, optional) – The source location of the literal.

class tvm.ir.StringImm(value: str, loc: Location = UnknownLoc())#

A string literal with StringType.

class tvm.ir.ExprOperand#

Python operator surface for anything that denotes an expression.

Methods:

expr_ty()

Return this expression's primitive result type.
expr_ty()#

Return this expression’s primitive result type.

class tvm.ir.ExprWithOp(loc=<MISSING>, ty=<MISSING>)#

Common type-directed operator behavior for core 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.LambdaExpr(parameter_types, function: Callable, *, ret_type=None)#

A typed staging expression representing a lambda computation.

LambdaExpr records computations such as reduction combiners and predication rules. Its body may describe computations on runtime values.

Parameters are bound within the expression body, which may produce a scalar or tuple result. The lambda has a FuncType describing its parameter and return types.

As a StagingExpr, LambdaExpr is eliminated during compilation and does not remain in executable IR.

Parameters:

  • parameter_types (list[Type]) – Explicit parameter types, in callable argument order. Primitive dtype strings and script scalar constructors are also accepted.

  • function (Callable) – A callable evaluated once with one fresh typed Var per supplied type. Tuple/list results become shared IR Tuple expressions.

  • ret_type (Type, optional) – An exact return-type check. No implicit conversion or cast is inserted.

Methods:

apply(arguments)

Substitute arguments simultaneously for the lambda's bound variables.
apply(arguments: list[Expr]) → Expr#

Substitute arguments simultaneously for the lambda’s bound variables.

class tvm.ir.StagingExpr(loc=<MISSING>, ty=<MISSING>)#

A traversable expression eliminated before executable IR.

class tvm.ir.OpaqueExpr(loc=<MISSING>, ty=<MISSING>)#

Base class for opaque values that must be removed from finished IR.

class tvm.ir.Range(begin: Expr, end: Expr | None = None, loc: Location = UnknownLoc())#

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 (Expr) – The begin value of the range when end is None. Otherwise it is the length of the range.

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

  • loc (Location, optional) – 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[, loc])

Construct a Range by min and extent.
static from_min_extent(min_value: Expr, extent: Expr, loc: Location = UnknownLoc()) → Range#

Construct a Range by min and extent.

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

Parameters:

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

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

  • loc (Location, optional) – The location of this node in the source code.

Returns:

rng – The constructed range.

Return type:

Range

class tvm.ir.TensorLoad(*args, **kwargs)#

An indexed load from an expression source.

TensorLoad objects are constructed by a dialect-specific helper that validates the source and derives the result type.

class tvm.ir.TensorRegion(source: Expr, region: list[Range], ty: Type, loc: Location = UnknownLoc())#

A region of an arbitrary tensor expression.

Parameters:

  • source (Expr) – The source expression.

  • region (list[Range]) – The ranges describing the region.

  • ty (tvm.ir.Type) – The result type, including any dialect-specific subscript semantics.

  • loc (Location, optional) – The location of the expression in the source code.

class tvm.ir.Tuple(fields: list[Expr] | tuple[Expr, ...], loc: Location = UnknownLoc())#

Tuple expression that groups several fields together.

Parameters:

  • fields (list[Expr] | tuple[Expr, ...]) – The fields in the tuple.

  • loc (Location, optional) – Location that points to the original source code.

class tvm.ir.TupleGetItem(tuple_value: Expr, index: int, loc: Location = UnknownLoc())#

Get the index-th item from a tuple.

Parameters:

  • tuple_value (Expr) – The input tuple expression.

  • index (int) – The field index.

  • loc (Location, optional) – Location that points to the original source code.

class tvm.ir.Var(name: str | None = None, ty: Type | str | None = None, loc: Location = UnknownLoc(), *, name_hint: str | None = None)#

A canonical local variable in the IR.

Parameters:

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

  • ty (Optional[Type or str]) – The exact type of the variable. A string denotes a primitive dtype.

  • loc (Location, optional) – Location that points to the original source code.

tvm.ir.is_prim_expr(value: object) → bool#

Return whether an expression has a primitive result type.

tvm.ir.is_prim_var(value: object) → bool#

Return whether a value is an ordinary variable with a primitive type.

tvm.ir.reinfer_type(call: Call) → Type#

Derive a Call’s result type from its current inputs without changing the Call.

The operator must register a fixed return type or a context-free inference rule. This does not invoke the operator’s validator.

class tvm.ir.BaseFunc(attrs, loc=<MISSING>, ty=<MISSING>)#

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

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

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

static from_expr(expr, functions=None)#

Construct a module from a standalone expression.

Parameters:

  • expr (Expr) – 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 a registered operator by name.

list_op_names()

List registered operator names in unspecified order.

set_attr(attr_name, value[, override])

Set an operator attribute.

set_signature([args, ty_args, var_args, ...])

Replace this Op's argument signature and validate Call arity.

add_arg(name, doc)

Append an argument's name and documentation.

add_ty_arg(name, doc)

Append a type-argument name and documentation without imposing a count rule.

get_attr(attr_name)

Return this Op's attribute or None; an unregistered column raises InternalError.

has_attr(attr_name)

Return whether the attribute column exists in the registry.

reset_attr(attr_name)

Remove this Op's current value; missing values are ignored and cached views observe removal.

set_attrs_type_key(key)

Resolve and set the attribute object type key and runtime index together; an unknown key raises before updating.
static get(op_name)#

Get a registered operator by name.

Parameters:

op_name (str) – The canonical operator name.

Returns:

A handle to the registered operator.

Return type:

Op

static list_op_names()#

List registered operator names in unspecified order.

Returns:

The registered operator names.

Return type:

list[str]

set_attr(attr_name, value, override=False)#

Set an operator attribute.

Parameters:

  • attr_name (str) – Attribute column name.

  • value (object) – Non-None attribute value.

  • override (bool, optional) – Replace an existing value if True. Duplicate registration otherwise raises ValueError. Cached views observe replacements; no history is kept.

Return type:

None

set_signature(args=(), *, ty_args=(), var_args=None, var_ty_args=None) → None#

Replace this Op’s argument signature and validate Call arity.

Each entry is a name string or a (name, doc) tuple of strings. Fixed entries form required prefixes; a variadic entry allows zero or more additional arguments. Calls must have exactly the required count without a tail, or at least that count with one. This method does not check argument types or Call attrs. It replaces a generated typed validator with a count-only one, while an existing custom validator keeps precedence.

Parameters:

  • args (sequence[str or tuple[str, str]], optional) – Fixed value arguments, in order.

  • ty_args (sequence[str or tuple[str, str]], optional) – Fixed type arguments, in order.

  • var_args (str or tuple[str, str], optional) – Variadic value-argument tail.

  • var_ty_args (str or tuple[str, str], optional) – Variadic type-argument tail.

Return type:

None

add_arg(name: str, doc: str) → None#

Append an argument’s name and documentation.

add_ty_arg(name: str, doc: str) → None#

Append a type-argument name and documentation without imposing a count rule.

get_attr(attr_name: str) → object#

Return this Op’s attribute or None; an unregistered column raises InternalError.

has_attr(attr_name: str) → bool#

Return whether the attribute column exists in the registry.

reset_attr(attr_name: str) → None#

Remove this Op’s current value; missing values are ignored and cached views observe removal.

set_attrs_type_key(key: str) → None#

Resolve and set the attribute object type key and runtime index together; an unknown key raises before updating.

tvm.ir.register_op_attr(op_name, attr_key, value=None, override=False)#

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

  • override (bool, optional) – Replace an existing value if True; otherwise duplicate registration raises ValueError. Cached views observe replacements; no priority history is kept.

Returns:

result – The registered value when supplied, or a decorator that registers and returns its argument. The named Op is created if it does not exist.

Return type:

object or function

class tvm.ir.Stmt(loc)#

Base class of all the statements.

class tvm.ir.SeqStmt(seq: Stmt | Sequence[Stmt], loc: Location | None = None)#

Sequence of statements.

Parameters:

  • seq (Stmt | Sequence[Stmt]) – The statements, flattened into one sequence. Empty and singleton sequences are valid.

  • loc (Location or None, optional) – None inherits the location when seq is an existing statement. Explicit UNKNOWN_LOC replaces that location with unknown.

class tvm.ir.Bind(var: Var, value: Expr, loc: Location = UnknownLoc())#

Bind node.

Bind a variable to a value in the enclosing scope. Bind has no body field. The bound variable is visible in all subsequent statements within the same enclosing scope (SeqStmt, ForNode.body, etc.).

Parameters:

  • var (Var) – The variable in the binding.

  • value (Expr) – The value to be bound.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.Evaluate(value: Expr, loc: Location = UnknownLoc())#

Evaluate node.

Parameters:

  • value (Expr) – The expression to be evaluated.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.Return(value: Expr, loc: Location = UnknownLoc())#

Return node.

Parameters:

  • value (Expr) – The value to return.

  • loc (Location, optional) – The location of this statement in the source code.

class tvm.ir.If(condition: Expr, then_case: Stmt | Sequence[Stmt], else_case: Stmt | Sequence[Stmt] | None, loc: Location = UnknownLoc())#

If node.

Parameters:

  • condition (Expr) – The expression

  • then_case (Stmt | Sequence[Stmt]) – The statement to execute if condition is true.

  • else_case (Stmt | Sequence[Stmt] | None) – The statement to execute if condition is false.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.ForKind(value)#

The kind of the for loop.

Note

ForKind can change the control flow semantics of the loop and need to be considered in all TIR passes.

class tvm.ir.For(loop_var: Var, min: Expr, extent: Expr, kind: ForKind, body: Stmt | Sequence[Stmt], annotations: Mapping[str, Object] | None = None, step: Expr | None = None, loc: Location = UnknownLoc())#

For node.

Parameters:

  • loop_var (Var) – The loop variable.

  • min (Expr) – The beginning value.

  • extent (Expr) – The length of the loop.

  • kind (ForKind) – The type of the for.

  • body (Stmt | Sequence[Stmt]) – The body statement.

  • annotations (Optional[Mapping[str, Object]]) – Additional loop annotations interpreted by the consuming dialect.

  • step (Expr) – The loop step. Defaults to None, which represents one.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.While(condition: Expr, body: Stmt | Sequence[Stmt], loc: Location = UnknownLoc())#

While node.

Parameters:

  • condition (Expr) – The termination condition.

  • body (Stmt | Sequence[Stmt]) – The body statement.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.Break(loc: Location = UnknownLoc())#

Break node.

class tvm.ir.Continue(loc: Location = UnknownLoc())#

Continue node.

class tvm.ir.AssertStmt(kind: StringImm, condition: Expr, message_parts: list | None = None, loc: Location = UnknownLoc())#

AssertStmt node.

Parameters:

  • kind (StringImm) – The error kind, e.g. “RuntimeError”, “TypeError”, “ValueError”.

  • condition (Expr) – The assert condition.

  • message_parts (list[StringImm]) – Error message fragments, concatenated at runtime when assertion fails.

  • loc (Location, optional) – The location of the stmt in the source code.

class tvm.ir.RegionStmt(op: Op | str, args: Sequence[Expr], body_params: Sequence[Var], attrs: DictAttrs | Mapping[str, Any] | None, body: Stmt | Sequence[Stmt], result_vars: Sequence[Var] | None = None, loc: Location = UnknownLoc())#

An operation with enclosing-scope operands and one lexical body.

body_params define variables visible only within body. Their count and types must match the operation’s required FRegionGetBodyParams hook. Zero-parameter regions register an empty-return hook; operations without a hook do not support region construction. Explicit parameter identities and their references in body are preserved. result_vars define variables after the region in the enclosing sequence. Attributes are evaluated outside the body-parameter scope. Direct construction and JSON serialization support result variables; structured script syntax currently supports only result-free regions.

class tvm.ir.TensorStore(dest: Expr, indices: list[Expr], value: Expr, loc: Location = UnknownLoc())#

Store a primitive value into a tensor destination.

Parameters:

  • dest (Expr) – The destination expression with a supported concrete store type.

  • indices (list[Expr]) – The indices location to be stored.

  • value (Expr) – The primitive value to be stored.

  • loc (Location, optional) – The location of the stmt in the source code.

tvm.ir.stmt_seq(*args: Expr | Stmt) → SeqStmt#

Make sequence of statements

Parameters:

*args (Union[Expr, Stmt]) – List of statements to be combined as sequence.

Returns:

stmt – The combined sequence.

Return type:

SeqStmt

tvm.ir.stmt_list(stmt: Stmt) → list[Stmt]#

Make list of stmt from blocks.

Parameters:

stmt (Stmt) – The input statement.

Returns:

stmt_list – The unpacked list of statements

Return type:

List[Stmt]

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

Contents