Function test_fp8_packing

tests/python/codegen/test_target_codegen_cuda_fp8.py:98–142

(dtype)

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
def test_fp8_packing(dtype):
    length = 64
    vector_length = 4
    native_dtype, packed_dtype = (f"{dtype}x{vector_length}", "uint32")

    def _create_mod(native_dtype, packed_dtype, length):
        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(
                A: T.Buffer((length,), native_dtype),
                R: T.Buffer((length,), packed_dtype),
                B: T.Buffer((length,), native_dtype),
            ):
                T.func_attr({"tirx.noalias": True})
                for i_0 in T.thread_binding(2, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(32, thread="threadIdx.x"):
                        with T.sblock("R"):
                            v_i = T.axis.spatial(length, i_0 * 32 + i_1)
                            T.reads(A[v_i])
                            T.writes(R[v_i])
                            R[v_i] = T.reinterpret(packed_dtype, A[v_i])
                for i_0 in T.thread_binding(2, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(32, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(length, i_0 * 32 + i_1)
                            T.reads(R[v_i])
                            T.writes(B[v_i])
                            B[v_i] = T.reinterpret(native_dtype, R[v_i])

        return Module

    mod = _create_mod(native_dtype, packed_dtype, length)
    target = "cuda"
    f = tvm.compile(mod, target=target)
    dev = tvm.device(target, 0)

    np_shape = (length, vector_length)
    a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(dtype)
    a = tvm.runtime.empty(shape=(length,), dtype=native_dtype, device=dev)
    r = tvm.runtime.empty(shape=(length,), dtype=packed_dtype, device=dev)
    b = tvm.runtime.empty(shape=(length,), dtype=native_dtype, device=dev)
    a.copyfrom(a_np)
    f(a, r, b)
    tvm.testing.assert_allclose(a.numpy().astype("float16"), b.numpy().astype("float16"))


native_dtype, promoted_dtype, numpytype = tvm.testing.parameters(