[GitHub] asmushetzel opened a new pull request #10997: speed up of topk-operator

[GitBox <gi...@apache.org>]

asmushetzel opened a new pull request #10997: speed up of topk-operator
URL: https://github.com/apache/incubator-mxnet/pull/10997
 
 
   ## Description ##
   Speedup of topk-operator. Specifically
   - on CPU, avoid doing 3 full sorts and instead do a single partial_sort resp a single full sort when K is reasonably large.
   - on CPU, parallelize over the batch size
   - on GPU, have a faster methods for very small K (which only works on shared memory). This is an important special case for example in beam search applications. 
   
   It is hard to put exact criteria in such code on when to switch between different versions of sort algorithms. The CPU criteria should be reasonable stable, the GPU one as well. Tests on Volta showed that a back2back full sort will start to outperform the special small-K version for some cases only (batch size 1, k > 500,000) and not by much. We can put more logic in there and try to be smarter for such cases, but not sure that it is worth it.
   
   The work here is driven by the fact that the performance in cases with small-K of the existing code is non-acceptable. We have an important application where we resorted actually to numpy to get around the current performance issues. 
   
   
   
   ## Checklist ##
   ### Essentials ###
   Please feel free to remove inapplicable items for your PR.
   - [ x] Changes are complete (i.e. I finished coding on this PR)
   - [x ] All changes have test coverage:
   - [x ] Code is well-documented: 
   - [ x] To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
   

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on 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


With regards,
Apache Git Services