code: def try_warp_memory():
      """Test using warp memory
      skip this in default test because it require higher arch"""
      arr_size = 128
      placeholder_a = te.placeholder((arr_size,), name="A")
      result_b = te.compute(
          (arr_size,), lambda i: placeholder_a[i] + 3, name="B"
      )
      warp_size = 32
      schedule = te.create_schedule(result_b.op)
      cache_read_aa = schedule.cache_read(
          placeholder_a, "warp", [result_b]
      )
      axis_x0, axis_xi = schedule[result_b].split(
          result_b.op.axis[0], warp_size * 2
      )
      _, axis_xi1 = schedule[result_b].split(axis_xi, factor=warp_size)
      thread_axis_tx = te.thread_axis("threadIdx.x")
      schedule[result_b].bind(axis_xi1, thread_axis_tx)
      schedule[result_b].bind(axis_x0, te.thread_axis("blockIdx.x"))
      schedule[cache_read_aa].compute_at(schedule[result_b], axis_x0)
      axis_x0, axis_xi = schedule[cache_read_aa].split(
          schedule[cache_read_aa].op.axis[0], warp_size
      )
      schedule[cache_read_aa].bind(axis_xi, thread_axis_tx)
  
      @tvm.register_func("tvm_callback_cuda_compile", override=True)
      def tvm_callback_cuda_compile(
          code, _
      ):  # pylint: disable=unused-variable
          ptx = nvcc.compile_cuda(code)
          return ptx
  
      # one line to build the function.
      def check_device(device):
          dev = tvm.device(device, 0)
          if not tvm.testing.device_enabled(device):
              print("skip because %s is not enabled.." % device)
              return
          myfunc = tvm.build(
              schedule, [placeholder_a, result_b], device
          )
          buff_a = tvm.nd.array(
              (np.random.uniform(size=arr_size) * 256).astype(
                  placeholder_a.dtype
              ),
              dev,
          )
          buff_b = tvm.nd.array(
              np.zeros(arr_size, dtype=result_b.dtype), dev
          )
          myfunc(buff_a, buff_b)
          tvm.testing.assert_allclose(
              buff_b.numpy(), buff_a.numpy() + 3, rtol=1e-6
          )
  
      check_device("cuda")
  
file_path: tvm/tests/python/integration/test_ewise.py
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.contrib import nvcc
label_desc: Use image node_id 7 for calling a global function with example usage: try_warp_memory() and returns: ptx