Add Python wrappers for c.parallel radix_sort API (NVIDIA#4353)

Author: NaderAlAwar

python/cuda_parallel/cuda/parallel/experimental/_bindings.pyi

@@ -1,5 +1,5 @@
 import ctypes
-from typing import Any
+from typing import Any, Optional
 
 from typing_extensions import Buffer
 
@@ -54,18 +54,26 @@ class Enumeration_IteratorKind:
     @property
     def ITERATOR(self) -> IntEnumerationMember: ...
 
+class Enumeration_SortOrder:
+    @property
+    def ASCENDING(self) -> IntEnumerationMember: ...
+    @property
+    def DESCENDING(self) -> IntEnumerationMember: ...
+
 TypeEnum: Enumeration_CCCLType
 OpKind: Enumeration_OpKind
 IteratorKind: Enumeration_IteratorKind
+SortOrder: Enumeration_SortOrder
 
 def is_TypeEnum(obj) -> bool: ...
 def is_OpKind(obj) -> bool: ...
 def is_IteratorKind(obj) -> bool: ...
+def is_SortOrder(obj) -> bool: ...
 
 class Op:
     def __init__(
         self,
-        name: str = ...,
+        name: Optional[str] = ...,
         operator_type: IntEnumerationMember = ...,
         ltoir=None,
         state=None,
@@ -264,13 +272,13 @@ class DeviceMergeSortBuildResult:
         d_in_keys: Iterator,
         d_in_items: Iterator,
         d_out_keys: Iterator,
-        d_out_itemss: Iterator,
+        d_out_items: Iterator,
         binary_op: Op,
         info: CommonData,
     ) -> int: ...
     def compute(
         self,
-        temp_storage_ptr: int,
+        temp_storage_ptr: int | None,
         temp_storage_nbytes: int,
         d_in_keys: Iterator,
         d_in_items: Iterator,
@@ -299,7 +307,7 @@ class DeviceUniqueByKeyBuildResult:
     ) -> int: ...
     def compute(
         self,
-        temp_storage_ptr: int,
+        temp_storage_ptr: int | None,
         temp_storage_nbytes: int,
         d_keys_in: Iterator,
         d_values_in: Iterator,
@@ -311,6 +319,29 @@ class DeviceUniqueByKeyBuildResult:
         stream,
     ) -> tuple[int, int]: ...
 
+# -----------------
+# DeviceRadixSort
+# -----------------
+
+class DeviceRadixSortBuildResult:
+    def __init__(self): ...
+    def compute(
+        self,
+        temp_storage_ptr: int | None,
+        temp_storage_nbytes: int,
+        d_keys_in: Iterator,
+        d_keys_out: Iterator,
+        d_values_in: Iterator,
+        d_values_out: Iterator,
+        decomposer_op: Op,
+        num_items: int,
+        begin_bit: int,
+        end_bit: int,
+        is_overwrite_okay: bool,
+        selector: int,
+        stream,
+    ) -> tuple[int, int]: ...
+
 # --------------------
 # DeviceUnaryTransform
 # --------------------

python/cuda_parallel/cuda/parallel/experimental/_bindings.pyx

@@ -77,6 +77,9 @@ cdef extern from "cccl/c/types.h":
         cccl_type_info value_type
         void *state
 
+    ctypedef enum cccl_sort_order_t:
+        CCCL_ASCENDING
+        CCCL_DESCENDING
 
 cdef void arg_type_check(
     str arg_name,
@@ -410,11 +413,47 @@ cdef class Enumeration_IteratorKind(IntEnumerationBase):
     def ITERATOR(self):
         return self._iterator
 
+cdef class Enumeration_SortOrder(IntEnumerationBase):
+    "Enumeration of sort orders (ascending or descending)"
+    cdef IntEnumerationMember _ascending
+    cdef IntEnumerationMember _descending
+
+    def __cinit__(self):
+        self.enum_name = "SortOrder"
+        self._ascending = self.make_ASCENDING()
+        self._descending = self.make_DESCENDING()
+
+    cdef IntEnumerationMember make_ASCENDING(self):
+        cdef str prop_name = "ASCENDING"
+        return IntEnumerationMember(
+            type(self),
+            self.enum_name,
+            prop_name,
+            cccl_sort_order_t.CCCL_ASCENDING
+        )
+
+    cdef IntEnumerationMember make_DESCENDING(self):
+        cdef str prop_name = "DESCENDING"
+        return IntEnumerationMember(
+            type(self),
+            self.enum_name,
+            prop_name,
+            cccl_sort_order_t.CCCL_DESCENDING
+        )
+
+    @property
+    def ASCENDING(self):
+        return self._ascending
+
+    @property
+    def DESCENDING(self):
+        return self._descending
+
 
 TypeEnum = Enumeration_CCCLType()
 OpKind = Enumeration_OpKind()
 IteratorKind = Enumeration_IteratorKind()
-
+SortOrder = Enumeration_SortOrder()
 
 cpdef bint is_TypeEnum(IntEnumerationMember attr):
     "Return True if attribute is a member of TypeEnum enumeration"
@@ -430,6 +469,10 @@ cpdef bint is_IteratorKind(IntEnumerationMember attr):
     "Return True if attribute is a member of IteratorKind enumeration"
     return attr.parent_class is Enumeration_IteratorKind
 
+cpdef bint is_SortOrder(IntEnumerationMember attr):
+    "Return True if attribute is a member of SortOrder enumeration"
+    return attr.parent_class is Enumeration_SortOrder
+
 
 cdef void _validate_alignment(int alignment) except *:
     cdef uint32_t val
@@ -1621,12 +1664,147 @@ cdef class DeviceUniqueByKeyBuildResult:
                 <uint64_t>num_items,
                 c_stream
             )
+
         if status != 0:
             raise RuntimeError(
                 f"Failed executing unique_by_key, error code: {status}"
             )
         return storage_sz
 
+# -----------------
+# DeviceRadixSort
+# -----------------
+
+cdef extern from "cccl/c/radix_sort.h":
+    cdef struct cccl_device_radix_sort_build_result_t 'cccl_device_radix_sort_build_result_t':
+        pass
+
+    cdef CUresult cccl_device_radix_sort_build(
+        cccl_device_radix_sort_build_result_t *build_ptr,
+        cccl_sort_order_t sort_order,
+        cccl_iterator_t d_keys_in,
+        cccl_iterator_t d_values_in,
+        cccl_op_t decomposer,
+        const char* decomposer_return_type,
+        int, int, const char *, const char *, const char *, const char *
+    ) nogil
+
+    cdef CUresult cccl_device_radix_sort(
+        cccl_device_radix_sort_build_result_t build,
+        void *d_storage_ptr,
+        size_t *d_storage_nbytes,
+        cccl_iterator_t d_keys_in,
+        cccl_iterator_t d_keys_out,
+        cccl_iterator_t d_values_in,
+        cccl_iterator_t d_values_out,
+        cccl_op_t decomposer,
+        size_t num_items,
+        int begin_bit,
+        int end_bit,
+        bint is_overwrite_okay,
+        int* selector,
+        CUstream stream
+    ) nogil
+
+    cdef CUresult cccl_device_radix_sort_cleanup(
+        cccl_device_radix_sort_build_result_t *build_ptr,
+    ) nogil
+
+
+cdef class DeviceRadixSortBuildResult:
+    cdef cccl_device_radix_sort_build_result_t build_data
+
+    def __dealloc__(DeviceRadixSortBuildResult self):
+        cdef CUresult status = -1
+        with nogil:
+            status = cccl_device_radix_sort_cleanup(&self.build_data)
+        if (status != 0):
+            print(f"Return code {status} encountered during radix_sort result cleanup")
+
+    def __cinit__(
+        DeviceRadixSortBuildResult self,
+        cccl_sort_order_t order,
+        Iterator d_keys_in,
+        Iterator d_values_in,
+        Op decomposer_op,
+        const char* decomposer_return_type,
+        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_radix_sort_build_result_t))
+        with nogil:
+            status = cccl_device_radix_sort_build(
+                &self.build_data,
+                order,
+                d_keys_in.iter_data,
+                d_values_in.iter_data,
+                decomposer_op.op_data,
+                decomposer_return_type,
+                cc_major,
+                cc_minor,
+                cub_path,
+                thrust_path,
+                libcudacxx_path,
+                ctk_path,
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed building radix_sort, error code: {status}"
+            )
+
+    cpdef tuple compute(
+        DeviceRadixSortBuildResult self,
+        temp_storage_ptr,
+        temp_storage_bytes,
+        Iterator d_keys_in,
+        Iterator d_keys_out,
+        Iterator d_values_in,
+        Iterator d_values_out,
+        Op decomposer_op,
+        size_t num_items,
+        int begin_bit,
+        int end_bit,
+        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_radix_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,
+                decomposer_op.op_data,
+                <uint64_t>num_items,
+                begin_bit,
+                end_bit,
+                is_overwrite_okay,
+                &selector_int,
+                c_stream
+            )
+
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing ascending radix_sort, error code: {status}"
+            )
+        return <object>storage_sz, <object>selector_int
+
 
     def _get_cubin(self):
         return self.build_data.cubin[:self.build_data.cubin_size]

python/cuda_parallel/cuda/parallel/experimental/_cccl_interop.py

@@ -178,7 +178,16 @@ def to_cccl_value(array_or_struct: np.ndarray | GpuStruct) -> Value:
         return to_cccl_value(array_or_struct._data)
 
 
-def to_cccl_op(op: Callable, sig) -> Op:
+def to_cccl_op(op: Callable | None, sig) -> Op:
+    if op is None:
+        return Op(
+            operator_type=OpKind.STATELESS,
+            name=None,
+            ltoir=None,
+            state_alignment=1,
+            state=None,
+        )
+
     ltoir, _ = cuda.compile(op, sig=sig, output="ltoir")
     return Op(
         operator_type=OpKind.STATELESS,

python/cuda_parallel/cuda/parallel/experimental/algorithms/__init__.py

@@ -4,6 +4,8 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._merge_sort import merge_sort as merge_sort
+from ._radix_sort import DoubleBuffer, SortOrder
+from ._radix_sort import radix_sort as radix_sort
 from ._reduce import reduce_into as reduce_into
 from ._scan import exclusive_scan as exclusive_scan
 from ._scan import inclusive_scan as inclusive_scan
@@ -18,6 +20,9 @@
     "inclusive_scan",
     "segmented_reduce",
     "unique_by_key",
+    "radix_sort",
+    "DoubleBuffer",
+    "SortOrder",
     "binary_transform",
     "unary_transform",
 ]

python/cuda_parallel/cuda/parallel/experimental/algorithms/_merge_sort.py

@@ -72,7 +72,9 @@ def __init__(
         d_out_items: DeviceArrayLike | None,
         op: Callable,
     ):
-        assert (d_in_items is None) == (d_out_items is None)
+        present_in_values = d_in_items is not None
+        present_out_values = d_out_items is not None
+        assert present_in_values == present_out_values
 
         self.d_in_keys_cccl = cccl.to_cccl_iter(d_in_keys)
         self.d_in_items_cccl = cccl.to_cccl_iter(d_in_items)