code: def register_topi_schedule(task_name, func=None):
      """Register a tunable template for a topi schedule function.
  
      The registration will wrap this topi schedule to take `cfg` as the first argument,
      followed by the original argument list.
  
      Note that this function will try to find "workload" from all the ComputeOp in the input.
      You can attach "workload" to your compute op by using :any:`register_topi_compute`.
  
      The task name has to be the same as that of the corresponding topi compute function.
  
      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_schedule):
          @functools.wraps(topi_schedule)
          @_register_task_schedule(task_name)
          def wrapper(outs, *args, **kwargs):
              """wrapper function for topi schedule"""
              workload = get_workload(outs, task_name)
              if workload is None:
                  raise RuntimeError(
                      f"Cannot find TOPI workload {task_name}. "
                      "Is it registered with `register_topi_compute`?"
                  )
              tgt = Target.current()
              cfg = DispatchContext.current.query(tgt, workload)
              return topi_schedule(cfg, outs, *args, **kwargs)
  
          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 2 for calling a global function with example usage: register_topi_schedule(task_name, func) and returns: _decorate, wrapper, _decorate, topi_schedule