Function test_cuda_vectorize_load

tests/python/codegen/test_target_codegen_cuda.py:217–250

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_load():
    num_thread = 8

    def check_cuda(dtype, n, lanes):
        dev = tvm.cuda(0)
        vec_dtype = f"{dtype}x{lanes}"
        num_blocks = n // num_thread

        @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(num_thread, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(n, i_0 * num_thread + i_1)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = A[v_i]

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

        np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
        a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_a)
        b = tvm.runtime.empty((n,), vec_dtype, dev)
        fun(a, b)
        tvm.testing.assert_allclose(a.numpy(), b.numpy())

    check_cuda("int8", 64, 2)
    check_cuda("int8", 64, 3)
    check_cuda("int8", 64, 4)
    check_cuda("int8", 64, 8)
    check_cuda("int8", 64, 16)


@pytest.mark.gpu