From 6375234d275911f1d5388a0a769ac36e8e083701 Mon Sep 17 00:00:00 2001 From: makihiro Date: Fri, 29 Mar 2019 17:20:07 +0900 Subject: [PATCH 1/9] [Relay][TOPI] add rsqrt operator --- docs/langref/relay_op.rst | 2 ++ python/tvm/relay/op/_tensor.py | 1 + python/tvm/relay/op/tensor.py | 16 ++++++++++++++++ src/relay/op/tensor/unary.cc | 11 ++++++++++- tests/python/relay/test_ir_op.py | 2 +- tests/python/relay/test_op_level1.py | 5 +++++ topi/include/topi/elemwise.h | 1 + topi/python/topi/math.py | 17 +++++++++++++++++ topi/src/topi.cc | 5 +++++ topi/tests/python/test_topi_basic.py | 1 + topi/tests/python/test_topi_math.py | 1 + 11 files changed, 60 insertions(+), 2 deletions(-) diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index bbb27ec83b48..d1485b316236 100644 --- a/docs/langref/relay_op.rst +++ b/docs/langref/relay_op.rst @@ -24,6 +24,7 @@ This level enables fully connected multi-layer perceptron. tvm.relay.log tvm.relay.sqrt + tvm.relay.rsqrt tvm.relay.exp tvm.relay.sigmoid tvm.relay.add @@ -168,6 +169,7 @@ Level 1 Definitions ------------------- .. autofunction:: tvm.relay.log .. autofunction:: tvm.relay.sqrt +.. autofunction:: tvm.relay.rsqrt .. autofunction:: tvm.relay.exp .. autofunction:: tvm.relay.sigmoid .. autofunction:: tvm.relay.add diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index 05d9acb27330..a4f3dc6dd246 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -11,6 +11,7 @@ register_schedule("log", schedule_broadcast) register_schedule("exp", schedule_broadcast) register_schedule("sqrt", schedule_broadcast) +register_schedule("rsqrt", schedule_broadcast) register_schedule("sigmoid", schedule_broadcast) register_schedule("floor", schedule_broadcast) register_schedule("ceil", schedule_broadcast) diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index d51208d478aa..0b146463a662 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -63,6 +63,22 @@ def sqrt(data): return _make.sqrt(data) +def rsqrt(data): + """Compute elementwise rsqrt of data. + + Parameters + ---------- + data : relay.Expr + The input data + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.rsqrt(data) + + def sigmoid(data): """Compute elementwise sigmoid of data. diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc index 4befc4b664a8..30326d78fe8b 100644 --- a/src/relay/op/tensor/unary.cc +++ b/src/relay/op/tensor/unary.cc @@ -45,7 +45,7 @@ RELAY_REGISTER_UNARY_OP("exp") .set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::exp)); RELAY_REGISTER_UNARY_OP("sqrt") -.describe(R"code(Returns the rsqrt input array, computed element-wise. +.describe(R"code(Returns the sqrt input array, computed element-wise. .. math:: sqrt(x) @@ -54,6 +54,15 @@ RELAY_REGISTER_UNARY_OP("sqrt") .set_support_level(1) .set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::sqrt)); +RELAY_REGISTER_UNARY_OP("rsqrt") +.describe(R"code(Returns the rsqrt input array, computed element-wise. + +.. math:: + 1/sqrt(x) + +)code" TVM_ADD_FILELINE) +.set_support_level(1) +.set_attr("FTVMCompute", RELAY_UNARY_COMPUTE(topi::rsqrt)); RELAY_REGISTER_UNARY_OP("zeros_like") .describe(R"code(Returns an array of zeros, with same type and shape as the input. diff --git a/tests/python/relay/test_ir_op.py b/tests/python/relay/test_ir_op.py index f1d835d2b43b..93d4bdd0b335 100644 --- a/tests/python/relay/test_ir_op.py +++ b/tests/python/relay/test_ir_op.py @@ -14,7 +14,7 @@ def test(x): def test_op_level1(): x = relay.Var("x") - for op_name in ["log", "exp", "sqrt", "tanh"]: + for op_name in ["log", "exp", "sqrt", "rsqrt","tanh"]: y = getattr(relay, op_name)(x) assert y.op.name == op_name assert y.op.support_level == 1 diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index b954e42bf1ab..6277835fa0f8 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -14,6 +14,10 @@ def relu(x): np.maximum(x_copy, 0, x_copy) return x_copy +def rsqrt(x): + one = np.ones_like(x) + return one / np.sqrt(x) + def test_unary_op(): def check_single_op(opfunc, ref): shape = (10, 4) @@ -41,6 +45,7 @@ def check_single_op(opfunc, ref): for opfunc, ref in [(tvm.relay.log, np.log), (tvm.relay.exp, np.exp), (tvm.relay.sqrt, np.sqrt), + (tvm.relay.rsqrt, rsqrt), (tvm.relay.sigmoid, sigmoid), (tvm.relay.tanh, np.tanh), (relay.nn.relu, relu)]: diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index e5d8778041b1..a1d7ee5cba8e 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -30,6 +30,7 @@ TOPI_DECLARE_UNARY_OP(exp); TOPI_DECLARE_UNARY_OP(tanh); TOPI_DECLARE_UNARY_OP(sigmoid); TOPI_DECLARE_UNARY_OP(sqrt); +TOPI_DECLARE_UNARY_OP(rsqrt); TOPI_DECLARE_UNARY_OP(log); TOPI_DECLARE_UNARY_OP(floor); TOPI_DECLARE_UNARY_OP(ceil); diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index faddc8ac8a90..ae9a736ca96b 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -208,6 +208,23 @@ def sqrt(x): return tvm.compute(x.shape, lambda *i: tvm.sqrt(x(*i))) +@tvm.tag_scope(tag=tag.ELEMWISE) +def rsqrt(x): + """Take inverse square root of input x. + + Parameters + ---------- + x : tvm.Tensor + Input argument. + + Returns + ------- + y : tvm.Tensor + The result. + """ + return tvm.compute(x.shape, lambda *i: 1/tvm.sqrt(x(*i))) + + @tvm.tag_scope(tag=tag.ELEMWISE) def sigmoid(x): """Take sigmoid tanh of input x. diff --git a/topi/src/topi.cc b/topi/src/topi.cc index aed2eab9c6bc..1d3e82ffb8d8 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -143,6 +143,11 @@ TVM_REGISTER_GLOBAL("topi.sqrt") *rv = sqrt(args[0]); }); +TVM_REGISTER_GLOBAL("topi.rsqrt") +.set_body([](TVMArgs args, TVMRetValue *rv) { +*rv = rsqrt(args[0]); +}); + TVM_REGISTER_GLOBAL("topi.log") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = log(args[0]); diff --git a/topi/tests/python/test_topi_basic.py b/topi/tests/python/test_topi_basic.py index 077331cd6d1a..d6c282164243 100644 --- a/topi/tests/python/test_topi_basic.py +++ b/topi/tests/python/test_topi_basic.py @@ -24,6 +24,7 @@ def test_apply(func, name): test_apply(topi.sigmoid, "sigmoid") test_apply(topi.log, "log") test_apply(topi.sqrt, "sqrt") + test_apply(topi.rsqrt, "rsqrt") if __name__ == "__main__": diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py index f2f2471c868f..dbc06d036676 100644 --- a/topi/tests/python/test_topi_math.py +++ b/topi/tests/python/test_topi_math.py @@ -59,6 +59,7 @@ def check_device(device): test_apply(topi.sigmoid, "sigmoid", lambda x:1/(1+np.exp(-x)), -1, 1) test_apply(topi.log, "log", np.log, 0, 100) test_apply(topi.sqrt, "sqrt", np.sqrt, 0, 100) + test_apply(topi.rsqrt, "rsqrt", lambda x:1/np.sqrt(x), 0, 100) if __name__ == "__main__": test_util() From c78977e883e38135f260df5eebf4f0bba0984d69 Mon Sep 17 00:00:00 2001 From: makihiro Date: Mon, 1 Apr 2019 18:08:43 +0900 Subject: [PATCH 2/9] [Relay][TOPI] add rsqrt operator --- docs/api/python/topi.rst | 2 ++ include/tvm/expr_operator.h | 1 + topi/python/topi/math.py | 2 +- 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/api/python/topi.rst b/docs/api/python/topi.rst index 84ded7b697d1..1be9bcc581ba 100644 --- a/docs/api/python/topi.rst +++ b/docs/api/python/topi.rst @@ -19,6 +19,7 @@ List of operators topi.tanh topi.log topi.sqrt + topi.rsqrt topi.sigmoid topi.clip topi.cast @@ -105,6 +106,7 @@ topi .. autofunction:: topi.tanh .. autofunction:: topi.log .. autofunction:: topi.sqrt +.. autofunction:: topi.rsqrt .. autofunction:: topi.sigmoid .. autofunction:: topi.clip .. autofunction:: topi.cast diff --git a/include/tvm/expr_operator.h b/include/tvm/expr_operator.h index c4d2d555f3a3..3647678e1ca8 100644 --- a/include/tvm/expr_operator.h +++ b/include/tvm/expr_operator.h @@ -468,6 +468,7 @@ TVM_DECLARE_INTRIN_UNARY(exp); TVM_DECLARE_INTRIN_UNARY(tanh); TVM_DECLARE_INTRIN_UNARY(sigmoid); TVM_DECLARE_INTRIN_UNARY(sqrt); +TVM_DECLARE_INTRIN_UNARY(rsqrt); TVM_DECLARE_INTRIN_UNARY(log); TVM_DECLARE_INTRIN_UNARY(popcount); diff --git a/topi/python/topi/math.py b/topi/python/topi/math.py index ae9a736ca96b..ca15858266b3 100644 --- a/topi/python/topi/math.py +++ b/topi/python/topi/math.py @@ -222,7 +222,7 @@ def rsqrt(x): y : tvm.Tensor The result. """ - return tvm.compute(x.shape, lambda *i: 1/tvm.sqrt(x(*i))) + return tvm.compute(x.shape, lambda *i: tvm.rsqrt(x(*i))) @tvm.tag_scope(tag=tag.ELEMWISE) From 7a5ff304fcb09fc6ca69e5947d1e669882c6f7a7 Mon Sep 17 00:00:00 2001 From: makihiro Date: Mon, 1 Apr 2019 18:40:49 +0900 Subject: [PATCH 3/9] [Relay][TOPI] add rsqrt operator --- topi/include/topi/elemwise.h | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index a1d7ee5cba8e..4560f2a69475 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -30,7 +30,6 @@ TOPI_DECLARE_UNARY_OP(exp); TOPI_DECLARE_UNARY_OP(tanh); TOPI_DECLARE_UNARY_OP(sigmoid); TOPI_DECLARE_UNARY_OP(sqrt); -TOPI_DECLARE_UNARY_OP(rsqrt); TOPI_DECLARE_UNARY_OP(log); TOPI_DECLARE_UNARY_OP(floor); TOPI_DECLARE_UNARY_OP(ceil); @@ -111,6 +110,23 @@ inline Tensor sign(const Tensor& x, }, name, tag); } +/*! +* \brief Creates an operation that returns rsqrt of a given tensor +* +* \param x The input tensor +* \param name The name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor whose op member is the rsqrt operation +*/ +inline Tensor rsqrt(const Tensor& x, + std::string name = "tensor", + std::string tag = kElementWise) { + return compute(x->shape, [&](const Array& i) { + return 1/tvm::sqrt(x(i)); + }, name, tag); +} + /*! * \brief Creates an operation that clips each element of a tensor to * the interval [a_min, a_max] From a80bb3fe43eee67523c251aabcbccba915414c05 Mon Sep 17 00:00:00 2001 From: makihiro Date: Mon, 1 Apr 2019 19:07:49 +0900 Subject: [PATCH 4/9] [Relay][TOPI] add rsqrt operator --- python/tvm/hybrid/runtime.py | 17 +++++++++++++++++ python/tvm/intrin.py | 18 +++++++++++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/tvm/hybrid/runtime.py b/python/tvm/hybrid/runtime.py index b3c744f42652..6ba12b1219c5 100644 --- a/python/tvm/hybrid/runtime.py +++ b/python/tvm/hybrid/runtime.py @@ -36,6 +36,22 @@ def allocate(shape, dtype='float32', scope='global'): #pylint: disable=unused-ar return numpy.zeros(shape).astype(dtype) +def rsqrt(x): + """ + Computes reciprocal of square root of x element-wise + + Parameters + ---------- + x: a real number + + Returns + ------- + res: a real number + The result of reciprocal of square root of x + """ + return numpy.ones_like(x) / numpy.sqrt(x) + + def popcount(x): """ Count ones in the binary representation of number x @@ -87,6 +103,7 @@ def max_num_threads(allow_none=True): 'allocate' : allocate, 'output_tensor' : allocate, 'sqrt' : numpy.sqrt, + 'rsqrt' : rsqrt, 'log' : numpy.log, 'tanh' : numpy.tanh, 'power' : numpy.power, diff --git a/python/tvm/intrin.py b/python/tvm/intrin.py index a0dabfc8a6e2..4c09f7882793 100644 --- a/python/tvm/intrin.py +++ b/python/tvm/intrin.py @@ -244,7 +244,7 @@ def log(x): def sqrt(x): - """Take log of input x. + """Take square root of input x. Parameters ---------- @@ -259,6 +259,22 @@ def sqrt(x): return call_pure_intrin(x.dtype, "sqrt", x) +def rsqrt(x): + """Take reciprocal of square root of input x. + + Parameters + ---------- + x : Expr + Input argument. + + Returns + ------- + y : Expr + The result. + """ + return call_pure_intrin(x.dtype, "rsqrt", x) + + def floor(x): """Take floor of float input x. From 790eb00410af172bf0ca26a5e9ec3df43672ace2 Mon Sep 17 00:00:00 2001 From: makihiro Date: Mon, 1 Apr 2019 19:28:51 +0900 Subject: [PATCH 5/9] [Relay][TOPI] add rsqrt operator --- src/codegen/intrin_rule.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/codegen/intrin_rule.cc b/src/codegen/intrin_rule.cc index f326fceb6ee8..d40429363e47 100644 --- a/src/codegen/intrin_rule.cc +++ b/src/codegen/intrin_rule.cc @@ -21,6 +21,9 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tanh") TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sqrt") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.rsqrt") +.set_body(DispatchExtern); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.pow") .set_body(DispatchExtern); From b789573961269308b239300ab01bf2bcbce5fbc8 Mon Sep 17 00:00:00 2001 From: makihiro Date: Tue, 2 Apr 2019 11:34:36 +0900 Subject: [PATCH 6/9] [Relay][TOPI] add rsqrt operator --- src/codegen/intrin_rule.cc | 3 --- topi/include/topi/elemwise.h | 3 ++- topi/tests/python/test_topi_math.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/codegen/intrin_rule.cc b/src/codegen/intrin_rule.cc index d40429363e47..f326fceb6ee8 100644 --- a/src/codegen/intrin_rule.cc +++ b/src/codegen/intrin_rule.cc @@ -21,9 +21,6 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tanh") TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sqrt") .set_body(DispatchExtern); -TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.rsqrt") -.set_body(DispatchExtern); - TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.pow") .set_body(DispatchExtern); diff --git a/topi/include/topi/elemwise.h b/topi/include/topi/elemwise.h index 4560f2a69475..7edc3e11da4b 100644 --- a/topi/include/topi/elemwise.h +++ b/topi/include/topi/elemwise.h @@ -123,7 +123,8 @@ inline Tensor rsqrt(const Tensor& x, std::string name = "tensor", std::string tag = kElementWise) { return compute(x->shape, [&](const Array& i) { - return 1/tvm::sqrt(x(i)); + Expr one = make_const(x->dtype, 1); + return one/tvm::sqrt(x(i)); }, name, tag); } diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py index dbc06d036676..78e8a6d04f9b 100644 --- a/topi/tests/python/test_topi_math.py +++ b/topi/tests/python/test_topi_math.py @@ -59,7 +59,7 @@ def check_device(device): test_apply(topi.sigmoid, "sigmoid", lambda x:1/(1+np.exp(-x)), -1, 1) test_apply(topi.log, "log", np.log, 0, 100) test_apply(topi.sqrt, "sqrt", np.sqrt, 0, 100) - test_apply(topi.rsqrt, "rsqrt", lambda x:1/np.sqrt(x), 0, 100) + test_apply(topi.rsqrt, "rsqrt", lambda x:np.ones_like(x)/np.sqrt(x), 0, 100, skip_name_check=True) if __name__ == "__main__": test_util() From aceaf900f18bf079e17d35886be061cad01efb41 Mon Sep 17 00:00:00 2001 From: makihiro Date: Wed, 3 Apr 2019 11:39:08 +0900 Subject: [PATCH 7/9] [Relay][TOPI] add rsqrt operator - fix comments --- python/tvm/hybrid/runtime.py | 4 ++-- python/tvm/relay/op/tensor.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python/tvm/hybrid/runtime.py b/python/tvm/hybrid/runtime.py index 6ba12b1219c5..e62bbf8fd627 100644 --- a/python/tvm/hybrid/runtime.py +++ b/python/tvm/hybrid/runtime.py @@ -42,11 +42,11 @@ def rsqrt(x): Parameters ---------- - x: a real number + x: Tensor Returns ------- - res: a real number + res: Tensor The result of reciprocal of square root of x """ return numpy.ones_like(x) / numpy.sqrt(x) diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index 0b146463a662..1a23b52f2bbc 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -66,6 +66,9 @@ def sqrt(data): def rsqrt(data): """Compute elementwise rsqrt of data. + .. math:: + 1/sqrt(x) + Parameters ---------- data : relay.Expr From 451bc061dd869381370a42604342ea30a87b0497 Mon Sep 17 00:00:00 2001 From: makihiro Date: Wed, 3 Apr 2019 16:31:47 +0900 Subject: [PATCH 8/9] [Relay][TOPI] add rsqrt operator - fix unresolved intrinsic error --- src/codegen/intrin_rule.cc | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/codegen/intrin_rule.cc b/src/codegen/intrin_rule.cc index f326fceb6ee8..0e1c3d9519d7 100644 --- a/src/codegen/intrin_rule.cc +++ b/src/codegen/intrin_rule.cc @@ -21,6 +21,16 @@ TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.tanh") TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.sqrt") .set_body(DispatchExtern); +TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.rsqrt") +.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 = one / sqrt(call->args[0]); + }); + TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.pow") .set_body(DispatchExtern); From 0681e8e87383e5c68326e3100ecef1304934ae14 Mon Sep 17 00:00:00 2001 From: makihiro Date: Fri, 12 Apr 2019 11:31:30 +0900 Subject: [PATCH 9/9] [Relay][TOPI] add rsqrt operator - fix math formula --- python/tvm/relay/op/tensor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index 1a23b52f2bbc..a2f5a300902f 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -67,6 +67,7 @@ def rsqrt(data): """Compute elementwise rsqrt of data. .. math:: + 1/sqrt(x) Parameters