[VM] Add set_input interface; Fix e2e tuning script. (tlc-pack#166)

* Add set_input interface.

* Address comments.

Author: YuchenJin

apps/relax_examples/e2e_auto_tir.py

@@ -129,15 +129,18 @@ def apply_opt_before_tuning(
     return relax_mod
 
 
-def f_measurement(rt_mod: runtime.Module, device: runtime.ndarray.Device, *input_data):
+def f_measurement(
+    rt_mod: runtime.Module, device: runtime.ndarray.Device, input_data: Dict[str, runtime.NDArray]
+):
     vm = relax.vm.VirtualMachine(exec=rt_mod, device=device)
+    vm.set_input("main", **input_data)
     evaluator = vm.module.time_evaluator(
         func_name="main",
         dev=device,
         repeat=5,
         min_repeat_ms=500,
     )
-    print(evaluator(*input_data))
+    print(evaluator())
 
 
 def get_runner():
@@ -166,10 +169,12 @@ def main():
         ARGS.input_shape,
         cache_dir=ARGS.cache_dir,
     )
-    print(f"Workload: {ARGS.workload}")
-    print(f"  input_name: {input_name}")
-    print(f"  input_shape: {input_shape}")
-    print(f"  input_dtype: {input_dtype}")
+    input_info = {input_name: input_shape}
+    input_data = {}
+    for input_name, input_shape in input_info.items():
+        print(f"  input_name: {input_name}")
+        print(f"  input_shape: {input_shape}")
+        print(f"  input_dtype: {input_dtype}")
 
     # translate the ResNet model from Relay to Relax
     relax_mod = apply_opt_before_tuning(relay_mod, params, target=ARGS.target)
@@ -189,10 +194,13 @@ def main():
         num_threads=os.cpu_count(),
     )
 
-    if input_dtype.startswith("float"):
-        input_data = [np.random.uniform(size=input_shape).astype(input_dtype)]
-    else:
-        input_data = [np.random.randint(low=0, high=10000, size=input_shape, dtype=input_dtype)]
+    for input_name, input_shape in input_info.items():
+        if input_dtype.startswith("float"):
+            input_data[input_name] = np.random.uniform(size=input_shape).astype(input_dtype)
+        else:
+            input_data[input_name] = np.random.randint(
+                low=0, high=10000, size=input_shape, dtype=input_dtype
+            )
 
     if ARGS.rpc_config:
         run_module_via_rpc(
@@ -204,8 +212,7 @@ def main():
         )
     else:
         dev = tvm.device(ARGS.target.kind.name)
-        input_data = [runtime.ndarray.array(arg, dev) for arg in input_data]
-        f_measurement(executable.mod, dev, *input_data)
+        f_measurement(executable.mod, dev, input_data)
 
 
 if __name__ == "__main__":

include/tvm/relax/exec_builder.h

@@ -52,8 +52,9 @@ class ExecBuilderNode : public Object {
    * \brief To annotate the start of a vm function.
    * \param func The function name.
    * \param num_inputs The number of inputs.
+   * \param param_names The function parameter names.
    */
-  void EmitFunction(std::string func, int64_t num_inputs);
+  void EmitFunction(std::string func, int64_t num_inputs, Array<String> param_names);
   /*!
    * \brief Emit a call instruction for a packed function.
    * \param func The packed function name.

include/tvm/runtime/relax_vm/executable.h

@@ -78,6 +78,8 @@ struct VMFunction {
   Index num_args;
   /*! \brief The register file size of the function. */
   Index register_file_size;
+  /*! \brief The function parameter names.*/
+  std::vector<std::string> param_names;
 };
 
 /*!

include/tvm/runtime/relax_vm/vm.h

@@ -167,6 +167,29 @@ class VirtualMachine : public runtime::ModuleNode {
    */
   inline void RunInstrCall(VMFrame* curr_frame, Instruction inst);
 
+  /*!
+   * \brief Set inputs to a function.
+   * \param func_name The function name.
+   * \param args args[offset:] are arguments to the function. If the arguments are not of the
+   * correct device for the function, they will be copied to the device.
+   * \param offset Starting offset of the arguments in \p args.
+   * \note This interface works when using VM over RPC by internally converting NDArray in
+   * the arguments to DLTensor, which is supported in RPC where remote could only have a minimal C
+   * runtime.
+   */
+  void SetInput(std::string func_name, TVMArgs args, int offset);
+
+  /*!
+   * \brief Set a function argument with a given index to an input tensor.
+   * \param func_args the function arguments.
+   * \param inp_tensor some input tensor (not necessarily DLTensor). When it's an NDArray or a list
+   * of NDArray, they will be converted.
+   * \param index The input tensor index in the function arguments.
+   * \param dev device to copy to if needed.
+   */
+  void SetInputTensorWithIndex(std::vector<RegType>& func_args, const TVMArgValue& inp_tensor,
+                               int index, Device dev);
+
  private:
   /*! \brief The loaded executable. */
   ObjectPtr<Executable> exec_;
@@ -189,6 +212,8 @@ class VirtualMachine : public runtime::ModuleNode {
   RegType return_value_;
   /*! \brief The global constant pool */
   std::vector<TVMRetValue> constants;
+  /*! \brief The function name to input register mapping. */
+  std::unordered_map<std::string, std::vector<RegType>> inputs_;
 };
 
 }  // namespace relax_vm

python/tvm/relax/exec_builder.py

@@ -71,9 +71,11 @@ def void_arg(self) -> int:
     def vm_state(self) -> int:
         return self.r(SpecialReg.VM_STATE)
 
-    def function(self, func_name: str, num_inputs: Optional[int] = 0) -> VMFuncScope:
+    def function(
+        self, func_name: str, num_inputs: Optional[int] = 0, param_names: List[str] = None
+    ) -> VMFuncScope:
         """annotate a VM function."""
-        _ffi_api.ExecBuilderFunction(self, func_name, num_inputs)
+        _ffi_api.ExecBuilderFunction(self, func_name, num_inputs, param_names)
         return VMFuncScope()
 
     def _check_scope(self) -> None:

python/tvm/relax/vm.py

@@ -14,15 +14,17 @@
 # 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
+# pylint: disable=invalid-name, redefined-builtin, no-else-return
 """The Relax virtual machine"""
 from typing import List, Optional, Union, Dict, Tuple
+from tvm._ffi import base as _base
+import numpy as np
 
 import tvm
 from tvm import relax
 from tvm.ir.module import IRModule
 from tvm.relay import Any
-from tvm.runtime import Device, Module, PackedFunc
+from tvm.runtime import Device, Module, PackedFunc, container
 from tvm.runtime.object import Object
 from tvm.tir.function import PrimFunc
 from . import _ffi_api
@@ -97,6 +99,8 @@ def __init__(
             else exec["vm_load_executable"]()
         )
         self._invoke_closure = self.module["invoke_closure"]
+        self._set_input = self.module["set_input"]
+        self._get_func_param_names = self.module["get_func_param_names"]
         self._setup_device(device, memory_cfg)
 
     def _setup_device(self, dev: Device, memory_cfg: Union[str, Dict[Device, str]]) -> None:
@@ -161,6 +165,79 @@ def invoke_closure(self, closure: Object, *args: Any) -> Object:
         """
         return self._invoke_closure(closure, *args)
 
+    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, (_base.integer_types, bool)):
+                return "int32"
+            else:
+                return "float32"
+
+        if isinstance(arg, Object):
+            cargs.append(arg)
+        elif isinstance(arg, np.ndarray):
+            nd_arr = tvm.nd.array(arg, device=tvm.cpu(0))
+            cargs.append(nd_arr)
+        elif isinstance(arg, tvm.runtime.NDArray):
+            cargs.append(arg)
+        elif isinstance(arg, (tuple, list)):
+            field_args = []
+            for field in arg:
+                self._convert(field, field_args)
+            cargs.append(container.tuple_object(field_args))
+        elif isinstance(arg, (_base.numeric_types, bool)):
+            dtype = _gettype(arg)
+            value = tvm.nd.array(np.array(arg, dtype=dtype), device=tvm.cpu(0))
+            cargs.append(value)
+        elif isinstance(arg, str):
+            cargs.append(arg)
+        else:
+            raise TypeError("Unsupported type: %s" % (type(arg)))
+
+    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 NDArray in
+        the arguments to DLTensor, which is supported in RPC where remote could only
+        have a minimal C runtime.
+
+        Parameters
+        ----------
+        func_name : str
+            The name of the function.
+        args: List[tvm.runtime.NDArray] or List[np.ndarray]
+            The arguments to the function.
+        kwargs: dict of str to tvm.runtime.NDArray or np.ndarray
+            Named arguments to the function.
+        """
+        cargs = []
+
+        if kwargs:
+            # kwargs can be a super set of the required function parameters.
+            # We only find the ones that are needed.
+            func_params = list(self._get_func_param_names(func_name))
+            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
+            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
+            args = new_args
+
+        for arg in args:
+            self._convert(arg, cargs)
+
+        self._set_input(func_name, *cargs)
+
 
 def build(mod: tvm.IRModule, target: tvm.target.Target) -> Executable:
     """

src/relax/backend/vm/codegen_vm.cc

@@ -66,15 +66,19 @@ class CodeGenVM : public ExprFunctor<Instruction::Arg(const Expr&)> {
   size_t NewRegister() { return registers_num_++; }
   Instruction::Arg VisitExpr_(const FunctionNode* func_node) {
     Optional<String> gsymbol = func_node->GetAttr<String>(tvm::attr::kGlobalSymbol);
+    Array<String> param_names;
+    for (Var param : func_node->params) {
+      param_names.push_back(param->name_hint());
+    }
     if (gsymbol.defined()) {
-      builder_->EmitFunction(gsymbol.value(), func_node->params.size());
+      builder_->EmitFunction(gsymbol.value(), func_node->params.size(), param_names);
     } else {
       // TODO(@yuchen): handle local functions that capture local vars outside the func
       // TODO(@yuchen): a renaming pass to resolve name conflicts, e.g. the input module has a
       // function named "local_funcN"
       // lift the local func to a global func and process it normally
       builder_->EmitFunction("local_func" + std::to_string(local_func_counter_++),
-                             func_node->params.size());
+                             func_node->params.size(), param_names);
     }
 
     for (Var param : func_node->params) {

src/relax/backend/vm/exec_builder.cc

@@ -43,13 +43,19 @@ vm::Index ExecBuilderNode::EmitConstant(TVMRetValue obj) {
   return vm::Instruction::Arg(vm::Instruction::kConstIdx, idx).data;
 }
 
-void ExecBuilderNode::EmitFunction(std::string func_name, int64_t num_inputs) {
+void ExecBuilderNode::EmitFunction(std::string func_name, int64_t num_inputs,
+                                   Array<String> param_names) {
   const auto& m = exec->global_map;
   ICHECK(m.find(func_name) == m.end());
   VMFunction vmfunc;
   vmfunc.name = func_name;
   vmfunc.start_instr = exec->instr_offset.size();
   vmfunc.num_args = num_inputs;
+  std::vector<std::string> names;
+  for (size_t i = 0; i < param_names.size(); ++i) {
+    names.push_back(param_names[i]);
+  }
+  vmfunc.param_names = names;
   exec->global_map[func_name] = exec->global_funcs.size();
   exec->global_funcs.push_back(vmfunc);
 }

src/runtime/relax_vm/executable.cc

@@ -290,6 +290,7 @@ void SerializeVMFunc(const VMFunction& func, dmlc::Stream* strm) {
   strm->Write(func.start_instr);
   strm->Write(func.num_args);
   strm->Write(func.register_file_size);
+  strm->Write(func.param_names);
 }
 
 VMFunction DeserializeVMFunc(dmlc::Stream* strm) {
@@ -298,6 +299,7 @@ VMFunction DeserializeVMFunc(dmlc::Stream* strm) {
   STREAM_CHECK(strm->Read(&func.start_instr), "vmfunc start_instr");
   STREAM_CHECK(strm->Read(&func.num_args), "vmfunc num_args");
   STREAM_CHECK(strm->Read(&func.register_file_size), "vmfunc register_file_size");
+  STREAM_CHECK(strm->Read(&func.param_names), "vmfunc params");
   return func;
 }