From 04464836e16e858fd99c4b5bd28b156584732e8e Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Fri, 12 Oct 2018 10:25:52 +0530 Subject: [PATCH 1/5] [RELAY][OP]Strided slice --- docs/langref/relay_op.rst | 2 + include/tvm/relay/attrs/transform.h | 13 +++ python/tvm/relay/op/transform.py | 27 +++++++ src/relay/op/tensor/transform.cc | 116 +++++++++++++++++++++++++++ tests/python/relay/test_op_level4.py | 21 +++++ 5 files changed, 179 insertions(+) diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index d1549cd8326e..cbccd76b284b 100644 --- a/docs/langref/relay_op.rst +++ b/docs/langref/relay_op.rst @@ -122,6 +122,7 @@ This level enables additional math and transform operators. tvm.relay.min tvm.relay.mean tvm.relay.prod + tvm.relay.strided_slice **Level 5: Vision/Image Operators** @@ -225,6 +226,7 @@ Level 4 Definitions .. autofunction:: tvm.relay.min .. autofunction:: tvm.relay.mean .. autofunction:: tvm.relay.prod +.. autofunction:: tvm.relay.strided_slice diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index dfad1013701f..d63cc1a6d2ff 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -119,6 +119,19 @@ struct SplitAttrs : public tvm::AttrsNode { "the entries indicate where along axis the array is split."); TVM_ATTR_FIELD(axis).set_default(0) .describe("the axis to be splitted."); +/*! \brief Attributes for StridedSlice operator */ +struct StridedSliceAttrs : public tvm::AttrsNode { + Array begin; + Array end; + Array stride; + + TVM_DECLARE_ATTRS(StridedSliceAttrs, "relay.attrs.StridedSliceAttrs") { + TVM_ATTR_FIELD(begin) + .describe("Indices for begin of slice"); + TVM_ATTR_FIELD(end) + .describe("Indices for end of the slice"); + TVM_ATTR_FIELD(stride).set_default(Array({})) + .describe("Stride values of the slice"); } }; diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 9d14463a530c..b14d841c2d57 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -336,3 +336,30 @@ def split(data, indices_or_sections, axis=0): else: ret_size = len(indices_or_sections) + 1 return TupleWrapper(_make.split(data, indices_or_sections, axis), ret_size) + +def strided_slice(data, begin, end, stride=None): + """Strided slice of an array.. + + Parameters + ---------- + data : relay.Expr + The source array to be sliced. + + begin: list of int + The indices to begin with in the slicing. + + end: list of int + Indicies indicating end of the slice. + + stride: list of int, optional + Specifies the stride values, it can be negative in that case, + the input tensor will be reversed in that particular axis. + + Returns + ------- + ret : relay.Expr + The computed result. + """ + stride = stride or [] + return _make.strided_slice(data, list(begin), list(end), list(stride)) +>>>>>>> [RELAY][OP]Strided slice diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 5faa0805426a..2456c0495d66 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -891,6 +891,122 @@ RELAY_REGISTER_OP("broadcast_to_like") .add_argument("broadcast_type", "Tensor", "Provide the type to broadcast to.") .set_support_level(10) .add_type_rel("BroadCastToLike", BroadCastToLikeRel); +// strided_slice +TVM_REGISTER_NODE_TYPE(StridedSliceAttrs); +bool StridedSliceRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 2); + const auto* data = types[0].as(); + CHECK(data != nullptr); + if (data->shape.size() == 0) return false; + + const StridedSliceAttrs *param = attrs.as(); + CHECK(param != nullptr); + + auto dshape = data->shape; + auto num_axis = dshape.size(); + + std::vector begin_vec; + for (auto i : param->begin) { + begin_vec.push_back(i); + } + for (auto i = begin_vec.size(); i < num_axis; ++i) { + begin_vec.push_back(0); + } + + std::vector end_vec; + for (auto i : param->end) { + end_vec.push_back(i); + } + for (auto i = end_vec.size(); i < num_axis; ++i) { + end_vec.push_back(dshape[i]); + } + + std::vector stride_vec; + for (auto i : param->stride) { + stride_vec.push_back(i); + } + for (auto i = stride_vec.size(); i < num_axis; ++i) { + stride_vec.push_back(1); + } + std::vector oshape(dshape.size()); + + #define MAX(a, b) (reporter->Assert((a) > (b)) ? (a) : (b)) + #define MIN(a, b) (reporter->Assert((a) < (b)) ? (a) : (b)) + + for (size_t i = 0; i < num_axis; ++i) { + auto begin_range = reporter->Assert(stride_vec[i] < 0) ? -1 : 0; + auto end_range = reporter->Assert(stride_vec[i] < 0) ? dshape[i] - 1 : dshape[i]; + auto begin = reporter->Assert(begin_vec[i] < 0) ? dshape[i] + begin_vec[i] : begin_vec[i]; + auto end = reporter->Assert(end_vec[i] < 0) ? dshape[i] + end_vec[i] : end_vec[i]; + + begin = MIN(MAX(begin, begin_range), end_range); + end = MIN(MAX(end, begin_range), end_range); + auto interval = abs((end - begin)); + auto slice_size = (interval + abs(stride_vec[i]) - 1) / abs(stride_vec[i]); + + CHECK(reporter->Assert(stride_vec[i] < 0) ? + reporter->Assert(end < begin) : reporter->Assert(begin < end)) + << ": Input [Begin=" << begin_vec[i] << ", End=" << end_vec[i] + << "] is invalid for axis=" << i; + oshape[i] = slice_size; + } + + reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype)); + return true; +} + + +// Positional relay function to create StridedSlice operator used by frontend FFI. +Expr MakeStridedSlice(Expr data, + Array begin, + Array end, + Array stride) { + auto attrs = make_node(); + attrs->begin = std::move(begin); + attrs->end = std::move(end); + attrs->stride = std::move(stride); + static const Op& op = Op::Get("strided_slice"); + return CallNode::make(op, {data}, Attrs(attrs), {}); +} + + +TVM_REGISTER_API("relay.op._make.strided_slice") + .set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeStridedSlice, args, rv); + }); + + +RELAY_REGISTER_OP("strided_slice") + .describe(R"code(Strided slice of an array. + +Examples:: + + x = [[ 1., 4., 7., 10.], + [ 2., 5., 8., 11.], + [ 3., 6., 9., 12.]] + + strided_slice(x, begin=[0, 1], end=[2, 4], stride=[1, 1]) = [[ 4., 7., 10.], + [ 5., 8., 11.]] + + x = [[[ 1., 2.], + [ 3., 4.]], + + [[ 5., 6.], + [ 7., 8.]]] + + strided_slice(x, begin=[0, 0], end=[2, 2]) = [[[ 1., 2.], + [ 3., 4.]], + + [[ 5., 6.], + [ 7., 8.]]] +)code" TVM_ADD_FILELINE) +.set_num_inputs(1) +.add_argument("data", "Tensor", "The input tensor.") +.set_support_level(4) +.add_type_rel("StridedSlice", StridedSliceRel); // Split TVM_REGISTER_NODE_TYPE(SplitAttrs); diff --git a/tests/python/relay/test_op_level4.py b/tests/python/relay/test_op_level4.py index 2dc643cfd7e4..b7505ca08c62 100644 --- a/tests/python/relay/test_op_level4.py +++ b/tests/python/relay/test_op_level4.py @@ -90,6 +90,26 @@ def test_reduce_functions(): verify_reduce(func, (128, 24, 128), (0, 2), False, False, (24,)) verify_reduce(func, (128, 24, 128), (0, 1), True, False, (1, 1, 128)) verify_reduce(func, (128, 24, 128), (0, 2), True, False, (1, 24, 1)) +def verify_strided_slice(data, begin, end, stride, output): + x = relay.var("x", relay.TensorType(data, "float32")) + z = relay.strided_slice(x, begin=begin, end=end, stride=stride) + zz = relay.ir_pass.infer_type(z) + assert "begin=" in z.astext() + assert "end=" in z.astext() + if stride: + assert "stride=" in z.astext() + assert zz.checked_type == relay.ty.TensorType(output, "float32") + +def test_strided_slice(): + verify_strided_slice((3, 4, 3), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2)) + verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 4, 3], [2, 1, 1], (1, 3, 3)) + verify_strided_slice((3, 4, 3), [1, -1, 0], [4, -5, 3], [2, -1, 1], (1, 4, 3)) + verify_strided_slice((3, 4, 3), [1, 0, 0], [2, 2, 3], [1, 1, 2], (1, 2, 2)) + verify_strided_slice((3, 4, 3), [1, -1, 0], [2, -3, 3], [1, -1, 1], (1, 2, 3)) + verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 4, 3], None, (2, 3, 3)) + verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 1000, 3], None, (2, 3, 3)) + verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 4], None, (2, 3, 3)) + verify_strided_slice((3, 4, 3), [1, 1], [4, 4, 3], None, (2, 3, 3)) if __name__ == "__main__": test_binary_op() @@ -97,3 +117,4 @@ def test_reduce_functions(): test_binary_int_broadcast() test_where() test_reduce_functions() + test_strided_slice() From 78749c11fd3c1132b130dbeefbfe5efee62760d2 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Mon, 15 Oct 2018 16:21:13 +0530 Subject: [PATCH 2/5] Review comments --- src/relay/op/tensor/transform.cc | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 2456c0495d66..0dadd6137b27 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -933,17 +933,14 @@ bool StridedSliceRel(const Array& types, } std::vector oshape(dshape.size()); - #define MAX(a, b) (reporter->Assert((a) > (b)) ? (a) : (b)) - #define MIN(a, b) (reporter->Assert((a) < (b)) ? (a) : (b)) - for (size_t i = 0; i < num_axis; ++i) { auto begin_range = reporter->Assert(stride_vec[i] < 0) ? -1 : 0; auto end_range = reporter->Assert(stride_vec[i] < 0) ? dshape[i] - 1 : dshape[i]; auto begin = reporter->Assert(begin_vec[i] < 0) ? dshape[i] + begin_vec[i] : begin_vec[i]; auto end = reporter->Assert(end_vec[i] < 0) ? dshape[i] + end_vec[i] : end_vec[i]; - begin = MIN(MAX(begin, begin_range), end_range); - end = MIN(MAX(end, begin_range), end_range); + begin = min(max(begin, begin_range), end_range); + end = min(max(end, begin_range), end_range); auto interval = abs((end - begin)); auto slice_size = (interval + abs(stride_vec[i]) - 1) / abs(stride_vec[i]); From 16c075145e02d66293f4cb4c2d34a76fd621abd1 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Mon, 15 Oct 2018 17:16:20 +0530 Subject: [PATCH 3/5] Description updated for begin&end strided slice --- include/tvm/relay/attrs/transform.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index d63cc1a6d2ff..417e3b64dd71 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -127,9 +127,9 @@ struct StridedSliceAttrs : public tvm::AttrsNode { TVM_DECLARE_ATTRS(StridedSliceAttrs, "relay.attrs.StridedSliceAttrs") { TVM_ATTR_FIELD(begin) - .describe("Indices for begin of slice"); + .describe("Indices for begin of slice, begin index is also inclusive"); TVM_ATTR_FIELD(end) - .describe("Indices for end of the slice"); + .describe("Indices for end of slice, end index is also inclusive"); TVM_ATTR_FIELD(stride).set_default(Array({})) .describe("Stride values of the slice"); } From 7ff403e96a00f35ccc23c336b2dd6b38187f7f85 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Sat, 20 Oct 2018 08:39:53 +0530 Subject: [PATCH 4/5] set_attrs_type_key and test_format testcase added --- include/tvm/relay/attrs/transform.h | 4 +++- python/tvm/relay/op/transform.py | 2 +- src/relay/op/tensor/transform.cc | 1 + 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index 417e3b64dd71..6228b6825283 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -119,6 +119,9 @@ struct SplitAttrs : public tvm::AttrsNode { "the entries indicate where along axis the array is split."); TVM_ATTR_FIELD(axis).set_default(0) .describe("the axis to be splitted."); + } +}; + /*! \brief Attributes for StridedSlice operator */ struct StridedSliceAttrs : public tvm::AttrsNode { Array begin; @@ -134,7 +137,6 @@ struct StridedSliceAttrs : public tvm::AttrsNode { .describe("Stride values of the slice"); } }; - } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_TRANSFORM_H_ diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index b14d841c2d57..9368befd7428 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -337,6 +337,7 @@ def split(data, indices_or_sections, axis=0): ret_size = len(indices_or_sections) + 1 return TupleWrapper(_make.split(data, indices_or_sections, axis), ret_size) + def strided_slice(data, begin, end, stride=None): """Strided slice of an array.. @@ -362,4 +363,3 @@ def strided_slice(data, begin, end, stride=None): """ stride = stride or [] return _make.strided_slice(data, list(begin), list(end), list(stride)) ->>>>>>> [RELAY][OP]Strided slice diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 0dadd6137b27..871ea67d4aa2 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1003,6 +1003,7 @@ Examples:: .set_num_inputs(1) .add_argument("data", "Tensor", "The input tensor.") .set_support_level(4) +.set_attrs_type_key("relay.attrs.StridedSliceAttrs") .add_type_rel("StridedSlice", StridedSliceRel); // Split From ba9864d7a1dc648525898185ea0f97b474cd7d19 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Sun, 28 Oct 2018 21:17:43 +0530 Subject: [PATCH 5/5] Review comment fixed --- include/tvm/relay/attrs/transform.h | 8 +++---- python/tvm/relay/op/transform.py | 8 +++---- src/relay/op/tensor/transform.cc | 36 +++++++++++++++++++--------- tests/python/relay/test_op_level4.py | 12 +++++++--- 4 files changed, 42 insertions(+), 22 deletions(-) diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index 6228b6825283..afb2030979ef 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -124,16 +124,16 @@ struct SplitAttrs : public tvm::AttrsNode { /*! \brief Attributes for StridedSlice operator */ struct StridedSliceAttrs : public tvm::AttrsNode { - Array begin; - Array end; - Array stride; + Array begin; + Array end; + Array strides; TVM_DECLARE_ATTRS(StridedSliceAttrs, "relay.attrs.StridedSliceAttrs") { TVM_ATTR_FIELD(begin) .describe("Indices for begin of slice, begin index is also inclusive"); TVM_ATTR_FIELD(end) .describe("Indices for end of slice, end index is also inclusive"); - TVM_ATTR_FIELD(stride).set_default(Array({})) + TVM_ATTR_FIELD(strides).set_default(Array({})) .describe("Stride values of the slice"); } }; diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 9368befd7428..e737ff9ed950 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -338,7 +338,7 @@ def split(data, indices_or_sections, axis=0): return TupleWrapper(_make.split(data, indices_or_sections, axis), ret_size) -def strided_slice(data, begin, end, stride=None): +def strided_slice(data, begin, end, strides=None): """Strided slice of an array.. Parameters @@ -352,7 +352,7 @@ def strided_slice(data, begin, end, stride=None): end: list of int Indicies indicating end of the slice. - stride: list of int, optional + strides: list of int, optional Specifies the stride values, it can be negative in that case, the input tensor will be reversed in that particular axis. @@ -361,5 +361,5 @@ def strided_slice(data, begin, end, stride=None): ret : relay.Expr The computed result. """ - stride = stride or [] - return _make.strided_slice(data, list(begin), list(end), list(stride)) + strides = strides or [] + return _make.strided_slice(data, list(begin), list(end), list(strides)) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 871ea67d4aa2..cbcee320b44d 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -891,6 +891,8 @@ RELAY_REGISTER_OP("broadcast_to_like") .add_argument("broadcast_type", "Tensor", "Provide the type to broadcast to.") .set_support_level(10) .add_type_rel("BroadCastToLike", BroadCastToLikeRel); + + // strided_slice TVM_REGISTER_NODE_TYPE(StridedSliceAttrs); bool StridedSliceRel(const Array& types, @@ -908,7 +910,7 @@ bool StridedSliceRel(const Array& types, auto dshape = data->shape; auto num_axis = dshape.size(); - std::vector begin_vec; + std::vector begin_vec; for (auto i : param->begin) { begin_vec.push_back(i); } @@ -924,8 +926,8 @@ bool StridedSliceRel(const Array& types, end_vec.push_back(dshape[i]); } - std::vector stride_vec; - for (auto i : param->stride) { + std::vector stride_vec; + for (auto i : param->strides) { stride_vec.push_back(i); } for (auto i = stride_vec.size(); i < num_axis; ++i) { @@ -934,10 +936,22 @@ bool StridedSliceRel(const Array& types, std::vector oshape(dshape.size()); for (size_t i = 0; i < num_axis; ++i) { - auto begin_range = reporter->Assert(stride_vec[i] < 0) ? -1 : 0; - auto end_range = reporter->Assert(stride_vec[i] < 0) ? dshape[i] - 1 : dshape[i]; - auto begin = reporter->Assert(begin_vec[i] < 0) ? dshape[i] + begin_vec[i] : begin_vec[i]; - auto end = reporter->Assert(end_vec[i] < 0) ? dshape[i] + end_vec[i] : end_vec[i]; + const int64_t* stride_t = as_const_int(stride_vec[i]); + CHECK(stride_t != nullptr) << "Stride cannot be symbolic."; + int64_t stride_v = stride_t[0]; + + const int64_t* begin_t = as_const_int(begin_vec[i]); + CHECK(begin_t != nullptr) << "Begin index cannot be symbolic."; + int64_t begin_v = begin_t[0]; + + const int64_t* end_t = as_const_int(end_vec[i]); + CHECK(end_t != nullptr) << "End index cannot be symbolic."; + int64_t end_v = end_t[0]; + + auto begin_range = make_const(Int(64), (stride_v < 0) ? -1 : 0); + auto end_range = (stride_v < 0) ? dshape[i] - 1 : dshape[i]; + auto begin = (begin_v < 0) ? dshape[i] + begin_vec[i] : begin_vec[i]; + auto end = (end_v < 0) ? dshape[i] + end_vec[i] : end_vec[i]; begin = min(max(begin, begin_range), end_range); end = min(max(end, begin_range), end_range); @@ -958,13 +972,13 @@ bool StridedSliceRel(const Array& types, // Positional relay function to create StridedSlice operator used by frontend FFI. Expr MakeStridedSlice(Expr data, - Array begin, - Array end, - Array stride) { + Array begin, + Array end, + Array strides) { auto attrs = make_node(); attrs->begin = std::move(begin); attrs->end = std::move(end); - attrs->stride = std::move(stride); + attrs->strides = std::move(strides); static const Op& op = Op::Get("strided_slice"); return CallNode::make(op, {data}, Attrs(attrs), {}); } diff --git a/tests/python/relay/test_op_level4.py b/tests/python/relay/test_op_level4.py index b7505ca08c62..bb6c28f0cf62 100644 --- a/tests/python/relay/test_op_level4.py +++ b/tests/python/relay/test_op_level4.py @@ -90,17 +90,22 @@ def test_reduce_functions(): verify_reduce(func, (128, 24, 128), (0, 2), False, False, (24,)) verify_reduce(func, (128, 24, 128), (0, 1), True, False, (1, 1, 128)) verify_reduce(func, (128, 24, 128), (0, 2), True, False, (1, 24, 1)) + + def verify_strided_slice(data, begin, end, stride, output): x = relay.var("x", relay.TensorType(data, "float32")) - z = relay.strided_slice(x, begin=begin, end=end, stride=stride) + z = relay.strided_slice(x, begin=begin, end=end, strides=stride) zz = relay.ir_pass.infer_type(z) assert "begin=" in z.astext() assert "end=" in z.astext() if stride: - assert "stride=" in z.astext() - assert zz.checked_type == relay.ty.TensorType(output, "float32") + assert "strides=" in z.astext() + if output: + assert zz.checked_type == relay.ty.TensorType(output, "float32") def test_strided_slice(): + d1, d2, d3, d4 = tvm.var("d1"), tvm.var("d2"), tvm.var("d3"), tvm.var("d4") + verify_strided_slice((d1, d2, d3), [0, 0, 0], [4, -5, 4], [1, -1, 2], None) verify_strided_slice((3, 4, 3), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2)) verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 4, 3], [2, 1, 1], (1, 3, 3)) verify_strided_slice((3, 4, 3), [1, -1, 0], [4, -5, 3], [2, -1, 1], (1, 4, 3)) @@ -111,6 +116,7 @@ def test_strided_slice(): verify_strided_slice((3, 4, 3), [1, 1, 0], [4, 4], None, (2, 3, 3)) verify_strided_slice((3, 4, 3), [1, 1], [4, 4, 3], None, (2, 3, 3)) + if __name__ == "__main__": test_binary_op() test_cmp_type()