tvm.ir

Contents

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

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[, span])

The type of tuple values.

Type([span])

The base class of all types.

Call(op, args[, attrs, ty_args, span, ret_ty])

Core function call node.

GlobalVar(name_hint)

A global variable in the IR.

Range(begin[, end, span])

Represent a range in TVM.

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

A canonical local variable in the IR.

BaseFunc(attrs[, span, ty])

Base class of all functions.

CallingConv(value)

Possible kinds of calling conventions.

GlobalInfo(*args, **kwargs)

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.

Array(input_list)

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

Map(input_dict)

Map container.

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.

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.

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

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.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.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.TupleType(fields, span=None)#

The type of tuple values.

Parameters:

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

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

The base class of all types.

Methods:

missing()

Return the sentinel for missing type information.

Missing()

Return the sentinel for missing type information.

is_missing()

Return whether this is the missing-type sentinel.

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()#

Return the sentinel for missing type information.

static Missing()#

Return the sentinel for missing type information.

is_missing()#

Return whether this is the missing-type sentinel.

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, span: Span | None = None, ret_ty: Type | str | None = None)#

Core function call node.

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.Range(begin: Expr, end: Expr | 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 (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.

  • 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: Expr, extent: Expr, 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 (Expr) – The minimum value of the range.

  • extent (Expr) – 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.Var(name: str | None = None, ty: Type | str | None = None, span: Span | None = None, *, 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.

  • span (Optional[Span]) – Span that points to the original source code.

Methods:

access_ptr(access_mask[, ptr_type, ...])

Get an access pointer to the head of buffer.

byte_offset_of(indices[, inner])

Get the byte offset of the buffer at the given indices.

chunk(spec)

Split dims into equal contiguous chunks and pick a chunk per dim — rank-preserving.

elem_offset_of(indices[, inner])

Get the element offset of the buffer at the given indices.

get_flattened_buffer()

Generate a Buffer that is a flattened version of this buffer.

is_scalar([alloc_or_decl])

Check if the buffer is a scalar.

local(*shape[, layout])

Create a thread-local view of this buffer.

offset_of(indices)

Determine the offset of the provided indices in the flattened buffer.

permute(*dims)

Permute the dimensions of the buffer.

ptr_to(indices)

Get the pointer to the buffer at the given indices (logical indices).

rearrange([pattern])

einops-style relayout in one line: buf.rearrange("b (2 r) -> 2 b r").

scope()

Return the storage scope associated with this buffer.

tile(*specs)

Chunk a dim: split it into factors, pick a chunk, keep the rest.

view(*args, **kwargs)

Creates a new view of the buffer.

vload(begin[, dtype, predicate])

Generate an Expr that loads dtype from begin index.

vstore(begin, value[, predicate])

Generate a Stmt that store value into begin index.

with_allocated_addr(allocated_addr)

Return a new buffer with the allocated address.

with_dtype(dtype)

Return a new buffer with the dtype.

Attributes:

byte_offset

Get the byte offset of the buffer.

sub

buf.sub[2, 4:8, ::4].
access_ptr(access_mask, ptr_type='handle', content_lanes=1, offset=0, extent=None)#

Get an access pointer to the head of buffer.

This is the recommended method to get buffer data ptress when interacting with external functions.

Parameters:

  • access_mask (int) – The access pattern MASK. Indicate whether the access will read or write to the data content.

  • ptr_type (str or tvm.ir.Type, optional) – The data type of the result pointer. Do not specify unless we want to cast pointer to specific type.

  • content_lanes (int, optional) – The number of lanes for the data type. This value is greater than one for vector types.

  • offset (Expr, optional) – The offset of pointer. We can use it to offset by the number of elements from the address of ptr.

  • extent (Expr, optional) – The extent of pointer.

Examples

# Get access ptr for read
buffer.access_ptr("r")
# Get access ptr for read/write with bitmask
buffer.access_ptr(BufferAccessKind.READ | BufferAccessKind.WRITE)
# Get access ptr for read/write with str flag
buffer.access_ptr("rw")
# Get access ptr for read with offset
buffer.access_ptr("r", offset = 100)
# Get access ptr for read with extent
buffer.access_ptr("r", extent = 100)
property byte_offset#

Get the byte offset of the buffer.

byte_offset_of(indices, inner=True)#

Get the byte offset of the buffer at the given indices. Note that indices subject to buffer’s layout mapping.

Parameters:

  • indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.

  • inner (bool, optional) – If False, the offset is relative to the original buffer. Default is True.

Returns:

offset – The byte offset of the buffer at the given indices.

Return type:

Expr

chunk(spec) → ChunkIndexer#

Split dims into equal contiguous chunks and pick a chunk per dim — rank-preserving. Index the result with [picks].

spec is a per-dim tuple (length = rank). Each entry is None (leave the dim) or a positive int n (split that dim, extent E with E % n == 0, into n equal chunks of E // n). Then chunk(spec)[picks] takes one entry per dim: a chunked dim’s pick is the chunk index (int / Expr) and narrows that dim to the chunk’s [c*E//n : (c+1)*E//n) range — the dim is kept at E // n, no dimension is added; an unchunked dim’s pick is a normal index (: / int / slice). The result is the same BufferRegion as the hand-written slice — one line instead of the c*k : (c+1)*k arithmetic:

X[.., c * k : (c + 1) * k, ..]        # before (k = E // n)
X.chunk((None, .., n, ..))[.., c, ..]  # after (k inferred)
elem_offset_of(indices, inner=True)#

Get the element offset of the buffer at the given indices. Note that indices subject to buffer’s layout mapping.

Parameters:

  • indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.

  • inner (bool, optional) – If False, the offset is relative to the original buffer. Default is True.

Returns:

offset – The element offset of the buffer at the given indices.

Return type:

Expr

get_flattened_buffer()#

Generate a Buffer that is a flattened version of this buffer.

Returns:

flattened – The corresponding flat buffer.

Return type:

Buffer

is_scalar(alloc_or_decl=True)#

Check if the buffer is a scalar.

Parameters:

alloc_or_decl (bool, optional) – Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar, False for decl_scalar.

Returns:

bool

Return type:

True if the buffer is a scalar, False otherwise.

local(*shape, layout=None) → Var#

Create a thread-local view of this buffer.

By default, both the inferred and explicit-shape forms address the raw physical storage span. local()[k] is the k-th physical storage element, including any gaps or layout offset, while local(d0, d1, ...) is a row-major reshape of that same span. Pass layout= to request a mediated view explicitly. This is an escape hatch whose shape is interpreted by the supplied layout. When that shape is explicit, the parent buffer does not need a layout.

When called with no shape arguments, auto-infers a 1D shape from the span of the parent layout’s non-thread component (i.e. self.layout.storage().span()). The explicit-layout= form instead infers the parent layout’s storage().size() for compatibility. Either inference requires the parent buffer to have a layout.

Parameters:

  • shape (tuple of Expr) – The shape of the local view for indexing. Without layout=, its product must equal the per-thread physical storage span. With an explicit layout, the shape is not constrained by the raw span. If omitted, a matching 1D shape is computed automatically.

  • layout (optional) – Override layout. If None, the default (identity) layout is used.

Returns:

local – The corresponding local buffer.

Return type:

DeclBufferFrame

offset_of(indices)#

Determine the offset of the provided indices in the flattened buffer.

Parameters:

indices (Union[Expr, List[Expr]]) – The indices of the element in the original buffer.

Returns:

flattened_indices – The offset indices of the element in the flattened buffer.

Return type:

List[Expr]

permute(*dims) → Var#

Permute the dimensions of the buffer.

Parameters:

dims (tuple of int) – The permutation of dimensions.

Returns:

permuted – The buffer with permuted dimensions.

Return type:

DeclBufferFrame

ptr_to(indices)#

Get the pointer to the buffer at the given indices (logical indices).

Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.

rearrange(pattern: str = <object object>, /, **sizes) → Var#

einops-style relayout in one line: buf.rearrange("b (2 r) -> 2 b r").

A pure reshape+permute+reshape over the SAME physical bytes, spelled as an einops pattern. Lowers to view (split lhs groups) → permute (reorder to rhs atom order) → view (merge rhs groups), so it inherits whatever the underlying axis machinery does: a plain (unswizzled) buffer collapses to a flat layout, a swizzled buffer keeps its swizzle, and a tmem buffer carries allocated_addr through. It therefore does NOT flatten a swizzle atom — the same pattern on a swizzled SMEM buffer vs an unswizzled TMEM buffer legitimately yields different physical layouts (that is the point: rearrange acts on the operand, not a string).

pattern is "lhs -> rhs"; each side is space-separated axis names, with (a b) grouping a product axis. Every lhs group’s product must equal that input dim; at most one axis per group may be unknown (inferred from the dim), the rest supplied via **sizes. Cannot express a replica (R[...]), a stride-fiction/padded view, or a reshape crossing a swizzle-atom boundary — keep those as explicit view(layout=...).

scope()#

Return the storage scope associated with this buffer. :returns: scope – The storage scope associated with this buffer. :rtype: str

property sub: SubIndexer#

buf.sub[2, 4:8, ::4].

Unlike plain buf[...] (BufferLoad for scalar indices, extent-1 BufferRegion dims for tile-primitive operands), sub follows numpy basic-indexing semantics as a view constructor: an integer index removes the dim (select), a:b narrows it, and a::s takes every s-th element (requires the extent divisible by s and a < s). Trailing dims are kept whole.

Type:

Numpy-style view indexer

tile(*specs) → TileIndexer#

Chunk a dim: split it into factors, pick a chunk, keep the rest.

Rank-preserving — the picked dim’s remaining factors merge back into that one dim, and every other dim is untouched, so N dims in gives N dims out. Chunk multiple dims by chaining (dims never shift): buf.tile(0, (nx, -1))[cx, :].tile(1, (-1, ny))[:, cy].

Call as tile(dim, factors) for one dim, or pass several (dim, factors) specs as sugar for a chain. factors is the tuple the dim splits into (row-major, like unflatten(); one -1 inferred). The indexer takes one entry per factor: an int / Expr picks it (fixing the chunk, dropping the axis, folding its offset) and : keeps it. At least one factor per dim must be picked — a pure keep-everything split is unflatten(), not a chunk:

# 64 rows split into (stripe, warp, row) = (-1, WARPS, 4); this
# warp's 16 interleaved rows (stripe x row merged):
buf.tile(1, (-1, WARPS, 4))[:, warp, :]

tile(d, (n, -1))[c, :]   # contiguous block c
tile(d, (-1, n))[:, c]   # round-robin chunk c

A picked index may be a dynamic Expr (e.g. a warp id); picking several factors of one dim is allowed.

view(*args, **kwargs) → Var#

Creates a new view of the buffer. (used by parser)

Supported signatures are view(*shape, layout=None), where shape can contain -1 to indicate that the dimension size is auto-inferred, and view(dtype: Union[str, tvm.DataType]).

Returns:

view – The corresponding view buffer.

Return type:

DeclBufferFrame

vload(begin, dtype=None, predicate=None)#

Generate an Expr that loads dtype from begin index.

Parameters:

  • begin (Array of Expr) – The beginning index in unit of Buffer.dtype

  • dtype (str) – The data type to be loaded, can be vector type which have lanes that is multiple of Buffer.dtype

  • predicate (Optional[Expr]) – A vector mask of boolean values indicating which lanes of a vector are to be loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

Returns:

load – The corresponding load expression.

Return type:

Expr

vstore(begin, value, predicate=None)#

Generate a Stmt that store value into begin index.

Parameters:

  • begin (Array of Expr) – The beginning index in unit of Buffer.dtype

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

  • predicate (Optional[Expr]) – A vector mask of boolean values indicating which lanes of a vector are to be stored. The number lanes of the mask must be equal to the number of lanes in value.

Returns:

store – The corresponding store stmt.

Return type:

Stmt

with_allocated_addr(allocated_addr)#

Return a new buffer with the allocated address.

with_dtype(dtype)#

Return a new buffer with the dtype.

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.

class tvm.ir.BaseFunc(attrs, span=<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(*args: Any, **kwargs: Any)#

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.

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