code: def test_cuda_shuffle():
      idxm = tvm.tir.indexmod
      a = te.placeholder((64,), "int32")
      b = te.placeholder((64,), "int32")
      c = te.compute(
          (64,), lambda x: a[x] + b[x - idxm(x, 4) + (3 - idxm(x, 4))]
      )
      sch = te.create_schedule(c.op)
      x = c.op.axis[0]
      xo, xi = sch[c].split(x, 4)
      thrx = te.thread_axis("threadIdx.x")
      sch[c].bind(xo, thrx)
      sch[c].vectorize(xi)
  
      def MyVectorize():
          def vectorizer(op):
              if op.kind == tvm.tir.ForKind.VECTORIZED:
                  idx = tvm.tir.Ramp(4 * thrx.var, 1, 4)
                  store = op.body
                  value = store.value
                  new_a = tvm.tir.BufferLoad(value.a.buffer, [idx])
                  bs, ids = [], []
                  for i in range(4):
                      bs.append(
                          tvm.tir.BufferLoad(
                              value.b.buffer, [4 * thrx.var + i]
                          )
                      )
                      ids.append(3 - i)
                  new_b = tvm.tir.Shuffle(bs, ids)
                  return tvm.tir.BufferStore(
                      store.buffer, new_a + new_b, [idx]
                  )
              return None
  
          def _transform(f, *_):
              return f.with_body(
                  tvm.tir.stmt_functor.ir_transform(
                      f.body, None, vectorizer, ["tir.For"]
                  )
              )
  
          return tvm.tir.transform.prim_func_pass(
              _transform, opt_level=0, name="MyVectorize"
          )
  
      with tvm.transform.PassContext(
          config={"tir.add_lower_pass": [(1, MyVectorize())]}
      ):
          module = tvm.build(sch, [a, b, c], target="cuda")
          a_ = np.array(list(range(64)), dtype="int32")
          b_ = np.array((list(range(4))[::-1]) * 16, dtype="int32")
          c_ = np.zeros((64,), dtype="int32")
          ref = a_ + np.array((list(range(4))) * 16, dtype="int32")
          nda, ndb, ndc = [
              tvm.nd.array(i, tvm.cuda(0)) for i in [a_, b_, c_]
          ]
          module(nda, ndb, ndc)
          tvm.testing.assert_allclose(ndc.numpy(), ref)
  
imports: import os
  import re
  import tvm
  from tvm import te
  import numpy
  from tvm import topi
  from tvm.contrib.nvcc import have_fp16, have_int8, have_bf16
  from tvm.contrib import utils
  from tvm.script import tir
  import tvm.testing
  import pytest
label_desc: Use image node_id 8 for calling a global function with example usage: test_cuda_shuffle() and returns: tvm, None, f, tvm