From d50092a7745a7f82902295ff672ca22ce544a247 Mon Sep 17 00:00:00 2001
From: Oleksandr Pavlyk <oleksandr.pavlyk@intel.com>
Date: Sun, 22 Dec 2024 18:48:25 -0600
Subject: [PATCH] Use radix sort when size is large enough so that one-group
 sorter wont be needed

---
 dpctl/tensor/libtensor/source/sorting/topk.cpp | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/dpctl/tensor/libtensor/source/sorting/topk.cpp b/dpctl/tensor/libtensor/source/sorting/topk.cpp
index dea20fd494..6f1821e45e 100644
--- a/dpctl/tensor/libtensor/source/sorting/topk.cpp
+++ b/dpctl/tensor/libtensor/source/sorting/topk.cpp
@@ -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 =