Add python wrappers for c.parallel segmented_sort

c/parallel/src/segmented_sort.cu

@@ -12,7 +12,7 @@
 #include <cub/detail/launcher/cuda_driver.cuh> // cub::detail::CudaDriverLauncherFactory
 #include <cub/detail/ptx-json-parser.h>
 #include <cub/device/dispatch/dispatch_segmented_sort.cuh> // cub::DispatchSegmentedSort
-#include <cub/device/dispatch/kernels/segmented_sort.cuh> // DeviceSegmentedSort kernels
+#include <cub/device/dispatch/kernels/kernel_segmented_sort.cuh> // DeviceSegmentedSort kernels
 #include <cub/device/dispatch/tuning/tuning_segmented_sort.cuh> // policy_hub
 #include <cub/thread/thread_load.cuh> // cub::LoadModifier
 
@@ -56,10 +56,10 @@ std::string get_device_segmented_sort_fallback_kernel_name(
   cccl_sort_order_t sort_order)
 {
   std::string chained_policy_t;
-  check(nvrtcGetTypeName<device_segmented_sort_policy>(&chained_policy_t));
+  check(cccl_type_name_from_nvrtc<device_segmented_sort_policy>(&chained_policy_t));
 
   std::string offset_t;
-  check(nvrtcGetTypeName<OffsetT>(&offset_t));
+  check(cccl_type_name_from_nvrtc<OffsetT>(&offset_t));
 
   /*
   template <SortOrder Order,             // 0 (ascending)
@@ -90,10 +90,10 @@ std::string get_device_segmented_sort_kernel_small_name(
   cccl_sort_order_t sort_order)
 {
   std::string chained_policy_t;
-  check(nvrtcGetTypeName<device_segmented_sort_policy>(&chained_policy_t));
+  check(cccl_type_name_from_nvrtc<device_segmented_sort_policy>(&chained_policy_t));
 
   std::string offset_t;
-  check(nvrtcGetTypeName<OffsetT>(&offset_t));
+  check(cccl_type_name_from_nvrtc<OffsetT>(&offset_t));
 
   /*
   template <SortOrder Order,             // 0 (ascending)
@@ -124,10 +124,10 @@ std::string get_device_segmented_sort_kernel_large_name(
   cccl_sort_order_t sort_order)
 {
   std::string chained_policy_t;
-  check(nvrtcGetTypeName<device_segmented_sort_policy>(&chained_policy_t));
+  check(cccl_type_name_from_nvrtc<device_segmented_sort_policy>(&chained_policy_t));
 
   std::string offset_t;
-  check(nvrtcGetTypeName<OffsetT>(&offset_t));
+  check(cccl_type_name_from_nvrtc<OffsetT>(&offset_t));
 
   /*
   template <SortOrder Order,             // 0 (ascending)
@@ -182,11 +182,11 @@ cccl_op_t make_segments_selector_op(
   cccl_op_t selector_op{};
   auto selector_op_state = std::make_unique<selector_state_t>();
   std::string offset_t;
-  check(nvrtcGetTypeName<OffsetT>(&offset_t));
+  check(cccl_type_name_from_nvrtc<OffsetT>(&offset_t));
 
   const std::string code = std::format(
     R"XXX(
-#include <cub/device/dispatch/kernels/segmented_sort.cuh>
+#include <cub/device/dispatch/kernels/kernel_segmented_sort.cuh>
 
 extern "C" __device__ void {0}(void* state_ptr, const void* arg_ptr, void* result_ptr)
 {{
@@ -297,7 +297,7 @@ std::string get_three_way_partition_init_kernel_name()
 std::string get_three_way_partition_kernel_name(std::string_view large_selector_t, std::string_view small_selector_t)
 {
   std::string chained_policy_t;
-  check(nvrtcGetTypeName<device_three_way_partition_policy>(&chained_policy_t));
+  check(cccl_type_name_from_nvrtc<device_three_way_partition_policy>(&chained_policy_t));
 
   static constexpr std::string_view input_it_t =
     "thrust::counting_iterator<cub::detail::segmented_sort::local_segment_index_t>";
@@ -308,7 +308,7 @@ std::string get_three_way_partition_kernel_name(std::string_view large_selector_
   static constexpr std::string_view num_selected_it_t = "cub::detail::segmented_sort::local_segment_index_t*";
   static constexpr std::string_view scan_tile_state_t = "cub::detail::three_way_partition::ScanTileStateT";
   std::string offset_t;
-  check(nvrtcGetTypeName<OffsetT>(&offset_t));
+  check(cccl_type_name_from_nvrtc<OffsetT>(&offset_t));
 
   static constexpr std::string_view per_partition_offset_t = "cub::detail::three_way_partition::per_partition_offset_t";
   static constexpr std::string_view streaming_context_t =
@@ -638,9 +638,9 @@ CUresult cccl_device_segmented_sort_build_ex(
 
     const std::string final_src = std::format(
       R"XXX(
-#include <cub/device/dispatch/kernels/segmented_sort.cuh>
+#include <cub/device/dispatch/kernels/kernel_segmented_sort.cuh>
 #include <cub/device/dispatch/tuning/tuning_segmented_sort.cuh>
-#include <cub/device/dispatch/kernels/three_way_partition.cuh>
+#include <cub/device/dispatch/kernels/kernel_three_way_partition.cuh>
 #include <cub/device/dispatch/tuning/tuning_three_way_partition.cuh>
 
 {0}

python/cuda_cccl/cuda/compute/__init__.py

@@ -17,13 +17,15 @@
     make_radix_sort,
     make_reduce_into,
     make_segmented_reduce,
+    make_segmented_sort,
     make_three_way_partition,
     make_unary_transform,
     make_unique_by_key,
     merge_sort,
     radix_sort,
     reduce_into,
     segmented_reduce,
+    segmented_sort,
     three_way_partition,
     unary_transform,
     unique_by_key,
@@ -59,6 +61,7 @@
     "make_radix_sort",
     "make_reduce_into",
     "make_segmented_reduce",
+    "make_segmented_sort",
     "make_three_way_partition",
     "make_unary_transform",
     "make_unique_by_key",
@@ -69,6 +72,7 @@
     "reduce_into",
     "ReverseIterator",
     "segmented_reduce",
+    "segmented_sort",
     "SortOrder",
     "TransformIterator",
     "TransformOutputIterator",

python/cuda_cccl/cuda/compute/_bindings.pyi

@@ -447,6 +447,29 @@ class DeviceHistogramBuildResult:
         stream,
     ) -> None: ...
 
+# -----------------
+# DeviceSegmentedSort
+# -----------------
+
+class DeviceSegmentedSortBuildResult:
+    def __init__(self): ...
+    def compute(
+        self,
+        temp_storage_ptr: int | None,
+        temp_storage_nbytes: int,
+        d_in_keys: Iterator,
+        d_out_keys: Iterator,
+        d_in_values: Iterator,
+        d_out_values: Iterator,
+        num_items: int,
+        num_segments: int,
+        d_begin_offsets: Iterator,
+        d_end_offsets: Iterator,
+        is_overwrite_okay: bool,
+        selector: int,
+        stream,
+    ) -> tuple[int, int]: ...
+
 # ---------------------
 # DeviceThreeWayPartition
 # ---------------------

python/cuda_cccl/cuda/compute/_bindings_impl.pyx

@@ -1779,7 +1779,7 @@ cdef class DeviceRadixSortBuildResult:
 
         if status != 0:
             raise RuntimeError(
-                f"Failed executing ascending radix_sort, error code: {status}"
+                f"Failed executing radix_sort, error code: {status}"
             )
         return <object>storage_sz, <object>selector_int
 
@@ -2271,3 +2271,145 @@ cdef class DeviceThreeWayPartitionBuildResult:
             <const char*>self.build_data.cubin,
             self.build_data.cubin_size
         )
+
+
+# -------------------
+# DeviceSegmentedSort
+# -------------------
+
+cdef extern from "cccl/c/segmented_sort.h":
+    cdef struct cccl_device_segmented_sort_build_result_t 'cccl_device_segmented_sort_build_result_t':
+        const char* cubin
+        size_t cubin_size
+
+    cdef CUresult cccl_device_segmented_sort_build(
+        cccl_device_segmented_sort_build_result_t *build_ptr,
+        cccl_sort_order_t sort_order,
+        cccl_iterator_t d_keys_in,
+        cccl_iterator_t d_keys_out,
+        cccl_iterator_t begin_offset_in,
+        cccl_iterator_t end_offset_in,
+        int, int, const char *, const char *, const char *, const char *
+    ) nogil
+
+    cdef CUresult cccl_device_segmented_sort(
+        cccl_device_segmented_sort_build_result_t build,
+        void* d_temp_storage,
+        size_t* temp_storage_bytes,
+        cccl_iterator_t d_keys_in,
+        cccl_iterator_t d_keys_out,
+        cccl_iterator_t d_values_in,
+        cccl_iterator_t d_values_out,
+        int64_t num_items,
+        int64_t num_segments,
+        cccl_iterator_t start_offset_in,
+        cccl_iterator_t end_offset_in,
+        bint is_overwrite_okay,
+        int* selector,
+        CUstream stream
+    ) nogil
+
+    cdef CUresult cccl_device_segmented_sort_cleanup(
+        cccl_device_segmented_sort_build_result_t* build_ptr
+    ) nogil
+
+cdef class DeviceSegmentedSortBuildResult:
+    cdef cccl_device_segmented_sort_build_result_t build_data
+
+    def __dealloc__(DeviceSegmentedSortBuildResult self):
+        cdef CUresult status = -1
+        with nogil:
+            status = cccl_device_segmented_sort_cleanup(&self.build_data)
+        if (status != 0):
+            print(f"Return code {status} encountered during segmented_sort result cleanup")
+
+    def __cinit__(
+        DeviceSegmentedSortBuildResult self,
+        cccl_sort_order_t order,
+        Iterator d_keys_in,
+        Iterator d_values_in,
+        Iterator begin_offset_in,
+        Iterator end_offset_in,
+        CommonData common_data,
+    ):
+        cdef CUresult status = -1
+        cdef int cc_major = common_data.get_cc_major()
+        cdef int cc_minor = common_data.get_cc_minor()
+        cdef const char *cub_path = common_data.cub_path_get_c_str()
+        cdef const char *thrust_path = common_data.thrust_path_get_c_str()
+        cdef const char *libcudacxx_path = common_data.libcudacxx_path_get_c_str()
+        cdef const char *ctk_path = common_data.ctk_path_get_c_str()
+
+        memset(&self.build_data, 0, sizeof(cccl_device_segmented_sort_build_result_t))
+        with nogil:
+            status = cccl_device_segmented_sort_build(
+                &self.build_data,
+                order,
+                d_keys_in.iter_data,
+                d_values_in.iter_data,
+                begin_offset_in.iter_data,
+                end_offset_in.iter_data,
+                cc_major,
+                cc_minor,
+                cub_path,
+                thrust_path,
+                libcudacxx_path,
+                ctk_path,
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed building segmented_sort, error code: {status}"
+            )
+
+    cpdef tuple compute(
+        DeviceSegmentedSortBuildResult self,
+        temp_storage_ptr,
+        temp_storage_bytes,
+        Iterator d_keys_in,
+        Iterator d_keys_out,
+        Iterator d_values_in,
+        Iterator d_values_out,
+        size_t num_items,
+        size_t num_segments,
+        Iterator start_offset_in,
+        Iterator end_offset_in,
+        bint is_overwrite_okay,
+        selector,
+        stream
+    ):
+        cdef CUresult status = -1
+        cdef void *storage_ptr = (<void *><size_t>temp_storage_ptr) if temp_storage_ptr else NULL
+        cdef size_t storage_sz = <size_t>temp_storage_bytes
+        cdef int selector_int = <int>selector
+        cdef CUstream c_stream = <CUstream><size_t>(stream) if stream else NULL
+
+        with nogil:
+            status = cccl_device_segmented_sort(
+                self.build_data,
+                storage_ptr,
+                &storage_sz,
+                d_keys_in.iter_data,
+                d_keys_out.iter_data,
+                d_values_in.iter_data,
+                d_values_out.iter_data,
+                <uint64_t>num_items,
+                <uint64_t>num_segments,
+                start_offset_in.iter_data,
+                end_offset_in.iter_data,
+                is_overwrite_okay,
+                &selector_int,
+                c_stream
+            )
+
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing segmented_sort, error code: {status}"
+            )
+        return <object>storage_sz, <object>selector_int
+
+
+    def _get_cubin(self):
+        return PyBytes_FromStringAndSize(
+            <const char*>self.build_data.cubin,
+            self.build_data.cubin_size
+        )

python/cuda_cccl/cuda/compute/algorithms/__init__.py

@@ -5,11 +5,6 @@
 
 from ._histogram import histogram_even as histogram_even
 from ._histogram import make_histogram_even as make_histogram_even
-from ._merge_sort import make_merge_sort as make_merge_sort
-from ._merge_sort import merge_sort as merge_sort
-from ._radix_sort import DoubleBuffer, SortOrder
-from ._radix_sort import make_radix_sort as make_radix_sort
-from ._radix_sort import radix_sort as radix_sort
 from ._reduce import make_reduce_into as make_reduce_into
 from ._reduce import reduce_into as reduce_into
 from ._scan import exclusive_scan as exclusive_scan
@@ -18,6 +13,13 @@
 from ._scan import make_inclusive_scan as make_inclusive_scan
 from ._segmented_reduce import make_segmented_reduce as make_segmented_reduce
 from ._segmented_reduce import segmented_reduce
+from ._sort import DoubleBuffer, SortOrder
+from ._sort import make_merge_sort as make_merge_sort
+from ._sort import make_radix_sort as make_radix_sort
+from ._sort import make_segmented_sort as make_segmented_sort
+from ._sort import merge_sort as merge_sort
+from ._sort import radix_sort as radix_sort
+from ._sort import segmented_sort as segmented_sort
 from ._three_way_partition import make_three_way_partition as make_three_way_partition
 from ._three_way_partition import three_way_partition as three_way_partition
 from ._transform import binary_transform, unary_transform
@@ -47,6 +49,8 @@
     "make_segmented_reduce",
     "unique_by_key",
     "make_unique_by_key",
+    "segmented_sort",
+    "make_segmented_sort",
     "three_way_partition",
     "make_three_way_partition",
     "DoubleBuffer",

python/cuda_cccl/cuda/compute/algorithms/_sort/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._merge_sort import make_merge_sort as make_merge_sort
+from ._merge_sort import merge_sort as merge_sort
+from ._radix_sort import make_radix_sort as make_radix_sort
+from ._radix_sort import radix_sort as radix_sort
+from ._segmented_sort import make_segmented_sort as make_segmented_sort
+from ._segmented_sort import segmented_sort as segmented_sort
+from ._sort_common import DoubleBuffer, SortOrder
+
+__all__ = [
+    "make_merge_sort",
+    "merge_sort",
+    "make_radix_sort",
+    "radix_sort",
+    "make_segmented_sort",
+    "segmented_sort",
+    "DoubleBuffer",
+    "SortOrder",
+]

python/cuda_cccl/cuda/compute/algorithms/_merge_sort.py → python/cuda_cccl/cuda/compute/algorithms/_sort/_merge_sort.py

@@ -7,19 +7,19 @@
 
 import numba
 
-from .. import _bindings
-from .. import _cccl_interop as cccl
-from .._caching import CachableFunction, cache_with_key
-from .._cccl_interop import call_build, set_cccl_iterator_state
-from .._utils import protocols
-from .._utils.protocols import (
+from ... import _bindings
+from ... import _cccl_interop as cccl
+from ..._caching import CachableFunction, cache_with_key
+from ..._cccl_interop import call_build, set_cccl_iterator_state
+from ..._utils import protocols
+from ..._utils.protocols import (
     get_data_pointer,
     validate_and_get_stream,
 )
-from .._utils.temp_storage_buffer import TempStorageBuffer
-from ..iterators._iterators import IteratorBase
-from ..op import OpKind
-from ..typing import DeviceArrayLike
+from ..._utils.temp_storage_buffer import TempStorageBuffer
+from ...iterators._iterators import IteratorBase
+from ...op import OpKind
+from ...typing import DeviceArrayLike
 
 
 def make_cache_key(