Function test_rocm_vectorize_add

tests/python/codegen/test_target_codegen_rocm.py:82–109

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
def test_rocm_vectorize_add():
    def check_rocm(dtype, n, lanes):
        vec_dtype = f"{dtype}x{lanes}"
        num_blocks = n // 4

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(A: T.Buffer((n,), vec_dtype), B: T.Buffer((n,), vec_dtype)):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(num_blocks, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(4, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(n, i_0 * 4 + i_1)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = A[v_i] + T.Broadcast(T.Cast(dtype, 1), lanes)

        fun = tvm.compile(Module, target="rocm")

        dev = tvm.rocm(0)
        a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np.random.uniform(size=(n, lanes)))
        c = tvm.runtime.empty((n,), vec_dtype, dev)
        fun(a, c)
        tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)

    check_rocm("float32", 64, 2)
    check_rocm("float16", 64, 2)


@pytest.mark.gpu