fix auto schedule

Author: YJ Shi

python/tvm/meta_schedule/testing/tune_relay_auto_scheduler.py

@@ -134,10 +134,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}")
     tasks, task_weights = auto_scheduler.extract_tasks(
         mod["main"],
         params,
@@ -170,19 +173,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,
@@ -196,17 +201,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))
@@ -219,7 +225,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,
     )