[Relay][Topi]Add Sort Op to Relay#6978
Conversation
|
cc @kevinthesun, PTAL when you have time |
|
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. |
|
I tried this PR, I see the following error when used with dynamic input shape |
| return _schedule_sort(outs) | ||
|
|
||
|
|
||
| def _dyn_topk_legalize(attrs, inputs, arg_types): |
There was a problem hiding this 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.
There was a problem hiding this 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.
There was a problem hiding this comment.
I also get segfault using current topi dynamic strided slice to cut topk result. I'll spend some time on it.
|
@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. |
|
#7018 |
|
closing in favor of #7018 |
|
@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? |
|
yes hummingbird people have asked for a converter for pytorch sort op, that would be exactly what I need. Thanks. cc @interesaaat |
|
Okay, I'll reopen and rip out the legalization. |
|
@kevinthesun Can you re-review without the topk legalization? |
| 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], |
There was a problem hiding this 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.
There was a problem hiding this comment.
Hmm, you're right, we don't strictly need that. Want me to add a todo?
masahi
left a comment
There was a problem hiding this comment.
Please fix remaining doc issues, otherwise LGTM.
|
Thanks @mbrookhart @zhiics @kevinthesun @jwfromm |
* Add sort op to relay * fix lint * fix sort docstring * fix docs * add TODO, shape_func, cleanup * add dynamic tests for sort and argsort
* Add sort op to relay * fix lint * fix sort docstring * fix docs * add TODO, shape_func, cleanup * add dynamic tests for sort and argsort
* Add sort op to relay * fix lint * fix sort docstring * fix docs * add TODO, shape_func, cleanup * add dynamic tests for sort and argsort
* Add sort op to relay * fix lint * fix sort docstring * fix docs * add TODO, shape_func, cleanup * add dynamic tests for sort and argsort
| Parameters | ||
| ---------- | ||
| outs: Array of Tensor | ||
| The computation graph description of argsort |
Re purposing PR for Relay Sort