speed up of topk-operator

[asmushetzel]

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

[marcoabreu]

Hi Asmus, thanks for the great work! Could you please elaborate on the test coverage? I don't think we have these exact test cases covered (the different values of N, K and batch sizes specifically). Would be great if you could add some tests for that

[asmushetzel]

test_order() in test_operator.py tests full sort, partial sort, different choice of axis, ascending/descending, different K (including cases where K is greater/bigger than the threshold in GPUs) etc. It is quite comprehensive and IMO sufficient.

[sxjscience]

Great!