code: def nms_ir(
      data,
      sorted_index,
      valid_count,
      indices,
      out_bboxes,
      out_scores,
      out_class_ids,
      out_features,
      box_indices,
      num_valid_boxes,
      max_output_size,
      iou_threshold,
      force_suppress,
      top_k,
      coord_start,
      id_index,
      score_index,
      return_indices,
  ):
      """Low level IR routing for transform location in multibox_detection operator.
  
      Parameters
      ----------
      data : Buffer
          Buffer of output boxes with class and score.
  
      sorted_index : Buffer
          Buffer of output box indexes sorted by score.
  
      valid_count : Buffer
          Buffer of number of valid output boxes.
  
      indices : Buffer
          indices in original tensor, with shape [batch_size, num_anchors],
          represents the index of box in original data. It could be the third
          output out_indices of get_valid_counts. The values in the second
          dimension are like the output of arange(num_anchors) if get_valid_counts
          is not used before non_max_suppression.
  
      out_bboxes : Buffer
          Output buffer, to be filled with sorted box coordinates.
  
      out_scores : Buffer
          Output buffer, to be filled with sorted scores.
  
      out_class_ids : Buffer
          Output buffer, to be filled with sorted class ids.
  
      box_indices : Buffer
          A indices tensor mapping sorted indices to original indices
          This is the first output of NMS when return_indices=True.
  
      num_valid_boxes : Buffer
          Record the number of boxes that have survived IOU tests.
          This is the second output of NMS when return_indices=True.
  
      max_output_size : int
          Max number of output valid boxes for each instance.
          By default all valid boxes are returned.
  
      iou_threshold : float
          Overlapping(IoU) threshold to suppress object with smaller score.
  
      force_suppress : boolean
          Whether to suppress all detections regardless of class_id.
  
      top_k : int
          Keep maximum top k detections before nms, -1 for no limit.
  
      coord_start : int
          Start index of the consecutive 4 coordinates.
  
      id_index : int
          index of the class categories, -1 to disable.
  
      score_index : optional, int
          Index of the scores/confidence of boxes.
  
      return_indices : boolean
          Whether to return box indices in input data.
  
      Returns
      -------
      stmt : Stmt
          The result IR statement.
      """
      batch_size = data.shape[0]
      num_anchors = data.shape[1]
      box_data_length = data.shape[2]
      num_features = out_features.shape[2]
  
      ib = tvm.tir.ir_builder.create()
  
      data = ib.buffer_ptr(data)
      sorted_index = ib.buffer_ptr(sorted_index)
      valid_count = ib.buffer_ptr(valid_count)
      indices = ib.buffer_ptr(indices)
  
      # outputs
      out_bboxes = ib.buffer_ptr(out_bboxes)
      out_scores = ib.buffer_ptr(out_scores)
      out_class_ids = ib.buffer_ptr(out_class_ids)
      out_features = ib.buffer_ptr(out_features)
      box_indices = ib.buffer_ptr(box_indices)
      num_valid_boxes = ib.buffer_ptr(num_valid_boxes)
  
      if isinstance(iou_threshold, float):
          iou_threshold = tvm.tir.FloatImm("float32", iou_threshold)
      top_k = tvm.tir.IntImm("int32", top_k)
      coord_start = tvm.tir.IntImm("int32", coord_start)
      id_index = tvm.tir.IntImm("int32", id_index)
      score_index = tvm.tir.IntImm("int32", score_index)
      force_suppress = tvm.tir.IntImm(
          "int32", 1 if force_suppress else 0
      )
  
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
  
      with ib.new_scope():
          nthread_tx = max_threads
          nthread_bx = ceil_div(num_anchors, max_threads)
          nthread_by = batch_size
          tx = te.thread_axis("threadIdx.x")
          bx = te.thread_axis("blockIdx.x")
          by = te.thread_axis("blockIdx.y")
          ib.scope_attr(by, "thread_extent", nthread_by)
          ib.scope_attr(tx, "thread_extent", nthread_tx)
          ib.scope_attr(bx, "thread_extent", nthread_bx)
          i = by
          base_src_idx = i * num_anchors * box_data_length
          base_bbox_idx = i * num_anchors * 4
          base_features_idx = i * num_anchors * num_features
  
          with ib.if_scope(
              tvm.tir.all(iou_threshold > 0, valid_count[i] > 0)
          ):
              # Reorder output
              nkeep = if_then_else(
                  tvm.tir.all(top_k > 0, top_k < valid_count[i]),
                  top_k,
                  valid_count[i],
              )
              j = bx * max_threads + tx
              with ib.if_scope(j < nkeep):
                  src_idx = (
                      base_src_idx
                      + sorted_index[i * num_anchors + j]
                      * box_data_length
                  )
                  with ib.for_range(0, 4, kind="unroll") as k:
                      out_bboxes[(base_bbox_idx + j * 4 + k)] = data[
                          src_idx + coord_start + k
                      ]
                  with ib.for_range(
                      0, num_features, kind="unroll"
                  ) as k:
                      out_features[
                          (base_features_idx + j * num_features + k)
                      ] = data[src_idx + coord_start + 4 + k]
  
                  out_scores[i * num_anchors + j] = data[
                      src_idx + score_index
                  ]
  
                  if id_index >= 0:
                      out_class_ids[i * num_anchors + j] = data[
                          src_idx + id_index
                      ]
  
              with ib.else_scope():
                  # Indices > nkeep are discarded
                  # Only needed for return_indices = False case
                  if return_indices is False:
                      with ib.if_scope(j < num_anchors):
                          with ib.for_range(0, 4, kind="unroll") as k:
                              out_bboxes[
                                  (base_bbox_idx + j * 4 + k)
                              ] = -1.0
                          with ib.for_range(
                              0, num_features, kind="unroll"
                          ) as k:
                              out_features[
                                  (
                                      base_features_idx
                                      + j * num_features
                                      + k
                                  )
                              ] = -1.0
  
                          out_scores[i, j] = -1.0
  
                          if id_index >= 0:
                              out_class_ids[i, j] = -1.0
  
              if return_indices:
                  with ib.if_scope(j < num_anchors):
                      box_indices[i * num_anchors + j] = -1
  
          with ib.else_scope():
              # Need to copy all boxes if not using return_indices
              bounds = valid_count[i] if return_indices else num_anchors
              with ib.if_scope(j < bounds):
                  src_offset = base_src_idx + j * box_data_length
  
                  with ib.for_range(0, 4, kind="unroll") as k:
                      out_bboxes[base_bbox_idx + j * 4 + k] = data[
                          src_offset + coord_start + k
                      ]
                  with ib.for_range(
                      0, num_features, kind="unroll"
                  ) as k:
                      out_features[
                          (base_features_idx + j * num_features + k)
                      ] = data[src_offset + coord_start + 4 + k]
                  out_scores[i * num_anchors + j] = data[
                      src_offset + score_index
                  ]
  
                  if id_index >= 0:
                      out_class_ids[i * num_anchors + j] = data[
                          src_offset + id_index
                      ]
  
                  box_indices[i * num_anchors + j] = j
  
      if isinstance(max_output_size, int):
          max_output_size = tvm.tir.const(max_output_size)
  
      def calc_overlap(i, j, k):
          offset_j = j * 4
          offset_k = k * 4
          base_bbox_idx = i * num_anchors * 4
          return calculate_overlap(
              out_bboxes,
              base_bbox_idx + offset_j,
              base_bbox_idx + offset_k,
          )
  
      def on_new_valid_box(ib, tid, num_current_valid_box, i, j):
          # When return_indices is False, no need to populate box_indices
          if return_indices:
              with ib.if_scope(tid + 0 == 0):
                  orig_idx = sorted_index[i * num_anchors + j]
                  box_indices[i, num_current_valid_box] = indices[
                      i, orig_idx
                  ]
  
      def on_new_invalidated_box(i, k):
          if return_indices is False and id_index >= 0:
              out_class_ids[i, k] = -1.0
  
      def needs_bbox_check(i, j, k):
          return tvm.tir.any(
              force_suppress > 0,
              id_index < 0,
              out_class_ids[i, k] == out_class_ids[i, j],
          )
  
      return _nms_loop(
          ib,
          batch_size,
          top_k,
          iou_threshold,
          max_output_size,
          valid_count,
          on_new_valid_box,
          on_new_invalidated_box,
          needs_bbox_check,
          calc_overlap,
          out_scores,
          num_valid_boxes,
      )
  
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 8 for calling a global function with example usage: nms_ir(data, sorted_index, valid_count, indices, out_bboxes, out_scores, out_class_ids, out_features, box_indices, num_valid_boxes, max_output_size, iou_threshold, force_suppress, top_k, coord_start, id_index, score_index, return_indices) and returns: _nms_loop, calculate_overlap, tvm