diff --git a/docs/api/python/te.rst b/docs/api/python/te.rst index dc3d3dacd2cae..7a0f4b7918a25 100644 --- a/docs/api/python/te.rst +++ b/docs/api/python/te.rst @@ -23,7 +23,7 @@ tvm.te :members: :imported-members: :exclude-members: - exp, erf, tanh, sigmoid, log, cos, sin, atan, sqrt, rsqrt, floor, ceil, + exp, erf, tanh, sigmoid, log, tan, cos, sin, atan, sqrt, rsqrt, floor, ceil, trunc, abs, round, nearbyint, isnan, power, popcount, fmod, if_then_else, div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, comm_reducer, min, max, sum diff --git a/docs/frontend/tensorflow.rst b/docs/frontend/tensorflow.rst index 87341ab6b7c60..8a54033a2b09c 100644 --- a/docs/frontend/tensorflow.rst +++ b/docs/frontend/tensorflow.rst @@ -135,6 +135,7 @@ Supported Ops - ConcatV2 - Conv2D - Cos +- Tan - CropAndResize - DecodeJpeg - DepthwiseConv2dNative diff --git a/include/tvm/tir/op.h b/include/tvm/tir/op.h index 5172b1496ad21..0a714d8fc8c85 100644 --- a/include/tvm/tir/op.h +++ b/include/tvm/tir/op.h @@ -515,6 +515,7 @@ TVM_DECLARE_INTRIN_UNARY(sqrt); TVM_DECLARE_INTRIN_UNARY(rsqrt); TVM_DECLARE_INTRIN_UNARY(log); TVM_DECLARE_INTRIN_UNARY(popcount); +TVM_DECLARE_INTRIN_UNARY(tan); TVM_DECLARE_INTRIN_UNARY(cos); TVM_DECLARE_INTRIN_UNARY(sin); TVM_DECLARE_INTRIN_UNARY(atan); diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index d74277bbe402f..e34bebe0ce5a3 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -1691,6 +1691,7 @@ def _get_bias_requantize_scale(_inputs, _data_scale, _kernel_scale): "ones_like", "where", "gather_nd", + "tan", "cos", "sin" ] diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 5532e3a5c1a47..837f430fec6c2 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1561,6 +1561,7 @@ def _impl(inputs, attr, params): 'LessEqual' : _broadcast('less_equal'), 'Log' : AttrCvt('log'), 'Log1p' : _log1p(), + 'Tan' : AttrCvt('tan'), 'Cos' : AttrCvt('cos'), 'Sin' : AttrCvt('sin'), 'LogicalAnd' : _logical('logical_and'), diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index 352bc6302ee04..501513ec41f1f 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -69,6 +69,7 @@ def __init__(self, model, subgraph, exp_tab): 'LOG': self.convert_log, 'SIN': self.convert_sin, 'COS': self.convert_cos, + 'TAN': self.convert_tan, 'SQRT': self.convert_sqrt, 'RSQRT': self.convert_rsqrt, 'NEG': self.convert_neg, @@ -608,6 +609,13 @@ def convert_sin(self, op): 'TFlite quantized SIN operator is not supported yet.') return self._convert_unary_elemwise(_op.sin, op) + def convert_tan(self, op): + """Convert TFLite TAN""" + if self.is_quantized(op): + raise tvm.error.OpNotImplemented( + 'TFlite quantized TAN operator is not supported yet.') + return self._convert_unary_elemwise(_op.tan, op) + def convert_cos(self, op): """Convert TFLite COS""" if self.is_quantized(op): diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index 7c8ccb7dd8277..190ed437636dc 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -27,6 +27,7 @@ register_broadcast_schedule("log") +register_broadcast_schedule("tan") register_broadcast_schedule("cos") register_broadcast_schedule("sin") register_broadcast_schedule("atan") diff --git a/python/tvm/relay/op/_tensor_grad.py b/python/tvm/relay/op/_tensor_grad.py index 944e51e636f5e..33a1937992888 100644 --- a/python/tvm/relay/op/_tensor_grad.py +++ b/python/tvm/relay/op/_tensor_grad.py @@ -61,6 +61,13 @@ def log_grad(orig, grad): return [grad * ones_like(x) / x] +@register_gradient("tan") +def tan_grad(orig, grad): + """Returns [grad / (cos^2(x))]""" + x = orig.args[0] + return [grad / (cos(x) * cos(x))] + + @register_gradient("cos") def cos_grad(orig, grad): """Returns [grad * (-sin(x))]""" diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index 898038dea0043..3f5276c748c0f 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -47,6 +47,21 @@ def log(data): """ return _make.log(data) +def tan(data): + """Compute elementwise tan of data. + + Parameters + ---------- + data : relay.Expr + The input data + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.tan(data) + def cos(data): """Compute elementwise cos of data. diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py index 5970315e854b0..34f5d24042bf3 100644 --- a/python/tvm/te/__init__.py +++ b/python/tvm/te/__init__.py @@ -18,7 +18,7 @@ """Namespace for Tensor Expression Language """ # expose all operators in tvm tir.op -from tvm.tir import exp, erf, tanh, sigmoid, log, cos, sin, atan, sqrt, rsqrt, floor, ceil +from tvm.tir import exp, erf, tanh, sigmoid, log, tan, cos, sin, atan, sqrt, rsqrt, floor, ceil from tvm.tir import trunc, abs, round, nearbyint, isnan, power, popcount, fmod, if_then_else from tvm.tir import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod from tvm.tir import comm_reducer, min, max, sum diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py index ab78ca6d6d63d..cf0ac8e6744c0 100644 --- a/python/tvm/tir/__init__.py +++ b/python/tvm/tir/__init__.py @@ -31,7 +31,7 @@ from .op import call_packed, call_pure_intrin, call_intrin, call_pure_extern, call_extern from .op import call_llvm_intrin, all, any, min_value, max_value -from .op import exp, erf, tanh, sigmoid, log, cos, sin, atan, sqrt, rsqrt, floor, ceil +from .op import exp, erf, tanh, sigmoid, log, tan, cos, sin, atan, sqrt, rsqrt, floor, ceil from .op import trunc, abs, round, nearbyint, isnan, power, popcount, fmod, if_then_else from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod from .op import comm_reducer, min, max, sum diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py index 66e70c508438d..716c5c1579670 100644 --- a/python/tvm/tir/op.py +++ b/python/tvm/tir/op.py @@ -393,6 +393,22 @@ def log(x): """ return call_pure_intrin(x.dtype, "log", x) +def tan(x): + """Take tan of input x. + + Parameters + ---------- + x : PrimExpr + Input argument. + + Returns + ------- + y : PrimExpr + The result. + """ + return call_pure_intrin(x.dtype, "tan", x) + + def cos(x): """Take cos of input x. diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc index caa6451542c98..9425a13fe03f5 100644 --- a/src/relay/op/tensor/unary.cc +++ b/src/relay/op/tensor/unary.cc @@ -51,6 +51,17 @@ RELAY_REGISTER_UNARY_OP("log") .set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::log)); +RELAY_REGISTER_UNARY_OP("tan") +.describe(R"code(Returns the tan of input array, computed element-wise. + +.. math:: + Y = tan(X) + +)code" TVM_ADD_FILELINE) +.set_support_level(1) +.set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::tan)); + + RELAY_REGISTER_UNARY_OP("cos") .describe(R"code(Returns the cos of input array, computed element-wise. diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc index 7e9ac71bb753d..ade3bbdf93042 100644 --- a/src/target/intrin_rule.cc +++ b/src/target/intrin_rule.cc @@ -40,6 +40,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.log") TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tanh") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tan") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.cos") .set_body(DispatchExtern); diff --git a/src/target/llvm/intrin_rule_llvm.cc b/src/target/llvm/intrin_rule_llvm.cc index 758b0afc22ff4..6c5a9cd079afd 100644 --- a/src/target/llvm/intrin_rule_llvm.cc +++ b/src/target/llvm/intrin_rule_llvm.cc @@ -91,6 +91,20 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow") TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount") .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tan") +.set_body([](const TVMArgs& targs, TVMRetValue* rv) { + PrimExpr e = targs[0]; + const tir::CallNode* call = e.as(); + CHECK(call != nullptr); + const PrimExpr& x = call->args[0]; + PrimExpr sin_x = tir::CallNode::make( + x.dtype(), "sin", {x}, tir::CallNode::PureIntrinsic); + PrimExpr cos_x = tir::CallNode::make( + x.dtype(), "cos", {x}, tir::CallNode::PureIntrinsic); + PrimExpr tan_x = sin_x / cos_x; + *rv = tan_x; +}); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cos") .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>); diff --git a/src/target/llvm/intrin_rule_nvptx.cc b/src/target/llvm/intrin_rule_nvptx.cc index 6f7a89cebd971..6d41d132724b2 100644 --- a/src/target/llvm/intrin_rule_nvptx.cc +++ b/src/target/llvm/intrin_rule_nvptx.cc @@ -81,6 +81,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.pow") TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.tanh") .set_body(DispatchExternLibDevice); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.tan") +.set_body(DispatchExternLibDevice); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.nvptx.cos") .set_body(DispatchExternLibDevice); diff --git a/src/target/llvm/intrin_rule_rocm.cc b/src/target/llvm/intrin_rule_rocm.cc index 31b7bf19419b5..4e6a661f298d0 100644 --- a/src/target/llvm/intrin_rule_rocm.cc +++ b/src/target/llvm/intrin_rule_rocm.cc @@ -80,6 +80,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.rocm.pow") TVM_REGISTER_GLOBAL("tvm.intrin.rule.rocm.tanh") .set_body(DispatchExternOCML); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.rocm.tan") +.set_body(DispatchExternOCML); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.rocm.cos") .set_body(DispatchExternOCML); diff --git a/src/target/source/intrin_rule_cuda.cc b/src/target/source/intrin_rule_cuda.cc index aed6c86e965ba..fc62dae7c9f9e 100644 --- a/src/target/source/intrin_rule_cuda.cc +++ b/src/target/source/intrin_rule_cuda.cc @@ -97,6 +97,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.erf") TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.log") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.tan") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.cos") .set_body(DispatchExtern); diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc index 07cae5e2c746f..7572f8d6bbb04 100644 --- a/src/tir/ir/expr.cc +++ b/src/tir/ir/expr.cc @@ -229,7 +229,7 @@ PrimExpr LetNode::make(Var var, PrimExpr value, PrimExpr body) { const char* CallNode::vectorizable_intrinsics[] = { "floor", "ceil", "sign", "trunc", "fabs", "round", "exp", "tanh", "sqrt", - "log", "sin", "cos", "pow", tir::CallNode::shift_left, tir::CallNode::shift_right, + "log", "sin", "cos", "pow", "tan", tir::CallNode::shift_left, tir::CallNode::shift_right, tir::CallNode::likely, tir::CallNode::popcount }; diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 2340bd4e6318e..63f96a98472cc 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -2570,6 +2570,15 @@ def test_forward_cos(): compare_tf_with_tvm([np_data], ['in_data:0'], 'cos:0') +def test_forward_tan(): + """test operator tan """ + np_data = np.random.uniform(1, 100, size=(2, 3, 5)).astype(np.float32) + tf.reset_default_graph() + in_data = tf.placeholder(tf.float32, (2, 3, 5), name="in_data") + tf.tan(in_data, name="tan") + compare_tf_with_tvm([np_data], ['in_data:0'], 'tan:0') + + def test_forward_sin(): """test operator sin """ np_data = np.random.uniform(1, 100, size=(2, 3, 5)).astype(np.float32) diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py index 427d4bfe28106..9302c20878b19 100644 --- a/tests/python/frontend/tflite/test_forward.py +++ b/tests/python/frontend/tflite/test_forward.py @@ -721,6 +721,13 @@ def _test_cos(data): """ One iteration of cos """ return _test_unary_elemwise(math_ops.cos, data) ####################################################################### +# Tan +# --- + +def _test_tan(data): + """ One iteration of tan """ + return _test_unary_elemwise(math_ops.tan, data) +####################################################################### # Sqrt # ---- @@ -763,6 +770,7 @@ def test_all_unary_elemwise(): if package_version.parse(tf.VERSION) >= package_version.parse('1.14.0'): _test_forward_unary_elemwise(_test_ceil) _test_forward_unary_elemwise(_test_cos) + _test_forward_unary_elemwise(_test_tan) ####################################################################### # Element-wise diff --git a/tests/python/relay/test_op_grad_level1.py b/tests/python/relay/test_op_grad_level1.py index 3be62a3170fb5..5a69d29d4df2e 100644 --- a/tests/python/relay/test_op_grad_level1.py +++ b/tests/python/relay/test_op_grad_level1.py @@ -63,6 +63,7 @@ def check_single_op(opfunc, ref): (relay.nn.relu, lambda x: np.where(x < 0, np.zeros_like(x), np.ones_like(x))), (tvm.relay.cos, lambda x: -1.0 * np.sin(x)), (tvm.relay.sin, lambda x: np.cos(x)), + (tvm.relay.tan, lambda x: 1.0 / (np.cos(x) ** 2)), (tvm.relay.atan, lambda x: 1 / (1 + np.power(x, 2.0)))]: check_single_op(opfunc, ref) diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index 194b095642882..dfa42ab9e8e54 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -75,6 +75,7 @@ def check_single_op(opfunc, ref, dtype): (relay.nn.relu, relu), (tvm.relay.cos, np.cos), (tvm.relay.sin, np.sin), + (tvm.relay.tan, np.tan), (tvm.relay.atan, np.arctan)]: for dtype in ['float16', 'float32']: check_single_op(opfunc, ref, dtype) diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index e35e3e424d6e9..e37683256e4f5 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -55,6 +55,7 @@ TOPI_DECLARE_UNARY_OP(round); TOPI_DECLARE_UNARY_OP(trunc); TOPI_DECLARE_UNARY_OP(abs); TOPI_DECLARE_UNARY_OP(cos); +TOPI_DECLARE_UNARY_OP(tan); TOPI_DECLARE_UNARY_OP(sin); TOPI_DECLARE_UNARY_OP(atan); TOPI_DECLARE_UNARY_OP(isnan); diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index 148d53a54cfe0..ed6bce0d11416 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -109,6 +109,23 @@ def tanh(x): return tvm.compute(x.shape, lambda *i: tvm.tanh(x(*i))) +@tvm.tag_scope(tag=tag.ELEMWISE) +def tan(x): + """Take tan of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + return tvm.compute(x.shape, lambda *i: tvm.tan(x(*i))) + + @tvm.tag_scope(tag=tag.ELEMWISE) def cos(x): """Take cos of input x. diff --git a/topi/src/topi.cc b/topi/src/topi.cc index 79e223c30975a..2fba6828eda1c 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -175,6 +175,11 @@ TVM_REGISTER_GLOBAL("topi.erf") *rv = erf(args[0]); }); +TVM_REGISTER_GLOBAL("topi.tan") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = tan(args[0]); + }); + TVM_REGISTER_GLOBAL("topi.cos") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = cos(args[0]); diff --git a/topi/tests/python/test_topi_basic.py b/topi/tests/python/test_topi_basic.py index 53b29df4f36d4..63dbdbe2f599f 100644 --- a/topi/tests/python/test_topi_basic.py +++ b/topi/tests/python/test_topi_basic.py @@ -44,6 +44,7 @@ def test_apply(func, name): test_apply(topi.rsqrt, "rsqrt") test_apply(topi.sin, "sin") test_apply(topi.cos, "cos") + test_apply(topi.tan, "tan") test_apply(topi.atan, "atan") diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py index debc3efe0d275..c9a0453f74d38 100644 --- a/topi/tests/python/test_topi_math.py +++ b/topi/tests/python/test_topi_math.py @@ -126,6 +126,7 @@ def check_device(device): test_apply(topi.sqrt, "sqrt", np.sqrt, 0, 100) test_apply(topi.rsqrt, "rsqrt", lambda x: np.ones_like(x) / np.sqrt(x), 0, 100, skip_name_check=True) test_apply(topi.cos, "cos", np.cos, -2.0*np.pi, 2.0*np.pi) + test_apply(topi.tan, "tan", np.tan, -2.0*np.pi, 2.0*np.pi) test_apply(topi.sin, "sin", np.sin, -2.0*np.pi, 2.0*np.pi) test_apply(topi.erf, "erf", scipy.special.erf, -.1, .1, dtype="float32") test_isnan(-100, 100)