Function test_half_broadcast

tests/python/codegen/test_target_codegen_cuda_fp8.py:226–254

(bcast_length)

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_half_broadcast(bcast_length):
    dtype = "float16"

    def _create_mod(bcast_length, dtype):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(a: T.Buffer((), dtype), vec: T.Buffer((bcast_length,), dtype)):
                for i_0 in T.thread_binding(1, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(1, thread="threadIdx.x"):
                        with T.sblock("broadcast"):
                            vec[0:bcast_length] = T.broadcast(a[()], bcast_length)

        return Module

    mod = _create_mod(bcast_length, dtype)
    target = "cuda"
    func = tvm.compile(mod, target=target)
    dev = tvm.device(target, 0)

    a_np = np.random.uniform(low=0, high=4, size=()).astype(dtype)
    a = tvm.runtime.tensor(a_np, device=dev)
    b = tvm.runtime.empty((bcast_length,), dtype=dtype, device=dev)

    func(a, b)

    b_np = np.full((bcast_length,), a_np)

    tvm.testing.assert_allclose(b.numpy(), b_np)


vector_length = tvm.testing.parameter(2, 4)