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/25 21:07:32 UTC
[GitHub] [tvm] mbrookhart opened a new pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
mbrookhart opened a new pull request #6978:
URL: https://github.com/apache/tvm/pull/6978
@zhiics @icemelon9 @jroesch
After several weeks of intermittently banging my head against dynamic topi composition and a direct tir topk op for GPU, I spent this morning banging out a legalization-based solution. What do you guys think?
Also, Happy Thanksgiving!
----------------------------------------------------------------
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] [tvm] zhiics commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r532856178
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ On GPU, we don't directly implement a topi kernel. Instead, we combine
+ topi sort and topi strided slice to implement topk. This topi-level composition
+ doesn't work with the dynamic op. To support the dynamic op on gpu, we instead
+ legalize it into the same logic (sort + strided slice) in relay to take advantage
+ of the VM's dynamic shape capabilities
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
+ The legalized expr
+ """
+ data = tvm.relay.var(
+ "dyn_topk_data", shape=inputs[0].checked_type.shape, dtype=inputs[0].checked_type.dtype
+ )
+ k = tvm.relay.var(
+ "dyn_topk_k", shape=inputs[1].checked_type.shape, dtype=inputs[1].checked_type.dtype
+ )
+ sort_out = tvm.relay.sort(data, axis=attrs.axis, is_ascend=attrs.is_ascend)
+ argsort_out = tvm.relay.argsort(
+ data, axis=attrs.axis, is_ascend=attrs.is_ascend, dtype=attrs.dtype
+ )
+ dshape = tvm.relay.shape_of(data)
+
+ axis = tvm.relay.const([attrs.axis])
+
+ zero = tvm.relay.const([0])
+ one = tvm.relay.const([1])
+ rank = tvm.relay.shape_of(dshape)
+
+ normalized_axis = tvm.relay.where(axis < zero, axis + rank, axis)
Review comment:
I didn't look into the impl much, but would these ops impact the performance?
----------------------------------------------------------------
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] [tvm] mbrookhart closed pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart closed pull request #6978:
URL: https://github.com/apache/tvm/pull/6978
----------------------------------------------------------------
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] [tvm] mbrookhart commented on pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-742003328
@kevinthesun Can you re-review without the topk legalization?
----------------------------------------------------------------
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] [tvm] zhiics commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
zhiics commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-736008580
cc @kevinthesun, PTAL when you have time
----------------------------------------------------------------
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] [tvm] kevinthesun commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
kevinthesun commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-737571529
https://github.com/apache/tvm/pull/7018
Fix for dynamic strided slice in topi to enable topk.
----------------------------------------------------------------
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] [tvm] masahi commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r544676738
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -316,6 +316,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with teh same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], axis, is_ascend, indices_out=outs[1]),
+ out_buffers=[value_buf, indices_buf],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+ return out
+
+
+def sort_thrust(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: tvm.tir.call_packed(
+ "tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend
+ ),
+ out_buffers=[value_buf, indices_buf],
Review comment:
Yeah that would be nice
----------------------------------------------------------------
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] [tvm] kevinthesun commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
kevinthesun commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-736129235
For dynamic topk, can we use thrust? IMO after that the only work left is dynamic strided_slice in topi. Not quite sure whether we need to do this in relay level.
----------------------------------------------------------------
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] [tvm] kevinthesun commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
kevinthesun commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r533639174
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
Review comment:
I also get segfault using current topi dynamic strided slice to cut topk result. I'll spend some time on it.
----------------------------------------------------------------
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] [tvm] mbrookhart commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-741025397
Okay, I'll reopen and rip out the legalization.
----------------------------------------------------------------
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] [tvm] masahi commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r544666115
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -316,6 +316,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with teh same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], axis, is_ascend, indices_out=outs[1]),
+ out_buffers=[value_buf, indices_buf],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+ return out
+
+
+def sort_thrust(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
Review comment:
This also needs updating
----------------------------------------------------------------
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] [tvm] mbrookhart commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r544675262
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -316,6 +316,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with teh same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], axis, is_ascend, indices_out=outs[1]),
+ out_buffers=[value_buf, indices_buf],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+ return out
+
+
+def sort_thrust(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: tvm.tir.call_packed(
+ "tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend
+ ),
+ out_buffers=[value_buf, indices_buf],
Review comment:
Hmm, you're right, we don't strictly need that. Want me to add a todo?
----------------------------------------------------------------
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] [tvm] masahi merged pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi merged pull request #6978:
URL: https://github.com/apache/tvm/pull/6978
----------------------------------------------------------------
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] [tvm] mbrookhart commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-740143976
closing in favor of #7018
----------------------------------------------------------------
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] [tvm] kevinthesun commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
kevinthesun commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r532998223
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
Review comment:
I suggest we improve gpu dynamic strided_slice and it helps both dynamic topk and argwhere. For sorting, we can still rely on existing thrust routine.
----------------------------------------------------------------
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] [tvm] masahi commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r543995039
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -315,6 +315,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
Review comment:
This seems the doc for argsort
----------------------------------------------------------------
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] [tvm] masahi commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r544668080
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -316,6 +316,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with teh same shape as the input data.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: sort_ir(ins[0], outs[0], axis, is_ascend, indices_out=outs[1]),
+ out_buffers=[value_buf, indices_buf],
+ name="sort_gpu",
+ tag="sort_gpu",
+ )[0]
+ return out
+
+
+def sort_thrust(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ dtype = "float32"
+
+ ndim = len(data.shape)
+ axis = ndim + axis if axis < 0 else axis
+
+ if axis != ndim - 1:
+ # Prepare for sorting along axis -1.
+ axes = swap(list(range(ndim)), axis)
+ data = transpose(data, axes)
+
+ value_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8)
+ indices_buf = tvm.tir.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8)
+ out = te.extern(
+ [data.shape, data.shape],
+ [data],
+ lambda ins, outs: tvm.tir.call_packed(
+ "tvm.contrib.thrust.sort", ins[0], outs[0], outs[1], is_ascend
+ ),
+ out_buffers=[value_buf, indices_buf],
Review comment:
There is no need to pass `indices_buf` for normal sort, but `tvm.contrib.thrust.sort` always expects `indices_buf` to be passed in... What we call `tvm.contrib.thrust.sort` is really argsort.
For optimal performance we should have the true `tvm.contrib.thrust.sort` someday.
----------------------------------------------------------------
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] [tvm] anijain2305 commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
anijain2305 commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-736133076
I tried this PR, I see the following error when used with dynamic input shape
~~~
File "/home/ubuntu/workplace/tvm/t1/tvm/src/relay/backend/compile_engine.cc", line 508
TVMError:
---------------------------------------------------------------
An internal invariant was violated during the execution of TVM.
Please read TVM's error reporting guidelines.
More details can be found here: https://discuss.tvm.ai/t/error-reporting/7793.
---------------------------------------------------------------
Check failed: fshape_func.count(op) > 0 (0 vs. 0) : Internal error, cannot find ShapeFunc for sort
~~~
----------------------------------------------------------------
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] [tvm] mbrookhart commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r533605262
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
Review comment:
Any tips for debugging generated code? I have a branch where I've done this a a topi composition, it passes if I compile with AddressSanitizer, but segfaults with a normal build. @zhiics attempted to do argwhere, but gets compilation errors with topi dropping variables.
----------------------------------------------------------------
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] [tvm] mbrookhart commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r533605262
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
Review comment:
Any tips for debugging generated code? I have a branch where I've done this a a topi composition, it passes if I compile with AddressSanitizer, but segfaults with a normal build in the generated code. @zhiics attempted to do argwhere, but gets compilation errors with topi dropping variables.
----------------------------------------------------------------
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] [tvm] zhiics commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r532856178
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ On GPU, we don't directly implement a topi kernel. Instead, we combine
+ topi sort and topi strided slice to implement topk. This topi-level composition
+ doesn't work with the dynamic op. To support the dynamic op on gpu, we instead
+ legalize it into the same logic (sort + strided slice) in relay to take advantage
+ of the VM's dynamic shape capabilities
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
+ The legalized expr
+ """
+ data = tvm.relay.var(
+ "dyn_topk_data", shape=inputs[0].checked_type.shape, dtype=inputs[0].checked_type.dtype
+ )
+ k = tvm.relay.var(
+ "dyn_topk_k", shape=inputs[1].checked_type.shape, dtype=inputs[1].checked_type.dtype
+ )
+ sort_out = tvm.relay.sort(data, axis=attrs.axis, is_ascend=attrs.is_ascend)
+ argsort_out = tvm.relay.argsort(
+ data, axis=attrs.axis, is_ascend=attrs.is_ascend, dtype=attrs.dtype
+ )
+ dshape = tvm.relay.shape_of(data)
+
+ axis = tvm.relay.const([attrs.axis])
+
+ zero = tvm.relay.const([0])
+ one = tvm.relay.const([1])
+ rank = tvm.relay.shape_of(dshape)
+
+ normalized_axis = tvm.relay.where(axis < zero, axis + rank, axis)
Review comment:
Would these ops impact the performance?
----------------------------------------------------------------
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] [tvm] mbrookhart commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-736714836
@anijain2305 I'll add a shape func for sort and a regression test
@kevinthesun I added a thrust implementation of sort, and tested dynamic topk with thrust enabled, it worked just fine.
----------------------------------------------------------------
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] [tvm] kevinthesun edited a comment on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
kevinthesun edited a comment on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-736129235
For dynamic topk, can we use Thrust? IMO after sorting with Thrust the only work left is dynamic strided_slice in topi. Not quite sure whether we need to do this in relay level.
----------------------------------------------------------------
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] [tvm] masahi commented on a change in pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r544665982
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -316,6 +316,89 @@ def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32")
return out[1]
+def sort(data, axis=-1, is_ascend=1):
+ """Performs sorting along the given axis and returns an array of
+ sorted values with teh same shape as the input data.
Review comment:
typo teh
----------------------------------------------------------------
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] [tvm] mbrookhart commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r532864974
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,82 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ On GPU, we don't directly implement a topi kernel. Instead, we combine
+ topi sort and topi strided slice to implement topk. This topi-level composition
+ doesn't work with the dynamic op. To support the dynamic op on gpu, we instead
+ legalize it into the same logic (sort + strided slice) in relay to take advantage
+ of the VM's dynamic shape capabilities
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
+ The legalized expr
+ """
+ data = tvm.relay.var(
+ "dyn_topk_data", shape=inputs[0].checked_type.shape, dtype=inputs[0].checked_type.dtype
+ )
+ k = tvm.relay.var(
+ "dyn_topk_k", shape=inputs[1].checked_type.shape, dtype=inputs[1].checked_type.dtype
+ )
+ sort_out = tvm.relay.sort(data, axis=attrs.axis, is_ascend=attrs.is_ascend)
+ argsort_out = tvm.relay.argsort(
+ data, axis=attrs.axis, is_ascend=attrs.is_ascend, dtype=attrs.dtype
+ )
+ dshape = tvm.relay.shape_of(data)
+
+ axis = tvm.relay.const([attrs.axis])
+
+ zero = tvm.relay.const([0])
+ one = tvm.relay.const([1])
+ rank = tvm.relay.shape_of(dshape)
+
+ normalized_axis = tvm.relay.where(axis < zero, axis + rank, axis)
Review comment:
Most of the small elementwise ops to get shapes will be fused, and this matches what I did in TOPI pretty closely. The only real performance concern I have with this is that it does sort and argsort separately. The way the cuda
kernels are written, I think that ends up sorting the input data twice.
----------------------------------------------------------------
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] [tvm] jwfromm commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
jwfromm commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r530650365
##########
File path: python/tvm/topi/sort.py
##########
@@ -161,3 +203,25 @@ def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"):
tag="topk_cpu",
)
return out
+
+
+@tvm.target.generic_func
+def dyn_topk_legalize(attrs, inputs, types):
+ """Legalizes dyn.topk op.
Review comment:
Can you add a brief description of why legalization is needed here and what transformation is applied?
----------------------------------------------------------------
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] [tvm] masahi commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-741024087
yes hummingbird people have asked for a converter for pytorch sort op, that would be exactly what I need. Thanks.
cc @interesaaat
----------------------------------------------------------------
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] [tvm] zhiics commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r530659622
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,76 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
Review comment:
I see. Thanks. I didn't realize this before.
----------------------------------------------------------------
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] [tvm] mbrookhart commented on pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-741008268
@masahi @zhiics @jwfromm
@kevinthesun found a fix for dynamic topk, so I don't think I need this legalization anymore, but do you think the sort op at the relay level is still useful?
----------------------------------------------------------------
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] [tvm] mbrookhart commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r530658429
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,76 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
Review comment:
This is how the other legalization functions are structured, which admittedly seems odd to me, but they have functions in topi implementing relay passes and calling relay, see cuda conv2d here: https://github.com/apache/tvm/blob/c58b20bbd8e6ac7128e82994624d922ee6b8a069/python/tvm/topi/cuda/conv2d_alter_op.py#L271-L348
I did attempt to do this more directly in the relay files, and got a lot of circular imports, so I decided to follow the example of the other ops.
----------------------------------------------------------------
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] [tvm] zhiics commented on a change in pull request #6978: [Relay][Topi][Dynamic] Add a Sort op and use the legalization pass to perform dynamic topk on GPU
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#discussion_r530653838
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -593,3 +694,76 @@ def schedule_topk(outs):
The computation schedule for the op.
"""
return _schedule_sort(outs)
+
+
+def _dyn_topk_legalize(attrs, inputs, arg_types):
+ """Legalizes dyn.topk op.
+
+ Parameters
+ ----------
+ attrs : tvm.ir.Attrs
+ Attributes of current convolution
+ inputs : list of tvm.relay.Expr
+ The args of the Relay expr to be legalized
+ types : list of types
+ List of input and output types
+
+ Returns
+ -------
+ result : tvm.relay.Expr
Review comment:
wouldn't be wired to access relay in topi since topi should conceptually lower level?
----------------------------------------------------------------
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] [tvm] masahi commented on pull request #6978: [Relay][Topi]Add Sort Op to Relay
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #6978:
URL: https://github.com/apache/tvm/pull/6978#issuecomment-747728047
Thanks @mbrookhart @zhiics @kevinthesun @jwfromm
----------------------------------------------------------------
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