| 112 | @pytest.mark.gpu |
| 113 | @pytest.mark.skipif(not env.has_rocm(), reason="need rocm") |
| 114 | def test_rocm_warp_shuffle(): |
| 115 | @T.prim_func(s_tir=True) |
| 116 | def func( |
| 117 | A_handle: T.handle, |
| 118 | ): |
| 119 | A = T.match_buffer(A_handle, (32,), dtype="float32") |
| 120 | |
| 121 | for bx in T.thread_binding(1, thread="blockIdx.x"): |
| 122 | for tx in T.thread_binding(32, thread="threadIdx.x"): |
| 123 | with T.sblock("test"): |
| 124 | A_local = T.sblock_alloc_buffer((1,), "float32", scope="local") |
| 125 | mask = T.sblock_alloc_buffer((1,), "uint32", scope="local") |
| 126 | t0 = T.sblock_alloc_buffer((1,), "float32", scope="local") |
| 127 | |
| 128 | A_local[0] = A[tx] |
| 129 | A_local[0] = T.tvm_warp_shuffle(mask[0], A_local[0], 0, 32, 32) |
| 130 | A[tx] = A_local[0] |
| 131 | |
| 132 | mod = tvm.compile(func, target="rocm") |
| 133 | dev = tvm.rocm(0) |
| 134 | a = tvm.runtime.tensor(np.random.uniform(size=(32,)).astype("float32"), dev) |
| 135 | mod(a) |
| 136 | tvm.testing.assert_allclose(a.numpy(), np.ones((32,)) * a.numpy()[0]) |
| 137 | |
| 138 | |
| 139 | @pytest.mark.gpu |