(t0, t1, factor)
| 557 | @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") |
| 558 | def test_vectorized_casts(): |
| 559 | def check(t0, t1, factor): |
| 560 | if (t0 == "float16" or t1 == "float16") and not have_fp16(tvm.cuda(0).compute_version): |
| 561 | print("Skip because gpu does not have fp16 support") |
| 562 | return |
| 563 | |
| 564 | n = 128 |
| 565 | num_thread = n // factor |
| 566 | |
| 567 | @I.ir_module(s_tir=True) |
| 568 | class Module: |
| 569 | @T.prim_func(s_tir=True) |
| 570 | def main(A: T.Buffer((n,), t0), B: T.Buffer((n,), t1), C: T.Buffer((n,), t0)): |
| 571 | T.func_attr({"tirx.noalias": True}) |
| 572 | for i_0 in T.thread_binding(num_thread, thread="threadIdx.x"): |
| 573 | for i_1 in T.vectorized(factor): |
| 574 | with T.sblock("C"): |
| 575 | v_i = T.axis.spatial(n, i_0 * factor + i_1) |
| 576 | T.reads(A[v_i], B[v_i]) |
| 577 | T.writes(C[v_i]) |
| 578 | C[v_i] = A[v_i] + T.Cast(t0, B[v_i]) |
| 579 | |
| 580 | func = tvm.compile(Module, target="cuda") |
| 581 | |
| 582 | # correctness |
| 583 | dev = tvm.cuda(0) |
| 584 | low, high = (0, 20) if t0.startswith("u") or t1.startswith("u") else (-10, 10) |
| 585 | a_np = np.random.randint(low, high, size=n).astype(t0) |
| 586 | b_np = np.random.randint(low, high, size=n).astype(t1) |
| 587 | c_np = (a_np + b_np).astype(t0) |
| 588 | a_nd = tvm.runtime.tensor(a_np, dev) |
| 589 | b_nd = tvm.runtime.tensor(b_np, dev) |
| 590 | c_nd = tvm.runtime.tensor(np.zeros(c_np.shape, dtype=c_np.dtype), dev) |
| 591 | func(a_nd, b_nd, c_nd) |
| 592 | tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-3) |
| 593 | |
| 594 | def skip(t0, t1): |
| 595 | if t0 == t1: |
no test coverage detected
searching dependent graphs…