code: def register_topi_compute(task_name, func=None):
      """Register a tunable template for a topi compute function.
  
      The registration will wrap this topi compute to take `cfg` as the first argument,
      followed by the original argument list. It uses all its argument as workload and
      stores this "workload" to its final ComputeOp, which can be used to reconstruct
      "workload" in the following topi_schedule call.
  
      Parameters
      ----------
      task_name: str
          The AutoTVM task name
  
      func: None or callable
          If it is None, return a decorator.
          If is callable, decorate this function.
  
      Returns
      -------
      decorator: callable
          A decorator
  
      Examples
      --------
      See tvm/topi/python/topi/arm_cpu/depthwise_conv2d.py for example usage.
      """
  
      def _decorate(topi_compute):
          @functools.wraps(topi_compute)
          @_register_task_compute(task_name)
          def wrapper(*args, **kwargs):
              """wrapper function for topi compute"""
              assert (
                  not kwargs
              ), "Do not support kwargs in template function call"
              task_env = TaskExtractEnv.current
              if task_env is not None and task_env.tracing:
                  task_env.add_task(task_name, args)
              workload = args_to_workload(args, task_name)
              tgt = Target.current()
              cfg = DispatchContext.current.query(tgt, workload)
              node = topi_compute(cfg, *args)
  
              # attach workload to return op
              op = node.op
              attrs = {}
              for k, v in node.op.attrs.items():
                  attrs[k] = v
              attrs["workload"] = workload
              if isinstance(op, tensor.ComputeOp):
                  op = tvm.te._ffi_api.ComputeOp(
                      op.name, op.tag, attrs, op.axis, op.body
                  )
              elif isinstance(op, tensor.ExternOp):
                  op = tvm.te._ffi_api.ExternOp(
                      op.name,
                      op.tag,
                      attrs,
                      op.inputs,
                      op.input_placeholders,
                      op.output_placeholders,
                      op.body,
                  )
              else:
                  raise RuntimeError(
                      "Unsupported op type: " + str(type(op))
                  )
  
              if isinstance(node, tensor.Tensor):
                  return op.output(0)
              return [op.output(i) for i in range(len(node))]
  
          return wrapper
  
      if func:
          return _decorate(func)
      return _decorate
  
file_path: tvm/python/tvm/autotvm/task/topi_integration.py
imports: import functools
  import tvm.te._ffi_api
  from tvm.target import Target
  from tvm.te import tensor
  from .task import args_to_workload, serialize_args, DispatchContext, _register_task_compute, _register_task_schedule
label_desc: Use image node_id 1 for calling a global function with example usage: register_topi_compute(task_name, func) and returns: _decorate, wrapper, _decorate, unknown, op