tvm_tirx.ir

Contents

tvm_tirx.ir#

Shared IR types, expressions, statements, source locations, and modules.

Functions:

const(value[, ty, loc, dtype])

Construct a constant with an explicit or inferred primitive type.

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

Classes:

Attrs()

Base class for structured operator and function attributes.

DictAttrs(dict)

A string-keyed attribute map for functions, modules, and calls.

EnvFunc(name, func)

A registered global function referenced and serialized by name.

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 source location together with its caller's location.

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.

NoneType()

The concrete type whose only value is None.

OpaqueType()

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

PtrType(element_type[, scope, loc])

A pointer to an element type in a specified memory scope.

PrimType(dtype)

A primitive scalar or vector data type, including the void sentinel.

StringType()

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

TensorRegionType()

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

TupleType(fields[, loc])

The type of tuple values.

Type()

The base class of all types.

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

Call a registered operator or another callable IR expression.

CallEffectKind(value)

Possible call effects, matching the shared C++ IR effect kinds.

IntImm(ty, value[, loc])

Int constant.

FloatImm(ty, value[, loc])

Float constant.

Constant([loc, ty])

Base class of literal constants.

GenericConst(ty, value[, 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[, loc])

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

A half-open interval stored as begin and extent.

TensorLoad(source, indices[, loc])

Load an element or vector from an indexed 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 for callable IR function definitions.

CallingConv(value)

Possible kinds of calling conventions.

GlobalInfo()

Base class for module-level information referenced by the IR.

IRModule([functions, attrs, global_infos])

A collection of global functions and module-level metadata.

Op()

A registered primitive operator used as a Call callee.

OpAttrMap(attr_name, column)

The values of one operator attribute, keyed by operator.

Stmt(loc)

Base class of all the statements.

SeqStmt(seq[, loc])

Sequence of statements.

Bind(var, value[, loc])

Bind node.

Evaluate(value[, loc])

Evaluate an expression for its effects and discard its result.

Return(value[, loc])

Exit the enclosing function with a return value.

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

Execute one branch according to a boolean condition.

ForKind(value)

Execution semantics of a For loop.

For(var, begin, end, step, kind, body[, ...])

Iterate from begin while below end, incrementing by step.

While(condition, body[, loc])

Repeat the body while the condition evaluates to true.

Break([loc])

Exit the innermost enclosing loop.

Continue([loc])

Skip to the next iteration of the innermost enclosing loop.

AssertStmt(condition, error_kind[, ...])

Check a condition and raise the specified error when it is false.

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.
tvm_tirx.ir.const(value, ty=None, loc: Location = UnknownLoc(), *, dtype=None) → Expr#

Construct a constant with an explicit or inferred primitive type.

The primitive type may be supplied positionally, as ty= or as dtype=. Supply at most one non-None type. When both are omitted or None, boolean, integer and floating-point scalars use bool, int32 and float32, respectively. The same defaults apply to NumPy scalars, regardless of their original width. Values outside the inferred type’s range raise an error; the type never widens automatically. NaN and infinity are accepted floating-point values. Specify ty or dtype explicitly to select another width or a vector type.

Parameters:

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

  • ty (str, DataType or tvm_tirx.ir.PrimType, optional) – The primitive type. None uses the scalar-kind defaults above.

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

  • dtype (str, DataType or tvm_tirx.ir.PrimType, optional) – Keyword-only alias for ty. None leaves ty unchanged. Supplying non-None values for both ty and dtype raises TypeError.

Returns:

const_val – The result expression.

Return type:

tvm_tirx.ir.Expr

Examples

ir.const(True)                 # bool
ir.const(1)                    # int32
ir.const(1.0)                  # float32
ir.const(1 << 31)              # raises ValueError: exceeds int32 range
ir.const(1 << 31, ty="int64")  # explicit wider type
ir.const(1 << 31, dtype="int64")  # equivalent type keyword
class tvm_tirx.ir.Attrs#

Base class for structured operator and function attributes.

Concrete subclasses declare named fields exposed through attribute access or attrs["field"]. Use make_node() to construct a registered attribute type by its type key and fields.

class tvm_tirx.ir.DictAttrs(dict)#

A string-keyed attribute map for functions, modules, and calls.

Construct with DictAttrs(dict={"key": value}) or make_node("ir.DictAttrs", key=value). Entries are available through indexing, get(), and attribute access; dict holds the underlying map.

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_tirx.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. For ir.DictAttrs, these are the dictionary entries.

Returns:

node – The corresponding IR Node

Return type:

Node

Note

If the created node is instance of AttrsObj, 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_tirx.ir.make_node(
    "ir.IntImm", ty=tvm_tirx.ir.PrimType("int32"), value=10, loc=tvm_tirx.ir.UnknownLoc()
)
assert isinstance(x, tvm_tirx.ir.IntImm)
assert x.value == 10
class tvm_tirx.ir.EnvFunc(name, func)#

A registered global function referenced and serialized by name.

Use EnvFunc.get(name) to look up a function and call the returned object with its arguments. Deserialization resolves the same name in the registry.

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_tirx.ir.Node#

Base class of all IR Nodes.

class tvm_tirx.ir.Scriptable#

Marker for display methods installed by the TVMScript owner.

Methods:

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

Print TVM IR into TVMScript text format

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

Print highlighted TVMScript.
script(*, name: str | None = None, show_meta: bool = False, 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, 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, extra_config: dict | 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

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

  • 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[tvm_ffi.Object]] = None) – Object to be underlined

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

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

Returns:

script – The TVM Script of the given TVM IR

Return type:

str

show(style: str | None = None, *, name: str | None = None, show_meta: bool = False, 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, 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, extra_config: dict | None = None) → None#

Print highlighted TVMScript.

Parameters:

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

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

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

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

  • 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[tvm_ffi.Object]] = None) – Object to be underlined

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

  • extra_config (Optional[dict] = None) – Dialect-specific configuration passed through to PrinterConfig.extra_config, including ir.prefix for the IR namespace alias (default "I").

tvm_tirx.ir.assert_structural_equal(lhs, rhs, map_free_vars=False)#

Assert lhs and rhs are structurally equal to each other.

Parameters:

  • lhs (tvm_ffi.Object) – The left operand.

  • rhs (tvm_ffi.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_tirx.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:

tvm_ffi.Object

tvm_tirx.ir.save_json(node) → str#

Save tvm object as json string.

Parameters:

node (tvm_ffi.Object) – A TVM object to be saved.

Returns:

json_str – Saved json string.

Return type:

str

class tvm_tirx.ir.SourceName(name)#

An identifier for a source location.

Parameters:

name (str) – The name of the source.

class tvm_tirx.ir.Location#

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

IR constructors default to the shared UNKNOWN_LOC.

class tvm_tirx.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_tirx.ir.UnknownLoc#

The canonical immutable location used when source information is unavailable.

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

A callee’s source location together with its caller’s location.

Parameters:

  • callee (Location) – The location of the called or inlined code.

  • caller (Location) – The calling location, which may itself be a CallSiteLoc.

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

The top type, which admits any value.

class tvm_tirx.ir.FuncType(arg_types, ret_type)#

Function type.

Records the argument types in order and a single return type.

Parameters:
class tvm_tirx.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_tirx.ir.NoneType#

The concrete type whose only value is None.

class tvm_tirx.ir.OpaqueType#

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

class tvm_tirx.ir.PtrType(element_type, scope=MemoryScope('global'), *, loc=UnknownLoc())#

A pointer to an element type in a specified memory scope.

Parameters:

  • element_type (tvm_tirx.ir.Type) – The type of pointer’s element.

  • scope (str or tvm_tirx.runtime.MemoryScope) – The storage scope into which the pointer addresses.

class tvm_tirx.ir.PrimType(dtype)#

A primitive scalar or vector data type, including the void sentinel.

Parameters:

dtype (str or tvm_ffi.dtype) – Element kind, bit width, and optional vector lane count, such as float32x4.

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_tirx.ir.StringType#

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

class tvm_tirx.ir.TensorRegionType#

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

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

The type of tuple values.

Parameters:

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

class tvm_tirx.ir.Type#

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.

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.

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_tirx.ir.Call(op: Expr | str, args: list[Expr] | tuple[Expr, ...], attrs: Attrs | dict | None = None, ty_args: list[tvm_tirx.ir.Type] | tuple[tvm_tirx.ir.Type, ...] | None = None, loc: Location = UnknownLoc(), *, ty: Type | str | None = None)#

Call a registered operator or another callable IR expression.

Omitted or None keyword-only 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. Constructing a Call never emits a builder statement, including calls to allocation operators.

For an Op callee, a registered fixed return type takes precedence over its type-inference callback. For another callee with a FuncType, inference uses the signature’s return type. Without an applicable rule, the result is MissingType. An explicit ty, including MissingType, bypasses inference.

Parameters:

  • op (Expr or str) – The callee expression, or a registered operator name resolved by Op.get.

  • args (sequence[Expr]) – Positional value arguments.

  • attrs (Attrs or dict, optional) – Operator-specific attributes. A dictionary becomes DictAttrs.

  • ty_args (sequence[Type], optional) – Explicit type arguments, empty by default.

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

  • ty (Type or str, optional) – Keyword-only result type; a primitive dtype string is accepted.

Methods:

validate()

Run the registered operator's validator without changing this Call.
validate() → None#

Run the registered operator’s validator without changing this Call.

An Op callee with a validator invokes that validator; other callees and operators without validators perform no checks. Success returns None, and failed checks raise a diagnostic exception. The checks depend on the registered validator and do not generally establish agreement between the stored result type and the type implied by the inputs.

class tvm_tirx.ir.CallEffectKind(value)#

Possible call effects, matching the shared C++ IR effect kinds.

class tvm_tirx.ir.IntImm(ty: str | dtype | PrimType, value: int, loc: Location = UnknownLoc())#

Int constant.

Parameters:

  • ty (str, tvm_ffi.dtype or PrimType) – The data type

  • value (int) – The constant value.

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

class tvm_tirx.ir.FloatImm(ty: str | dtype | PrimType, value: float, loc: Location = UnknownLoc())#

Float constant.

Parameters:

  • ty (str, tvm_ffi.dtype or PrimType) – The data type

  • value (float) – The constant value.

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

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

Base class of literal constants.

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

A literal payload with an explicit expression type.

Parameters:

  • ty (Type) – The type assigned to the literal.

  • value (object) – The payload, convertible to an FFI value.

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

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

A data type literal whose expression type is AnyType.

Parameters:

  • value (str or tvm_ffi.dtype) – The represented data type.

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

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

A string literal with StringType.

Parameters:

  • value (str) – The literal text.

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

class tvm_tirx.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_tirx.ir.ExprWithOp(loc=<MISSING>, ty=<MISSING>)#

Common type-directed operator behavior for core expressions.

class tvm_tirx.ir.GlobalVar(name_hint: str, loc: Location = UnknownLoc())#

A global variable in the IR.

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

Each construction creates a distinct reference, even when names match. Use module.get_global_var(name) to retrieve the stored reference when constructing calls or replacing definitions.

Parameters:

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

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

class tvm_tirx.ir.LambdaExpr(parameter_types, function: Callable, *, ret_type=None, loc: Location = UnknownLoc())#

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. In TVMScript, use T.Lambda to construct it.

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.

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

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_tirx.ir.StagingExpr(loc=<MISSING>, ty=<MISSING>)#

A traversable expression eliminated before executable IR.

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

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

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

A half-open interval stored as begin and extent.

Use from_min_extent() when the extent is already known.

Parameters:

  • begin (Expr) – The lower bound when end is provided; otherwise the extent of [0, begin).

  • end (Expr, optional) – The exclusive upper bound. With two arguments, extent is end - begin.

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

Methods:

from_min_extent(min_value, extent[, loc])

Construct a Range by min and extent.

get_extent_from_begin_end(begin, end)

Extract the width from endpoints, reusing a shallow addition operand.
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

static get_extent_from_begin_end(begin: Expr, end: Expr) → Expr#

Extract the width from endpoints, reusing a shallow addition operand.

class tvm_tirx.ir.TensorLoad(source: Expr, indices: list[Expr], loc: Location = UnknownLoc())#

Load an element or vector from an indexed expression source.

The source type’s registered inference hook validates the inputs and derives the primitive result type.

Parameters:

  • source (Expr) – The tensor or other expression supporting indexed loads.

  • indices (list[Expr]) – Indices interpreted by the source type, in dimension order.

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

class tvm_tirx.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_tirx.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_tirx.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_tirx.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_tirx.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, optional) – The display name. Variables with the same name remain distinct bindings.

  • ty (Optional[Type or str]) – The variable type, or MissingType when omitted. A string denotes a primitive dtype; "handle" denotes a pointer to void.

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

  • name_hint (str, optional) – Keyword-only alternative to name; supply exactly one name spelling.

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

Return whether an expression has a primitive result type.

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

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

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

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

For an Op callee, a registered TFixedReturnType takes precedence over FInferType. For another callee whose type is a FuncType, the result is that function type’s return type. Without an applicable rule, this returns MissingType. Inference errors propagate.

This does not invoke the operator’s validator or overwrite call.ty.

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

Base class for callable IR function definitions.

attrs stores function metadata as DictAttrs. Attribute update methods return a new function, preserving the original. Concrete function classes define the parameters, body, and function type in ty.

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 function and add the given attribute map to it.

without_attr(attr_key)

Copy the function and remove the given attribute key.
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 (tvm_ffi.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 function and add the given attribute map to it. :param attr_map: The attribute map :type attr_map: Union[DictAttrs, Dict[str, tvm_ffi.Object]]

Returns:

func – A new copy of the function

Return type:

BaseFunc

without_attr(attr_key: str) → BaseFunc#

Copy the function and remove the given attribute key.

Parameters:

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

Returns:

func – A new copy of the function

Return type:

BaseFunc

class tvm_tirx.ir.CallingConv(value)#

Possible kinds of calling conventions.

class tvm_tirx.ir.GlobalInfo#

Base class for module-level information referenced by the IR.

Concrete subclasses are stored in IRModule.global_infos, a map from names to arrays of global information objects.

Methods:

same_as(other)

Overload with structural equality.
same_as(other)#

Overload with structural equality.

class tvm_tirx.ir.IRModule(functions=None, attrs=None, global_infos=None)#

A collection of global functions and module-level metadata.

Modules are the input and output of compiler passes. Functions can be looked up by their GlobalVar or name.

Parameters:

  • functions (dict[GlobalVar or str, BaseFunc], optional) – Global function definitions. String keys are converted to GlobalVar objects.

  • attrs (dict, optional) – Module attributes, stored as DictAttrs.

  • global_infos (dict[str, list[GlobalInfo]], optional) – Named groups of global information referenced by the IR.

Methods:

clone()

Create a module object with the same functions and metadata.

functions_items()

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

update(other)

Insert function definitions from another module or dictionary in place.

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.
clone() → IRModule#

Create a module object with the same functions and metadata.

The clone initially shares its function objects and other contents with the original. Adding, replacing, or removing function-map entries on the clone preserves the original module’s entries. Contents are not recursively copied.

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 function definitions from another module or dictionary in place.

Definitions are merged using their GlobalVar keys. A dictionary’s string keys create fresh GlobalVar objects; a fresh key with the same name as an existing key raises RuntimeError. To replace an existing definition, reuse its stored GlobalVar or use assignment by name.

Parameters:

other (IRModule or dict[GlobalVar or str, BaseFunc]) – Function definitions to merge into this module. Module attributes and global_infos are not merged.

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 (tvm_ffi.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, tvm_ffi.Object]]

Returns:

mod – A new copy of the IRModule with the attribute

Return type:

IRModule

class tvm_tirx.ir.Op#

A registered primitive operator used as a Call callee.

Retrieve the shared operator with Op.get(name). Its registration supplies argument signatures, attributes, validation, and result-type inference; constructing a Call does not itself execute the operator.

Methods:

get(op_name)

Get a registered operator by name.

list_op_names()

List registered operator names in unspecified order.

list_attr_names()

List registered attribute column names in unspecified order.

get_attr_map(attr_name)

Get the map from operators to their values of an attribute.

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.

expose_script()

Expose the completed operator definition in TVMScript.

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]

static list_attr_names()#

List registered attribute column names in unspecified order.

Returns:

The attribute column names, such as "TCallEffectKind".

Return type:

list[str]

static get_attr_map(attr_name)#

Get the map from operators to their values of an attribute.

This is the Python counterpart of C++ Op::GetAttrMap. The map reads the registry’s storage directly, so later registrations, overrides and removals are visible through it, and it can be cached.

Parameters:

attr_name (str) – Attribute name, such as "TCallEffectKind".

Returns:

The attribute map, or None when no operator has registered the attribute.

Return type:

OpAttrMap or None

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

expose_script() → None#

Expose the completed operator definition in TVMScript.

Call after setting the signature and attributes. Early requests wait until the script builder is initialized; explicit builders take precedence.

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.

class tvm_tirx.ir.OpAttrMap(attr_name, column)#

The values of one operator attribute, keyed by operator.

Mirrors C++ OpAttrMap. Obtain one with Op.get_attr_map(). Lookups see later registrations, overrides and removals; register values with Op.set_attr(). Like the registry itself, the map is not synchronized with concurrent registration.

Attributes:

attr_name

The attribute name.

Methods:

get(expr[, default])

Get the attribute of an operator with a fallback.
property attr_name#

The attribute name.

get(expr, default=None)#

Get the attribute of an operator with a fallback.

Parameters:

  • expr (Expr) – The operator. Any other expression yields default.

  • default (object) – The value to return when expr is not an Op or has no value.

Returns:

The registered value or default.

Return type:

object

tvm_tirx.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 must already be registered.

Return type:

object or function

class tvm_tirx.ir.Stmt(loc)#

Base class of all the statements.

class tvm_tirx.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_tirx.ir.Bind(var: Var, value: Expr, loc: Location = UnknownLoc())#

Bind node.

Bind a variable to a value in the enclosing scope. Parsed TVMScript normally creates bindings with ordinary assignments. Use explicit T.Bind(var, value) when constructing a binding for an existing variable. Bind has no body field. The bound variable is visible in all subsequent statements within the same enclosing scope (SeqStmt, ForObj.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_tirx.ir.Evaluate(value: Expr, loc: Location = UnknownLoc())#

Evaluate an expression for its effects and discard its result.

Parameters:

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

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

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

Exit the enclosing function with a return value.

Parameters:

  • value (Expr) – The value to return.

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

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

Execute one branch according to a boolean condition.

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_tirx.ir.ForKind(value)#

Execution semantics of a For loop.

  • DEFAULT executes iterations in order, one after another.

  • PARALLEL allows different iterations to run concurrently.

  • VECTORIZED asks the compiler to combine iterations into vector operations that process multiple elements at once.

  • UNROLLED asks the compiler to expand the loop into repeated copies of its body.

Transformations must preserve the selected execution semantics.

class tvm_tirx.ir.For(var: Var, begin: Expr, end: Expr, step: Expr, kind: ForKind, body: Stmt | Sequence[Stmt], annotations: Mapping[str, Object] | None = None, loc: Location = UnknownLoc())#

Iterate from begin while below end, incrementing by step.

Parameters:

  • var (Var) – The loop variable.

  • begin (Expr) – The beginning value.

  • end (Expr) – The exclusive end of the loop.

  • step (Expr) – The required loop step, including a typed one for unit stride.

  • kind (ForKind) – Sequential, parallel, vectorized, or unrolled execution semantics.

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

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

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

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

Repeat the body while the condition evaluates to true.

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_tirx.ir.Break(loc: Location = UnknownLoc())#

Exit the innermost enclosing loop.

Parameters:

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

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

Skip to the next iteration of the innermost enclosing loop.

Parameters:

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

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

Check a condition and raise the specified error when it is false.

Parameters:

  • condition (Expr) – The assert condition.

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

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

Parameters:

  • op (Op or str) – Registered region operator or its name.

  • args (sequence[Expr]) – Operands evaluated in the enclosing scope.

  • body_params (sequence[Var]) – Variables bound for the body according to the operator’s contract.

  • attrs (DictAttrs or Mapping or None) – Operator-specific attributes; None gives an empty map.

  • body (Stmt or Sequence[Stmt]) – Region body, normalized to SeqStmt.

  • result_vars (sequence[Var], optional) – Variables defined for subsequent statements, empty by default.

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

class tvm_tirx.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_tirx.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

Contents