code: def _collect_selected_indices_ir(
      num_class, selected_indices, num_detections, row_offsets, out
  ):
      batch_classes, num_boxes = selected_indices.shape
  
      ib = tvm.tir.ir_builder.create()
  
      selected_indices = ib.buffer_ptr(selected_indices)
      num_detections = ib.buffer_ptr(num_detections)
      row_offsets = ib.buffer_ptr(row_offsets)
      out = ib.buffer_ptr(out)
  
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
      nthread_tx = max_threads
      nthread_bx = ceil_div(num_boxes, nthread_tx)
      nthread_by = batch_classes
      tx = te.thread_axis("threadIdx.x")
      bx = te.thread_axis("blockIdx.x")
      by = te.thread_axis("blockIdx.y")
      ib.scope_attr(tx, "thread_extent", nthread_tx)
      ib.scope_attr(bx, "thread_extent", nthread_bx)
      ib.scope_attr(by, "thread_extent", nthread_by)
  
      with ib.new_scope():
          idx = bx * nthread_tx + tx
          idy = cast(by, "int64")
          batch_id = idy // num_class
          class_id = idy % num_class
          with ib.if_scope(idx < num_detections[idy]):
              out[row_offsets[idy] + idx, 0] = batch_id
              out[row_offsets[idy] + idx, 1] = class_id
              out[row_offsets[idy] + idx, 2] = cast(
                  selected_indices[idy, idx], "int64"
              )
  
      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 16 for calling a global function with example usage: _collect_selected_indices_ir(num_class, selected_indices, num_detections, row_offsets, out) and returns: ib