apache--tvm/python/tvm/runtime/vm.py

# 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}")