| 651 | @pytest.mark.gpu |
| 652 | @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") |
| 653 | def test_vectorized_intrin1(): |
| 654 | test_funcs = [ |
| 655 | (tvm.tirx.floor, lambda x: np.floor(x)), |
| 656 | (tvm.tirx.ceil, lambda x: np.ceil(x)), |
| 657 | (tvm.tirx.trunc, lambda x: np.trunc(x)), |
| 658 | (tvm.tirx.abs, lambda x: np.fabs(x)), |
| 659 | (tvm.tirx.round, lambda x: np.round(x)), |
| 660 | (tvm.tirx.exp, lambda x: np.exp(x)), |
| 661 | (tvm.tirx.exp2, lambda x: np.exp2(x)), |
| 662 | (tvm.tirx.exp10, lambda x: np.power(10, x)), |
| 663 | (tvm.tirx.log, lambda x: np.log(x)), |
| 664 | (tvm.tirx.log2, lambda x: np.log2(x)), |
| 665 | (tvm.tirx.log10, lambda x: np.log10(x)), |
| 666 | (tvm.tirx.tan, lambda x: np.tan(x)), |
| 667 | (tvm.tirx.cos, lambda x: np.cos(x)), |
| 668 | (tvm.tirx.cosh, lambda x: np.cosh(x)), |
| 669 | (tvm.tirx.sin, lambda x: np.sin(x)), |
| 670 | (tvm.tirx.sinh, lambda x: np.sinh(x)), |
| 671 | (tvm.tirx.atan, lambda x: np.arctan(x)), |
| 672 | (tvm.tirx.tanh, lambda x: np.tanh(x)), |
| 673 | (tvm.tirx.sqrt, lambda x: np.sqrt(x)), |
| 674 | ] |
| 675 | |
| 676 | def run_test(tvm_intrin, np_func, dtype): |
| 677 | if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version): |
| 678 | print("Skip because gpu does not have fp16 support") |
| 679 | return |
| 680 | # set of intrinsics does not support fp16 yet. |
| 681 | skip_set = { |
| 682 | tvm.tirx.abs, |
| 683 | tvm.tirx.round, |
| 684 | tvm.tirx.tan, |
| 685 | tvm.tirx.atan, |
| 686 | tvm.tirx.tanh, |
| 687 | tvm.tirx.cosh, |
| 688 | tvm.tirx.sinh, |
| 689 | } |
| 690 | if dtype == "float16" and tvm_intrin in skip_set: |
| 691 | print(f"Skip because '{tvm_intrin.__name__}' does not support fp16 yet") |
| 692 | return |
| 693 | |
| 694 | n = 128 |
| 695 | f = sched(tvm_intrin, dtype, n) |
| 696 | dev = tvm.cuda(0) |
| 697 | a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev) |
| 698 | b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev) |
| 699 | f(a, b) |
| 700 | tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3) |
| 701 | |
| 702 | for func in test_funcs: |
| 703 | run_test(*func, "float32") |
| 704 | run_test(*func, "float16") |
| 705 | |
| 706 | |
| 707 | @pytest.mark.gpu |