Function test_cuda_bf16_vectorize_add

tests/python/codegen/test_target_codegen_cuda.py:112–165

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_bf16_vectorize_add():
    if not have_bf16(tvm.cuda(0).compute_version):
        print("skip because gpu does not support bf16")
        return
    num_thread = 8

    def np_float2np_bf16(arr):
        """Convert a numpy array of float to a numpy array
        of bf16 in uint16"""
        orig = arr.view("<u4")
        bias = np.bitwise_and(np.right_shift(orig, 16), 1) + 0x7FFF
        return np.right_shift(orig + bias, 16).astype("uint16")

    def np_bf162np_float(arr):
        """Convert a numpy array of bf16 (uint16) to a numpy array
        of float"""
        u32 = np.left_shift(arr.astype("uint32"), 16)
        return u32.view("<f4")

    def check_cuda(n, lanes):
        vec_dtype = f"bfloat16x{lanes}"
        num_blocks = n // num_thread
        one = tvm.tirx.Broadcast(tvm.tirx.const(1, "bfloat16"), lanes)

        @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] + one

        with tvm.transform.PassContext(
            disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
        ):
            fun = tvm.compile(Module, target="cuda")
        dev = tvm.cuda(0)
        np_a = np.random.uniform(size=(n, lanes)).astype("float32")
        np_a = np_bf162np_float(np_float2np_bf16(np_a))
        a = tvm.runtime.empty((n,), vec_dtype, dev).copyfrom(np_float2np_bf16(np_a))
        c = tvm.runtime.empty((n,), vec_dtype, dev)
        fun(a, c)
        c = tvm.runtime.empty((n, lanes), "uint16", dev).copyfrom(c)
        tvm.testing.assert_allclose(c.numpy(), np_float2np_bf16(np_a + 1))

    check_cuda(64, 2)
    check_cuda(64, 4)
    check_cuda(64, 6)
    check_cuda(64, 8)


@pytest.mark.gpu