Fix vm usage.

python/tvm/meta_schedule/testing/custom_builder_runner.py

@@ -17,14 +17,15 @@
 """Customized builder and runner methods"""
 # pylint: disable=import-outside-toplevel
 
-from typing import TYPE_CHECKING, Callable, Dict, List
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 
 if TYPE_CHECKING:
     import numpy as np  # type: ignore
     from tvm.ir import IRModule
     from tvm.meta_schedule.runner import EvaluatorConfig, RPCConfig
     from tvm.runtime import Device, Module, NDArray
     from tvm.target import Target
+    from tvm.runtime.vm import Executable
 
 
 def build_relay(
@@ -143,10 +144,11 @@ def run_with_graph_executor(
 
 def run_module_via_rpc(
     rpc_config: "RPCConfig",
-    lib: "Module",
+    lib: Union["Module", "Executable"],
     dev_type: str,
     args: Dict[str, "np.ndarray"],
     continuation: Callable,
+    use_vm: Optional[bool] = False,
 ):
     """Execute a tvm.runtime.Module on RPC remote"""
     # pylint: disable=import-outside-toplevel
@@ -160,13 +162,15 @@ def run_module_via_rpc(
 
     with tempfile.TemporaryDirectory() as tmp_dir:
         filename = os.path.join(tmp_dir, "tvm_tmp_mod." + tar.output_format)
+        if use_vm:
+            code, lib = lib.save()
         lib.export_library(filename, tar)
         session = rpc_config.connect_server()
         session.upload(filename)
         _, filename = os.path.split(filename)
         rt_mod = session.load_module(filename)
+        if use_vm:
+            rt_mod = session.get_function("runtime.Load_Executable")(code, rt_mod)
         dev = session.device(dev_type=dev_type, dev_id=0)
-        nd_args = {}
-        for arg_key, arg_value in args.items():
-            nd_args[arg_key] = ndarray.array(arg_value, dev)
+        nd_args = {k: ndarray.array(v, dev) for k, v in args.items()}
         return continuation(rt_mod, dev, nd_args)