Function local_run

python/tvm/testing/runner.py:80–149

Run a TVM module on a local device. Parameters ---------- mod : Module The TVM module to run. device_type : str The device type to run the module on. args : List[Union[np.ndarray, Tensor, int, float]] The arguments to be fed to the module. evaluator_c

(  # pylint: disable=too-many-arguments,too-many-locals
    mod: "Module",
    device_type: str,
    args: list[Union["np.ndarray", "Tensor", int, float]],
    evaluator_config: Optional["EvaluatorConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
)

def local_run(  # pylint: disable=too-many-arguments,too-many-locals
    mod: "Module",
    device_type: str,
    args: list[Union["np.ndarray", "Tensor", int, float]],
    evaluator_config: Optional["EvaluatorConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
):
    """Run a TVM module on a local device.

    Parameters
    ----------
    mod : Module
        The TVM module to run.
    device_type : str
        The device type to run the module on.
    args : List[Union[np.ndarray, Tensor, int, float]]
        The arguments to be fed to the module.
    evaluator_config : Optional[EvaluatorConfig]
        The evaluator configuration to use.
    export_func : Union[Callable[Module, str], Literal["tar", "ndk"]]
        The function to export the module to a file.
        If callable, it must be a function that takes two arguments: the module to export and the
        path to export to.
        If "tar", the module will be exported to a tar file.
        If "ndk", the module will be exported to a shared library.
    output_format : Optional[str]
        The format of the exported module.
        If not specified, it will be inferred from the `export_func` argument.

    Returns
    -------
    args : List[Union[np.ndarray, Tensor, int, float]]
        The results of running the module.
    profile_result : tvm.runtime.BenchmarkResult
        The profiling result of running the module.
    """
    import os.path as osp
    import tempfile

    from tvm.runtime import device, load_module
    from tvm.s_tir.meta_schedule.runner import EvaluatorConfig

    evaluator_config = EvaluatorConfig._normalized(evaluator_config)
    export_func, output_format = _normalize_export_func(export_func, output_format)

    with tempfile.TemporaryDirectory() as tmp_dir:
        artifact_path = osp.join(tmp_dir, "tvm_tmp_mod." + output_format)
        export_func(mod, artifact_path)
        device: Device = device(device_type, 0)

        try:
            args = _args_to_device(args, device)
            remote_mod = load_module(artifact_path)
            profile_result = remote_mod.time_evaluator(
                func_name=remote_mod.entry_name,
                dev=device,
                number=evaluator_config.number,