(input)
| 46 | @pytest.mark.gpu |
| 47 | @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") |
| 48 | def test_fp8_conversions(input): |
| 49 | dtype, nv_dtype = input |
| 50 | |
| 51 | def _create_mod(dtype): |
| 52 | @I.ir_module(s_tir=True) |
| 53 | class Module: |
| 54 | @T.prim_func(s_tir=True) |
| 55 | def main( |
| 56 | A: T.Buffer((64,), dtype), |
| 57 | B: T.Buffer((64,), dtype), |
| 58 | C: T.Buffer((64,), dtype), |
| 59 | ): |
| 60 | T.func_attr({"tirx.noalias": True}) |
| 61 | for i_0 in T.thread_binding(2, thread="blockIdx.x"): |
| 62 | for i_1 in T.thread_binding(32, thread="threadIdx.x"): |
| 63 | with T.sblock("C"): |
| 64 | v_i = T.axis.spatial(64, i_0 * 32 + i_1) |
| 65 | T.reads(A[v_i], B[v_i]) |
| 66 | T.writes(C[v_i]) |
| 67 | C[v_i] = T.Cast( |
| 68 | dtype, T.Cast("float16", A[v_i]) + T.Cast("float16", B[v_i]) |
| 69 | ) |
| 70 | |
| 71 | return Module |
| 72 | |
| 73 | mod = _create_mod(dtype) |
| 74 | target = "cuda" |
| 75 | fadd = tvm.tirx.build(mod, target=target) |
| 76 | |
| 77 | cuda_src = fadd.imports[0].inspect_source() |
| 78 | assert nv_dtype in cuda_src, f"{nv_dtype} datatype not found in generated CUDA" |
| 79 | |
| 80 | dev = tvm.device(target, 0) |
| 81 | |
| 82 | a = tvm.runtime.tensor(np.random.uniform(low=0, high=5, size=64).astype(dtype), dev) |
| 83 | b = tvm.runtime.tensor(np.random.uniform(low=0, high=5, size=64).astype(dtype), dev) |
| 84 | c = tvm.runtime.tensor(np.zeros(64, dtype=dtype), dev) |
| 85 | fadd(a, b, c) |
| 86 | |
| 87 | tvm.testing.assert_allclose( |
| 88 | c.numpy().astype("float16"), (a.numpy() + b.numpy()).astype("float16") |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected
searching dependent graphs…