| 707 | @pytest.mark.gpu |
| 708 | @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") |
| 709 | def test_vectorized_intrin2(dtype="float32"): |
| 710 | c2 = tvm.tirx.const(2, dtype=dtype) |
| 711 | test_funcs = [ |
| 712 | (tvm.tirx.power, lambda x: np.power(x, 2.0)), |
| 713 | (tvm.tirx.fmod, lambda x: np.fmod(x, 2.0)), |
| 714 | ] |
| 715 | |
| 716 | def run_test(tvm_intrin, np_func): |
| 717 | n = 128 |
| 718 | f = sched(lambda x: tvm_intrin(x, c2), dtype, n) |
| 719 | dev = tvm.cuda(0) |
| 720 | a = tvm.runtime.tensor(np.random.uniform(0, 1, size=n).astype(dtype), dev) |
| 721 | b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev) |
| 722 | f(a, b) |
| 723 | tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3) |
| 724 | |
| 725 | for func in test_funcs: |
| 726 | run_test(*func) |
| 727 | |
| 728 | |
| 729 | @pytest.mark.gpu |