Function rpc_run

python/tvm/testing/runner.py:152–235

Run a TVM module on a remote 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_

(  # 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,
    rpc_config: Optional["RPCConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
)

def rpc_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,
    rpc_config: Optional["RPCConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
):
    """Run a TVM module on a remote 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.
    rpc_config : Optional[RPCConfig]
        The RPC configuration to connect to the remote device.
        If not specified, the default RPC configuration will be used, which reads the following
        environment variables:
        - TVM_TRACKER_HOST
        - TVM_TRACKER_PORT
        - TVM_TRACKER_KEY
    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.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig

    evaluator_config = EvaluatorConfig._normalized(evaluator_config)
    rpc_config = RPCConfig._normalized(rpc_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)
        _, remote_path = osp.split(artifact_path)