From 274dd80932f01f18d707894ffad27cb701c52a6f Mon Sep 17 00:00:00 2001 From: xp56 Date: Tue, 17 Nov 2020 10:23:14 +0800 Subject: [PATCH 1/5] [TF frontend] add some Segment and UnsortedSegment ops * segment_max, segment_min, segment_mean, segment_sum, segment_prod * unsorted_segment_max, unsorted_segment_min, unsorted_segment_mean * unsorted_segment_prod, unsorted_segment_sum --- include/tvm/relay/attrs/transform.h | 10 + python/tvm/relay/frontend/tensorflow.py | 43 +++ python/tvm/relay/op/_tensor.py | 23 +- python/tvm/relay/op/op_attrs.py | 5 + python/tvm/relay/op/tensor.py | 110 ++++++ python/tvm/topi/generic/vision.py | 85 +++++ python/tvm/topi/tensor.py | 340 ++++++++++++++++++ src/relay/op/tensor/binary.cc | 110 ++++++ .../frontend/tensorflow/test_forward.py | 89 +++++ tests/python/topi/python/test_topi_math.py | 85 +++++ 10 files changed, 899 insertions(+), 1 deletion(-) diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index cc97a94a1406..44d9b1a8ebec 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -478,6 +478,16 @@ struct UniqueAttrs : public tvm::AttrsNode { } }; // struct UniqueAttrs +/*! \brief Attributes used in segment_max, segment_min, + segment_mean, segment_sum, segment_prod operator */ +struct SegmentAttrs : public tvm::AttrsNode { + int num_segments; + + TVM_DECLARE_ATTRS(SegmentAttrs, "relay.attrs.SegmentAttrs") { + TVM_ATTR_FIELD(num_segments).set_default(0).describe("The maximum of segment_ids."); + } +}; // struct SegmentAttrs + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_TRANSFORM_H_ diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 4af73702ad9c..9f2dde317166 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -912,6 +912,22 @@ def _impl(inputs, attr, params, mod): return _impl +def _unsorted_segment(name): + def _impl(inputs, attr, params, mod): + # op description: https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_max + try: + num_segments = _infer_value(inputs[2], params).asnumpy().tolist() + except Exception: + raise tvm.error.OpAttributeInvalid("Can't find num_segments.") + return AttrCvt( + op_name="segment_" + name, + ignores=["Tdim", "Tidx", "Tindices", "Tnumsegments"], + extras={"num_segments": num_segments}, + )([inputs[0], inputs[1]], attr) + + return _impl + + def _crop_and_resize(): def _impl(inputs, attr, params, mod): # input image is a 4-D tensor of shape [batch, image_height, image_width, depth] @@ -2617,6 +2633,23 @@ def _impl(inputs, attr, params, mod): return _impl +def _segment(opname): + def _impl(inputs, attr, params, mod): + # op description: https://www.tensorflow.org/api_docs/python/tf/math/segment_max + try: + segment_ids = _infer_value(inputs[1], params) + except Exception: + raise tvm.error.OpAttributeInvalid("Can't get value of segment_ids.") + + num_out = segment_ids.asnumpy().max() + 1 + out = AttrCvt(op_name=opname, ignores=["T", "Tindices"], extras={"num_segments": num_out})( + inputs, attr + ) + return out + + return _impl + + def _size(): def _impl(inputs, attr, params, mod): new_attr = attr @@ -2864,6 +2897,11 @@ def _impl(inputs, attr, params, mod): "SelectV2": _where(), "Selu": _selu(), "Shape": _shape(), + "SegmentMax": _segment("segment_max"), + "SegmentMean": _segment("segment_mean"), + "SegmentMin": _segment("segment_min"), + "SegmentProd": _segment("segment_prod"), + "SegmentSum": _segment("segment_sum"), "Sigmoid": AttrCvt("sigmoid"), "Sign": AttrCvt("sign"), "Sin": AttrCvt("sin"), @@ -2915,6 +2953,11 @@ def _impl(inputs, attr, params, mod): "UniqueWithCounts": _unique(True), "Unpack": _unpack(), "UnravelIndex": _unravel_index(), + "UnsortedSegmentMax": _unsorted_segment("max"), + "UnsortedSegmentMin": _unsorted_segment("min"), + "UnsortedSegmentMean": _unsorted_segment("mean"), + "UnsortedSegmentProd": _unsorted_segment("prod"), + "UnsortedSegmentSum": _unsorted_segment("sum"), "Where": _where(), "ZerosLike": AttrCvt("zeros_like"), } diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index d7d99c017b2b..88869abf9e22 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -21,7 +21,8 @@ from tvm import topi from tvm.runtime import convert -from .op import register_compute, register_shape_func +from . import strategy +from .op import register_compute, register_shape_func, register_strategy from .op import register_broadcast_schedule, register_injective_schedule from .op import register_pattern, OpPattern @@ -283,3 +284,23 @@ def elemwise_shape_func(attrs, inputs, _): register_shape_func("sigmoid", False, elemwise_shape_func) register_shape_func("tanh", False, elemwise_shape_func) register_shape_func("logical_not", False, elemwise_shape_func) + +# segment_max +register_strategy("segment_max", strategy.segment_max_strategy) +register_pattern("segment_max", OpPattern.OPAQUE) + +# segment_min +register_strategy("segment_min", strategy.segment_min_strategy) +register_pattern("segment_min", OpPattern.OPAQUE) + +# segment_mean +register_strategy("segment_mean", strategy.segment_mean_strategy) +register_pattern("segment_mean", OpPattern.OPAQUE) + +# segment_sum +register_strategy("segment_sum", strategy.segment_sum_strategy) +register_pattern("segment_sum", OpPattern.OPAQUE) + +# segment_prod +register_strategy("segment_prod", strategy.segment_prod_strategy) +register_pattern("segment_prod", OpPattern.OPAQUE) diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py index 6844d133a77e..8deafd6af598 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -329,6 +329,11 @@ class ProposalAttrs(Attrs): """Attributes used in proposal operators""" +@tvm._ffi.register_object("relay.attrs.SegmentAttrs") +class SegmentAttrs(Attrs): + """Attributes used in segment operators""" + + @tvm._ffi.register_object("relay.attrs.MaxPool2DAttrs") class MaxPool2DAttrs(Attrs): """Attributes used in max_pool2d operators""" diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index a38a23064d6f..259bd2a0d032 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -1285,3 +1285,113 @@ def isinf(data): The computed result. """ return _make.isinf(data) + + +def segment_max(data, segment_ids, num_segments): + """Computes the maximum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_max(data, segment_ids, num_segments) + + +def segment_min(data, segment_ids, num_segments): + """Computes the minimum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_min(data, segment_ids, num_segments) + + +def segment_mean(data, segment_ids, num_segments): + """Computes the mean along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_mean(data, segment_ids, num_segments) + + +def segment_sum(data, segment_ids, num_segments): + """Computes the sum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_sum(data, segment_ids, num_segments) + + +def segment_prod(data, segment_ids, num_segments): + """Computes the prod along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_prod(data, segment_ids, num_segments) diff --git a/python/tvm/topi/generic/vision.py b/python/tvm/topi/generic/vision.py index e7518b1110a1..b5d44ed8ced1 100644 --- a/python/tvm/topi/generic/vision.py +++ b/python/tvm/topi/generic/vision.py @@ -176,3 +176,88 @@ def schedule_proposal(outs): The computation schedule for the op. """ return _default_schedule(outs, False) + + +def schedule_segment_max(outs): + """Schedule for segment_max operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_max + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) + + +def schedule_segment_min(outs): + """Schedule for segment_min operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_min + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) + + +def schedule_segment_mean(outs): + """Schedule for segment_mean operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_mean + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) + + +def schedule_segment_sum(outs): + """Schedule for segment_sum operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_sum + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) + + +def schedule_segment_prod(outs): + """Schedule for segment_prod operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_prod + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) diff --git a/python/tvm/topi/tensor.py b/python/tvm/topi/tensor.py index 31ebe86760cb..01794930a8fd 100644 --- a/python/tvm/topi/tensor.py +++ b/python/tvm/topi/tensor.py @@ -18,6 +18,8 @@ """Elementwise operators""" from __future__ import absolute_import as _abs from . import cpp +from .. import te, tir +from .util import get_const_tuple def elemwise_sum(xs): @@ -73,3 +75,341 @@ def full_like(x, fill_value): The result. """ return cpp.full_like(x, fill_value) + + +def segment_max(data, segment_ids, num_out): + """segment_max operator. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _segment_max(data, segment_ids, out_buf): + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = -3.4028235e38 + + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + in_index = k * inner_size + l + out[out_index] = te.max(input_data[in_index], out[out_index]) + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_max(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out + + +def segment_min(data, segment_ids, num_out): + """segment_min operator. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _segment_min(data, segment_ids, out_buf): + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = 3.4028235e38 + + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + in_index = k * inner_size + l + out[out_index] = te.min(input_data[in_index], out[out_index]) + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_min(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out + + +def segment_mean(data, segment_ids, num_out): + """segment_mean operator. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _segment_mean(data, segment_ids, out_buf): + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + + temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") + num = ib.allocate("int32", (1), name="num", scope="local") + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = 0.0 + + num[0] = 0 + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + temp_index[num[0]] = k + num[0] += 1 + + with ib.if_scope(num[0] > 0): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + with ib.for_range(0, num[0]) as k: + in_index = temp_index[k] * inner_size + l + out[out_index] += input_data[in_index] + out[out_index] = out[out_index] / num[0] + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_mean(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out + + +def segment_sum(data, segment_ids, num_out): + """segment_sum operator. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _segment_sum(data, segment_ids, out_buf): + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + + temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") + num = ib.allocate("int32", (1), name="num", scope="local") + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = 0.0 + + num[0] = 0 + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + temp_index[num[0]] = k + num[0] += 1 + + with ib.if_scope(num[0] > 0): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + with ib.for_range(0, num[0]) as k: + in_index = temp_index[k] * inner_size + l + out[out_index] += input_data[in_index] + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_sum(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out + + +def segment_prod(data, segment_ids, num_out): + """segment_prod operator. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _segment_prod(data, segment_ids, out_buf): + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + + temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") + num = ib.allocate("int32", (1), name="num", scope="local") + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = 1.0 + + num[0] = 0 + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + temp_index[num[0]] = k + num[0] += 1 + + with ib.if_scope(num[0] > 0): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + with ib.for_range(0, num[0]) as k: + in_index = temp_index[k] * inner_size + l + out[out_index] *= input_data[in_index] + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_prod(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out diff --git a/src/relay/op/tensor/binary.cc b/src/relay/op/tensor/binary.cc index aafd4492fec4..1513b063c9ea 100644 --- a/src/relay/op/tensor/binary.cc +++ b/src/relay/op/tensor/binary.cc @@ -161,5 +161,115 @@ RELAY_REGISTER_CMP_OP("greater_equal") .set_support_level(4) .set_attr("FTVMCompute", RELAY_BINARY_COMPUTE(topi::greater_equal)); +// segment_max +TVM_REGISTER_NODE_TYPE(SegmentAttrs); + +bool SegmentRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 3); + + auto segment_attrs = attrs.as(); + int num_segments = segment_attrs->num_segments; + + const auto* data = types[0].as(); + const auto& dshape = data->shape; + Array oshape(dshape); + oshape.Set(0, num_segments); + + // assign output type + reporter->Assign(types[2], TensorType(oshape, data->dtype)); + return true; +} + +Expr MakeSegmentMax(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_max"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_max").set_body_typed(MakeSegmentMax); + +RELAY_REGISTER_OP("segment_max") + .describe(R"doc(Computes the maximum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMax", SegmentRel); + +Expr MakeSegmentMin(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_min"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_min").set_body_typed(MakeSegmentMin); + +RELAY_REGISTER_OP("segment_min") + .describe(R"doc(Computes the minimum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMin", SegmentRel); + +Expr MakeSegmentMean(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_mean"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_mean").set_body_typed(MakeSegmentMean); + +RELAY_REGISTER_OP("segment_mean") + .describe(R"doc(Computes the mean along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMean", SegmentRel); + +Expr MakeSegmentSum(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_sum"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_sum").set_body_typed(MakeSegmentSum); + +RELAY_REGISTER_OP("segment_sum") + .describe(R"doc(Computes the sum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentSum", SegmentRel); + +Expr MakeSegmentProd(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_prod"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_prod").set_body_typed(MakeSegmentProd); + +RELAY_REGISTER_OP("segment_prod") + .describe(R"doc(Computes the prod along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentProd", SegmentRel); + } // namespace relay } // namespace tvm diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index f29450dbb604..eb4e6aba05ce 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -4398,6 +4398,95 @@ def _test_forward_rint(shape): _test_forward_rint([2, 5, 2, 5]) +def test_forwars_segment(): + """test operator segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod """ + + def _get_segment_ids(length, size): + segment_ids = [0] + for i in range(size): + if np.array(segment_ids).sum() < length: + for _ in range(i): + segment_ids.append(i) + else: + break + length = length - len(segment_ids) + for i in range(length): + segment_ids.append(segment_ids[-1]) + return np.array(segment_ids).astype("int32") + + def _test_forwars_segment(name, data_shape, size): + np_data = np.random.uniform(0, 100, size=data_shape).astype(np.float32) + tf.reset_default_graph() + in_data = tf.placeholder(tf.float32, data_shape, name="in_data") + segment_ids = _get_segment_ids(data_shape[0], size) + tf_segment_ids = tf.constant(segment_ids) + segment_op = getattr(tf.math, name) + result = segment_op(in_data, tf_segment_ids, name=name) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, True) + + _test_forwars_segment("segment_max", [100], 20) + _test_forwars_segment("segment_max", [10, 3, 4], 5) + _test_forwars_segment("segment_max", [10, 3, 4, 4], 10) + _test_forwars_segment("segment_min", [100], 20) + _test_forwars_segment("segment_min", [10, 3, 4], 5) + _test_forwars_segment("segment_min", [10, 3, 4, 4], 10) + _test_forwars_segment("segment_mean", [100], 20) + _test_forwars_segment("segment_mean", [10, 3, 4], 5) + _test_forwars_segment("segment_mean", [10, 3, 4, 4], 10) + _test_forwars_segment("segment_sum", [100], 20) + _test_forwars_segment("segment_sum", [10, 3, 4], 5) + _test_forwars_segment("segment_sum", [10, 3, 4, 4], 10) + _test_forwars_segment("segment_prod", [100], 20) + _test_forwars_segment("segment_prod", [10, 3, 4], 5) + _test_forwars_segment("segment_prod", [10, 3, 4, 4], 10) + + +def test_forwars_unsorted_segment(): + """test operator unsorted_segment_max, unsorted_segmet_min, + unsorted_segmet_mean, unsorted_segmet_sum, unsorted_segmet_prod""" + + def _get_segment_ids(length, size): + segment_ids = [0] + for i in range(size): + if np.array(segment_ids).sum() < length: + for _ in range(i): + segment_ids.append(i) + else: + break + length = length - len(segment_ids) + for i in range(length): + segment_ids.append(segment_ids[-1]) + np.random.shuffle(segment_ids) + segment_ids = np.array(segment_ids).astype("int32") + return segment_ids, segment_ids.max() + 1 + + def _test_forwars_unsorted_segment(name, data_shape, size): + np_data = np.random.uniform(0, 100, size=data_shape).astype(np.float32) + tf.reset_default_graph() + in_data = tf.placeholder(tf.float32, data_shape, name="in_data") + segment_ids, num_segments = _get_segment_ids(data_shape[0], size) + tf_segment_ids = tf.constant(segment_ids) + segment_op = getattr(tf.math, name) + result = segment_op(in_data, tf_segment_ids, num_segments, name=name) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, True) + + _test_forwars_unsorted_segment("unsorted_segment_max", [100], 20) + _test_forwars_unsorted_segment("unsorted_segment_max", [10, 3, 4], 5) + _test_forwars_unsorted_segment("unsorted_segment_max", [10, 3, 4, 4], 10) + _test_forwars_unsorted_segment("unsorted_segment_min", [100], 20) + _test_forwars_unsorted_segment("unsorted_segment_min", [10, 3, 4], 5) + _test_forwars_unsorted_segment("unsorted_segment_min", [10, 3, 4, 4], 10) + _test_forwars_unsorted_segment("unsorted_segment_mean", [100], 20) + _test_forwars_unsorted_segment("unsorted_segment_mean", [10, 3, 4], 5) + _test_forwars_unsorted_segment("unsorted_segment_mean", [10, 3, 4, 4], 10) + _test_forwars_unsorted_segment("unsorted_segment_sum", [100], 20) + _test_forwars_unsorted_segment("unsorted_segment_sum", [10, 3, 4], 5) + _test_forwars_unsorted_segment("unsorted_segment_sum", [10, 3, 4, 4], 10) + _test_forwars_unsorted_segment("unsorted_segment_prod", [100], 20) + _test_forwars_unsorted_segment("unsorted_segment_prod", [10, 3, 4], 5) + _test_forwars_unsorted_segment("unsorted_segment_prod", [10, 3, 4, 4], 10) + + def test_forward_negative(): """test tf operator Neg""" np_data = np.random.uniform(-100, 255, size=(224, 224, 3)).astype(np.float32) diff --git a/tests/python/topi/python/test_topi_math.py b/tests/python/topi/python/test_topi_math.py index c7f80033bdf3..17e33fa8ad0b 100644 --- a/tests/python/topi/python/test_topi_math.py +++ b/tests/python/topi/python/test_topi_math.py @@ -23,6 +23,8 @@ import tvm.testing import tvm.topi.testing from tvm.topi import utils +from tvm.topi import util +import tensorflow as tf def test_util(): @@ -234,8 +236,91 @@ def check_target(target): ) +_segment_implement = { + "segment_max": { + "generic": (topi.segment_max, topi.generic.schedule_segment_max), + }, + "segment_min": { + "generic": (topi.segment_min, topi.generic.schedule_segment_min), + }, + "segment_mean": { + "generic": (topi.segment_mean, topi.generic.schedule_segment_mean), + }, + "segment_sum": { + "generic": (topi.segment_sum, topi.generic.schedule_segment_sum), + }, + "segment_prod": { + "generic": (topi.segment_prod, topi.generic.schedule_segment_prod), + }, +} + + +def verify_segmet(name, data_shape, segmnet_size): + def get_segment_ids(length, size): + segment_ids = [0] + for i in range(size): + if np.array(segment_ids).sum() < length: + for _ in range(i): + segment_ids.append(i) + else: + break + length = length - len(segment_ids) + for i in range(length): + segment_ids.append(segment_ids[-1]) + return np.array(segment_ids).astype("int32") + + def get_ref_data(): + tf.reset_default_graph() + np_data = np.random.uniform(0, 100, size=data_shape).astype("float32") + segment_ids = get_segment_ids(data_shape[0], segmnet_size) + num_out = segment_ids[-1] + 1 + segment_op = getattr(tf.math, name) + out_calcu = segment_op(np_data, segment_ids) + with tf.Session() as sess: + tf_out = sess.run(out_calcu) + return np_data, segment_ids, num_out, tf_out + + device = "llvm" + ctx = tvm.context(device, 0) + print("Running on target: %s" % device) + data = te.placeholder(data_shape) + segment_ids = te.placeholder([data_shape[0]], dtype="int32") + np_data, np_segment_ids, num_out, np_out = get_ref_data() + with tvm.target.create(device): + fcompute, fschedule = tvm.topi.testing.dispatch(device, _segment_implement[name]) + out = fcompute(data, segment_ids, num_out=num_out) + s = fschedule(out) + f = tvm.build(s, [data, segment_ids, out], device) + tvm_data = tvm.nd.array(np_data, ctx=ctx) + tvm_segment_ids = tvm.nd.array(np_segment_ids, ctx=ctx) + tvm_out = tvm.nd.empty(ctx=ctx, shape=out.shape, dtype=out.dtype) + f(tvm_data, tvm_segment_ids, tvm_out) + tvm.testing.assert_allclose(tvm_out.asnumpy(), np_out, rtol=1e-4) + + +def test_segment(): + # segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod + + verify_segmet("segment_max", [100], 20) + verify_segmet("segment_max", [10, 3, 4], 5) + verify_segmet("segment_max", [10, 3, 4, 4], 10) + verify_segmet("segment_min", [100], 20) + verify_segmet("segment_min", [10, 3, 4], 5) + verify_segmet("segment_min", [10, 3, 4, 4], 10) + verify_segmet("segment_mean", [100], 20) + verify_segmet("segment_mean", [10, 3, 4], 5) + verify_segmet("segment_mean", [10, 3, 4, 4], 10) + verify_segmet("segment_sum", [100], 20) + verify_segmet("segment_sum", [10, 3, 4], 5) + verify_segmet("segment_sum", [10, 3, 4, 4], 10) + verify_segmet("segment_prod", [100], 20) + verify_segmet("segment_prod", [10, 3, 4], 5) + verify_segmet("segment_prod", [10, 3, 4, 4], 10) + + if __name__ == "__main__": test_util() test_ewise() test_cast() test_fastmath() + test_segment() From 60fc80810ac7696e1234e6fde272475d275f8f61 Mon Sep 17 00:00:00 2001 From: xp56 Date: Tue, 17 Nov 2020 14:20:42 +0800 Subject: [PATCH 2/5] fix bug --- python/tvm/topi/tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/topi/tensor.py b/python/tvm/topi/tensor.py index 01794930a8fd..d70c2b4baf03 100644 --- a/python/tvm/topi/tensor.py +++ b/python/tvm/topi/tensor.py @@ -19,7 +19,7 @@ from __future__ import absolute_import as _abs from . import cpp from .. import te, tir -from .util import get_const_tuple +from .utils import get_const_tuple def elemwise_sum(xs): From e352601e07b8918e53ca3f07c69179e9a16d39ac Mon Sep 17 00:00:00 2001 From: xp56 Date: Thu, 19 Nov 2020 16:49:58 +0800 Subject: [PATCH 3/5] Fixed suggested changes --- python/tvm/topi/generic/vision.py | 70 +--- python/tvm/topi/tensor.py | 315 +++--------------- .../frontend/tensorflow/test_forward.py | 84 ++--- tests/python/topi/python/test_topi_math.py | 199 +++++++---- 4 files changed, 213 insertions(+), 455 deletions(-) diff --git a/python/tvm/topi/generic/vision.py b/python/tvm/topi/generic/vision.py index b5d44ed8ced1..751fb78a31a4 100644 --- a/python/tvm/topi/generic/vision.py +++ b/python/tvm/topi/generic/vision.py @@ -178,7 +178,7 @@ def schedule_proposal(outs): return _default_schedule(outs, False) -def schedule_segment_max(outs): +def schedule_segment_op(outs): """Schedule for segment_max operator. Parameters @@ -193,71 +193,3 @@ def schedule_segment_max(outs): The computation schedule for the op. """ return _default_schedule(outs, False) - - -def schedule_segment_min(outs): - """Schedule for segment_min operator. - - Parameters - ---------- - outs: Array of Tensor - The computation graph description of segment_min - in the format of an array of tensors. - - Returns - ------- - s: Schedule - The computation schedule for the op. - """ - return _default_schedule(outs, False) - - -def schedule_segment_mean(outs): - """Schedule for segment_mean operator. - - Parameters - ---------- - outs: Array of Tensor - The computation graph description of segment_mean - in the format of an array of tensors. - - Returns - ------- - s: Schedule - The computation schedule for the op. - """ - return _default_schedule(outs, False) - - -def schedule_segment_sum(outs): - """Schedule for segment_sum operator. - - Parameters - ---------- - outs: Array of Tensor - The computation graph description of segment_sum - in the format of an array of tensors. - - Returns - ------- - s: Schedule - The computation schedule for the op. - """ - return _default_schedule(outs, False) - - -def schedule_segment_prod(outs): - """Schedule for segment_prod operator. - - Parameters - ---------- - outs: Array of Tensor - The computation graph description of segment_prod - in the format of an array of tensors. - - Returns - ------- - s: Schedule - The computation schedule for the op. - """ - return _default_schedule(outs, False) diff --git a/python/tvm/topi/tensor.py b/python/tvm/topi/tensor.py index d70c2b4baf03..1886554e064f 100644 --- a/python/tvm/topi/tensor.py +++ b/python/tvm/topi/tensor.py @@ -77,8 +77,10 @@ def full_like(x, fill_value): return cpp.full_like(x, fill_value) -def segment_max(data, segment_ids, num_out): - """segment_max operator. +def segment_op(data, segment_ids, num_out, opname): + """segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod + unsorted_segment_max, unsorted_segmet_min, unsorted_segmet_mean, + unsorted_segmet_sum, unsorted_segmet_prod operators. Parameters ---------- @@ -91,67 +93,8 @@ def segment_max(data, segment_ids, num_out): num_out : int number of output - Returns - ------- - out : tvm.te.Tensor - Tensor with shape determined by the segment ids. - """ - - def _segment_max(data, segment_ids, out_buf): - - ib = tir.ir_builder.create() - input_data = ib.buffer_ptr(data) - seg_ids = ib.buffer_ptr(segment_ids) - out = ib.buffer_ptr(out_buf) - - shape = get_const_tuple(data.shape) - num_segment = get_const_tuple(out_buf.shape)[0] - inner_size = 1 - for s in range(1, len(shape)): - inner_size = inner_size * shape[s] - - with ib.for_range(0, num_segment) as n: - with ib.for_range(0, inner_size) as j: - out_index = n * inner_size + j - out[out_index] = -3.4028235e38 - - with ib.for_range(0, shape[0]) as k: - with ib.if_scope(seg_ids[k] == n): - with ib.for_range(0, inner_size) as l: - out_index = n * inner_size + l - in_index = k * inner_size + l - out[out_index] = te.max(input_data[in_index], out[out_index]) - - return ib.get() - - assert len(segment_ids.shape) == 1 - - out_shape = list(get_const_tuple(data.shape)) - out_shape[0] = num_out - - out = te.extern( - out_shape, - [data, segment_ids], - lambda ins, outs: _segment_max(ins[0], ins[1], outs[0]), - dtype=data.dtype, - ) - - return out - - -def segment_min(data, segment_ids, num_out): - """segment_min operator. - - Parameters - ---------- - data : tvm.te.Tensor - input data - - segment_ids : tvm.te.Tensor - input segment ids - - num_out : int - number of output + opname : str + name of target op Returns ------- @@ -159,244 +102,64 @@ def segment_min(data, segment_ids, num_out): Tensor with shape determined by the segment ids. """ - def _segment_min(data, segment_ids, out_buf): - - ib = tir.ir_builder.create() - input_data = ib.buffer_ptr(data) - seg_ids = ib.buffer_ptr(segment_ids) - out = ib.buffer_ptr(out_buf) - - shape = get_const_tuple(data.shape) - num_segment = get_const_tuple(out_buf.shape)[0] - inner_size = 1 - for s in range(1, len(shape)): - inner_size = inner_size * shape[s] + def _max(in_data, out): + return te.max(in_data, out) - with ib.for_range(0, num_segment) as n: - with ib.for_range(0, inner_size) as j: - out_index = n * inner_size + j - out[out_index] = 3.4028235e38 + def _min(in_data, out): + return te.min(in_data, out) - with ib.for_range(0, shape[0]) as k: - with ib.if_scope(seg_ids[k] == n): - with ib.for_range(0, inner_size) as l: - out_index = n * inner_size + l - in_index = k * inner_size + l - out[out_index] = te.min(input_data[in_index], out[out_index]) - - return ib.get() + def _add(in_data, out): + return in_data + out - assert len(segment_ids.shape) == 1 - - out_shape = list(get_const_tuple(data.shape)) - out_shape[0] = num_out + def _prod(in_data, out): + return in_data * out - out = te.extern( - out_shape, - [data, segment_ids], - lambda ins, outs: _segment_min(ins[0], ins[1], outs[0]), - dtype=data.dtype, - ) + func_dict = {"max": _max, "min": _min, "mean": _add, "sum": _add, "prod": _prod} + init_dict = { + "max": -float("inf"), + "min": float("inf"), + "mean": 0.0, + "sum": 0.0, + "prod": 1.0, + } - return out - - -def segment_mean(data, segment_ids, num_out): - """segment_mean operator. - - Parameters - ---------- - data : tvm.te.Tensor - input data - - segment_ids : tvm.te.Tensor - input segment ids - - num_out : int - number of output - - Returns - ------- - out : tvm.te.Tensor - Tensor with shape determined by the segment ids. - """ - - def _segment_mean(data, segment_ids, out_buf): + def _segment_op(data, segment_ids, out_buf): + func = func_dict[opname] + init_data = init_dict[opname] ib = tir.ir_builder.create() input_data = ib.buffer_ptr(data) seg_ids = ib.buffer_ptr(segment_ids) out = ib.buffer_ptr(out_buf) - - temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") num = ib.allocate("int32", (1), name="num", scope="local") shape = get_const_tuple(data.shape) num_segment = get_const_tuple(out_buf.shape)[0] + # The number of data to be calculated for each output inner_size = 1 for s in range(1, len(shape)): inner_size = inner_size * shape[s] + # Init output tensor with ib.for_range(0, num_segment) as n: with ib.for_range(0, inner_size) as j: out_index = n * inner_size + j - out[out_index] = 0.0 + out[out_index] = init_data num[0] = 0 + # Operate on numbers with the same id with ib.for_range(0, shape[0]) as k: with ib.if_scope(seg_ids[k] == n): - temp_index[num[0]] = k - num[0] += 1 - - with ib.if_scope(num[0] > 0): - with ib.for_range(0, inner_size) as l: - out_index = n * inner_size + l - with ib.for_range(0, num[0]) as k: - in_index = temp_index[k] * inner_size + l - out[out_index] += input_data[in_index] - out[out_index] = out[out_index] / num[0] - - return ib.get() - - assert len(segment_ids.shape) == 1 - - out_shape = list(get_const_tuple(data.shape)) - out_shape[0] = num_out - - out = te.extern( - out_shape, - [data, segment_ids], - lambda ins, outs: _segment_mean(ins[0], ins[1], outs[0]), - dtype=data.dtype, - ) - - return out - - -def segment_sum(data, segment_ids, num_out): - """segment_sum operator. - - Parameters - ---------- - data : tvm.te.Tensor - input data - - segment_ids : tvm.te.Tensor - input segment ids - - num_out : int - number of output - - Returns - ------- - out : tvm.te.Tensor - Tensor with shape determined by the segment ids. - """ - - def _segment_sum(data, segment_ids, out_buf): - - ib = tir.ir_builder.create() - input_data = ib.buffer_ptr(data) - seg_ids = ib.buffer_ptr(segment_ids) - out = ib.buffer_ptr(out_buf) - - temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") - num = ib.allocate("int32", (1), name="num", scope="local") - - shape = get_const_tuple(data.shape) - num_segment = get_const_tuple(out_buf.shape)[0] - inner_size = 1 - for s in range(1, len(shape)): - inner_size = inner_size * shape[s] - - with ib.for_range(0, num_segment) as n: - with ib.for_range(0, inner_size) as j: - out_index = n * inner_size + j - out[out_index] = 0.0 - - num[0] = 0 - with ib.for_range(0, shape[0]) as k: - with ib.if_scope(seg_ids[k] == n): - temp_index[num[0]] = k - num[0] += 1 - - with ib.if_scope(num[0] > 0): - with ib.for_range(0, inner_size) as l: - out_index = n * inner_size + l - with ib.for_range(0, num[0]) as k: - in_index = temp_index[k] * inner_size + l - out[out_index] += input_data[in_index] - - return ib.get() - - assert len(segment_ids.shape) == 1 - - out_shape = list(get_const_tuple(data.shape)) - out_shape[0] = num_out - - out = te.extern( - out_shape, - [data, segment_ids], - lambda ins, outs: _segment_sum(ins[0], ins[1], outs[0]), - dtype=data.dtype, - ) - - return out - - -def segment_prod(data, segment_ids, num_out): - """segment_prod operator. - - Parameters - ---------- - data : tvm.te.Tensor - input data - - segment_ids : tvm.te.Tensor - input segment ids - - num_out : int - number of output - - Returns - ------- - out : tvm.te.Tensor - Tensor with shape determined by the segment ids. - """ - - def _segment_prod(data, segment_ids, out_buf): - - ib = tir.ir_builder.create() - input_data = ib.buffer_ptr(data) - seg_ids = ib.buffer_ptr(segment_ids) - out = ib.buffer_ptr(out_buf) - - temp_index = ib.allocate("int32", (data.shape[0]), name="temp_index", scope="local") - num = ib.allocate("int32", (1), name="num", scope="local") - - shape = get_const_tuple(data.shape) - num_segment = get_const_tuple(out_buf.shape)[0] - inner_size = 1 - for s in range(1, len(shape)): - inner_size = inner_size * shape[s] - - with ib.for_range(0, num_segment) as n: - with ib.for_range(0, inner_size) as j: - out_index = n * inner_size + j - out[out_index] = 1.0 - - num[0] = 0 - with ib.for_range(0, shape[0]) as k: - with ib.if_scope(seg_ids[k] == n): - temp_index[num[0]] = k + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + in_index = k * inner_size + l + out[out_index] = func(input_data[in_index], out[out_index]) num[0] += 1 - - with ib.if_scope(num[0] > 0): - with ib.for_range(0, inner_size) as l: - out_index = n * inner_size + l - with ib.for_range(0, num[0]) as k: - in_index = temp_index[k] * inner_size + l - out[out_index] *= input_data[in_index] + if opname == "mean": + with ib.if_scope(num[0] > 0): + with ib.for_range(0, inner_size) as s: + out_index = n * inner_size + s + out[out_index] = out[out_index] / num[0] return ib.get() @@ -408,7 +171,7 @@ def _segment_prod(data, segment_ids, out_buf): out = te.extern( out_shape, [data, segment_ids], - lambda ins, outs: _segment_prod(ins[0], ins[1], outs[0]), + lambda ins, outs: _segment_op(ins[0], ins[1], outs[0]), dtype=data.dtype, ) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index eb4e6aba05ce..faf15ed75e0e 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -4398,7 +4398,7 @@ def _test_forward_rint(shape): _test_forward_rint([2, 5, 2, 5]) -def test_forwars_segment(): +def test_forward_segment(): """test operator segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod """ def _get_segment_ids(length, size): @@ -4414,34 +4414,34 @@ def _get_segment_ids(length, size): segment_ids.append(segment_ids[-1]) return np.array(segment_ids).astype("int32") - def _test_forwars_segment(name, data_shape, size): - np_data = np.random.uniform(0, 100, size=data_shape).astype(np.float32) + def _test_forward_segment(name, data_shape, size): + np_data = np.random.uniform(-100, 100, size=data_shape).astype(np.float32) tf.reset_default_graph() in_data = tf.placeholder(tf.float32, data_shape, name="in_data") segment_ids = _get_segment_ids(data_shape[0], size) tf_segment_ids = tf.constant(segment_ids) segment_op = getattr(tf.math, name) result = segment_op(in_data, tf_segment_ids, name=name) - compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, True) - - _test_forwars_segment("segment_max", [100], 20) - _test_forwars_segment("segment_max", [10, 3, 4], 5) - _test_forwars_segment("segment_max", [10, 3, 4, 4], 10) - _test_forwars_segment("segment_min", [100], 20) - _test_forwars_segment("segment_min", [10, 3, 4], 5) - _test_forwars_segment("segment_min", [10, 3, 4, 4], 10) - _test_forwars_segment("segment_mean", [100], 20) - _test_forwars_segment("segment_mean", [10, 3, 4], 5) - _test_forwars_segment("segment_mean", [10, 3, 4, 4], 10) - _test_forwars_segment("segment_sum", [100], 20) - _test_forwars_segment("segment_sum", [10, 3, 4], 5) - _test_forwars_segment("segment_sum", [10, 3, 4, 4], 10) - _test_forwars_segment("segment_prod", [100], 20) - _test_forwars_segment("segment_prod", [10, 3, 4], 5) - _test_forwars_segment("segment_prod", [10, 3, 4, 4], 10) - - -def test_forwars_unsorted_segment(): + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, False) + + _test_forward_segment("segment_max", [100], 20) + _test_forward_segment("segment_max", [10, 3, 4], 5) + _test_forward_segment("segment_max", [10, 3, 4, 4], 10) + _test_forward_segment("segment_min", [100], 20) + _test_forward_segment("segment_min", [10, 3, 4], 5) + _test_forward_segment("segment_min", [10, 3, 4, 4], 10) + _test_forward_segment("segment_mean", [100], 20) + _test_forward_segment("segment_mean", [10, 3, 4], 5) + _test_forward_segment("segment_mean", [10, 3, 4, 4], 10) + _test_forward_segment("segment_sum", [100], 20) + _test_forward_segment("segment_sum", [10, 3, 4], 5) + _test_forward_segment("segment_sum", [10, 3, 4, 4], 10) + _test_forward_segment("segment_prod", [100], 20) + _test_forward_segment("segment_prod", [10, 3, 4], 5) + _test_forward_segment("segment_prod", [10, 3, 4, 4], 10) + + +def test_forward_unsorted_segment(): """test operator unsorted_segment_max, unsorted_segmet_min, unsorted_segmet_mean, unsorted_segmet_sum, unsorted_segmet_prod""" @@ -4460,31 +4460,31 @@ def _get_segment_ids(length, size): segment_ids = np.array(segment_ids).astype("int32") return segment_ids, segment_ids.max() + 1 - def _test_forwars_unsorted_segment(name, data_shape, size): - np_data = np.random.uniform(0, 100, size=data_shape).astype(np.float32) + def _test_forward_unsorted_segment(name, data_shape, size): + np_data = np.random.uniform(-100, 100, size=data_shape).astype(np.float32) tf.reset_default_graph() in_data = tf.placeholder(tf.float32, data_shape, name="in_data") segment_ids, num_segments = _get_segment_ids(data_shape[0], size) tf_segment_ids = tf.constant(segment_ids) segment_op = getattr(tf.math, name) result = segment_op(in_data, tf_segment_ids, num_segments, name=name) - compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, True) - - _test_forwars_unsorted_segment("unsorted_segment_max", [100], 20) - _test_forwars_unsorted_segment("unsorted_segment_max", [10, 3, 4], 5) - _test_forwars_unsorted_segment("unsorted_segment_max", [10, 3, 4, 4], 10) - _test_forwars_unsorted_segment("unsorted_segment_min", [100], 20) - _test_forwars_unsorted_segment("unsorted_segment_min", [10, 3, 4], 5) - _test_forwars_unsorted_segment("unsorted_segment_min", [10, 3, 4, 4], 10) - _test_forwars_unsorted_segment("unsorted_segment_mean", [100], 20) - _test_forwars_unsorted_segment("unsorted_segment_mean", [10, 3, 4], 5) - _test_forwars_unsorted_segment("unsorted_segment_mean", [10, 3, 4, 4], 10) - _test_forwars_unsorted_segment("unsorted_segment_sum", [100], 20) - _test_forwars_unsorted_segment("unsorted_segment_sum", [10, 3, 4], 5) - _test_forwars_unsorted_segment("unsorted_segment_sum", [10, 3, 4, 4], 10) - _test_forwars_unsorted_segment("unsorted_segment_prod", [100], 20) - _test_forwars_unsorted_segment("unsorted_segment_prod", [10, 3, 4], 5) - _test_forwars_unsorted_segment("unsorted_segment_prod", [10, 3, 4, 4], 10) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, False) + + _test_forward_unsorted_segment("unsorted_segment_max", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_max", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_max", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_min", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_min", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_min", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_mean", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_mean", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_mean", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_sum", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_sum", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_sum", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_prod", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_prod", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_prod", [10, 3, 4, 4], 10) def test_forward_negative(): diff --git a/tests/python/topi/python/test_topi_math.py b/tests/python/topi/python/test_topi_math.py index 17e33fa8ad0b..5ba95b3c5295 100644 --- a/tests/python/topi/python/test_topi_math.py +++ b/tests/python/topi/python/test_topi_math.py @@ -23,8 +23,6 @@ import tvm.testing import tvm.topi.testing from tvm.topi import utils -from tvm.topi import util -import tensorflow as tf def test_util(): @@ -236,59 +234,21 @@ def check_target(target): ) -_segment_implement = { - "segment_max": { - "generic": (topi.segment_max, topi.generic.schedule_segment_max), - }, - "segment_min": { - "generic": (topi.segment_min, topi.generic.schedule_segment_min), - }, - "segment_mean": { - "generic": (topi.segment_mean, topi.generic.schedule_segment_mean), - }, - "segment_sum": { - "generic": (topi.segment_sum, topi.generic.schedule_segment_sum), - }, - "segment_prod": { - "generic": (topi.segment_prod, topi.generic.schedule_segment_prod), - }, -} - - -def verify_segmet(name, data_shape, segmnet_size): - def get_segment_ids(length, size): - segment_ids = [0] - for i in range(size): - if np.array(segment_ids).sum() < length: - for _ in range(i): - segment_ids.append(i) - else: - break - length = length - len(segment_ids) - for i in range(length): - segment_ids.append(segment_ids[-1]) - return np.array(segment_ids).astype("int32") - - def get_ref_data(): - tf.reset_default_graph() - np_data = np.random.uniform(0, 100, size=data_shape).astype("float32") - segment_ids = get_segment_ids(data_shape[0], segmnet_size) - num_out = segment_ids[-1] + 1 - segment_op = getattr(tf.math, name) - out_calcu = segment_op(np_data, segment_ids) - with tf.Session() as sess: - tf_out = sess.run(out_calcu) - return np_data, segment_ids, num_out, tf_out - +def verify_segment(name, np_data, np_segment_ids, np_out): device = "llvm" ctx = tvm.context(device, 0) print("Running on target: %s" % device) - data = te.placeholder(data_shape) - segment_ids = te.placeholder([data_shape[0]], dtype="int32") - np_data, np_segment_ids, num_out, np_out = get_ref_data() - with tvm.target.create(device): - fcompute, fschedule = tvm.topi.testing.dispatch(device, _segment_implement[name]) - out = fcompute(data, segment_ids, num_out=num_out) + data = te.placeholder(np_data.shape) + segment_ids = te.placeholder(np_segment_ids.shape, dtype="int32") + num_out = np_segment_ids[-1] + 1 + with tvm.target.Target(device): + fcompute, fschedule = tvm.topi.testing.dispatch( + device, + { + "generic": (topi.segment_op, topi.generic.schedule_segment_op), + }, + ) + out = fcompute(data, segment_ids, num_out, name) s = fschedule(out) f = tvm.build(s, [data, segment_ids, out], device) tvm_data = tvm.nd.array(np_data, ctx=ctx) @@ -300,22 +260,125 @@ def get_ref_data(): def test_segment(): # segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod - - verify_segmet("segment_max", [100], 20) - verify_segmet("segment_max", [10, 3, 4], 5) - verify_segmet("segment_max", [10, 3, 4, 4], 10) - verify_segmet("segment_min", [100], 20) - verify_segmet("segment_min", [10, 3, 4], 5) - verify_segmet("segment_min", [10, 3, 4, 4], 10) - verify_segmet("segment_mean", [100], 20) - verify_segmet("segment_mean", [10, 3, 4], 5) - verify_segmet("segment_mean", [10, 3, 4, 4], 10) - verify_segmet("segment_sum", [100], 20) - verify_segmet("segment_sum", [10, 3, 4], 5) - verify_segmet("segment_sum", [10, 3, 4, 4], 10) - verify_segmet("segment_prod", [100], 20) - verify_segmet("segment_prod", [10, 3, 4], 5) - verify_segmet("segment_prod", [10, 3, 4, 4], 10) + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0, 1, 45, 1, 0.7, 60, 80]) + verify_segment("max", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [1, 0.8, 30, 60, 50, 80], + [2, 0.4, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ) + verify_segment("max", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0, 0.8, -25, 1, 0.7, -30, 80]) + verify_segment("min", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.0, 0.7, 1.0, 20.0, -25.0, 45.0], + [0.0, -0.9, 4.0, 21.0, 19.0, 40.0], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("min", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 0.9, 13.333333, 1.0, 0.7, 26.666666, 80.0]) + verify_segment("mean", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.5, 0.75, 15.5, 40.0, 12.5, 62.5], + [1.0, -0.25, 19.5, 41.0, 35.5, 59.5], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("mean", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 1.8, 40.0, 1.0, 0.7, 80.0, 80.0]) + verify_segment("sum", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [1.0, 1.5, 31.0, 80.0, 25.0, 125.0], + [2.0, -0.5, 39.0, 82.0, 71.0, 119.0], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("sum", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 0.8, -22500.0, 1.0, 0.7, -90000, 80]) + verify_segment("prod", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.0, 0.56, 30.0, 1200.0, -1250.0, 3600.0], + [0.0, -0.36, 140.0, 1281.0, 988.0, 3160.0], + [1.0, 0.5, 100.0, 60.0, 70, 110], + ] + ) + verify_segment("prod", np_data, segment_ids, np_result) if __name__ == "__main__": From d048e98a70285c58f706950d85b5a6fc3e668494 Mon Sep 17 00:00:00 2001 From: xp56 Date: Thu, 19 Nov 2020 20:10:27 +0800 Subject: [PATCH 4/5] set no_gpu True --- tests/python/frontend/tensorflow/test_forward.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index faf15ed75e0e..203d5ba479ee 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -4422,7 +4422,7 @@ def _test_forward_segment(name, data_shape, size): tf_segment_ids = tf.constant(segment_ids) segment_op = getattr(tf.math, name) result = segment_op(in_data, tf_segment_ids, name=name) - compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, False) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, no_gpu=True) _test_forward_segment("segment_max", [100], 20) _test_forward_segment("segment_max", [10, 3, 4], 5) @@ -4468,7 +4468,7 @@ def _test_forward_unsorted_segment(name, data_shape, size): tf_segment_ids = tf.constant(segment_ids) segment_op = getattr(tf.math, name) result = segment_op(in_data, tf_segment_ids, num_segments, name=name) - compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, False) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, no_gpu=True) _test_forward_unsorted_segment("unsorted_segment_max", [100], 20) _test_forward_unsorted_segment("unsorted_segment_max", [10, 3, 4], 5) From cfc0ccb7d1863d2cfb081d1b3d7ec68d95f637d7 Mon Sep 17 00:00:00 2001 From: xp56 Date: Thu, 14 Jan 2021 11:33:14 +0800 Subject: [PATCH 5/5] fix conflicting bug --- python/tvm/relay/op/strategy/generic.py | 115 ++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index a6ad06e544a6..8ea5c95752e5 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -1540,6 +1540,29 @@ def uniform_strategy(attrs, inputs, out_type, target): return strategy +# segment_max +def wrap_compute_segment_max(topi_compute): + """wrap segment_max topi compute""" + + def _compute_segment_max(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "max")] + + return _compute_segment_max + + +@override_native_generic_func("segment_max_strategy") +def segment_max_strategy(attrs, inputs, out_type, target): + """segment_max generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_max(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_max.generic", + ) + return strategy + + def wrap_compute_scanop(topi_compute): """Wrap scanop style topi compute""" @@ -1561,6 +1584,29 @@ def cumsum_strategy(attrs, inputs, out_type, target): return strategy +# segment_min +def wrap_compute_segment_min(topi_compute): + """wrap segment_min topi compute""" + + def _compute_segment_min(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "min")] + + return _compute_segment_min + + +@override_native_generic_func("segment_min_strategy") +def segment_min_strategy(attrs, inputs, out_type, target): + """segment_min generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_min(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_min.generic", + ) + return strategy + + @override_native_generic_func("cumprod_strategy") def cumprod_strategy(attrs, inputs, out_type, target): """cumprod generic strategy""" @@ -1573,6 +1619,29 @@ def cumprod_strategy(attrs, inputs, out_type, target): return strategy +# segment_mean +def wrap_compute_segment_mean(topi_compute): + """wrap segment_mean topi compute""" + + def _compute_segment_mean(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "mean")] + + return _compute_segment_mean + + +@override_native_generic_func("segment_mean_strategy") +def segment_mean_strategy(attrs, inputs, out_type, target): + """segment_mean generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_mean(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_mean.generic", + ) + return strategy + + def wrap_compute_unique(topi_compute): """Wrap unique topi compute""" @@ -1594,8 +1663,54 @@ def unique_strategy(attrs, inputs, out_type, target): return strategy +# segment_sum +def wrap_compute_segment_sum(topi_compute): + """wrap segment_sum topi compute""" + + def _compute_segment_sum(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "sum")] + + return _compute_segment_sum + + +@override_native_generic_func("segment_sum_strategy") +def segment_sum_strategy(attrs, inputs, out_type, target): + """segment_sum generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_sum(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_sum.generic", + ) + return strategy + + @generic_func def schedule_transpose(attrs, outs, target): """schedule transpose""" with target: return schedule_injective(attrs, outs, target) + + +# segment_prod +def wrap_compute_segment_prod(topi_compute): + """wrap segment_prod topi compute""" + + def _compute_segment_prod(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "prod")] + + return _compute_segment_prod + + +@override_native_generic_func("segment_prod_strategy") +def segment_prod_strategy(attrs, inputs, out_type, target): + """segment_prod generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_prod(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_prod.generic", + ) + return strategy