tvm.relax¶

The Relax IR namespace containing the IR, type, operator, builder, vm, etc.

class tvm.relax.Binding¶: The base class of a binding in Relax.

class tvm.relax.BindingBlock(bindings: List[Binding], span: Span | None = None)¶: base class of binding block, bindings inside can be impure (with side effect or control flow)

Function call node in Relax.

relax.Call node corresponds the operator application node in computational graph terminology.

Parameters:

op (tvm.ir.Op or any tvm.relax.Expr with function type.) – The operation to be called.
args (Union[List[Expr], Tuple[Expr, ...]]) – The arguments to the call.
attrs (Optional[tvm.ir.Attrs]) – Attributes to the call, can be None
sinfo_args (Optional[Union[List[StructInfo], Tuple[StructInfo, ...]]]) – The structure info arguments of a CallNode. sinfo_args is designed to be non-empty only for intrinsic op (e.g., call_tir, call_builtin_with_ctx, etc.) and calls to ExternFuncs, with the main usage of structure info inference.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.Constant(data: NDArray, struct_info: StructInfo | None = None, span: Span | None = None)¶

Constant Tensor

Parameters:

data (tvm.nd.NDArray) – The data of the constant tensor.
struct_info (Optional[StructInfo]) – The struct info of the constant tensor. If not specified, infer it from data.
span (Optional[Span]) – Span that points to original source code

Note

Scalar constants are represented by ndim-0 constant tensors.

class tvm.relax.DataTypeImm(value: DataType | str, span: Span | None = None)¶: Represent a data type constant.

class tvm.relax.DataflowBlock(bindings: List[Binding], span: Span | None = None)¶: dataflow block, bindings inside are pure (no side effect and no control flow)

class tvm.relax.DataflowBlockRewrite(dfb: DataflowBlock, root_fn: Function)¶

A binding/statement-level dataflow block rewriter.

Notes

Due to the immutable and copy-on-write nature of TVM AST nodes, the rewriting is not done in place. Instead, a new DataflowBlock is created and returned with mutated_dfb. Similarly, its new root Function is created and returned by mutated_root_fn. To apply this change for an IRModule, use mutate_irmodule which rewrites the old function that registered in the constructor.

add(expr: Expr, name: str | None = None, is_dfvar: bool = False) → None¶

Add a new statement to the DataflowBlock with an automatically generated variable name.

Parameters:

expr (Expr) – The expression to add.
name (Optional[str], optional) – Variable name, by default None
is_dfvar (bool, optional) – The variable type, by default False

Notes

If the variable name is not given, it will be automatically generated in a form of “tmp${COUNTER}”. The variable type will be DataflowVar if is_dfvar is True, otherwise it will be Var. Being relax.Var means the variables are output variables of the DataflowBlock. While being DataflowVar means the variables are internal variables of the DataflowBlock.

mutate_irmodule(irmodule: IRModule) → IRModule¶

Return an updated IRModule by replacing the old function with the mutated root function.

Parameters:: irmodule (tvm.IRModule) – The base IRModule to update.
Returns:: The updated IRModule.
Return type:: tvm.IRModule

mutated_dfb() → DataflowBlock¶: Returns the mutated DataflowBlock.

mutated_root_fn() → Function¶: Returns the mutated root function.

remove_all_unused() → None¶

Remove all unused variables.

Notes

This could remove unused variables in other DataflowBlocks as well.

remove_unused(var: Var, allow_undef=False) → None¶

Remove a statement by its variable definition if and only if it is unused.

Parameters:

var (relax.Var) – The unused variable definition.
allow_undef (bool, optional) – Whether to allow var being undefined variable, by default False

Raises:

TVMError if the variable is used or undefined (allow_undef=False). –

replace_all_uses(old_var: Var, new_var: Var) → None¶

Replace all uses of old_var with new_var.

Parameters:

old_var (relax.Var) – The old variable to replace.
new_var (relax.Var) – The new variable to replace with.

class tvm.relax.DataflowVar(name_hint: str | Id, struct_info: StructInfo | None = None, span: Span | None = None)¶

A sub-type of the variable node used to mark dataflow variables from normal visible “function local” bindings.

Parameters:

name_hint (Union[str, Id]) – The name hint of the variable.
struct_info (Optional[StructInfo]) – The struct info annotation of the variable.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.DynTensorType(ndim=-1, dtype='float32', span: Span | None = None)¶

A dynamic tensor type in Relax.

This is the type assigned to tensors with a known dtype and unknown shape.

Parameters:

ndim (Optional[int]) – The ndim of the Tensor
dtype (Optional[str]) – The content data type.

class tvm.relax.ExecBuilder¶

A builder to emit instructions and build executable for the virtual machine.

c(idx: int) → int¶: set instruction’s argument as a constant.

declare_function(func_name: str, kind: VMFuncKind = VMFuncKind.PACKED_FUNC) → None¶: Declare a function

emit_call(name: str, args: List[NDArray | DataType] | None = None, dst: int | None = None) → None¶: emit a call instruction which calls a packed function.

emit_goto(pc_offset)¶: emit a goto instruction

emit_if(cond, false_offset)¶: emit an if instruction

emit_ret(result: int) → None¶: emit a return instruction

f(name: str) → int¶: set instruction’s argument as a function.

function(func_name: str, num_inputs: int | None = 0, param_names: List[str] | None = None) → VMFuncScope¶: annotate a VM function.

get() → Executable¶: return the executable

imm(value: int) → int¶: set instruction’s argument as an immediate.

r(idx: int) → int¶: set instruction’s argument as a register.

class tvm.relax.Executable(mod: Module)¶

The executable object emitted by the VM compiler or the ExecBuilder.

as_python() → str¶: print the instructions as python program.

as_text() → str¶: print the instructions as text format.

export_library(file_name: str, fcompile: str | callable | None = None, workspace_dir: str | None = None, **kwargs) → Any¶

Export the executable to a library which can then be loaded back.

Parameters:

file_name (str) – The name of the shared library.
fcompile (function(target, file_list, kwargs), optional) – The compilation function to use create the final library object during
workspace_dir (str, optional) – The path of the directory used to create the intermediate artifacts when exporting the module. If this is not provided a temporary dir will be created.
kwargs (dict, optional) – Additional arguments passed to fcompile

Returns:

result of fcompile() – If the compilation function returns an artifact it would be returned via export_library, if any.

Return type:

unknown, optional

Examples

ex = relax.build(mod, target)
# export the library
ex.export_library("exported.so")

# load it back for future uses.
rt_mod = tvm.runtime.load_module("exported.so")
vm = tvm.relax.VirtualMachine(rt_mod, tvm.cuda())

jit(fcompile=None, addons=None, **kwargs) → Module¶

Just-in-time compile and link the modules.

The Executable returned by relax.build may not be directly runnable as they may contain cuda source files and objects that are yet to be compiled and linked. This function helps to create a runtime.Module for these cases.

Parameters:

fcompile (function(target, file_list, kwargs), optional) – The compilation function to use create the final library object during
kwargs (dict, optional) – Additional arguments passed to fcompile

Returns:

rt_mod – A runnable runtime module that can be passed to VirtualMachine.

Return type:

tvm.runtime.Module

Examples

ex = relax.build(mod, target)
# build a runnable module using nvcc to link everything
rt_mod = ex.jit()
vm = tvm.relax.VirtualMachine(rt_mod, tvm.cuda())

stats() → str¶: print the detailed statistics of the executable.

tvm.relax.Expr¶: alias of RelayExpr

class tvm.relax.ExprFunctor¶

An abstract visitor defined over Expr. Defines the default dispatch over expressions, and implements memoization.

visit_expr(expr: Expr) → Expr¶: Apply the visitor to an expression.

class tvm.relax.ExternFunc(global_symbol: String, span: Span | None = None)¶: extern function, which represents a PackedFunc.

class tvm.relax.FuncStructInfo(params: List[StructInfo], ret: StructInfo, purity: bool = True, span: Span | None = None)¶

StructInfo of a function value.

Parameters:

params (List[StructInfo]) – The struct info of the fields.
ret (StructInfo) – The struct info of return value
purity (bool) – Whether the function is pure (has no visible side effects). Note: We consider a function to be pure only if it is pure on all inputs. If a function can have visible side effects only in some cases, we still consider it impure.

static opaque_func(*, ret: StructInfo | None = None, derive_func: EnvFunc | None = None, purity: bool = False, span: Span | None = None) → FuncStructInfo¶

Create an opaque FuncStructInfo.

The opaque function takes either a ret that specificies the struct info of the return value or a derive_func that provides a customized derivation rule.

Parameters:

ret (Optional[StructInfo]) – The struct info of the function return value.
derive_func (Optional[EnvFunc]) – The environment function used for derivation
purity (bool) – Whether the function is pure (false by default, as most opaque functions are not pure)
span (Optional[Span]) – Optional span information of the ast.

Returns:

info

Return type:

FuncStructInfo

Note

We cannot specify ret and derive_func simultaneously.

class tvm.relax.Function(params: List[Var], body: Expr, ret_struct_info: StructInfo | None = None, is_pure: bool | None = True, attrs: DictAttrs | None = None, span: Span | None = None)¶

A Relax function.

Return a new function with updated symbolic variable

Parameters:

binding_map (Mapping[) – Union[str, relax.Var], Union[int, float, PrimExpr, tvm.runtime.NDArray, _np.ndarray, Expr],
] –
The mapping of values to be replaced.

Keys may be either a relax.Var or a string name of the Relax variable. If the variables are referred to by name, the name must uniquely identify a parameter in the function.

Values must be a relax expression, or a value that is convertible into a relax expression. The value must be compatible with the variable being replaced.

Returns:

func – The updated function

Return type:

Function

bind_symbolic_vars(binding_map: Mapping[str | Var, PrimExpr]) → Function¶

Return a new function with updated symbolic variable

Parameters:: binding_map (Mapping[Union[str, tvm.tir.Var], PrimExpr]) – The mapping of values to be replaced. Keys may be either a tir.Var or a string name of the variable. If the variables are referred to by name, the name must uniquely identify a symbolic variable in the function.
Returns:: func – The updated function
Return type:: Function

static create_empty(params: List[Var], ret_struct_info: StructInfo, is_pure: bool | None = True, attrs: DictAttrs | None = None, span: Span | None = None)¶: Construct a relax.Function but without body

class tvm.relax.Id¶: Unique identifier(name) used in Var. Guaranteed to be stable across all passes.

class tvm.relax.If(cond: Expr, true_branch: Expr, false_branch: Expr, span: Span | None = None)¶

A conditional expression in Relax.

Parameters:

cond (Expr) – The condition.
true_branch (Expr) – The expression evaluated when condition is true.
false_branch (Expr) – The expression evaluated when condition is false.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.MatchCast(var: Var, value: Expr, struct_info: StructInfo, span: Span | None = None)¶

Runtime-match the value to the struct info.

This operation does runtime check, populates the un-defined symbolic shape vars and vars in struct_info in the first occurrence, and insert equality assertions in other cases.

Parameters:

var (relax.Var) – The return variable that the match cast bind to.
value (Expr) – The input value expression.
struct_info (tvm.relax.StructInfo) – The struct info to match cast to.

class tvm.relax.ObjectStructInfo(span: Span | None = None)¶: StructInfo of an Object.

class tvm.relax.ObjectType(span: Span | None = None)¶: A type that corresponds to tvm::runtime::Object, is base of all possible object values in TVM.

class tvm.relax.PackedFuncType(span: Span | None = None)¶: The type of ExternFunc in Relax.

StructInfo of a primitive POD value.

Parameters:: dtype_or_expr (Union[str, DataType, PrimExpr]) – The data type of the prim value, or a known expression for the prim value.

class tvm.relax.PrimValue(value: PrimExpr | int, span: Span | None = None)¶: The prim expr representing the value.

class tvm.relax.PyExprMutator(mod: IRModule | None = None)¶

An abstract ExprMutator with customized methods on the python-side. This is the user facing class for method overwriting inheritance. _tvm_metadata discribes the class to inherit(“cls”), the methods that users can overwrite(“methods”), the constructor’s parameters(“fields”)

Note: @relax.expr_functor.mutator is required for proper usage of any inherited class.

See also: visitor, _PyExprVisitor

Example:

@relax.expr_functor.visitor
def MyExprVisitor(PyExprVisitor):
    ...

visit_binding(binding: Binding) → None¶

Generic dispatcher for Binding. Users can customized this function to overwrite VisitBinding(const Binding& binding) on the C++ side.

Parameters:: binding (Binding) – The binding to be visited.

visit_binding_block(block: BindingBlock) → None¶

Generic dispatcher for BindingBlock. Users can customized this function to overwrite VisitBindingBlock(const BindingBlock& block) on the C++ side.

Parameters:: block (BindingBlock) – The block to be visited.

visit_binding_block_(block: BindingBlock) → None¶

Visit BindingBlock. Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode* block) on the C++ side.

Parameters:: block (BindingBlock) – The BindingBlock to be visited.

visit_call_(op: Call) → None¶

Visit Call. Users can customized this function to overwrite VisitExpr_(const CallNode* op) on the C++ side.

Parameters:: op (relax.Call) – The relax.Call to be visited.

visit_constant_(op: Constant) → None¶

Visit Constant. Users can customized this function to overwrite VisitExpr_(const ConstantNode* op) on the C++ side.

Parameters:: op (Constant) – The Constant to be visited.

visit_data_type_imm_(op: DataTypeImm) → None¶

Visit DataTypeImm. Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op) on the C++ side.

Parameters:: op (DataTypeImm) – The DataTypeImm to be visited.

visit_dataflow_block_(block: DataflowBlock) → None¶

Visit DataflowBlock. Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode* block) on the C++ side.

Parameters:: block (DataflowBlock) – The DataflowBlock to be visited.

visit_dataflow_var_(op: DataflowVar) → None¶

Visit DataflowVar. Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op) on the C++ side.

Parameters:: op (DataflowVar) – The DataflowVar to be visited.

visit_dataflow_var_def_(var: DataflowVar) → None¶

Visit the DataflowVar definition site. Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var) on the C++ side.

Parameters:: var (DataflowVar) – The DataflowVar to be visited.

visit_expr(expr: Expr) → None¶

Generic dispatcher for Expr. Users can customized this function to overwrite VisitExpr(const Expr& expr) on the C++ side.

Parameters:: expr (Expr) – The expr to be visited.

visit_extern_func_(op: ExternFunc) → None¶

Visit ExternFunc. Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op) on the C++ side.

Parameters:: op (ExternFunc) – The ExternFunc to be visited.

visit_function_(op: Function) → None¶

Visit Function. Users can customized this function to overwrite VisitExpr_(const FunctionNode* op) on the C++ side.

Parameters:: op (Function) – The Function to be visited.

visit_global_var_(op: GlobalVar) → None¶

Visit GlobalVar. Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op) on the C++ side.

Parameters:: op (GlobalVar) – The GlobalVar to be visited.

visit_if_(op: If) → None¶

Visit If. Users can customized this function to overwrite VisitExpr_(const IfNode* op) on the C++ side.

Parameters:: op (If) – The If to be visited.

visit_match_cast_(binding: MatchCast) → None¶

Visit MatchCast. Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding) on the C++ side.

Parameters:: binding (MatchCast) – The MatchCast to be visited.

visit_op_(op: Op) → None¶

Visit Op. Users can customized this function to overwrite VisitExpr_(const OpNode* op) on the C++ side.

Parameters:: op (Op) – The Op to be visited.

visit_prim_value_(op: PrimValue) → None¶

Visit PrimValue. Users can customized this function to overwrite VisitExpr_(const PrimValueNode* op) on the C++ side.

Parameters:: op (PrimValue) – The PrimValue to be visited.

visit_seq_expr_(op: SeqExpr) → None¶

Visit SeqExpr. Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op) on the C++ side.

Parameters:: op (SeqExpr) – The SeqExpr to be visited.

visit_shape_expr_(op: ShapeExpr) → None¶

Visit ShapeExpr. Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op) on the C++ side.

Parameters:: op (ShapeExpr) – The ShapeExpr to be visited.

visit_span(span: Span) → None¶

Visit Span. Users can customized this function to overwrite VisitSpan(const Span& span) on the C++ side.

Parameters:: span (Span) – The Span to be visited.

visit_string_imm_(op: StringImm) → None¶

Visit StringImm. Users can customized this function to overwrite VisitExpr_(const StringImmNode* op) on the C++ side.

Parameters:: op (StringImm) – The StringImm to be visited.

visit_tuple_(op: Tuple) → None¶

Visit Tuple. Users can customized this function to overwrite VisitExpr_(const TupleNode* op) on the C++ side.

Parameters:: op (Tuple) – The Tuple to be visited.

visit_tuple_getitem_(op: TupleGetItem) → None¶

Visit TupleGetItem. Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op) on the C++ side.

Parameters:: op (TupleGetItem) – The TupleGetItem to be visited.

visit_var_(op: Var) → None¶

Visit Var. Users can customized this function to overwrite VisitExpr_(const VarNode* op) on the C++ side.

Parameters:: op (relax.Var) – The relax.Var to be visited.

visit_var_binding_(binding: VarBinding) → None¶

Visit VarBinding. Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding) on the C++ side.

Parameters:: binding (VarBinding) – The VarBinding to be visited.

visit_var_def(var: Var) → None¶

Generic dispatcher for visiting the var definition site. Users can customized this function to overwrite VisitVarDef(const relax.Var& var) on the C++ side. Note that visit_var_() will only visit the usage site of an Var.

Parameters:: var (relax.Var) – The var to be visited.

visit_var_def_(var: Var) → None¶

Visit the relax.Var definition site. Users can customized this function to overwrite VisitVarDef_(const VarNode* var) on the C++ side.

Parameters:: var (relax.Var) – The relax.Var to be visited.

class tvm.relax.SeqExpr(blocks: List[BindingBlock], body: Expr, span: Span | None = None)¶: A sequence of binding blocks followed by an expression.

class tvm.relax.ShapeExpr(values: List[PrimExpr] | Tuple[PrimExpr, ...] | Array, span: Span | None = None)¶

A shape expression which allows users to construct a shape containing PrimExpr.

Parameters:

values (Union[List[PrimExpr], Tuple[PrimExpr, ...], tvm.ir.Array]) – The values of the shape expression.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.ShapeStructInfo(values: List[PrimExpr] | None = None, ndim: int = -1, span: Span | None = None)¶

StructInfo of a shape value.

Parameters:

values (Optional[List[PrimExpr]]) – The symbolic shape values if known.
ndim (Optional[int]) – The size of the shape.

Note

Do not specify values and ndim at the same time.

class tvm.relax.ShapeType(ndim: int = -1, span: Span | None = None)¶

The type of shape in Relax.

Parameters:: ndim (Optional[int]) – The size of the shape.

class tvm.relax.StringImm(value: str, span: Span | None = None)¶: Represent a string literal constant.

class tvm.relax.StructInfo¶

The base class of all StructInfo.

StructInfo contains both the static type and runtime structural information.

is_base_of(derived: StructInfo) → bool¶

Check if self is base of another derived struct info.

Parameters:: derived (StructInfo) – The derived struct info to be checked.
Returns:: result – The check result.
Return type:: bool

same_as(other)¶: Overload with structural equality.

StructInfo of a Tensor value.

Parameters:

shape (Optional[Expr]) – The shape expression.
dtype (Optional[str]) – The content data type.
vdevice (Optional[Vdevice]) – The virtual device.
ndim (Optional[int]) – The number of dimensions of the tensor.

Note

Do not specify shape and ndim at the same time.

class tvm.relax.Tuple(fields: List[Expr] | Tuple[Expr, ...], span: Span | None = None)¶

Tuple expression that groups several fields together.

Parameters:

fields (Union[List[Expr], Tuple[Expr, ...]]) – The fields in the tuple.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.TupleGetItem(tuple_value: Expr, index: int, span: Span | None = None)¶

Get index-th item from a tuple.

Parameters:

tuple_value (Expr) – The input tuple expression.
index (int) – The index.
span (Optional[Span]) – Span that points to original source code

class tvm.relax.TupleStructInfo(fields: List[StructInfo], span: Span | None = None)¶

StructInfo of a Tuple value.

Parameters:: fields (List[StructInfo]) – The struct info of the fields.

class tvm.relax.VMInstrumentReturnKind(value)¶: An enumeration.

class tvm.relax.Var(name_hint: str | Id, struct_info: StructInfo | None = None, span: Span | None = None)¶

The variable class for all Relax bindings.

Parameters:

name_hint (Union[str, Id]) – The name hint of the variable.
struct_info (Optional[StructInfo]) – The struct info annotation of the variable.
span (Optional[Span]) – Span that points to original source code

property name_hint: str¶: Get name hint of the current var.

class tvm.relax.VarBinding(var: Var, value: Expr, span: Span | None = None)¶

Variable binding, bind he variable of the lhs with the rhs.

Parameters:

var (relax.Var) – The return variable that the match cast bind to.
value (Expr) – The input value expression.

class tvm.relax.VirtualMachine(rt_mod: Module | Executable, device: Device | List[Device], memory_cfg: str | Dict[Device, str] | None = None, profile: bool = False)¶

Relax VM runtime.

get_outputs(func_name: str) → Object | Tuple[Any]¶

Get the value output by the function by the given name after a call of invoke_stateful.

It is an error to call this function without first calling invoke_stateful.

Parameters:: func_name (str) – The name of the function whose output should be fetched.
Returns:: ret – The result of the earlier call to the function via invoke_stateful. If the result is a tuple, it returns a list of the fields. The fields are potentially also tuples, so these can be arbitrily nested.
Return type:: Union[tvm.Object, Tuple[Any]]

invoke_closure(closure: Object, *args: Any) → Object¶

Invoke a closure.

Parameters:

closure (Object) – The VMClosure Object.
args (list[tvm.runtime.NDArray] or list[np.ndarray]) – The arguments to the closure.

Returns:

result – The output.

Return type:

Object

invoke_stateful(func_name: str) → None¶

relax.Call the named function from the VM module using the arguments set using set_input. It is an error to call invoke_stateful without using set_input first (even if it’s to set 0 inputs); conversely, if set_input has been called, it is an error to call the function without using invoke_stateful.

The results of the call can be obtained by calling get_outputs.

Parameters:: func_name (str) – The name of the function to call.

profile(func_name: str, *args)¶

Profile a function call.

Parameters:

func_name (str) – The name of the function.
args (List of NDArray or other objects supported by PackedFunc.) – The arguments to the function.

Returns:

report – The formatted profiling result, showing per-op timing measurements.

Return type:

tvm.runtime.profiling.Report

save_function(func_name: str, saved_name: str, *args: List[Any], include_return: bool = True, **kwargs: Dict[str, Any]) → None¶

Convenience function. Takes a function from the module and saves a PackedFunc that, when called, will invoke the function with the given arguments. The PackedFunc can be accessed from the module using saved_name. This is included to facilitate timing trials: Invoking the returned PackedFunc will have less overhead from dictionary lookups than normally running through the VM.

If the saved name is taken, it can be overridden, though it cannot override the name of a function defined in the Relax source.

This is really creating a closure, but the function has a different name to avoid confusion with invoke_closure (they are not meant to be used together).

Parameters:

func_name (str) – The function that should be packaged up.
saved_name (str) – The name that the resulting closure should be saved under.
include_return (bool) – Whether the saved PackedFunc should return its output. If timing over RPC, it may not be desirable to send output between machines.
args (List[Any]) – The arguments to package up with the function.
kwargs (Dict[str, Any]) – Any named arguments to package up with the function

set_input(func_name: str, *args: Any, **kwargs: Any) → None¶

Set the inputs to a function. This interface works when using VM over RPC by internally converting NDArray in the arguments to DLTensor, which is supported in RPC where remote could only have a minimal C runtime.

Note: If set_input is used, the function must be called using invoke_stateful and the results must be obtained using get_outputs.

Parameters:

func_name (str) – The name of the function.
args (List[tvm.runtime.NDArray] or List[np.ndarray]) – The arguments to the function.
kwargs (dict of str to tvm.runtime.NDArray or np.ndarray) – Named arguments to the function.

set_instrument(instrument: PackedFunc) → None¶

Set an instrumentation function.

If instrument is present, the function will be called before/after each relax.Call instruction. The function have the following signature:

def instrument(
    func: Union[VMClosure, PackedFunc],
    func_symbol: str,
    before_run: bool,
    ret_value: any,
    *args) -> bool:
    pass

The instrument takes the following parameters: - func: function object to be called. - func_symbol: the symbol name of the function. - before_run: whether it is before or after call. - ret_value: the return value of the call, only valid after run. - args: the arguments being passed to call.

The instrument function can choose an integer, which corresponds to action direction for the following run. See VMInstrumentReturnKind for more details.

Parameters:: instrument (tvm.runtime.PackedFunc) – A instrumentation function that get invoked every VM call instr.

See also

VMInstrumentReturnKind: the possible return values in VM.

time_evaluator(func_name: str, dev: Device, number: int = 10, repeat: int = 1, min_repeat_ms: int = 0, cooldown_interval_ms: int = 0, repeats_to_cooldown: int = 1, f_preproc: str = '') → Callable[[...], BenchmarkResult]¶

Returns an evaluator that times a function in the module. This follows the same convention as time_evaluator in tvm.runtime.module. This can be used in combination with save_function() so that the timings avoid extra dictionary lookups.

Parameters:

func_name (str) – The name of the function in the module.
dev (Device) – The device we should run this function on.
number (int) – The number of times to run this function for taking average. We call these runs as one repeat of measurement.
repeat (int, optional) – The number of times to repeat the measurement. In total, the function will be invoked (1 + number x repeat) times, where the first one is warm up and will be discarded. The returned result contains repeat costs, each of which is an average of number costs.
min_repeat_ms (int, optional) – The minimum duration of one repeat in milliseconds. By default, one repeat contains number runs. If this parameter is set, the parameters number will be dynamically adjusted to meet the minimum duration requirement of one repeat. i.e., When the run time of one repeat falls below this time, the number parameter will be automatically increased.
cooldown_interval_ms (int, optional) – The cooldown interval in milliseconds between the number of repeats defined by repeats_to_cooldown.
repeats_to_cooldown (int, optional) – The number of repeats before the cooldown is activated.
f_preproc (str, optional) – The preprocess function name we want to execute before executing the time evaluator.

Note

The function will be invoked (1 + number x repeat) times, with the first call discarded in case there is lazy initialization.

Example

Normal use with a VM function (may not work over RPC if the function returns a tuple):

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)

Use with the stateful API:

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.set_input("func_name", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")

With saved closures via save_function (this results in fewer dictionary lookups in the timed portion):

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()

Returns:: ftimer – The function that takes same argument as func and returns a BenchmarkResult. The ProfileResult reports repeat time costs in seconds.
Return type:: function

Build an IRModule to VM executable.

Parameters:

mod (IRModule) – The input IRModule to be built.
target (Optional[Union[str, tvm.target.Target]]) –
A build target which can have optional host side compilation target.

When TVM compiles device specific program such as CUDA, we also need host(CPU) side code to interact with the driver to setup the dimensions and parameters correctly. host is used to specify the host side codegen target. By default, llvm is used if it is enabled, otherwise a stackvm interpreter is used.
params (Optional[Dict[str, list]]) – Parameters for the input IRModule that will be bound.
pipeline (str = "default_build") – The compilation pipeline to use.
exec_mode ({"bytecode", "compiled"}) – The execution mode.
system_lib (Optional[bool]) – Whether to build system lib that is being packed statically and auto registers generated functions to the system. By default auto detects based on the target.

Returns:

ex – An executable that can be loaded by virtual machine.

Return type:

tvm.relax.Executable

Example

class InputModule:
    @R.function
    def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
        z = R.add(x, y)
        return z

mod = InputModule
target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(mod, target)

tvm.relax.call_dps_packed(func: str | Expr, args: Expr, out_sinfo: TensorStructInfo | List[TensorStructInfo]) → Call¶

relax.Call a destination-passing-style packed function and return the output.

Note: The called function is assumed to be _pure_ (other than modifying the designated output arguments). If the function _does_ result in other side effects, then the compiler may end up removing, reordering, or repeating those effects–no guarantees can be made.

Parameters:

func (Union[str, Expr]) – The destination-passing-style function, can be ExternFunc.
args (Expr) – The input arguments.
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – The structure info of the call_dps_packed output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.

Returns:

ret – A call node for the call_dps_packed operator.

Return type:

relax.Call

tvm.relax.call_pure_packed(func: str | ExternFunc | GlobalVar, *args: Expr, sinfo_args: StructInfo | List[StructInfo]) → Expr¶

Construct a call to a packed function that should be treated as pure, even though packed calls are normally not treated as pure.

The resulting call will have the same semantics as calling the packed function directly.

Note: This should be used for cases when the user knows that calling the packed function with these arguments will in reality not cause any side effects. If it is used for a call that does result in side effects, then the compiler may end up removing, reordering, or repeating that call, with no guarantees made about any side effects from the callee.

Parameters:

func (Union[str, ExternFunc]) – The name (global symbol) for a PackedFunc or an ExternFunc node.
args (Expr) – The arguments for the PackedFunc.
sinfo_args (Union[StructInfo, List[StructInfo]]) – The list of structure info arguments (giving the structural info for the returned value).

Returns:

result – A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)

Return type:

Expr

tvm.relax.call_tir(gvar: GlobalVar, args: Expr, out_sinfo: TensorStructInfo | List[TensorStructInfo], tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call¶

relax.Call a tir.prim_func and return the output.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a tir PrimFunc.
args (Expr) – The input arguments.
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – The structure info of the call_tir output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns:

ret – A call node for the call_tir operator.

Return type:

relax.Call

relax.Call a TIR PrimFunc and return the result, doing the specified computations in-place (based on the inplace_indices argument; outputs will alias the inputs selected by in-place indices).

Warning: This operator is considered pure by the type system but actually mutates the arguments specified by inplace_indices. This operator should not be used directly, but rather should be inserted by passes that have checked whether it is safe to perform operations in-place (i.e., none of the arguments specified as an output is aliased or is live after calling call_tir_inplace).

Direct calls to this operator should be done for testing purposes only.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a TIR PrimFunc.
args (Expr) – The input arguments.
input_indices (Union[int, List[int]]) – Specify which arguments should be used for in-place computations. If input_indices is a single integer, it will be made into a singleton list. Suppose input_indices[i] = j, where j >= 0. Then the i`th output will be an alias of `args[j]. If input_indices[i] = -1, then the i`th output will be a freshly allocated tensor. At least one member of `input_indices must not be -1.
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – The structure info of the call_tir_inplace output. It should be a single TensorStructInfo or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor. If a list of TensorStructInfo is given, the result will be a tuple of TensorStructInfo.
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns:

ret – A call node for the call_tir operator.

Return type:

relax.Call

tvm.relax.call_tir_with_grad(gvar: GlobalVar, args: Expr, out_sinfo: TensorStructInfo | List[TensorStructInfo], te_grad_name: str, te_grad_kwargs: Dict[str, Object] = None, tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call¶

relax.Call a tir.prim_func and return the output. This intrinsic will bind a te gradient function (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be called by the Gradient pass.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a tir PrimFunc.
args (Expr) – The input arguments.
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – The structure info of the call_tir_with_grad output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.
te_grad_name (str) – The registered name of the te gradient function associated with the call_tir_with_grad node. Must be provided as a keyword argument.
te_grad_kwargs (Dict[str, Object], optional) – The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}.
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns:

ret – A call node for the call_tir_with_grad operator.

Return type:

relax.Call

Create a constant value.

Parameters:

value (Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]) – The constant value.
dtype (Optional[str]) – The data type of the resulting constant.

Note

When dtype is None, we use the following rule:

int maps to “int32”
float maps to “float32”
bool maps to “bool”
other using the same default rule as numpy.

tvm.relax.extern(name: str, span: Span | None = None)¶: Create extern function.

tvm.relax.get_pipeline(name: str = 'zero', **kwargs) → Pass¶

Get pre-build pipeline by name

Parameters:

name (Optional[str]) – Name of the pipeline
kwargs (Dict[str, object]) – Keyword args for configuring the pipeline.

Returns:

pipeline – The transformation pipeline.

Return type:

tvm.transform.Pass

tvm.relax.get_shape_of(expr: Expr) → Expr¶

Get shape of expr.

Parameters:: expr (Expr) – The input expr.
Returns:: shape – The shape expression
Return type:: Expr

Note

This function requires expr to be normalized. The function will report an error if expr’s StructInfo is not TensorStructInfo. It will try to return symbolic function when possible. If the tensor do not have a compile-time symbolic shape, the function will then choose to return relax.Call(relax.op.shape_of, [expr]).

tvm.relax.register_pipeline(name: str)¶

Register a new pipeline

Parameters:: name (str) – Name of the pipeline