You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/11/17 02:52:17 UTC
[GitHub] [incubator-tvm] alter-xp opened a new pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
alter-xp opened a new pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928
* 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
@giuseros @siju-samuel
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#issuecomment-731956750
when testing for unstored_segmnet_mean, a single constant data will as input for strideslice op in tf graph. But tf frontend in tvm not support this situation. fix this in pr https://github.com/apache/incubator-tvm/pull/6949.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526696625
##########
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:
👌
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] giuseros commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
giuseros commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r525237522
##########
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):
Review comment:
Could you add a bit more comments through this file? This would make it easier to read and also more future proof
##########
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
Review comment:
Is there a better way to do this? In theory you could pre-compute also the segment sizes. Or you could calculate the sizes on the fly. At least, I would put something like: `sys.float_info.min`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526692645
##########
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:
At that time, it was to ensure that the division by 0 would not occur during the calculation process. Now modified to only be used in the `mean`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526695990
##########
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:
Okay, thanks for all your comments. All suggestions have been revised.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526691076
##########
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:
👌
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp edited a comment on pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp edited a comment on pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#issuecomment-731956750
when testing for unsorted_segment_mean, a single constant data will as input for strideslice op in tf graph. But tf frontend in tvm not support this situation. fix this in pr https://github.com/apache/incubator-tvm/pull/6949.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] giuseros commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526690293
##########
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
Review comment:
`sys.float_info.min` is a number like 2.2250738585072014e-308, , which is not very suitable here, because it is always greater than 0. I replaced it with `float("inf")`
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops
Posted by GitBox <gi...@apache.org>.
alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526690254
##########
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):
Review comment:
👌
##########
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
Review comment:
`sys.float_info.min` is a number like 2.2250738585072014e-308, , which is not very suitable here, because it is always greater than 0. I replaced it with float("int")
----------------------------------------------------------------
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:
users@infra.apache.org