tvm.runtime.relax_vm¶

The Relax virtual machine.

class tvm.runtime.relax_vm.VMInstrumentReturnKind(value)¶: An enumeration.

class tvm.runtime.relax_vm.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¶

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