Skip to content

Commit

Permalink
Use radix sort when size is large enough so that one-group sorter won…
Browse files Browse the repository at this point in the history
…t be needed
  • Loading branch information
oleksandr-pavlyk committed Dec 23, 2024
1 parent 51ead2b commit 3e5e303
Showing 1 changed file with 10 additions and 6 deletions.
16 changes: 10 additions & 6 deletions dpctl/tensor/libtensor/source/sorting/topk.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -112,13 +112,17 @@ sycl::event topk_caller(sycl::queue &exec_q,
{
if constexpr (use_radix_sort<argTy>::value) {
using dpctl::tensor::kernels::topk_radix_impl;
auto ascending = !largest;
return topk_radix_impl<argTy, IndexTy>(
exec_q, iter_nelems, axis_nelems, k, ascending, arg_cp, vals_cp,
inds_cp, iter_arg_offset, iter_vals_offset, iter_inds_offset,
axis_arg_offset, axis_vals_offset, axis_inds_offset, depends);
const auto ascending = !largest;

if (axis_nelems >= 16384) {
return topk_radix_impl<argTy, IndexTy>(
exec_q, iter_nelems, axis_nelems, k, ascending, arg_cp, vals_cp,
inds_cp, iter_arg_offset, iter_vals_offset, iter_inds_offset,
axis_arg_offset, axis_vals_offset, axis_inds_offset, depends);
}
}
else {

{
using dpctl::tensor::kernels::topk_merge_impl;
if (largest) {
using CompTy =
Expand Down

0 comments on commit 3e5e303

Please sign in to comment.