wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/runtime/vm.py
+477
View File
@@ -0,0 +1,477 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, redefined-builtin, no-else-return, consider-using-dict-items
# ruff: noqa: RUF005
"""The Relax virtual machine."""
from collections.abc import Callable
from enum import IntEnum
from numbers import Integral, Number
from typing import Any
import numpy as np # type: ignore
from tvm_ffi import Function, register_global_func
import tvm
from tvm.runtime import Device, Object
from ..rpc.base import RPC_SESS_MASK
class VMInstrumentReturnKind(IntEnum):
NO_OP = 0
# skip the following call, only valid in before
SKIP_RUN = 1
class VirtualMachine:
"""Relax VM runtime."""
NAIVE_ALLOCATOR = 1
POOLED_ALLOCATOR = 2
def __init__(
self,
rt_mod: tvm.runtime.Module | tvm.runtime.Executable,
device: Device | list[Device],
memory_cfg: str | dict[Device, str] | None = None,
) -> None:
"""
Construct a VirtualMachine wrapper object.
Parameters
----------
rt_mod: Union[tvm.runtime.Module, tvm.runtime.Executable]
Runtime module exported by the result of build.
device : Union[Device, List[Device]]
The device to deploy the module.
memory_cfg : Optional[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.
"""
if not isinstance(rt_mod, tvm.runtime.Module):
if isinstance(rt_mod, tvm.runtime.Executable):
rt_mod = rt_mod.jit()
else:
raise ValueError("Expect the rt_mod to be an runtime.Module")
self.module = rt_mod["vm_load_executable"]()
self._invoke_closure = self.module["invoke_closure"]
self._save_function = self.module["save_function"]
self._set_input = self.module["set_input"]
self._invoke_stateful = self.module["invoke_stateful"]
self._get_output = self.module["get_output"]
self._get_output_arity = self.module["get_output_arity"]
self._get_function_arity = self.module["get_function_arity"]
self._get_function_param_name = self.module["get_function_param_name"]
self._set_instrument = self.module["set_instrument"]
self._setup_device(device, memory_cfg)
def _setup_device(self, dev: Device, memory_cfg: str | dict[Device, str]) -> None:
"""init devices and allocators."""
devs = dev
if not isinstance(dev, list | tuple):
if not isinstance(dev, tvm.runtime.Device):
raise TypeError("dev is expected to be Device or List[Device]")
devs = [dev]
# CPU is required for executing shape functions
if devs[-1].dlpack_device_type() % RPC_SESS_MASK != tvm.cpu().dlpack_device_type():
devs.append(tvm.cpu())
default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
if memory_cfg is None:
memory_cfg = {}
elif isinstance(memory_cfg, str):
assert memory_cfg in ["naive", "pooled"]
if memory_cfg == "naive":
default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR
memory_cfg = {}
elif not isinstance(memory_cfg, dict):
raise TypeError(
"memory_cfg is expected be string or dictionary, "
+ f"but received {type(memory_cfg)}"
)
init_args = []
for device in devs:
init_args.append(device.dlpack_device_type() % RPC_SESS_MASK)
init_args.append(device.index)
alloc_type = memory_cfg[device] if device in memory_cfg else default_alloc_type
init_args.append(alloc_type)
self.module["vm_initialization"](*init_args)
def __getitem__(self, key: str) -> Function:
return self.module[key]
def invoke_closure(self, 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 : Object
The output.
"""
return self._invoke_closure(closure, *args)
def save_function(
self,
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
"""
cargs: list[Any] = []
if kwargs:
args = self._convert_func_named_args(func_name, args, **kwargs)
for arg in args:
self._convert(arg, cargs)
self._save_function(func_name, saved_name, int(include_return), *cargs)
def _convert(self, arg: Any, cargs: list) -> None:
"""helper function to convert arguments to vm function."""
def _gettype(arg):
if isinstance(arg, np.float16):
return "float16"
elif isinstance(arg, Integral | bool):
return "int32"
else:
return "float32"
if isinstance(arg, Object):
cargs.append(arg)
elif isinstance(arg, np.ndarray):
nd_arr = tvm.runtime.tensor(arg, device=tvm.cpu(0))
cargs.append(nd_arr)
elif isinstance(arg, tvm.runtime.Tensor):
cargs.append(arg)
elif isinstance(arg, tuple | list):
field_args: list[Any] = []
for field in arg:
self._convert(field, field_args)
cargs.append(tuple(field_args))
elif isinstance(arg, Number | bool):
dtype = _gettype(arg)
value = tvm.runtime.tensor(np.array(arg, dtype=dtype), device=tvm.cpu(0))
cargs.append(value)
elif isinstance(arg, str):
cargs.append(arg)
else:
raise TypeError(f"Unsupported type: {type(arg)}")
def _convert_func_named_args(self, func_name: str, args: Any, **kwargs: Any) -> Any:
"""
Takes named function parameters and returns a list of those needed,
in the order they should appear
"""
# kwargs can be a super set of the required function parameters.
# We only find the ones that are needed.
func_arity = self._get_function_arity(func_name)
func_params = [self._get_function_param_name(func_name, i) for i in range(func_arity)]
new_args = [None] * len(func_params)
cnt = 0
for k in kwargs:
if k in func_params:
idx = func_params.index(k)
new_args[idx] = kwargs[k]
cnt += 1
else:
print(f'Warning: Keyword argument "{k}" is unused in {func_name}')
assert len(args) + cnt == len(func_params)
idx = 0
for i, arg in enumerate(new_args):
if arg is None:
new_args[i] = args[idx]
idx += 1
return new_args
def set_input(self, 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.
"""
cargs: list[Any] = []
if kwargs:
args = self._convert_func_named_args(func_name, args, **kwargs)
for arg in args:
self._convert(arg, cargs)
self._set_input(func_name, *cargs)
def invoke_stateful(self, 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.
"""
self._invoke_stateful(func_name)
def get_outputs(self, func_name: str) -> tvm.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: 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.
"""
# to deal with potentially nested tuples, we need to query for arity recursively
def get_output_rec(func_name, *idx):
arity = self._get_output_arity(func_name, *idx)
if arity == -1:
return self._get_output(func_name, *idx)
# otherwise we need to specify more indices
idx_list = list(idx)
return tuple(get_output_rec(func_name, *(idx_list + [i])) for i in range(arity))
return get_output_rec(func_name)
def set_instrument(self, instrument: Function) -> 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:
.. code:: python
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.
"""
self._set_instrument(instrument)
def time_evaluator(
self,
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[..., 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):
.. code-block:: python
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:
.. code-block:: python
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):
.. code-block:: python
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 : function
The function that takes same argument as func and returns a BenchmarkResult.
The ProfileResult reports `repeat` time costs in seconds.
"""
return self.module.time_evaluator(
func_name,
dev,
number=number,
repeat=repeat,
min_repeat_ms=min_repeat_ms,
cooldown_interval_ms=cooldown_interval_ms,
repeats_to_cooldown=repeats_to_cooldown,
f_preproc=f_preproc,
)
@register_global_func("vm.builtin.debug_print")
def _print(lineo: str, array) -> None:
print(f"{lineo}: shape = {array.shape}, dtype = {array.dtype}, data =\n{array}")