code: def sort_ir(
      data,
      values_out,
      values_out_swap,
      axis,
      is_ascend,
      indices_out=None,
      indices_out_swap=None,
  ):
      """Low level IR to do sorting on the GPU, same usage as tvm.contrib.sort.argsort on the CPU.
  
      Parameters
      ----------
      data: Buffer
          Buffer of input data. Data will be sorted in place.
  
      values_out : Buffer
          Output buffer of values of sorted tensor with same shape as data.
  
      values_out_swap : Buffer
          Output buffer of values with same shape as data to use as swap.
  
      axis : Int
          Axis long which to sort the input tensor.
  
      is_ascend : Boolean
          Whether to sort in ascending or descending order.
  
      indicess_out : Buffer
          Output buffer of indices of sorted tensor with same shape as data.
  
      indices_out_swap : Buffer
          Output buffer of indices with same shape as data to use as swap.
  
      Returns
      -------
      stmt : Stmt
          The result IR statement.
      """
      ib = tvm.tir.ir_builder.create()
      shape = data.shape
  
      data = ib.buffer_ptr(data)
      values_out = ib.buffer_ptr(values_out)
      values_out_swap = ib.buffer_ptr(values_out_swap)
      if indices_out is not None:
          indices_out = ib.buffer_ptr(indices_out)
          assert indices_out_swap is not None
          indices_out_swap = ib.buffer_ptr(indices_out_swap)
  
      with ib.if_scope(shape[axis] > 0):
          axis_mul_before, axis_mul_after = _sort_init(
              ib,
              shape,
              axis,
              data,
              values_out,
              indices_out,
              value_init_func=lambda _, tid: tvm.tir.generic.cast(
                  tid, indices_out.dtype
              ),
          )
  
          _sort_common(
              ib,
              shape[axis],
              axis_mul_before,
              axis_mul_after,
              is_ascend,
              values_out,
              values_out_swap,
              values=indices_out,
              values_swap=indices_out_swap,
          )
  
      return ib.get()
  
imports: import tvm
  from tvm import te
  from .injective import schedule_injective_from_existing
  from ..transform import strided_slice, transpose
  from ..None import tag
  from ..utils import ceil_div, swap
  from ..math import cast, ceil_log2
label_desc: Use image node_id 6 for calling a global function with example usage: sort_ir(data, values_out, values_out_swap, axis, is_ascend, indices_out, indices_out_swap) and returns: ib