Feature/radix sort

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Improved performance of copying operation to C-/F-contig array, with optimization for batch of square matrices [gh-1850](https://github.com/IntelPython/dpctl/pull/1850)
 * Improved performance of `tensor.argsort` function for all types [gh-1859](https://github.com/IntelPython/dpctl/pull/1859)
 * Improved performance of `tensor.sort` and `tensor.argsort` for short arrays in the range [16, 64] elements [gh-1866](https://github.com/IntelPython/dpctl/pull/1866)
+* Implement radix sort algorithm to be used in `dpt.sort` and `dpt.argsort` [gh-1867](https://github.com/IntelPython/dpctl/pull/1867)
 
 ### Fixed
 

dpctl/tensor/CMakeLists.txt

@@ -116,6 +116,10 @@ set(_sorting_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/argsort.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/searchsorted.cpp
 )
+set(_sorting_radix_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/radix_sort.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/radix_argsort.cpp
+)
 set(_static_lib_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/simplify_iteration_space.cpp
 )
@@ -151,6 +155,10 @@ set(_tensor_sorting_impl_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_sorting.cpp
     ${_sorting_sources}
 )
+set(_tensor_sorting_radix_impl_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_sorting_radix.cpp
+    ${_sorting_radix_sources}
+)
 set(_linalg_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/elementwise_functions_type_utils.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/linalg_functions/dot.cpp
@@ -160,10 +168,10 @@ set(_tensor_linalg_impl_sources
     ${_linalg_sources}
 )
 set(_accumulator_sources
-${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/accumulators_common.cpp
-${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_logsumexp.cpp
-${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_prod.cpp
-${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_sum.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/accumulators_common.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_logsumexp.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_prod.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/accumulators/cumulative_sum.cpp
 )
 set(_tensor_accumulation_impl_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/tensor_accumulation.cpp
@@ -205,6 +213,12 @@ add_sycl_to_target(TARGET ${python_module_name} SOURCES ${_tensor_sorting_impl_s
 target_link_libraries(${python_module_name} PRIVATE ${_static_lib_trgt})
 list(APPEND _py_trgts ${python_module_name})
 
+set(python_module_name _tensor_sorting_radix_impl)
+pybind11_add_module(${python_module_name} MODULE ${_tensor_sorting_radix_impl_sources})
+add_sycl_to_target(TARGET ${python_module_name} SOURCES ${_tensor_sorting_radix_impl_sources})
+target_link_libraries(${python_module_name} PRIVATE ${_static_lib_trgt})
+list(APPEND _py_trgts ${python_module_name})
+
 set(python_module_name _tensor_linalg_impl)
 pybind11_add_module(${python_module_name} MODULE ${_tensor_linalg_impl_sources})
 add_sycl_to_target(TARGET ${python_module_name} SOURCES ${_tensor_linalg_impl_sources})

dpctl/tensor/_sorting.py

@@ -25,11 +25,26 @@
     _sort_ascending,
     _sort_descending,
 )
+from ._tensor_sorting_radix_impl import (
+    _radix_argsort_ascending,
+    _radix_argsort_descending,
+    _radix_sort_ascending,
+    _radix_sort_descending,
+    _radix_sort_dtype_supported,
+)
 
 __all__ = ["sort", "argsort"]
 
 
-def sort(x, /, *, axis=-1, descending=False, stable=True):
+def _get_mergesort_impl_fn(descending):
+    return _sort_descending if descending else _sort_ascending
+
+
+def _get_radixsort_impl_fn(descending):
+    return _radix_sort_descending if descending else _radix_sort_ascending
+
+
+def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):
     """sort(x, axis=-1, descending=False, stable=True)
 
     Returns a sorted copy of an input array `x`.
@@ -49,7 +64,10 @@ def sort(x, /, *, axis=-1, descending=False, stable=True):
             relative order of `x` values which compare as equal. If `False`,
             the returned array may or may not maintain the relative order of
             `x` values which compare as equal. Default: `True`.
-
+        kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
+            Sorting algorithm. The default is `"stable"`, which uses parallel
+            merge-sort or parallel radix-sort algorithms depending on the
+            array data type.
     Returns:
         usm_ndarray:
             a sorted array. The returned array has the same data type and
@@ -74,10 +92,33 @@ def sort(x, /, *, axis=-1, descending=False, stable=True):
             axis,
         ]
         arr = dpt.permute_dims(x, perm)
+    if kind is None:
+        kind = "stable"
+    if not isinstance(kind, str) or kind not in [
+        "stable",
+        "radixsort",
+        "mergesort",
+    ]:
+        raise ValueError(
+            "Unsupported kind value. Expected 'stable', 'mergesort', "
+            f"or 'radixsort', but got '{kind}'"
+        )
+    if kind == "mergesort":
+        impl_fn = _get_mergesort_impl_fn(descending)
+    elif kind == "radixsort":
+        if _radix_sort_dtype_supported(x.dtype.num):
+            impl_fn = _get_radixsort_impl_fn(descending)
+        else:
+            raise ValueError(f"Radix sort is not supported for {x.dtype}")
+    else:
+        dt = x.dtype
+        if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:
+            impl_fn = _get_radixsort_impl_fn(descending)
+        else:
+            impl_fn = _get_mergesort_impl_fn(descending)
     exec_q = x.sycl_queue
     _manager = du.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
-    impl_fn = _sort_descending if descending else _sort_ascending
     if arr.flags.c_contiguous:
         res = dpt.empty_like(arr, order="C")
         ht_ev, impl_ev = impl_fn(
@@ -109,7 +150,15 @@ def sort(x, /, *, axis=-1, descending=False, stable=True):
     return res
 
 
-def argsort(x, axis=-1, descending=False, stable=True):
+def _get_mergeargsort_impl_fn(descending):
+    return _argsort_descending if descending else _argsort_ascending
+
+
+def _get_radixargsort_impl_fn(descending):
+    return _radix_argsort_descending if descending else _radix_argsort_ascending
+
+
+def argsort(x, axis=-1, descending=False, stable=True, kind=None):
     """argsort(x, axis=-1, descending=False, stable=True)
 
     Returns the indices that sort an array `x` along a specified axis.
@@ -129,6 +178,10 @@ def argsort(x, axis=-1, descending=False, stable=True):
             relative order of `x` values which compare as equal. If `False`,
             the returned array may or may not maintain the relative order of
             `x` values which compare as equal. Default: `True`.
+        kind (Optional[Literal["stable", "mergesort", "radixsort"]]):
+            Sorting algorithm. The default is `"stable"`, which uses parallel
+            merge-sort or parallel radix-sort algorithms depending on the
+            array data type.
 
     Returns:
         usm_ndarray:
@@ -157,10 +210,33 @@ def argsort(x, axis=-1, descending=False, stable=True):
             axis,
         ]
         arr = dpt.permute_dims(x, perm)
+    if kind is None:
+        kind = "stable"
+    if not isinstance(kind, str) or kind not in [
+        "stable",
+        "radixsort",
+        "mergesort",
+    ]:
+        raise ValueError(
+            "Unsupported kind value. Expected 'stable', 'mergesort', "
+            f"or 'radixsort', but got '{kind}'"
+        )
+    if kind == "mergesort":
+        impl_fn = _get_mergeargsort_impl_fn(descending)
+    elif kind == "radixsort":
+        if _radix_sort_dtype_supported(x.dtype.num):
+            impl_fn = _get_radixargsort_impl_fn(descending)
+        else:
+            raise ValueError(f"Radix sort is not supported for {x.dtype}")
+    else:
+        dt = x.dtype
+        if dt in [dpt.bool, dpt.uint8, dpt.int8, dpt.int16, dpt.uint16]:
+            impl_fn = _get_radixargsort_impl_fn(descending)
+        else:
+            impl_fn = _get_mergeargsort_impl_fn(descending)
     exec_q = x.sycl_queue
     _manager = du.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
-    impl_fn = _argsort_descending if descending else _argsort_ascending
     index_dt = ti.default_device_index_type(exec_q)
     if arr.flags.c_contiguous:
         res = dpt.empty_like(arr, dtype=index_dt, order="C")

dpctl/tensor/libtensor/include/kernels/sorting/sort.hpp → dpctl/tensor/libtensor/include/kernels/sorting/merge_sort.hpp

@@ -32,7 +32,7 @@
 #include <vector>
 
 #include "kernels/dpctl_tensor_types.hpp"
-#include "kernels/sorting/sort_detail.hpp"
+#include "kernels/sorting/search_sorted_detail.hpp"
 
 namespace dpctl
 {
@@ -41,9 +41,11 @@ namespace tensor
 namespace kernels
 {
 
-namespace sort_detail
+namespace merge_sort_detail
 {
 
+using namespace dpctl::tensor::kernels::search_sorted_detail;
+
 /*! @brief Merge two contiguous sorted segments */
 template <typename InAcc, typename OutAcc, typename Compare>
 void merge_impl(const std::size_t offset,
@@ -699,18 +701,7 @@ merge_sorted_block_contig_impl(sycl::queue &q,
     return dep_ev;
 }
 
-} // end of namespace sort_detail
-
-typedef sycl::event (*sort_contig_fn_ptr_t)(sycl::queue &,
-                                            size_t,
-                                            size_t,
-                                            const char *,
-                                            char *,
-                                            ssize_t,
-                                            ssize_t,
-                                            ssize_t,
-                                            ssize_t,
-                                            const std::vector<sycl::event> &);
+} // end of namespace merge_sort_detail
 
 template <typename argTy, typename Comp = std::less<argTy>>
 sycl::event stable_sort_axis1_contig_impl(
@@ -741,8 +732,8 @@ sycl::event stable_sort_axis1_contig_impl(
     if (sort_nelems < sequential_sorting_threshold) {
         // equal work-item sorts entire row
         sycl::event sequential_sorting_ev =
-            sort_detail::sort_base_step_contig_impl<const argTy *, argTy *,
-                                                    Comp>(
+            merge_sort_detail::sort_base_step_contig_impl<const argTy *,
+                                                          argTy *, Comp>(
                 exec_q, iter_nelems, sort_nelems, arg_tp, res_tp, comp,
                 sort_nelems, depends);
 
@@ -753,16 +744,16 @@ sycl::event stable_sort_axis1_contig_impl(
 
         // Sort segments of the array
         sycl::event base_sort_ev =
-            sort_detail::sort_over_work_group_contig_impl<const argTy *,
-                                                          argTy *, Comp>(
+            merge_sort_detail::sort_over_work_group_contig_impl<const argTy *,
+                                                                argTy *, Comp>(
                 exec_q, iter_nelems, sort_nelems, arg_tp, res_tp, comp,
                 sorted_block_size, // modified in place with size of sorted
                                    // block size
                 depends);
 
         // Merge segments in parallel until all elements are sorted
         sycl::event merges_ev =
-            sort_detail::merge_sorted_block_contig_impl<argTy *, Comp>(
+            merge_sort_detail::merge_sorted_block_contig_impl<argTy *, Comp>(
                 exec_q, iter_nelems, sort_nelems, res_tp, comp,
                 sorted_block_size, {base_sort_ev});
 
@@ -816,8 +807,7 @@ sycl::event stable_argsort_axis1_contig_impl(
     const IndexComp<IndexTy, argTy, ValueComp> index_comp{arg_tp, ValueComp{}};
 
     static constexpr size_t determine_automatically = 0;
-    size_t sorted_block_size =
-        (sort_nelems >= 512) ? 512 : determine_automatically;
+    size_t sorted_block_size = determine_automatically;
 
     const size_t total_nelems = iter_nelems * sort_nelems;
 
@@ -837,21 +827,24 @@ sycl::event stable_argsort_axis1_contig_impl(
         });
 
     // Sort segments of the array
-    sycl::event base_sort_ev = sort_detail::sort_over_work_group_contig_impl(
-        exec_q, iter_nelems, sort_nelems, res_tp, res_tp, index_comp,
-        sorted_block_size, // modified in place with size of sorted block size
-        {populate_indexed_data_ev});
+    sycl::event base_sort_ev =
+        merge_sort_detail::sort_over_work_group_contig_impl(
+            exec_q, iter_nelems, sort_nelems, res_tp, res_tp, index_comp,
+            sorted_block_size, // modified in place with size of sorted block
+                               // size
+            {populate_indexed_data_ev});
 
     // Merge segments in parallel until all elements are sorted
-    sycl::event merges_ev = sort_detail::merge_sorted_block_contig_impl(
+    sycl::event merges_ev = merge_sort_detail::merge_sorted_block_contig_impl(
         exec_q, iter_nelems, sort_nelems, res_tp, index_comp, sorted_block_size,
         {base_sort_ev});
 
     sycl::event write_out_ev = exec_q.submit([&](sycl::handler &cgh) {
         cgh.depends_on(merges_ev);
 
         auto temp_acc =
-            sort_detail::GetReadOnlyAccess<decltype(res_tp)>{}(res_tp, cgh);
+            merge_sort_detail::GetReadOnlyAccess<decltype(res_tp)>{}(res_tp,
+                                                                     cgh);
 
         using KernelName = index_map_to_rows_krn<argTy, IndexTy, ValueComp>;