tvm.relax¶

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

class tvm.relax.VirtualMachine(rt_mod: Union[tvm.runtime.module.Module, tvm.relax.vm_build.Executable], device: Union[tvm._ffi.runtime_ctypes.Device, List[tvm._ffi.runtime_ctypes.Device]], memory_cfg: Optional[Union[str, Dict[tvm._ffi.runtime_ctypes.Device, str]]] = None, profile: bool = False)¶

Relax VM runtime.

invoke_closure(closure: tvm.runtime.object.Object, *args: Any) → tvm.runtime.object.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

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.

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.

get_outputs(func_name: str) → Union[tvm.runtime.object.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]]

set_instrument(instrument: tvm.runtime.packed_func.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: tvm._ffi.runtime_ctypes.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[[...], tvm.runtime.module.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

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

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

tvm.relax.Expr¶: alias of tvm.ir.expr.RelayExpr

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

class tvm.relax.Var(name_hint: Union[str, tvm.relax.expr.Id], struct_info: Optional[tvm.relax.expr.StructInfo] = None, span: Optional[tvm.ir.base.Span] = 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.DataflowVar(name_hint: Union[str, tvm.relax.expr.Id], struct_info: Optional[tvm.relax.expr.StructInfo] = None, span: Optional[tvm.ir.base.Span] = 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.Binding¶: The base class of a binding in Relax.

class tvm.relax.MatchCast(var: tvm.relax.expr.Var, value: tvm.ir.expr.RelayExpr, struct_info: tvm.relax.expr.StructInfo, span: Optional[tvm.ir.base.Span] = 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.VarBinding(var: tvm.relax.expr.Var, value: tvm.ir.expr.RelayExpr, span: Optional[tvm.ir.base.Span] = 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.BindingBlock(bindings: List[tvm.relax.expr.Binding], span: Optional[tvm.ir.base.Span] = None)¶: base class of binding block, bindings inside can be impure (with side effect or control flow)

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

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

class tvm.relax.ShapeExpr(values: Union[List[tvm.ir.expr.PrimExpr], Tuple[tvm.ir.expr.PrimExpr, ...], tvm.ir.container.Array], span: Optional[tvm.ir.base.Span] = None)¶

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

Parameters

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

class tvm.relax.Tuple(fields: Union[List[tvm.ir.expr.RelayExpr], Tuple[tvm.ir.expr.RelayExpr, ...]], span: Optional[tvm.ir.base.Span] = None)¶

Tuple expression that groups several fields together.

Parameters

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

class tvm.relax.TupleGetItem(tuple_value: tvm.ir.expr.RelayExpr, index: int, span: Optional[tvm.ir.base.Span] = 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.Function(params: List[tvm.relax.expr.Var], body: tvm.ir.expr.RelayExpr, ret_struct_info: Optional[tvm.relax.expr.StructInfo] = None, is_pure: Optional[bool] = True, attrs: Optional[tvm.ir.attrs.DictAttrs] = None, span: Optional[tvm.ir.base.Span] = None)¶

A Relax function.

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

bind_symbolic_vars(binding_map: Mapping[Union[str, tvm.tir.expr.Var], tvm.ir.expr.PrimExpr]) → tvm.relax.expr.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

bind_params(binding_map: Mapping[Union[str, tvm.relax.expr.Var], Union[int, float, tvm.ir.expr.PrimExpr, tvm.runtime.ndarray.NDArray, numpy.ndarray, tvm.ir.expr.RelayExpr]]) → tvm.relax.expr.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

class tvm.relax.ExternFunc(global_symbol: tvm.runtime.container.String, struct_info: Optional[tvm.relax.expr.StructInfo] = None, span: Optional[tvm.ir.base.Span] = None)¶: extern function, which represents a PackedFunc.

class tvm.relax.Call(op: Union[tvm.ir.expr.RelayExpr, tvm.ir.op.Op], args: Union[List[tvm.ir.expr.RelayExpr], Tuple[tvm.ir.expr.RelayExpr, ...]], attrs: Optional[tvm.ir.attrs.Attrs] = None, sinfo_args: Optional[Union[List[tvm.relax.expr.StructInfo], Tuple[tvm.relax.expr.StructInfo, ...]]] = None, span: Optional[tvm.ir.base.Span] = None)¶

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], typing.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], typing.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.If(cond: tvm.ir.expr.RelayExpr, true_branch: tvm.ir.expr.RelayExpr, false_branch: tvm.ir.expr.RelayExpr, span: Optional[tvm.ir.base.Span] = 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.Constant(data: tvm.runtime.ndarray.NDArray, struct_info: Optional[tvm.relax.expr.StructInfo] = None, span: Optional[tvm.ir.base.Span] = 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.PrimValue(value: Union[tvm.ir.expr.PrimExpr, int], span: Optional[tvm.ir.base.Span] = None)¶: The prim expr representing the value.

class tvm.relax.DataTypeImm(value: Union[tvm._ffi.runtime_ctypes.DataType, str], span: Optional[tvm.ir.base.Span] = None)¶: Represent a data type constant.

class tvm.relax.StringImm(value: str, span: Optional[tvm.ir.base.Span] = None)¶: Represent a string literal constant.

tvm.relax.const(value: Union[bool, int, float, numpy.ndarray, tvm.runtime.ndarray.NDArray], dtype: Optional[str] = None) → tvm.relax.expr.Constant¶

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, struct_info: Optional[tvm.relax.expr.StructInfo] = None, span: Optional[tvm.ir.base.Span] = None)¶: Create extern function.

tvm.relax.get_shape_of(expr: tvm.ir.expr.RelayExpr) → tvm.ir.expr.RelayExpr¶

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

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

class tvm.relax.ShapeType(ndim: int = - 1, span: Optional[tvm.ir.base.Span] = None)¶

The type of shape in Relax.

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

class tvm.relax.DynTensorType(ndim=- 1, dtype='float32', span: Optional[tvm.ir.base.Span] = 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.PackedFuncType(span: Optional[tvm.ir.base.Span] = None)¶: The type of ExternFunc in Relax.

class tvm.relax.ExecBuilder¶

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

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

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

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

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

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

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

emit_call(name: str, args: Optional[List[Union[tvm.runtime.ndarray.NDArray, tvm._ffi.runtime_ctypes.DataType]]] = None, dst: Optional[int] = None) → None¶: emit a call instruction which calls a packed function.

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

emit_goto(pc_offset)¶: emit a goto instruction

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

get() → tvm.relax.vm_build.Executable¶: return the executable

tvm.relax.call_tir(gvar: tvm.ir.expr.GlobalVar, args: tvm.ir.expr.RelayExpr, out_sinfo: Union[tvm.relax.struct_info.TensorStructInfo, List[tvm.relax.struct_info.TensorStructInfo]], tir_vars: Optional[Union[tvm.relax.expr.ShapeExpr, Tuple[tvm.ir.expr.PrimExpr], List[tvm.ir.expr.PrimExpr]]] = None) → tvm.relax.expr.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

tvm.relax.call_tir_inplace(gvar: tvm.ir.expr.GlobalVar, args: tvm.ir.expr.RelayExpr, inplace_indices: Union[int, List[int]], out_sinfo: Union[tvm.relax.struct_info.TensorStructInfo, List[tvm.relax.struct_info.TensorStructInfo]], tir_vars: Optional[Union[tvm.relax.expr.ShapeExpr, Tuple[tvm.ir.expr.PrimExpr], List[tvm.ir.expr.PrimExpr]]] = None) → tvm.relax.expr.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.
inplace_indices (Union[int, List[int]]) – Specify which arguments should be used for in-place computations. If inplace_indices is a single integer, it will be made into a singleton list. Suppose inplace_indices[i] = j, where j >= 0. Then the i`th output will be an alias of `args[j]. If inplace_indices[i] = -1, then the i`th output will be a freshly allocated tensor. At least one member of `inplace_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_pure_packed(func: Union[str, tvm.relax.expr.ExternFunc, tvm.ir.expr.GlobalVar], *args: tvm.ir.expr.RelayExpr, sinfo_args: Union[tvm.relax.expr.StructInfo, List[tvm.relax.expr.StructInfo]]) → tvm.ir.expr.RelayExpr¶

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_dps_packed(func: Union[str, tvm.ir.expr.RelayExpr], args: tvm.ir.expr.RelayExpr, out_sinfo: Union[tvm.relax.struct_info.TensorStructInfo, List[tvm.relax.struct_info.TensorStructInfo]]) → tvm.relax.expr.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_tir_with_grad(gvar: tvm.ir.expr.GlobalVar, args: tvm.ir.expr.RelayExpr, out_sinfo: Union[tvm.relax.struct_info.TensorStructInfo, List[tvm.relax.struct_info.TensorStructInfo]], te_grad_name: str, te_grad_kwargs: Dict[str, tvm.runtime.object.Object] = None, tir_vars: Optional[Union[tvm.relax.expr.ShapeExpr, Tuple[tvm.ir.expr.PrimExpr], List[tvm.ir.expr.PrimExpr]]] = None) → tvm.relax.expr.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

class tvm.relax.ExprFunctor¶

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

visit_expr(expr: tvm.ir.expr.RelayExpr) → tvm.ir.expr.RelayExpr¶: Apply the visitor to an expression.

class tvm.relax.PyExprVisitor¶

An abstract ExprVisitor 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”).

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

See also: visitor, _PyExprVisitor

Example:

@relax.expr_functor.mutator
def MyExprMutator(PyExprMutator):
    ...

visit_expr(expr: tvm.ir.expr.RelayExpr) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation.
Return type: Expr

visit_binding(binding: tvm.relax.expr.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: tvm.relax.expr.BindingBlock) → tvm.relax.expr.BindingBlock¶

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.
Returns: result – The binding block after transformation.
Return type: BindingBlock

visit_var_def(var: tvm.relax.expr.Var) → tvm.relax.expr.Var¶

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.
Returns: result – The var after post-order rewritten.
Return type: relax.Var

visit_constant_(op: tvm.relax.expr.Constant) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_tuple_(op: tvm.relax.expr.Tuple) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_var_(op: tvm.relax.expr.Var) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_dataflow_var_(op: tvm.relax.expr.DataflowVar) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_shape_expr_(op: tvm.relax.expr.ShapeExpr) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_extern_func_(op: tvm.relax.expr.ExternFunc) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_global_var_(op: tvm.ir.expr.GlobalVar) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_function_(op: tvm.relax.expr.Function) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_call_(op: tvm.relax.expr.Call) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_seq_expr_(op: tvm.relax.expr.SeqExpr) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_if_(op: tvm.relax.expr.If) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_op_(op: tvm.ir.op.Op) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_tuple_getitem_(op: tvm.relax.expr.TupleGetItem) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_prim_value_(op: tvm.relax.expr.PrimValue) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_string_imm_(op: tvm.relax.expr.StringImm) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_data_type_imm_(op: tvm.relax.expr.DataTypeImm) → tvm.ir.expr.RelayExpr¶

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.
Returns: result – The Expr after transformation
Return type: Expr

visit_var_binding_(binding: tvm.relax.expr.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_match_cast_(binding: tvm.relax.expr.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_binding_block_(block: tvm.relax.expr.BindingBlock) → tvm.relax.expr.BindingBlock¶

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.
Returns: result – The binding block after transformation
Return type: BindingBlock

visit_dataflow_block_(block: tvm.relax.expr.DataflowBlock) → tvm.relax.expr.BindingBlock¶

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.
Returns: result – The binding block after transformation
Return type: BindingBlock

visit_var_def_(var: tvm.relax.expr.Var) → tvm.relax.expr.Var¶

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.
Returns: result – The var after post-order rewritten.
Return type: relax.Var

visit_dataflow_var_def_(var: tvm.relax.expr.DataflowVar) → tvm.relax.expr.Var¶

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.
Returns: result – The var after post-order rewritten.
Return type: relax.Var

visit_span(span: tvm.ir.base.Span) → tvm.ir.base.Span¶

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.
Returns: result – The span after transformation.
Return type: Span

visit_expr_post_order(expr: tvm.ir.expr.RelayExpr) → tvm.ir.expr.RelayExpr¶

Post-order rewrite an Expr and normalize.

Parameters: expr (Expr) – The Expr to be rewritten.
Returns: result – The Expr after post-order rewritten.
Return type: Expr

set_var_remap(vid: tvm.relax.expr.Id, var: tvm.relax.expr.Var) → None¶

Remap a var to a new var in use-site.

Parameters

vid (Id) – The vid of the old var.
var (relax.Var) – The new var.

get_var_remap(vid: tvm.relax.expr.Id) → tvm.relax.expr.Var¶

Remap a var to a new var in use-site.

Parameters: vid (Id) – The vid of the old var
Returns: var – The remapped var.
Return type: relax.Var

visit_with_new_scope(expr: tvm.ir.expr.RelayExpr) → tvm.ir.expr.RelayExpr¶

Rewrite the expr with a new scope, used in a Function’s body and the branches of If.

Parameters: expr (Expr) – The expr to be visited.
Returns: var – The expr after visiting.
Return type: relax.Var

lookup_binding(var: tvm.relax.expr.Var) → Optional[tvm.ir.expr.RelayExpr]¶

Look up the value bound to a variable. Note: For function parameters, this function returns NullOpt.

Parameters: var (relax.Var) – The var to be looked up.
Returns: var – The value bound to the input var.
Return type: relax.Var

with_struct_info(var: tvm.relax.expr.Var, struct_info: tvm.relax.expr.StructInfo) → tvm.relax.expr.Var¶

Create a new var with specified shape and type if the original var’s shape or type does not match with the specified ones.

Parameters

var (relax.Var) – The var to be updated.
struct_info (StructInfo) – The struct info.

Returns

var – The var filled with shape and type.

Return type

relax.Var

class tvm.relax.StructInfo¶

The base class of all StructInfo.

StructInfo contains both the static type and runtime structural information.

same_as(other)¶: Overload with structural equality.

is_base_of(derived: tvm.relax.expr.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

class tvm.relax.ObjectStructInfo(span: Optional[tvm.ir.base.Span] = None)¶: StructInfo of an Object.

class tvm.relax.PrimStructInfo(dtype: Optional[Union[str, tvm._ffi.runtime_ctypes.DataType]] = None, value: Optional[Union[int, float, tvm.ir.expr.PrimExpr]] = None, span: Optional[tvm.ir.base.Span] = None)¶

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.ShapeStructInfo(values: Optional[List[tvm.ir.expr.PrimExpr]] = None, ndim: int = - 1, span: Optional[tvm.ir.base.Span] = 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.TensorStructInfo(shape: Union[tvm.ir.expr.RelayExpr, None, List[tvm.ir.expr.PrimExpr]] = None, dtype: str = 'float32', vdevice: Union[tvm.ir.global_info.VDevice, None, str] = None, ndim: int = - 1, span: Optional[tvm.ir.base.Span] = None)¶

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.TupleStructInfo(fields: List[tvm.relax.expr.StructInfo], span: Optional[tvm.ir.base.Span] = None)¶

StructInfo of a Tuple value.

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

class tvm.relax.FuncStructInfo(params: List[tvm.relax.expr.StructInfo], ret: tvm.relax.expr.StructInfo, purity: bool = True, span: Optional[tvm.ir.base.Span] = 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: Optional[tvm.relax.expr.StructInfo] = None, derive_func: Optional[Union[str, tvm.ir.base.EnvFunc]] = None, purity: bool = False, span: Optional[tvm.ir.base.Span] = None) → tvm.relax.struct_info.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[Union[str,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.

tvm.relax.get_pipeline(name: str = 'zero', **kwargs) → tvm.ir.transform.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.register_pipeline(name: str)¶

Register a new pipeline

Parameters: name (str) – Name of the pipeline

tvm.relax.convert_to_expr(value: Any) → tvm.ir.expr.RelayExpr¶

Helper function to convert the input to Expr, which follows the rules: 1. Return the input itself if it’s already a relax.Expr; 2. Return relax.PrimValue if the input is a PrimExpr; 3. Return relax.StringImm if the input is tvm.String or str; 4. Return relax.Tuple if the input is a tuple/list of Expr.

Notes

tvm.tir.StringImm is not allowed because of ambiguity, which can be either relax.StringImm or relax.PrimValue.

tvm.relax.build(mod: tvm.ir.module.IRModule, target: Optional[Union[str, tvm.target.target.Target]] = None, params: Optional[Dict[str, list]] = None, pipeline: Union[None, str, tvm.ir.transform.Pass] = 'default_build', exec_mode: str = 'bytecode', *, system_lib: Optional[bool] = None) → tvm.relax.vm_build.Executable¶

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)

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

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

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

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

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

jit(fcompile=None, addons=None, **kwargs) → tvm.runtime.module.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())

export_library(file_name: str, fcompile: Optional[Union[str, callable]] = None, workspace_dir: Optional[str] = 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())

class tvm.relax.DataflowBlockRewrite(dfb: tvm.relax.expr.DataflowBlock, root_fn: tvm.relax.expr.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.

replace_all_uses(old_var: tvm.relax.expr.Var, new_var: tvm.relax.expr.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.

add(expr: tvm.ir.expr.RelayExpr, name: Optional[str] = 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.

remove_unused(var: tvm.relax.expr.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) –

remove_all_unused() → None¶

Remove all unused variables.

Notes

This could remove unused variables in other DataflowBlocks as well.

mutated_dfb() → tvm.relax.expr.DataflowBlock¶: Returns the mutated DataflowBlock.

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

mutate_irmodule(irmodule: tvm.ir.module.IRModule) → tvm.ir.module.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