Function test_unary

tests/python/codegen/test_target_codegen_vulkan.py:511–565

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_vulkan(), reason="need vulkan")
def test_unary():
    test_funcs = [
        (tvm.tirx.sin, lambda x: np.sin(x)),
        (tvm.tirx.cos, lambda x: np.cos(x)),
        (tvm.tirx.tan, lambda x: np.tan(x)),
        (tvm.tirx.sinh, lambda x: np.sinh(x)),
        (tvm.tirx.cosh, lambda x: np.cosh(x)),
        (tvm.tirx.tanh, lambda x: np.tanh(x)),
        (tvm.tirx.asin, lambda x: np.arcsin(x)),
        (tvm.tirx.acos, lambda x: np.arccos(x)),
        (tvm.tirx.atan, lambda x: np.arctan(x)),
        (tvm.tirx.asinh, lambda x: np.arcsinh(x)),
        (tvm.tirx.acosh, lambda x: np.arccosh(x)),
        (tvm.tirx.atanh, lambda x: np.arctanh(x)),
    ]

    def run_test(tvm_intrin, np_func):
        n = 16

        @I.ir_module(s_tir=True)
        class Module:
            @T.prim_func(s_tir=True)
            def main(var_A: T.handle, var_B: T.handle):
                m = T.int32(is_size_var=True)
                A = T.match_buffer(var_A, (m,), "float32")
                B = T.match_buffer(var_B, (m,), "float32")
                for i_0 in T.thread_binding((m + 63) // 64, thread="blockIdx.x"):
                    for i_1 in T.thread_binding(64, thread="threadIdx.x"):
                        with T.sblock("B"):
                            v_i = T.axis.spatial(m, i_0 * 64 + i_1)
                            T.where(i_0 * 64 + i_1 < m)
                            T.reads(A[v_i])
                            T.writes(B[v_i])
                            B[v_i] = tvm_intrin(A[v_i])

        target = tvm.target.Target("vulkan")
        dev = tvm.device(target.kind.name, 0)
        func = tvm.compile(Module, target=target)

        if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
            data = np.random.uniform(-1.0, 1.0, size=n)
        elif tvm_intrin == tvm.tirx.atanh:
            data = np.random.uniform(-0.999, 0.999, size=n)
        elif tvm_intrin == tvm.tirx.acosh:
            data = np.random.uniform(1.0, 5.0, size=n)
        else:
            data = np.random.uniform(0.1, 0.9, size=n)

        a = tvm.runtime.tensor(data.astype("float32"), dev)
        b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
        func(a, b)
        tvm.testing.assert_allclose(b.numpy(), np_func(a.numpy()), atol=1e-3, rtol=1e-3)

    for func in test_funcs:
        run_test(*func)


@pytest.mark.gpu