fix issues for models with more than one inputs

Author: YJ Shi

python/tvm/meta_schedule/testing/custom_builder_runner.py

@@ -167,4 +167,4 @@ def run_module_via_rpc(
         rt_mod = session.load_module(filename)
         dev = session.device(dev_type=dev_type, dev_id=0)
         args = [ndarray.array(arg, dev) for arg in args]
-        return continuation(rt_mod, dev, *args)
+        return continuation(rt_mod, dev, args)

python/tvm/meta_schedule/testing/tune_relay_meta_schedule.py

@@ -103,10 +103,13 @@ def main():
         ARGS.input_shape,
         cache_dir=ARGS.cache_dir,
     )
+    input_info = {input_name: input_shape}
+    inputs = []
     print(f"Workload: {ARGS.workload}")
-    print(f"  input_name: {input_name}")
-    print(f"  input_shape: {input_shape}")
-    print(f"  input_dtype: {input_dtype}")
+    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}")
     alloc_repeat = 1
     runner = ms.runner.RPCRunner(
         rpc_config=ARGS.rpc_config,
@@ -133,19 +136,21 @@ def main():
         params=params,
     )
     graph, rt_mod, params = lib.graph_json, lib.lib, lib.params
-    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"):
+            inputs.append(np.random.uniform(size=input_shape).astype(input_dtype))
+        else:
+            inputs.append(np.random.randint(low=0, high=10000, size=input_shape, dtype=input_dtype))
 
-    def f_timer(rt_mod, dev, input_data):
+    def f_timer(rt_mod, dev, inputs):
         # pylint: disable=import-outside-toplevel
         from tvm.contrib.graph_executor import GraphModule
 
         # pylint: enable=import-outside-toplevel
 
         mod = GraphModule(rt_mod["default"](dev))
-        mod.set_input(input_name, input_data)
+        for index, (input_name, _) in enumerate(input_info.items()):
+            mod.set_input(input_name, inputs[index])
         ftimer = mod.module.time_evaluator(
             "run",
             dev,
@@ -159,17 +164,18 @@ def f_timer(rt_mod, dev, input_data):
         rpc_config=ARGS.rpc_config,
         lib=lib,
         dev_type=ARGS.target.kind.name,
-        args=[input_data],
+        args=inputs,
         continuation=f_timer,
     )
 
-    def f_per_layer(rt_mod, dev, input_data):
+    def f_per_layer(rt_mod, dev, inputs):
         # pylint: disable=import-outside-toplevel
         from tvm.contrib.debugger.debug_executor import create
 
         # pylint: enable=import-outside-toplevel
         mod = create(graph, rt_mod, dev)
-        mod.set_input(input_name, input_data)
+        for index, (input_name, _) in enumerate(input_info.items()):
+            mod.set_input(input_name, inputs[index])
         graph_nodes = [n["name"] for n in json.loads(graph)["nodes"]]
         graph_time = mod.run_individual(number=10, repeat=1, min_repeat_ms=5000)
         print("|graph_nodes| = ", len(graph_nodes))
@@ -182,7 +188,7 @@ def f_per_layer(rt_mod, dev, input_data):
         rpc_config=ARGS.rpc_config,
         lib=rt_mod,
         dev_type=ARGS.target.kind.name,
-        args=[input_data],
+        args=inputs,
         continuation=f_per_layer,
     )