code: def main():
      if ARGS.target.kind.name != "llvm" and ARGS.graph_tuner:
          raise ValueError("GraphTuner only supports llvm target")
      if ARGS.target.kind.name != "llvm" and ARGS.cpu_flush:
          raise ValueError("cpu_flush only supports llvm target")
      if ARGS.target.kind.name == "llvm" and not ARGS.cpu_flush:
          warnings.warn("cpu_flush is not enabled for llvm target")
  
      log_file = os.path.join(ARGS.work_dir, f"{ARGS.workload}.json")
      graph_opt_sch_file = os.path.join(
          ARGS.work_dir, f"{ARGS.workload}_graph_opt.log"
      )
      measure_option = autotvm.measure_option(
          builder=autotvm.LocalBuilder(),
          runner=autotvm.RPCRunner(
              key=ARGS.rpc_key,
              host=ARGS.rpc_host,
              port=ARGS.rpc_port,
              number=ARGS.number,
              repeat=ARGS.repeat,
              min_repeat_ms=ARGS.min_repeat_ms,
              enable_cpu_cache_flush=ARGS.cpu_flush,
          ),
      )
      describe()
      print(f"Workload: {ARGS.workload}")
      mod, params, (input_name, input_shape, input_dtype) = get_network(
          ARGS.workload,
          ARGS.input_shape,
          layout=ARGS.layout,
          cache_dir=ARGS.cache_dir,
      )
      input_info = [
          {
              "name": input_name,
              "shape": input_shape,
              "dtype": input_dtype,
          }
      ]
      input_data = {
          item["name"]: generate_input_data(
              item["shape"], item["dtype"]
          )
          for item in input_info
      }
      for item in input_info:
          print(f"  input_name : {item['name']}")
          print(f"  input_shape: {item['shape']}")
          print(f"  input_dtype: {item['dtype']}")
  
      with ms.Profiler() as profiler:
          with ms.Profiler.timeit("TaskExtraction"):
              # extract workloads from relay program
              tasks = autotvm.task.extract_from_program(
                  mod["main"],
                  target=ARGS.target,
                  params=params,
                  ops=(
                      relay.op.get("nn.conv2d"),
                      relay.op.get("nn.conv3d"),
                      relay.op.get("nn.conv2d_transpose"),
                      relay.op.get("nn.dense"),
                      relay.op.get("nn.batch_matmul"),
                  ),
              )
              for i, task in enumerate(tasks):
                  print(f"Task {i} {task.name}: {task}")
  
          with ms.Profiler.timeit("Tuning"):
              if ARGS.num_trials > 0:
                  for i, task in enumerate(tasks):
                      prefix = f"[Task {i + 1:2d}/{len(tasks):2d}] "
                      tuner_obj = XGBTuner(task, loss_type="reg")
                      n_trial = min(
                          len(task.config_space), ARGS.num_trials
                      )
                      tuner_obj.tune(
                          n_trial=n_trial,
                          early_stopping=800,
                          measure_option=measure_option,
                          callbacks=[
                              autotvm.callback.progress_bar(
                                  n_trial, prefix=prefix
                              ),
                              autotvm.callback.log_to_file(log_file),
                          ],
                      )
                  if ARGS.graph_tuner:
                      executor = DPTuner(
                          graph=mod["main"],
                          input_shapes={input_name: input_shape},
                          records=log_file,
                          target_ops=[relay.op.get("nn.conv2d")],
                          target=ARGS.target,
                      )
                      executor.benchmark_layout_transform(
                          min_exec_num=1000
                      )
                      executor.run()
                      executor.write_opt_sch2record_file(
                          graph_opt_sch_file
                      )
  
          relay_build = {"graph": relay.build, "vm": relay.vm.compile}[
              ARGS.backend
          ]
          with ms.Profiler.timeit("PostTuningCompilation"):
              if ARGS.graph_tuner:
                  ctx = autotvm.apply_graph_best(graph_opt_sch_file)
              else:
                  ctx = autotvm.apply_history_best(log_file)
              with ctx:
                  print("compile...")
                  with tvm.transform.PassContext(opt_level=3):
                      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/autotvm/testing/tune_relay.py
imports: import argparse
  import json
  import os
  import warnings
  from distutils.util import strtobool
  import tvm
  from tvm import autotvm
  from tvm import meta_schedule
  from tvm import relay
  from tvm.autotvm.graph_tuner import DPTuner
  from tvm.autotvm.tuner import XGBTuner
  from tvm.meta_schedule.testing.custom_builder_runner import run_module_via_rpc
  from tvm.meta_schedule.testing.relay_workload import get_network
  from tvm.meta_schedule.testing.tune_utils import create_timer, generate_input_data
  from tvm.support import describe
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types