Use radix sort when size is large enough so that one-group sorter won…

…t be needed
IntelPython · Dec 23, 2024 · de53055 · de53055
1 parent 51ead2b
commit de53055
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dpctl/tensor/libtensor/source/sorting/topk.cpp b/dpctl/tensor/libtensor/source/sorting/topk.cpp
@@ -110,9 +110,9 @@ sycl::event topk_caller(sycl::queue &exec_q,
                         py::ssize_t axis_inds_offset,
                         const std::vector<sycl::event> &depends)
 {
-    if constexpr (use_radix_sort<argTy>::value) {
+    if (use_radix_sort<argTy>::value && (axis_nelems >= 16384)) {
         using dpctl::tensor::kernels::topk_radix_impl;
-        auto ascending = !largest;
+        const 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,