tvm.relax

tvm.relax#

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

class tvm.relax.VirtualMachine(rt_mod: Module | Executable, device: Device | list[Device], memory_cfg: str | dict[Device, str] | None = None)#

Relax VM runtime.

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

Invoke a closure.

Parameters:

closure (Object) – The VMClosure Object.
args (list[tvm.runtime.Tensor] 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 Function that, when called, will invoke the function with the given arguments. The Function can be accessed from the module using saved_name. This is included to facilitate timing trials: Invoking the returned Function 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 Function 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 Tensor 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.Tensor] or List[np.ndarray]) – The arguments to the function.
kwargs (dict of str to tvm.runtime.Tensor 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) → 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: Function) → 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, Function],
    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_ffi.Function) – 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 = tvm.compile(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 = tvm.compile(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 = tvm.compile(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

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

class tvm.relax.DummyGlobalInfo#

class tvm.relax.VDevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global')#

class tvm.relax.Expr(loc=<MISSING>, ty=<MISSING>)#: Base class of all the expressions.

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

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

Parameters:

name (str) – The name of the variable.
ty (Optional[Type]) – The type annotation of the variable.
loc (Location) – Location that points to original source code

class tvm.relax.Binding(loc, var)#: The base class of a binding in Relax.

class tvm.relax.MatchCast(var: Var, value: Expr, ty: Type, loc: Location = UnknownLoc())#

Runtime-match the value to the type.

This operation does runtime check, populates the un-defined symbolic shape vars and vars in ty 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.
ty (tvm.relax.Type) – The type to match cast to.

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

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[Binding], loc: Location = UnknownLoc())#: base class of binding block, bindings inside can be impure (with side effect or control flow)

class tvm.relax.DataflowBlock(bindings: list[Binding], loc: Location = UnknownLoc())#: dataflow block, bindings inside are pure (no side effect and no control flow)

class tvm.relax.SeqExpr(blocks: list[BindingBlock], body: Expr, loc: Location = UnknownLoc())#: A sequence of binding blocks followed by an expression.

class tvm.relax.ShapeExpr(values: list[Expr] | tuple[Expr, ...] | Array, loc: Location = UnknownLoc())#

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

Parameters:

values (Union[List[Expr], Tuple[Expr, ...], tvm_ffi.Array]) – The values of the shape expression.
loc (Location) – Location that points to original source code

class tvm.relax.Function(params: list[Var], body: Expr, ret_ty: Type | None = None, is_pure: bool | None = True, attrs: DictAttrs | None = None, loc: Location = UnknownLoc())#

A Relax function.

static create_empty(params: list[Var], ret_ty: Type, is_pure: bool | None = True, attrs: DictAttrs | None = None, loc: Location = UnknownLoc())#: Construct a relax.Function but without body

bind_symbolic_vars(binding_map: Mapping[str | Var, Expr]) → Function#

Return a new function with updated symbolic variable

Parameters:: binding_map (Mapping[str | tvm.tirx.Var, Expr]) – The mapping of values to be replaced. Keys may be either a tirx.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

Return a new function with updated symbolic variable

Parameters:

binding_map (Mapping[) – str | relax.Var, int | float | Expr | tvm.runtime.Tensor | _np.ndarray,
] –
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: String, ty: Type | None = None, loc: Location = UnknownLoc())#: extern function, which represents a PackedFunc.

class tvm.relax.IfExpr(cond: Expr, true_branch: Expr, false_branch: Expr, loc: Location = UnknownLoc())#

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.
loc (Location) – Location that points to original source code

tvm.relax.prim_value(value: Expr | int | float, dtype: str | None = None) → Expr#

Convert a Python scalar or primitive expression to Expr.

Parameters:

value (Expr | int | float) – The value to convert.
dtype (Optional[str]) – The dtype to use when converting Python numeric values.

Returns:

result – The converted primitive expression. Existing primitive expression inputs are returned unchanged, even when dtype is supplied.

Return type:

Expr

Notes

Existing primitive expressions can be used directly in TVMScript. Python scalars are context dependent: an unannotated Relax binding or return converts them to tensor constants, while this helper constructs primitive expressions. dtype selects the type of Python numeric inputs; omitting it uses bool, int64, or float64. The script alias R.prim_value retains these programmatic conversion semantics.

Create a constant value.

Parameters:

value (bool | int | float | numpy.ndarray | tvm.runtime.Tensor) – 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, ty: Type | None = None, loc: Location = UnknownLoc())#: Create extern function.

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 Type is not TensorType. 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.ShapeType(values: list[Expr] | None = None, ndim: int = -1, loc: Location = UnknownLoc())#

Type of a shape value.

Parameters:

values (Optional[List[Expr]]) – 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.

Type 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.PackedFuncType(loc: Location = UnknownLoc())#: 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: VMFuncKind = VMFuncKind.PACKED_FUNC) → None#: Declare a function

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

emit_call(name: str, args: list[Tensor | dtype] | None = None, dst: int | None = 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() → VMExecutable#: return the executable

tvm.relax.call_tir(func, args, *, ty_args, attrs=None, ty=None, loc: Location = UnknownLoc(), **kwargs) → Call#

relax.Call a destination-passing TIR function and allocate its output tensors.

func is the function’s GlobalVar and args contains its ordered inputs. ty_args contains one output type: a TensorType or a TupleType of tensor results. Omitted ty uses the registered result inference; explicit types and locs are forwarded unchanged.

tvm.relax.call_tir_inplace(func, args, *, ty_args, attrs=None, ty=None, loc: Location = UnknownLoc(), **kwargs) → Call#

relax.Call a TIR function whose selected outputs alias its input tensors.

ty_args contains one output type. In the inplace_indices attribute, entry i >= 0 makes the corresponding output alias input i; -1 allocates a fresh output. At least one output must alias an input.

Although classified as pure, this operation mutates the selected inputs. Only use it after proving there are no live uses or aliases that could observe those mutations. Direct construction is intended for testing; optimization passes normally establish these preconditions.

tvm.relax.call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, loc: Location = UnknownLoc()) → Call#

relax.Call a TIRx Function through its native packed-call contract.

Every native parameter is supplied explicitly, in order. Unlike call_tir(), this operator does not allocate destination tensors or interpret destination parameters as results. It performs no implicit copies, casts, layout conversions, device transfers, or redistribution.

The result follows the native declared return type: a supported scalar keeps its exact primitive type, void becomes the empty tuple, and a pointer becomes an Any carrying an opaque pointer. A pointer result does not imply tensor ownership or a lifetime guarantee. Native tensor, nonempty tuple, callable, vector, and other unsupported return types are rejected. There is no out_ty argument.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a TIRx Function with its native tvm.ir.FuncType signature.
args (Expr) –
The ordered arguments, supplied as an inline Relax tuple, a Python tuple or list, or a single expression. Tensor parameters accept compatible Relax tensors, with dtype, rank, and known shapes checked against the native signature. Remaining supported native tensor constraints are checked by the packed ABI at runtime. Specialized global storage scopes require matching VDevice.memory_scope; non-default layouts, allocated-address contracts, unsupported storage scopes, and unlowered distributed tensors are rejected.

Scalar parameters require the exact primitive dtype: scalar bool, signed or unsigned integers up to 64 bits, or float16/32/64. These same scalar types are supported as direct results. The existing packed integer carrier is signed 64-bit: uint64 values must be in [0, 2**63 - 1]; larger unsigned values are not representable.

Pointer parameters accept Any or a handle-compatible object, including a runtime tensor. A tensor passed to a pointer parameter supplies its DLTensor header handle, not its data pointer. This erased carrier does not prove pointee type, address space, ownership, or lifetime compatibility. Known scalar values cannot serve as pointer arguments. The runtime carrier must satisfy the existing packed ABI’s null, opaque-pointer, DLTensor-pointer, or object-handle check.
to (The call is effectful and may mutate its arguments. For a call known)
effects (have no observable)
wrapper (use an explicit call_pure_packed)
inferred (around the relax.call_tir_packed operator. Purity is never)
ABI. (from the native signature or packed)

Returns:

ret – A call whose result type is derived from the native declared return during Relax type inference.

Return type:

relax.Call

Examples

A native (int64, int64) -> int64 function returns its scalar directly:

result = relax.call_tir_packed(add_scalar, (a, b))

A caller that knows the scalar function has no effects may assert purity:

result = relax.call_pure_packed(
    tvm.ir.Op.get("relax.call_tir_packed"), add_scalar, (a, b)
)

A native copy function with a void return writes to a caller-owned tensor:

relax.call_tir_packed(copy, (source, destination))

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, Op]) – The name (global symbol) for a PackedFunc or an ExternFunc node. The explicit relax.call_tir_packed Op is also accepted; its native callee and argument tuple follow as arguments to this wrapper.
args (Expr) – The arguments for the PackedFunc.
ty_args (Union[Type, List[Type]]) – The list of type information arguments (giving the type information for the returned value). Omit this for the native bridge, whose result follows the native signature.

Returns:

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

Return type:

Expr

tvm.relax.call_dps_packed(func, args, *, ty_args, attrs=None, ty=None, loc: Location = UnknownLoc(), **kwargs) → Call#

relax.Call a destination-passing packed function and allocate its outputs.

Python string callees become ExternFunc; explicit Expr callees retain their identity. ty_args contains one output type, including a TupleType for multiple results.

The function must be pure apart from writing its designated outputs. Other effects may be removed, reordered or repeated by the compiler.

tvm.relax.call_tir_with_grad(func, args, *, ty_args, attrs=None, ty=None, loc: Location = UnknownLoc(), **kwargs) → Call#

relax.Call a TIR function with a registered TE gradient rule.

ty_args contains the single output type, including a TupleType for multiple results. The te_grad_name and optional te_grad_kwargs attribute keywords select the rule used by the Gradient pass. Attributes may instead be supplied through attrs.

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.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: Expr) → Expr#

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: 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) → 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: Var) → Var#

Parameters:: var (relax.Var) – The var to be visited.
Returns:: result – The var after post-order rewritten.
Return type:: relax.Var

visit_generic_const_(op: GenericConst) → Expr#

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

Parameters:: op (GenericConst) – The GenericConst to be visited.
Returns:: result – The Expr after transformation
Return type:: Expr

visit_tuple_(op: Tuple) → Expr#

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: Var) → Expr#

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: DataflowVar) → Expr#

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: ShapeExpr) → Expr#

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: ExternFunc) → Expr#

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: GlobalVar) → Expr#

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: Function) → Expr#

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: Call) → Expr#

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: SeqExpr) → Expr#

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_expr_(op: IfExpr) → Expr#

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

Parameters:: op (IfExpr) – The IfExpr to be visited.
Returns:: result – The Expr after transformation
Return type:: Expr

visit_op_(op: Op) → Expr#

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: TupleGetItem) → Expr#

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_expr_fallback_(op: Expr) → Expr#

Visit an expression handled by the C++ fallback.

Parameters:: op (Expr) – The expression to be visited.
Returns:: result – The Expr after transformation
Return type:: Expr

visit_string_imm_(op: StringImm) → Expr#

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

Parameters:: op (relax.StringImm) – The relax.StringImm to be visited.
Returns:: result – The Expr after transformation
Return type:: Expr

visit_data_type_imm_(op: DataTypeImm) → Expr#

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: 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: 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: BindingBlock) → 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: DataflowBlock) → 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: Var) → 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: DataflowVar) → 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_loc(loc: Location) → Location#

Visit Location. Override this function to customize VisitLoc(const Location& loc) on the C++ side.

Parameters:: loc (Location) – The Location to be visited.
Returns:: result – The loc after transformation.
Return type:: Location

visit_expr_post_order(expr: Expr) → Expr#

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(old_var: Var, new_var: Var) → None#

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

Parameters:

old_var (relax.Var) – The old var.
new_var (relax.Var) – The new var.

get_var_remap(var: Var) → Var | None#

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

Parameters:: var (relax.Var) – The old var.
Returns:: var – The remapped var, or None if no mapping exists.
Return type:: relax.Var

visit_with_new_scope(expr: Expr) → Expr#

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

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

lookup_binding(var: Var) → Expr | None#

Look up the value bound to a variable. Note: For function parameters, this function returns std::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_type(var: Var, ty: Type) → 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.
ty (Type) – The type.

Returns:

var – The var filled with shape and type.

Return type:

relax.Var

tvm.relax.get_default_pipeline(target: Target)#: Get the default Relax compilation pipeline for the given target.

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.register_pipeline(name: str)#

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

tvm.relax.convert_to_expr(value: Any) → Expr#: 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 Expr if the input is a primitive scalar; 3. Return tvm.ir.StringImm if the input is tvm.String or str; 4. Return relax.Tuple if the input is a tuple/list of Expr.

class tvm.relax.BasePyModule(ir_mod: IRModule, device: Device, target: Target | None = None)#

Base class that allows Python functions in IRModule with DLPack conversion.

This class provides the infrastructure for: 1. JIT compilation of TIR and Relax functions. 2. DLPack-based conversion between PyTorch tensors and TVM Tensors. 3. Wrapping Relax functions for easy Python calling. 4. Cross-function calls between Python, TIR, and Relax functions.

Shared IRModules collect Python functions in __pyfuncs__. Decorating a subclass with R.py_module adds this executable runtime interface.

call_tir(tir_func, args, out_ty)#: relax.Call a TIR function with PyTorch tensors.

call_dps_packed(func_name: str, args, out_ty)#: relax.Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack.

call_py_func(func_name: str, args)#: relax.Call a Python function stored in the module’s pyfuncs.

get_function(name: str) → Function | None#: Get a compiled function by name.

list_functions() → dict[str, list[str]]#: List all available functions.

add_python_function(name: str, func: callable)#: Add a Python function to the module.

script(*, name: str | None = None, show_meta: bool = False, ir_prefix: str = 'I', module_alias: str = 'cls', int_dtype: str = 'int32', float_dtype: str = 'void', verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, show_all_ty: bool = True, extra_config: dict | None = None) → str#

Print TVM IR into TVMScript text format with Python function support.

This method extends the standard IRModule script() method to handle Python functions stored in the IRModule’s __pyfuncs__ attribute.

show(style: str | None = None, black_format: bool | None = None, **kwargs) → None#

A sugar for print highlighted TVM script with Python function support.

This method extends the standard IRModule show() method to handle Python functions stored in the IRModule’s pyfuncs attribute.

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 c backend is used.
params (Optional[Dict[str, list]]) – Parameters for the input IRModule that will be bound.
relax_pipeline (str = "default") – The Relax compilation pipeline to use.
tir_pipelinie (str = "default") – The TIR compilation pipeline to use.
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 = tvm.compile(mod, target)

class tvm.relax.VMExecutable(mod: Module)#

The virtual machine 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.

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.

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.

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.

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:

RuntimeError 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() → DataflowBlock#: Returns the mutated DataflowBlock.

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

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

tvm.relax

Contents

tvm.relax#