code: def _nms_loop(
      ib,
      batch_size,
      top_k,
      iou_threshold,
      max_output_size,
      valid_count,
      on_new_valid_box_func,
      on_new_invalidated_box_func,
      needs_bbox_check_func,
      calc_overlap_func,
      out_scores,
      num_valid_boxes,
  ):
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
  
      with ib.new_scope():
          nthread_by = batch_size
          nthread_tx = max_threads
  
          # Some cuda architectures have smaller limit of 32K for cudaDevAttrMaxRegistersPerBlock
          # vs 64K for most GPUs. Since this kernel uses many registers (around 35), the limit will
          # be exceeded with 1024 threads.
          target = tvm.target.Target.current(allow_none=False)
          if target.kind.name == "cuda":
              if nvcc.get_target_compute_version(target) in [
                  "3.2",
                  "5.3",
                  "6.2",
              ]:
                  nthread_tx = 512
  
          by = te.thread_axis("blockIdx.y")
          tx = te.thread_axis("threadIdx.x")
          ib.scope_attr(by, "thread_extent", nthread_by)
          ib.scope_attr(tx, "thread_extent", nthread_tx)
  
          num_valid_boxes_local = ib.allocate(
              "int32", (1,), name="num_valid_boxes_local", scope="local"
          )
          num_valid_boxes_local[0] = 0
  
          def nms_inner_loop(ib, i, j, nkeep):
              # The box j is valid, invalidate other boxes that overlap with j above iou_threshold
              on_new_valid_box_func(
                  ib, tx, num_valid_boxes_local[0], i, j
              )
              num_valid_boxes_local[0] += 1
  
              num_iter_per_thread = ceil_div(
                  nkeep - (j + 1), nthread_tx
              )
  
              with ib.for_range(
                  0, num_iter_per_thread, name="_k"
              ) as _k:
                  k = j + 1 + _k * nthread_tx + tx
  
                  with ib.if_scope(
                      tvm.tir.all(
                          k < nkeep,
                          out_scores[i, k]
                          > 0,  # is the box k still valid?
                          needs_bbox_check_func(i, j, k),
                      )
                  ):
                      iou = calc_overlap_func(i, j, k)
  
                      with ib.if_scope(iou >= iou_threshold):
                          # invalidate the box k
                          out_scores[i, k] = -1.0
                          on_new_invalidated_box_func(i, k)
  
                  ib.emit(
                      tvm.tir.Call(
                          None,
                          "tir.tvm_storage_sync",
                          tvm.runtime.convert(["shared"]),
                      )
                  )
  
          i = by
  
          nkeep = if_then_else(
              tvm.tir.all(top_k > 0, top_k < valid_count[i]),
              top_k,
              valid_count[i],
          )
          max_output_size = if_then_else(
              max_output_size > 0, max_output_size, nkeep
          )
  
          with ib.if_scope(
              tvm.tir.all(iou_threshold > 0, valid_count[i] > 0)
          ):
              # Apply nms
              # No need to do more iteration if we have already reached max_output_size boxes
              box_idx = ib.allocate(
                  "int32", (1,), name="box_idx", scope="local"
              )
              box_idx[0] = 0
              with ib.while_loop(
                  tvm.tir.all(
                      box_idx[0] < nkeep,
                      num_valid_boxes_local[0] < max_output_size,
                  )
              ):
                  # Proceed to the inner loop if the box with id box_idx is still valid
                  with ib.if_scope(out_scores[i, box_idx[0]] > -1.0):
                      nms_inner_loop(ib, i, box_idx[0], nkeep)
                  box_idx[0] += 1
  
              with ib.if_scope(tx + 0 == 0):
                  num_valid_boxes[i] = num_valid_boxes_local[0]
  
          with ib.else_scope():
              num_valid_boxes[i] = 0
  
      return ib.get()
  
imports: import tvm
  from tvm import te
  from tvm.contrib import nvcc
  from tvm.contrib.thrust import can_use_thrust, can_use_rocthrust
  from tvm.ir import register_intrin_lowering
  from tvm.tir import if_then_else
  from .sort import argsort, argsort_thrust
  from .scan import exclusive_scan
  from ..utils import ceil_div
  from ..math import cast
  from ..transform import reshape
  from ..vision.nms_util import calculate_overlap, binary_search, collect_selected_indices, collect_selected_indices_and_scores, run_all_class_nms
label_desc: Use image node_id 7 for calling a global function with example usage: _nms_loop(ib, batch_size, top_k, iou_threshold, max_output_size, valid_count, on_new_valid_box_func, on_new_invalidated_box_func, needs_bbox_check_func, calc_overlap_func, out_scores, num_valid_boxes) and returns: ib