code: def search_op_config(code_only=False):
      tvm_target = "cuda"
      logging.getLogger("autotvm").setLevel(logging.DEBUG)
      logging.getLogger("autotvm").addHandler(
          logging.StreamHandler(sys.stdout)
      )
  
      default_tune_op = importlib.import_module(
          "templates." + (os.environ["OP"])
      )
      print(
          "  >> Backend = %s, Python PID = %s, Task = %s;"
          % (backend, os.getpid(), default_tune_op.__name__)
      )
  
      task = autotvm.task.create(
          default_tune_op.get_template_op, args=(), target=tvm_target
      )
      op_attributes = default_tune_op.op_attributes
      op_summary = "_".join(
          [k + str(op_attributes[k]) for k in op_attributes]
      )
  
      def json_to_config(json_dict):
          config = ConfigEntity.from_json_dict(
              {"i": -1, "t": "", "c": None, "e": json_dict}
          )
          return config
  
      def config_to_json(config):
          jobj = config.to_json_dict()["e"]
          json_dict = dict()
          for i in range(len(jobj)):
              assert jobj[i][1] in ["sp", "ot"]
              json_dict[jobj[i][0]] = jobj[i][2]
          return json_dict
  
      num_trials = (
          int(os.environ["STEP"]) if "STEP" in os.environ else 0
      )
  
      if "CONFIG" in os.environ:
          params_given = json.loads(os.environ["CONFIG"])
          print("====>> [Current Config Option]", os.environ["CONFIG"])
  
          trial_config = []
          for key in params_given:
              trial_config.append(
                  [
                      key,
                      "sp" if type(params_given[key]) is list else "ot",
                      params_given[key],
                  ]
              )
          best_config = json_to_config(trial_config)
  
      elif "NNI_TRIAL_JOB_ID" in os.environ:
          show_search_space(
              task.config_space, os.environ["NNI_TRIAL_JOB_ID"] == "@"
          )
          import nni
  
          params_given = nni.get_next_parameter()
          if params_given is None:
              raise
          local_dir_id = os.environ["NNI_TRIAL_JOB_ID"]
          t = run_config_entity(params_given, local_dir_id)
          gflops = compute_gflops(task.flop, t)
          print("[TVM-engine] Final entity result is: %g" % gflops)
          try:
              nni.report_final_result(gflops)
          except:
              print("[TVM-engine] (not reporting final result to NNI.)")
          exit(0)
  
      elif num_trials > 0:
          n_parallel = (
              16
              if "BATCH" not in os.environ
              else int(os.environ["BATCH"])
          )
          measure_option = autotvm.measure_option(
              builder=autotvm.LocalBuilder(n_parallel=n_parallel),
              runner=autotvm.LocalRunner(
                  repeat=3, min_repeat_ms=100, timeout=4
              ),
          )
          # if DO_TUNING:
          tuner = autotvm.tuner.XGBTuner(task, num_threads=8)
  
          from concurrent.futures import ThreadPoolExecutor
  
          thread_pool = ThreadPoolExecutor(max_workers=n_parallel)
  
          dev_num = get_tuning_parallism()
  
          def parse_configs(task, configs):
              results = []
              futures = []
              expected_timecost = "inf"
              for i in range(len(configs)):
                  futures.append(
                      thread_pool.submit(
                          run_config_entity,
                          config_to_json(configs[i]),
                          i,
                          expected_timecost,
                          i % dev_num,
                      )
                  )
              for i in range(len(configs)):
                  t = futures[i].result()
                  if t < tuner.task.best_config[0]:
                      tuner.task.best_config = (t, configs[i])
                  results.append(
                      autotvm.measure.MeasureResult(
                          costs=(t,),
                          error_no=0,
                          all_cost=i,
                          timestamp=time.time(),
                      )
                  )
              return results
  
          tuner.task.best_config = (float("inf"), None)
          tuner.parse_configs = parse_configs
  
          tuner.tune(
              n_trial=num_trials,
              measure_option=measure_option,
              callbacks=[],
          )
          assert not math.isinf(tuner.task.best_config[0])
          best_config = tuner.task.best_config[1]
          print(
              "\n[Best Config]", json.dumps(config_to_json(best_config))
          )
      else:
          best_config = task.config_space
  
      with ApplyConfig(best_config):
          with tvm.target.create(tvm_target):
              s, arg_bufs = default_tune_op.get_template_op()
              lower_source = str(
                  tvm.lower(s, arg_bufs, simple_mode=True)
              )
  
              # Verify Source Code
              assert len(("\n" + lower_source).split("\nproduce ")) == 2
              lower_file = local_get_dir_file("my_kernel.lower")
              with open(lower_file, "w") as fp:
                  fp.write(lower_source)
  
              max_threads_per_block = tvm.ndarray.gpu(
                  0
              ).max_threads_per_block
              max_shared_memory_per_block = tvm.ndarray.gpu(
                  0
              ).max_shared_memory_per_block
  
              thread_extents = subprocess.getoutput(
                  "cat '%s' | grep '^ *// attr.*iter_var.*thread_extent'"
                  % (lower_file)
              ).split("\n")
              reserved_axes = dict(
                  {
                      "threadIdx.x": None,
                      "threadIdx.y": None,
                      "threadIdx.z": None,
                      "blockIdx.x": None,
                      "blockIdx.y": None,
                      "blockIdx.z": None,
                  }
              )
              for line in thread_extents:
                  thread_name = line.split("[iter_var(")[-1].split(",")[
                      0
                  ]
                  if thread_name in reserved_axes:
                      thread_val = int(
                          line.split("thread_extent = ")[-1]
                      )
                      if reserved_axes[thread_name] is not None:
                          if reserved_axes[thread_name] != thread_val:
                              assert False
                      else:
                          reserved_axes[thread_name] = thread_val
                  else:
                      raise Exception(
                          "Invalid thread_axis name: %s" % thread_name
                      )
  
              num_threads = 1
              for thread_name in [
                  "threadIdx.x",
                  "threadIdx.y",
                  "threadIdx.z",
              ]:
                  if reserved_axes[thread_name] is not None:
                      num_threads *= reserved_axes[thread_name]
              if num_threads > max_threads_per_block:
                  raise Exception(
                      "Invalid kernel code: using num_threads %d > max_threads_per_block %d"
                      % (num_threads, max_threads_per_block)
                  )
  
              allocate_shared = subprocess.getoutput(
                  "cat '%s' | grep 'allocate .*shared\[.*\]'"
                  % (lower_file)
              ).split("\n")
              shared_memory_in_bytes = 0
              for line in allocate_shared:
                  if not line:
                      continue
                  parts = line.split("[")
                  assert len(parts) == 2
                  parts = parts[1].split(" * ")
                  assert len(parts) == 2
                  assert parts[1][-1] == "]"
                  allocate_type = parts[0]
                  allocate_val = int(parts[1][:-1])
                  if allocate_type in ["float32"]:
                      shared_memory_in_bytes += allocate_val * 4
                  else:
                      raise Exception(
                          "Unrecognized shared memory data type: %s"
                          % allocate_type
                      )
              if shared_memory_in_bytes > max_shared_memory_per_block:
                  raise Exception(
                      "Invalid kernel code: using shared_memory_in_bytes %d > max_shared_memory_per_block %d"
                      % (
                          shared_memory_in_bytes,
                          max_shared_memory_per_block,
                      )
                  )
  
              func = tvm.build(
                  s, arg_bufs, tvm_target, name="template_op"
              )
  
      assert len(func.imported_modules) == 1
      device_source = translate_code(
          func.imported_modules[0].get_source()
      )
  
      if code_only:
          return device_source
  
      if lower_source and device_source:
          tune_slot_id = (
              0
              if "CUDA_VISIBLE_DEVICES" not in os.environ
              else int(os.environ["CUDA_VISIBLE_DEVICES"])
          )
          exec_fd, _ = system_lock([tune_slot_id])
          gpu_id = 0
          ctx = tvm.context(tvm_target, gpu_id)
          tensors, outs = [], []
          for arg in arg_bufs:
              shape = [int(x) for x in arg.shape]
              is_output = arg.op.__class__ != tvm.tensor.PlaceholderOp
              from tvm._ffi.ndarray import empty
  
              td = empty(shape, arg.dtype, ctx)
              if is_output:
                  outs.append(td)
              tensors.append(td)
  
          def timeout_handler():
              print("Error: Timeout during Kernel warmup")
              os._exit(1)
  
          my_timer = Timer(10, timeout_handler, [])
          my_timer.start()
          # Warmup
          func(*tensors)
          tvm.ndarray.gpu(gpu_id).sync()
          # Estimate
          t_start = time.time()
          func(*tensors)
          tvm.ndarray.gpu(gpu_id).sync()
          t_diff = time.time() - t_start
          my_timer.cancel()
          del my_timer
  
          num_runs = max(3, min(100, math.floor(1.0 / t_diff)))
          timeout_seconds = math.ceil((num_runs + 5) * t_diff)
          my_timer = Timer(timeout_seconds, timeout_handler, [])
          my_timer.start()
          timer_f = func.time_evaluator(
              func.entry_name, ctx, number=num_runs
          )
          t = timer_f(*tensors).mean
          my_timer.cancel()
          exec_fd()
  
          gflops = compute_gflops(task.flop, t)
          print(
              "[TVM-engine] Average time cost of %d runs = %g ms, %g gflops."
              % (num_runs, t * 1e3, gflops)
          )
  
          with open(local_get_dir_file("result.txt"), "w") as fp:
              fp.write(str(t))
  
imports: import numpy
  import tvm
  import logging
  import math
  import re
  import sys, time, subprocess, os, random, hashlib
  from tvm import autotvm
  import topi
  import json
  from topi.util import get_const_tuple
  import importlib
  from tvm.autotvm.task.dispatcher import ApplyConfig
  from tvm.autotvm.task import ConfigEntity
  from threading import Timer
label_desc: Use image node_id 11 for calling a global function with example usage: search_op_config(code_only) and returns: config, json_dict, device_source, results