giuseros commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526179923
##########
File path: python/tvm/topi/tensor.py
##########
@@ -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()
Review comment:
It looks like these function share a common implementation. Could you
write a single `_segment_op` with a string `op` parameters and add some logic
inside this function?
##########
File path: python/tvm/topi/tensor.py
##########
@@ -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):
Review comment:
So, if the segment is not present we omit it? Why we don't do this for
max and min? Some explanation of this could would be useful, I think
##########
File path: tests/python/topi/python/test_topi_math.py
##########
@@ -226,8 +228,91 @@ def check_device(device):
test_apply(topi.fast_tanh, "fast_tanh", np.tanh, low=-10, high=10,
step=0.01)
+_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):
Review comment:
Typo verify_segmet -> verify_segment
##########
File path: tests/python/topi/python/test_topi_math.py
##########
@@ -226,8 +228,91 @@ def check_device(device):
test_apply(topi.fast_tanh, "fast_tanh", np.tanh, low=-10, high=10,
step=0.01)
+_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():
Review comment:
Instead of using tflite to get the reference data, could you add some
reference data manually? In this way tests can be independent and also show
what the goal of each function is (also, you already added reference tests in
`test_forward.py`). What do you think?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
