code: def gen_ir(data, indices, updates, out, axis, reduce_func):
      """Generate ir for scatter elements
  
      Parameters
      ----------
      data : tir.Tensor
          The input data to the operator.
  
      indices : tir.Tensor
          The index locations to update.
  
      updates : tir.Tensor
          The values to update.
  
      out : tir.Tensor
          The output tensor.
  
      axis : int
          The axis to scatter on
  
      reduce_func : Any
          The function reduced update and output to output
  
      Returns
      -------
      ret : tir
          The computational ir.
      """
      ib = tir.ir_builder.create()
  
      data_ptr = ib.buffer_ptr(data)
      indices_ptr = ib.buffer_ptr(indices)
      updates_ptr = ib.buffer_ptr(updates)
      out_ptr = ib.buffer_ptr(out)
  
      # Prepare ranges and strides
      shape = data.shape
      if axis < 0:
          axis = len(shape) + axis
      axis_range = cast(shape[axis], indices.dtype)
  
      before_axis_range = 1
      after_axis_range = 1
      for i, value in enumerate(shape, 0):
          if i < axis:
              before_axis_range *= value
          elif i > axis:
              after_axis_range *= value
      before_axis_stride = axis_range * after_axis_range
      full_range = before_axis_range * before_axis_stride
  
      ind_shape = indices.shape
      ind_axis_range = ind_shape[axis]
  
      ind_before_axis_range = 1
      ind_after_axis_range = 1
      for i, value in enumerate(ind_shape, 0):
          if i < axis:
              ind_before_axis_range *= value
          elif i > axis:
              ind_after_axis_range *= value
      ind_before_axis_stride = ind_axis_range * ind_after_axis_range
      ind_full_range_excl_axis = (
          ind_before_axis_range * ind_after_axis_range
      )
  
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
      # Copy initial input data to output
      with ib.new_scope():
          num_blocks = ceil_div(full_range, max_threads)
          bx = te.thread_axis("blockIdx.x")
          tx = te.thread_axis("threadIdx.x")
          ib.scope_attr(bx, "thread_extent", num_blocks)
          ib.scope_attr(tx, "thread_extent", max_threads)
  
          index = bx * max_threads + tx
          with ib.if_scope(index < full_range):
              out_ptr[index] = data_ptr[index]
  
      with ib.new_scope():
          num_blocks_2 = ceil_div(ind_full_range_excl_axis, max_threads)
          bx2 = te.thread_axis("blockIdx.x")
          tx2 = te.thread_axis("threadIdx.x")
          ib.scope_attr(bx2, "thread_extent", num_blocks_2)
          ib.scope_attr(tx2, "thread_extent", max_threads)
  
          ind_fused = bx2 * max_threads + tx2
          with ib.if_scope(ind_fused < ind_full_range_excl_axis):
              i = ind_fused // ind_after_axis_range
              j = ind_fused % ind_after_axis_range
              pre_index1 = i * ind_before_axis_stride + j
              pre_index2 = i * before_axis_stride + j
              with ib.for_range(0, ind_axis_range, "k") as k:
                  # Offset along indices or updates
                  index1 = pre_index1 + k * ind_after_axis_range
                  # Get index and shift to positive side if need
                  new_index = indices_ptr[index1]
                  shifted_index = (
                      new_index + (new_index < 0) * axis_range
                  )
                  # Offset along data
                  index2 = pre_index2 + shifted_index * after_axis_range
                  reduce_func(out_ptr, index2, updates_ptr[index1])
  
      return ib.get()
  
file_path: tvm/python/tvm/topi/cuda/scatter_elements.py
imports: import tvm
  from tvm import te, tir
  from ..utils import ceil_div, get_const_int
  from ..math import cast
  from .nms import atomic_add
label_desc: Use image node_id 2 for calling a global function with example usage: gen_ir(data, indices, updates, out, axis, reduce_func) and returns: ib