code: def main():
      log_file = os.path.join(ARGS.work_dir, f"{ARGS.model_name}.json")
  
      runner = auto_scheduler.RPCRunner(
          key=ARGS.rpc_key,
          host=ARGS.rpc_host,
          port=ARGS.rpc_port,
          n_parallel=cpu_count(logical=True),
          number=ARGS.number,
          repeat=ARGS.repeat,
          min_repeat_ms=ARGS.min_repeat_ms,
          enable_cpu_cache_flush=ARGS.cpu_flush,
          timeout=ARGS.rpc_config.session_timeout_sec,
      )
  
      if ARGS.target.kind.name == "llvm":
          hardware_params = auto_scheduler.HardwareParams(
              num_cores=int(ARGS.target.attrs["num-cores"]),
              target=ARGS.target,
          )
      elif ARGS.target.kind.name == "cuda":
          hardware_params = auto_scheduler.HardwareParams(
              num_cores=-1,
              vector_unit_bytes=16,
              cache_line_bytes=64,
              max_shared_memory_per_block=int(
                  ARGS.target.attrs["max_shared_memory_per_block"]
              ),
              max_threads_per_block=int(
                  ARGS.target.attrs["max_threads_per_block"]
              ),
              # The value `max_local_memory_per_block` is not used in AutoScheduler,
              # but is required by the API.
              max_local_memory_per_block=12345678,
              max_vthread_extent=8,
              warp_size=32,
          )
      else:
          raise NotImplementedError(f"Unsupported target {ARGS.target}")
  
      describe()
      print(f"Workload: {ARGS.model_name}")
      onnx_model = onnx.load(ARGS.onnx_path)
      shape_dict = {}
      for item in ARGS.input_shape:
          print(f"  input_name : {item['name']}")
          print(f"  input_shape: {item['shape']}")
          print(f"  input_dtype: {item['dtype']}")
          shape_dict[item["name"]] = item["shape"]
      mod, params = from_onnx(
          onnx_model, shape_dict, freeze_params=True
      )
      input_data = {
          item["name"]: generate_input_data(
              item["shape"], item["dtype"]
          )
          for item in ARGS.input_shape
      }
  
      with ms.Profiler() as profiler:
          tasks, task_weights = auto_scheduler.extract_tasks(
              mod["main"],
              params,
              target=ARGS.target,
              hardware_params=hardware_params,
          )
          for idx, (task, task_weight) in enumerate(
              zip(tasks, task_weights)
          ):
              print(
                  f"==== Task {idx}: {task.desc} "
                  f"(weight {task_weight} key: {task.workload_key}) ====="
              )
              print(task.compute_dag)
  
          if ARGS.num_trials > 0:
              tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
              tuner.tune(
                  auto_scheduler.TuningOptions(
                      num_measure_trials=ARGS.num_trials,
                      runner=runner,
                      measure_callbacks=[
                          auto_scheduler.RecordToFile(log_file),
                      ],
                  ),
                  adaptive_training=ARGS.adaptive_training,
              )
  
          relay_build = {"graph": relay.build, "vm": relay.vm.compile}[
              ARGS.backend
          ]
          with auto_scheduler.ApplyHistoryBest(log_file):
              with tvm.transform.PassContext(
                  opt_level=3,
                  config={"relay.backend.use_auto_scheduler": True},
              ):
                  lib = relay_build(
                      mod,
                      target=ARGS.target,
                      params=params,
                  )
      print("Tuning Time:")
      print(profiler.table())
  
      run_module_via_rpc(
          rpc_config=ARGS.rpc_config,
          lib=lib,
          dev_type=ARGS.target.kind.name,
          args=input_data,
          continuation=create_timer(ARGS.backend),
          backend=ARGS.backend,
      )
  
file_path: tvm/python/tvm/auto_scheduler/testing/tune_onnx.py
imports: from distutils.util import strtobool
  import argparse
  import json
  import os
  import onnx
  import tvm
  from tvm import auto_scheduler
  from tvm import meta_schedule
  from tvm import relay
  from tvm.meta_schedule.testing.custom_builder_runner import run_module_via_rpc
  from tvm.meta_schedule.testing.tune_utils import generate_input_data, create_timer
  from tvm.meta_schedule.utils import cpu_count
  from tvm.relay.frontend import from_onnx
  from tvm.support import describe
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types