tvm.runtime.relax_vm

The Relax virtual machine.

class tvm.runtime.relax_vm.VMInstrumentReturnKind(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

class tvm.runtime.relax_vm.VirtualMachine(rt_mod, device, memory_cfg=None, profile=False)

Relax VM runtime.

参数:

• rt_mod (Module | tvm.relax.Executable)

• device (Device | List[Device])

• memory_cfg (str | Dict[Device, str] | None)

• profile (bool)

__init__(rt_mod, device, memory_cfg=None, profile=False)

Construct a VirtualMachine wrapper object.

Parameters

rt_mod: Union[tvm.runtime.Module, tvm.relax.Executable]

Runtime module exported by the result of build.

deviceUnion[Device, List[Device]]

The device to deploy the module.

memory_cfgOptional[Union[str, Dict[Device, str]]]

Config the type of memory allocator. The allocator type can be ["naive", "pooled"]. If memory_cfg is None, all devices will use pooled allocator by default. If memory_cfg is string, all devices will use the specified allocator type. If memory_cfg is a dict, each device uses the allocator type specified in the dict, or pooled allocator if not specified in the dict.

profileOptional[bool]

Whether or not to enable profiling.

参数:

• rt_mod (Module | Executable)

• device (Device | List[Device])

• memory_cfg (str | Dict[Device, str] | None)

• profile (bool)

返回类型:

None

_convert(arg, cargs)

helper function to convert arguments to vm function.

参数:

• arg (Any)

• cargs (List)

返回类型:

None

_convert_func_named_args(func_name, args, **kwargs)

Takes named function parameters and returns a list of those needed, in the order they should appear

参数:

• func_name (str)

• args (Any)

• kwargs (Any)

返回类型:

Any

_setup_device(dev, memory_cfg)

init devices and allocators.

参数:

• dev (Device)

• memory_cfg (str | Dict[Device, str])

返回类型:

None

get_outputs(func_name)

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: Union[tvm.Object, Tuple[Any]]

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.

参数:

func_name (str)

返回类型:

Object | Tuple[Any]

invoke_closure(closure, *args)

Invoke a closure.

Parameters

closureObject

The VMClosure Object.

argslist[tvm.runtime.NDArray] or list[np.ndarray]

The arguments to the closure.

Returns

resultObject

The output.

参数:

• closure (Object)

• args (Any)

返回类型:

Object

invoke_stateful(func_name)

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.

参数:

func_name (str)

返回类型:

None

profile(func_name, *args)

Profile a function call.

Parameters

func_namestr

The name of the function.

args: List of NDArray or other objects supported by PackedFunc.

The arguments to the function.

Returns

report: tvm.runtime.profiling.Report

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

参数:

func_name (str)

save_function(func_name, saved_name, *args, include_return=True, **kwargs)

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_namestr

The function that should be packaged up.

saved_namestr

The name that the resulting closure should be saved under.

include_returnbool

Whether the saved PackedFunc should return its output. If timing over RPC, it may not be desirable to send output between machines.

argsList[Any]

The arguments to package up with the function.

kwargsDict[str, Any]

Any named arguments to package up with the function

参数:

• func_name (str)

• saved_name (str)

• args (List[Any])

• include_return (bool)

• kwargs (Dict[str, Any])

返回类型:

None

set_input(func_name, *args, **kwargs)

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_namestr

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.

参数:

• func_name (str)

• args (Any)

• kwargs (Any)

返回类型:

None

set_instrument(instrument)

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.

参数:

instrument (PackedFunc)

返回类型:

None

time_evaluator(func_name, dev, number=10, repeat=1, min_repeat_ms=0, cooldown_interval_ms=0, repeats_to_cooldown=1, f_preproc='')

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

ftimerfunction

The function that takes same argument as func and returns a BenchmarkResult. The ProfileResult reports repeat time costs in seconds.

参数:

• func_name (str)

• dev (Device)

• number (int)

• repeat (int)

• min_repeat_ms (int)

• cooldown_interval_ms (int)

• repeats_to_cooldown (int)

• f_preproc (str)

返回类型:

Callable[[...], BenchmarkResult]