From 8233b0153c726df0947648544df8dd9333807d6b Mon Sep 17 00:00:00 2001 From: Alex Gladkov Date: Wed, 24 Jul 2019 19:25:04 +0000 Subject: [PATCH] Add support for Tensorflow operators log1p and cos The patch adds support for Tensorflow operators log1p and cos Tensorflow log1p is described at https://www.tensorflow.org/api_docs/python/tf/math/log1p Tensorflow cos is described at https://www.tensorflow.org/api_docs/python/tf/math/cos --- include/tvm/expr_operator.h | 2 +- python/tvm/intrin.py | 29 ++++++++++++++++++ python/tvm/relay/frontend/tensorflow.py | 2 ++ python/tvm/relay/op/_tensor.py | 2 ++ python/tvm/relay/op/_tensor_grad.py | 12 ++++++++ python/tvm/relay/op/tensor.py | 29 ++++++++++++++++++ src/codegen/intrin_rule.cc | 14 +++++++++ src/codegen/intrin_rule_cuda.cc | 3 ++ src/relay/op/tensor/unary.cc | 22 ++++++++++++++ .../frontend/tensorflow/test_forward.py | 19 ++++++++++++ topi/include/topi/elemwise.h | 10 +++++++ topi/python/topi/math.py | 30 +++++++++++++++++++ topi/src/topi.cc | 10 +++++++ 13 files changed, 183 insertions(+), 1 deletion(-) diff --git a/include/tvm/expr_operator.h b/include/tvm/expr_operator.h index d90ccf0871b13..2ad91296beecd 100644 --- a/include/tvm/expr_operator.h +++ b/include/tvm/expr_operator.h @@ -517,7 +517,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(cos); // Implementation details after this inline bool is_const(const Expr& x) { diff --git a/python/tvm/intrin.py b/python/tvm/intrin.py index df854e270e9de..49aad454e8286 100644 --- a/python/tvm/intrin.py +++ b/python/tvm/intrin.py @@ -258,6 +258,35 @@ def log(x): """ return call_pure_intrin(x.dtype, "log", x) +def log1p(x): + """Take log of input x. + + Parameters + ---------- + x : Expr + Input argument. + + Returns + ------- + y : Expr + The result. + """ + return call_pure_intrin(x.dtype, "log1p", x) + +def cos(x): + """Take cos of input x. + + Parameters + ---------- + x : Expr + Input argument. + + Returns + ------- + y : Expr + The result. + """ + return call_pure_intrin(x.dtype, "cos", x) def sqrt(x): """Take square root of input x. diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 8605edf1d4a62..b609b26d720cf 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1344,6 +1344,8 @@ def _impl(inputs, attr, params): 'Less' : _broadcast('less'), 'LessEqual' : _broadcast('less_equal'), 'Log' : AttrCvt('log'), + 'Log1p' : AttrCvt('log1p'), + 'Cos' : AttrCvt('cos'), 'LogicalAnd' : _logical('logical_and'), 'LogicalOr' : _logical('logical_or'), 'LogicalNot' : _logical('logical_not'), diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index 8b3dd72db043c..baecc70c45011 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -25,6 +25,8 @@ schedule_elemwise = schedule_injective register_schedule("log", schedule_broadcast) +register_schedule("log1p", schedule_broadcast) +register_schedule("cos", schedule_broadcast) register_schedule("exp", schedule_broadcast) register_schedule("sqrt", schedule_broadcast) register_schedule("rsqrt", schedule_broadcast) diff --git a/python/tvm/relay/op/_tensor_grad.py b/python/tvm/relay/op/_tensor_grad.py index 24da24cacfdae..f8aa51708c363 100644 --- a/python/tvm/relay/op/_tensor_grad.py +++ b/python/tvm/relay/op/_tensor_grad.py @@ -31,6 +31,18 @@ def log_grad(orig, grad): x = orig.args[0] return [grad * ones_like(x) / x] +@register_gradient("log1p") +def log1p_grad(orig, grad): + """Returns [grad * (1 / (x+1))]""" + x = orig.args[0] + return [grad * ones_like(x) / (x + ones_like(x))] + +@register_gradient("cos") +def cos_grad(orig, grad): + """Returns [grad * (-sin(x))]""" + x = orig.args[0] + ones = ones_like(x) + return [grad * (-ones * sin(x))] @register_gradient("exp") def exp_grad(orig, grad): diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index 3795e6598cabe..a9c821b4ea89e 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -46,6 +46,35 @@ def log(data): """ return _make.log(data) +def log1p(data): + """Compute elementwise log of data. + + Parameters + ---------- + data : relay.Expr + The input data + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.log1p(data) + +def cos(data): + """Compute elementwise cos of data. + + Parameters + ---------- + data : relay.Expr + The input data + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.cos(data) def exp(data): """Compute elementwise exp of data. diff --git a/src/codegen/intrin_rule.cc b/src/codegen/intrin_rule.cc index 230a3bccab3c1..90e5605b8f4ef 100644 --- a/src/codegen/intrin_rule.cc +++ b/src/codegen/intrin_rule.cc @@ -34,9 +34,23 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.exp") TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.log") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.log1p") +.set_body([](const TVMArgs& args, TVMRetValue* rv){ + Expr e = args[0]; + const Call* call = e.as(); + CHECK(call != nullptr); + + auto one = make_const(call->args[0].type(), 1); + *rv = log(call->args[0] + one); + }); + + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tanh") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.cos") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sqrt") .set_body(DispatchExtern); diff --git a/src/codegen/intrin_rule_cuda.cc b/src/codegen/intrin_rule_cuda.cc index 22648386f1509..89710276055d7 100644 --- a/src/codegen/intrin_rule_cuda.cc +++ b/src/codegen/intrin_rule_cuda.cc @@ -95,6 +95,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.exp") TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.log") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.cos") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.cuda.tanh") .set_body(DispatchExtern); diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc index 60e53784649bb..0813bc8543cf9 100644 --- a/src/relay/op/tensor/unary.cc +++ b/src/relay/op/tensor/unary.cc @@ -53,6 +53,28 @@ RELAY_REGISTER_UNARY_OP("log") .set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::log)); +RELAY_REGISTER_UNARY_OP("log1p") +.describe(R"code(Returns the log1p input array, computed element-wise. + +.. math:: + log(x+1) + +)code" TVM_ADD_FILELINE) +.set_support_level(1) +.set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::log1p)); + + +RELAY_REGISTER_UNARY_OP("cos") +.describe(R"code(Returns the cos of input array, computed element-wise. + +.. math:: + Y = cos(X) + +)code" TVM_ADD_FILELINE) +.set_support_level(1) +.set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::cos)); + + RELAY_REGISTER_UNARY_OP("exp") .describe(R"code(Returns the exp input array, computed element-wise. diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 6c9824e4ed13b..4c2fe4efb479c 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1850,6 +1850,22 @@ def test_forward_log(): tf.log(in_data, name="log") compare_tf_with_tvm([np_data], ['in_data:0'], 'log:0') +def test_forward_log1p(): + """test operator Log1p """ + 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.log1p(in_data, name="log1p") + compare_tf_with_tvm([np_data], ['in_data:0'], 'log1p:0') + +def test_forward_cos(): + """test operator cos """ + 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.cos(in_data, name="cos") + compare_tf_with_tvm([np_data], ['in_data:0'], 'cos:0') + def test_forward_negative(): """test tf operator Neg """ np_data = np.random.uniform(-100, 255, size=(224, 224, 3)).astype(np.float32) @@ -2097,6 +2113,7 @@ def test_placeholder(): # ---- if __name__ == '__main__': + # Transforms test_forward_transpose() test_forward_reshape() @@ -2140,6 +2157,8 @@ def test_placeholder(): test_forward_pow_exp() test_forward_sign() test_forward_log() + test_forward_log1p() + test_forward_cos() test_forward_negative() test_forward_divide() test_forward_abs() diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index 000567eeae140..6fb20c2e325a6 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -54,6 +54,7 @@ TOPI_DECLARE_UNARY_OP(ceil); TOPI_DECLARE_UNARY_OP(round); TOPI_DECLARE_UNARY_OP(trunc); TOPI_DECLARE_UNARY_OP(abs); +TOPI_DECLARE_UNARY_OP(cos); /* * \brief Fast_tanh_float implementation from Eigen @@ -214,6 +215,15 @@ inline Tensor rsqrt(const Tensor& x, }, name, tag); } +inline Tensor log1p(const Tensor& x, + std::string name = "tensor", + std::string tag = kElementWise) { + return compute(x->shape, [&](const Array& i) { + Expr one = make_const(x->dtype, 1); + return tvm::log(x(i)+one); + }, name, tag); +} + /*! * \brief Creates an operation that clips each element of a tensor to * the interval [a_min, a_max] diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index 87ac06c76c751..4e0b862edb7ae 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -90,6 +90,21 @@ def tanh(x): """ return tvm.compute(x.shape, lambda *i: tvm.tanh(x(*i))) +@tvm.tag_scope(tag=tag.ELEMWISE) +def cos(x): + """Take cos of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + return tvm.compute(x.shape, lambda *i: tvm.cos(x(*i))) @tvm.tag_scope(tag=tag.ELEMWISE) def floor(x): @@ -206,6 +221,21 @@ def log(x): """ return tvm.compute(x.shape, lambda *i: tvm.log(x(*i))) +@tvm.tag_scope(tag=tag.ELEMWISE) +def log1p(x): + """Take logarithm of input (x+1). + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + return tvm.compute(x.shape, lambda *i: tvm.log1p(x(*i))) @tvm.tag_scope(tag=tag.ELEMWISE) def sqrt(x): diff --git a/topi/src/topi.cc b/topi/src/topi.cc index bd2910d39032e..5dd98f5c7d5c9 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -148,6 +148,11 @@ TVM_REGISTER_GLOBAL("topi.exp") *rv = exp(args[0]); }); +TVM_REGISTER_GLOBAL("topi.cos") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = cos(args[0]); + }); + TVM_REGISTER_GLOBAL("topi.tanh") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = tanh(args[0]); @@ -173,6 +178,11 @@ TVM_REGISTER_GLOBAL("topi.log") *rv = log(args[0]); }); +TVM_REGISTER_GLOBAL("topi.log1p") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = log1p(args[0]); + }); + TVM_REGISTER_GLOBAL("topi.identity") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = identity(args[0]);