diff --git a/src/target/spirv/intrin_rule_spirv.cc b/src/target/spirv/intrin_rule_spirv.cc index e5f869de1718..ccb8d131c9d1 100644 --- a/src/target/spirv/intrin_rule_spirv.cc +++ b/src/target/spirv/intrin_rule_spirv.cc @@ -91,6 +91,39 @@ TVM_REGISTER_OP("tir.sin").set_attr("vulkan.FLowerIntrinsic", TVM_REGISTER_OP("tir.cos").set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); +TVM_REGISTER_OP("tir.tan").set_attr("vulkan.FLowerIntrinsic", + DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.asin") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.acos") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.atan") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.sinh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.cosh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.tanh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.asinh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.acosh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.atanh") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + +TVM_REGISTER_OP("tir.atan2") + .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); + TVM_REGISTER_OP("tir.log").set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); @@ -103,9 +136,6 @@ TVM_REGISTER_OP("tir.sqrt") TVM_REGISTER_OP("tir.pow").set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); -TVM_REGISTER_OP("tir.tanh") - .set_attr("vulkan.FLowerIntrinsic", DispatchGLSLPureIntrin); - TVM_REGISTER_OP("tir.erf").set_attr("vulkan.FLowerIntrinsic", codegen::intrin ::DispatchFastErf); } // namespace intrin diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py index b661ce486981..89acf598d6e3 100644 --- a/tests/python/codegen/test_target_codegen_vulkan.py +++ b/tests/python/codegen/test_target_codegen_vulkan.py @@ -568,5 +568,60 @@ def kernel(): vulkan_codegen(mod, target) +@tvm.testing.requires_gpu +@tvm.testing.requires_vulkan +def test_unary(): + test_funcs = [ + (tvm.tir.sin, lambda x: np.sin(x)), + (tvm.tir.cos, lambda x: np.cos(x)), + (tvm.tir.tan, lambda x: np.tan(x)), + (tvm.tir.sinh, lambda x: np.sinh(x)), + (tvm.tir.cosh, lambda x: np.cosh(x)), + (tvm.tir.tanh, lambda x: np.tanh(x)), + (tvm.tir.asin, lambda x: np.arcsin(x)), + (tvm.tir.acos, lambda x: np.arccos(x)), + (tvm.tir.atan, lambda x: np.arctan(x)), + (tvm.tir.asinh, lambda x: np.arcsinh(x)), + (tvm.tir.acosh, lambda x: np.arccosh(x)), + (tvm.tir.atanh, lambda x: np.arctanh(x)), + ] + + def run_test(tvm_intrin, np_func): + m = te.var("m") + A = te.placeholder((m,), name="A", dtype="float32") + B = te.compute((m,), lambda *i: tvm_intrin(A(*i)), name="B") + + mod = te.create_prim_func([A, B]) + sch = tir.Schedule(mod) + + block = sch.get_block("B") + loop = sch.get_loops(block)[0] + bx, tx = sch.split(loop, factors=[None, 64]) + sch.bind(bx, "blockIdx.x") + sch.bind(tx, "threadIdx.x") + + target = tvm.target.Target("vulkan") + dev = tvm.device(target.kind.name, 0) + func = tvm.compile(sch.mod, target=target) + + n = 16 + if tvm_intrin in [tvm.tir.asin, tvm.tir.acos]: + data = np.random.uniform(-1.0, 1.0, size=n) + elif tvm_intrin == tvm.tir.atanh: + data = np.random.uniform(-0.999, 0.999, size=n) + elif tvm_intrin == tvm.tir.acosh: + data = np.random.uniform(1.0, 5.0, size=n) + else: + data = np.random.uniform(0.1, 0.9, size=n) + + a = tvm.nd.array(data.astype(A.dtype), dev) + b = tvm.nd.array(np.zeros(n, dtype=A.dtype), 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) + + if __name__ == "__main__": tvm.testing.main()