| 556 | @pytest.mark.gpu |
| 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: |
| 596 | return True |
| 597 | # CUDA does support cast between {u}int8 and fp16. |
| 598 | skip_set = {"float16", "uint8", "int8"} |
| 599 | if t0 in skip_set and t1 in skip_set: |
| 600 | return True |
| 601 | return False |
| 602 | |
| 603 | types_4 = [ |
| 604 | "float16", |
| 605 | "float32", |
| 606 | "int8", |
| 607 | "uint8", |
| 608 | "int16", |
| 609 | "uint16", |
| 610 | "int32", |
| 611 | "uint32", |
| 612 | "float64", |
| 613 | "int64", |
| 614 | "uint64", |
| 615 | ] |