cuda-cccl 0.3.2__cp313-cp313-manylinux_2_24_aarch64.whl → 0.3.4__cp313-cp313-manylinux_2_24_aarch64.whl

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
+        )

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_exclusive_scan as make_exclusive_scan
 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 @@ __all__ = [
     "make_segmented_reduce",
     "unique_by_key",
     "make_unique_by_key",
+    "segmented_sort",
+    "make_segmented_sort",
     "three_way_partition",
     "make_three_way_partition",
     "DoubleBuffer",

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",
+]

cuda/compute/algorithms/{_merge_sort.py → _sort/_merge_sort.py}

@@ -7,19 +7,19 @@ from typing import Callable, Union
 
 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(

cuda/compute/algorithms/{_radix_sort.py → _sort/_radix_sort.py}

@@ -3,37 +3,18 @@
 #
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from enum import Enum
-from typing import Tuple
-
-from .. import _bindings
-from .. import _cccl_interop as cccl
-from .._caching import cache_with_key
-from .._cccl_interop import call_build, set_cccl_iterator_state
-from .._utils.protocols import (
+from ... import _bindings
+from ... import _cccl_interop as cccl
+from ..._caching import cache_with_key
+from ..._cccl_interop import call_build, set_cccl_iterator_state
+from ..._utils.protocols import (
     get_data_pointer,
     get_dtype,
     validate_and_get_stream,
 )
-from .._utils.temp_storage_buffer import TempStorageBuffer
-from ..typing import DeviceArrayLike
-
-
-class SortOrder(Enum):
-    ASCENDING = 0
-    DESCENDING = 1
-
-
-class DoubleBuffer:
-    def __init__(self, d_current: DeviceArrayLike, d_alternate: DeviceArrayLike):
-        self.d_buffers = [d_current, d_alternate]
-        self.selector = 0
-
-    def current(self):
-        return self.d_buffers[self.selector]
-
-    def alternate(self):
-        return self.d_buffers[1 - self.selector]
+from ..._utils.temp_storage_buffer import TempStorageBuffer
+from ...typing import DeviceArrayLike
+from ._sort_common import DoubleBuffer, SortOrder, _get_arrays
 
 
 def make_cache_key(
@@ -65,32 +46,6 @@ def make_cache_key(
     )
 
 
-def _get_arrays(
-    d_in_keys: DeviceArrayLike | DoubleBuffer,
-    d_out_keys: DeviceArrayLike | DoubleBuffer | None,
-    d_in_values: DeviceArrayLike | DoubleBuffer | None,
-    d_out_values: DeviceArrayLike | None,
-) -> Tuple:
-    if isinstance(d_in_keys, DoubleBuffer):
-        d_in_keys_array = d_in_keys.current()
-        d_out_keys_array = d_in_keys.alternate()
-
-        if d_in_values is not None:
-            assert isinstance(d_in_values, DoubleBuffer)
-            d_in_values_array = d_in_values.current()
-            d_out_values_array = d_in_values.alternate()
-        else:
-            d_in_values_array = None
-            d_out_values_array = None
-    else:
-        d_in_keys_array = d_in_keys
-        d_in_values_array = d_in_values
-        d_out_keys_array = d_out_keys
-        d_out_values_array = d_out_values
-
-    return d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array
-
-
 class _RadixSort:
     __slots__ = [
         "d_in_keys_cccl",
@@ -104,7 +59,7 @@ class _RadixSort:
     def __init__(
         self,
         d_in_keys: DeviceArrayLike | DoubleBuffer,
-        d_out_keys: DeviceArrayLike | DoubleBuffer | None,
+        d_out_keys: DeviceArrayLike | None,
         d_in_values: DeviceArrayLike | DoubleBuffer | None,
         d_out_values: DeviceArrayLike | None,
         order: SortOrder,
@@ -294,10 +249,6 @@ def radix_sort(
         end_bit,
         stream,
     )
-    # Use the appropriate array for namespace - prefer d_out_keys, fallback to d_in_keys
-    ref_array = d_out_keys if d_out_keys is not None else d_in_keys
-    if isinstance(ref_array, DoubleBuffer):
-        ref_array = ref_array.current()
     tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
     sorter(
         tmp_storage,

cuda/compute/algorithms/_sort/_segmented_sort.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+from ... import _bindings
+from ... import _cccl_interop as cccl
+from ..._caching import cache_with_key
+from ..._cccl_interop import call_build, set_cccl_iterator_state
+from ..._utils.protocols import (
+    get_data_pointer,
+    get_dtype,
+    validate_and_get_stream,
+)
+from ..._utils.temp_storage_buffer import TempStorageBuffer
+from ...typing import DeviceArrayLike
+from ._sort_common import DoubleBuffer, SortOrder, _get_arrays
+
+
+class _SegmentedSort:
+    __slots__ = [
+        "build_result",
+        "d_in_keys_cccl",
+        "d_out_keys_cccl",
+        "d_in_values_cccl",
+        "d_out_values_cccl",
+        "start_offsets_in_cccl",
+        "end_offsets_in_cccl",
+    ]
+
+    def __init__(
+        self,
+        d_in_keys: DeviceArrayLike | DoubleBuffer,
+        d_out_keys: DeviceArrayLike | None,
+        d_in_values: DeviceArrayLike | DoubleBuffer | None,
+        d_out_values: DeviceArrayLike | None,
+        start_offsets_in: DeviceArrayLike,
+        end_offsets_in: DeviceArrayLike,
+        order: SortOrder,
+    ):
+        d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
+            _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
+        )
+
+        self.d_in_keys_cccl = cccl.to_cccl_input_iter(d_in_keys_array)
+        self.d_out_keys_cccl = cccl.to_cccl_output_iter(d_out_keys_array)
+        self.d_in_values_cccl = cccl.to_cccl_input_iter(d_in_values_array)
+        self.d_out_values_cccl = cccl.to_cccl_output_iter(d_out_values_array)
+        self.start_offsets_in_cccl = cccl.to_cccl_input_iter(start_offsets_in)
+        self.end_offsets_in_cccl = cccl.to_cccl_input_iter(end_offsets_in)
+
+        cccl.cccl_iterator_set_host_advance(
+            self.start_offsets_in_cccl, start_offsets_in
+        )
+        cccl.cccl_iterator_set_host_advance(self.end_offsets_in_cccl, end_offsets_in)
+
+        self.build_result = call_build(
+            _bindings.DeviceSegmentedSortBuildResult,
+            _bindings.SortOrder.ASCENDING
+            if order is SortOrder.ASCENDING
+            else _bindings.SortOrder.DESCENDING,
+            self.d_in_keys_cccl,
+            self.d_in_values_cccl,
+            self.start_offsets_in_cccl,
+            self.end_offsets_in_cccl,
+        )
+
+    def __call__(
+        self,
+        temp_storage,
+        d_in_keys,
+        d_out_keys,
+        d_in_values,
+        d_out_values,
+        num_items,
+        num_segments,
+        start_offsets_in,
+        end_offsets_in,
+        stream=None,
+    ):
+        d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
+            _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
+        )
+
+        set_cccl_iterator_state(self.d_in_keys_cccl, d_in_keys_array)
+        set_cccl_iterator_state(self.d_out_keys_cccl, d_out_keys_array)
+        if d_in_values_array is not None:
+            set_cccl_iterator_state(self.d_in_values_cccl, d_in_values_array)
+        if d_out_values_array is not None:
+            set_cccl_iterator_state(self.d_out_values_cccl, d_out_values_array)
+        set_cccl_iterator_state(self.start_offsets_in_cccl, start_offsets_in)
+        set_cccl_iterator_state(self.end_offsets_in_cccl, end_offsets_in)
+
+        stream_handle = validate_and_get_stream(stream)
+        if temp_storage is None:
+            temp_storage_bytes = 0
+            d_temp_storage = 0
+        else:
+            temp_storage_bytes = temp_storage.nbytes
+            d_temp_storage = get_data_pointer(temp_storage)
+
+        # Detect overwrite mode and selector, similar to radix sort
+        is_overwrite_okay = isinstance(d_in_keys, DoubleBuffer)
+        selector = -1
+
+        temp_storage_bytes, selector = self.build_result.compute(
+            d_temp_storage,
+            temp_storage_bytes,
+            self.d_in_keys_cccl,
+            self.d_out_keys_cccl,
+            self.d_in_values_cccl,
+            self.d_out_values_cccl,
+            num_items,
+            num_segments,
+            self.start_offsets_in_cccl,
+            self.end_offsets_in_cccl,
+            is_overwrite_okay,
+            selector,
+            stream_handle,
+        )
+
+        if is_overwrite_okay and temp_storage is not None:
+            assert selector in (0, 1)
+            assert isinstance(d_in_keys, DoubleBuffer)
+            d_in_keys.selector = selector
+            if d_in_values is not None:
+                assert isinstance(d_in_values, DoubleBuffer)
+                d_in_values.selector = selector
+
+        return temp_storage_bytes
+
+
+def make_cache_key(
+    d_in_keys: DeviceArrayLike | DoubleBuffer,
+    d_out_keys: DeviceArrayLike | None,
+    d_in_values: DeviceArrayLike | DoubleBuffer | None,
+    d_out_values: DeviceArrayLike | None,
+    start_offsets_in: DeviceArrayLike,
+    end_offsets_in: DeviceArrayLike,
+    order: SortOrder,
+):
+    d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array = (
+        _get_arrays(d_in_keys, d_out_keys, d_in_values, d_out_values)
+    )
+
+    d_in_keys_key = get_dtype(d_in_keys_array)
+    d_out_keys_key = None if d_out_keys_array is None else get_dtype(d_out_keys_array)
+    d_in_values_key = (
+        None if d_in_values_array is None else get_dtype(d_in_values_array)
+    )
+    d_out_values_key = (
+        None if d_out_values_array is None else get_dtype(d_out_values_array)
+    )
+    start_offsets_in_key = get_dtype(start_offsets_in)
+    end_offsets_in_key = get_dtype(end_offsets_in)
+
+    return (
+        d_in_keys_key,
+        d_out_keys_key,
+        d_in_values_key,
+        d_out_values_key,
+        start_offsets_in_key,
+        end_offsets_in_key,
+        order,
+    )
+
+
+@cache_with_key(make_cache_key)
+def make_segmented_sort(
+    d_in_keys: DeviceArrayLike | DoubleBuffer,
+    d_out_keys: DeviceArrayLike | None,
+    d_in_values: DeviceArrayLike | DoubleBuffer | None,
+    d_out_values: DeviceArrayLike | None,
+    start_offsets_in: DeviceArrayLike,
+    end_offsets_in: DeviceArrayLike,
+    order: SortOrder,
+):
+    """
+    Performs a device-wide segmented sort using the specified keys and values.
+
+    Example:
+        Below, ``make_segmented_sort`` is used to create a segmented sort object that can be reused.
+
+        .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_object.py
+            :language: python
+            :start-after: # example-begin
+
+    Args:
+        d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted
+        d_out_keys: Device array to store the sorted keys
+        d_in_values: Optional Device array or DoubleBuffer containing the input values to be sorted
+        d_out_values: Device array to store the sorted values
+        start_offsets_in: Device array or iterator containing the sequence of beginning offsets
+        end_offsets_in: Device array or iterator containing the sequence of ending offsets
+        order: SortOrder specifying the order of the sort
+
+    Returns:
+        A callable object that can be used to perform the segmented sort
+    """
+    return _SegmentedSort(
+        d_in_keys,
+        d_out_keys,
+        d_in_values,
+        d_out_values,
+        start_offsets_in,
+        end_offsets_in,
+        order,
+    )
+
+
+def segmented_sort(
+    d_in_keys: DeviceArrayLike | DoubleBuffer,
+    d_out_keys: DeviceArrayLike | None,
+    d_in_values: DeviceArrayLike | DoubleBuffer | None,
+    d_out_values: DeviceArrayLike | None,
+    num_items: int,
+    num_segments: int,
+    start_offsets_in: DeviceArrayLike,
+    end_offsets_in: DeviceArrayLike,
+    order: SortOrder,
+    stream=None,
+):
+    """
+    Performs device-wide segmented sort.
+
+    This function automatically handles temporary storage allocation and execution.
+
+    Example:
+        Below, ``segmented_sort`` is used to perform a segmented sort. It also rearranges the values according to the keys' order.
+
+        .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_basic.py
+            :language: python
+            :start-after: # example-begin
+
+
+        In the following example, ``segmented_sort`` is used to perform a segmented sort with a ``DoubleBuffer` for reduced temporary storage.
+
+        .. literalinclude:: ../../python/cuda_cccl/tests/parallel/examples/sort/segmented_sort_buffer.py
+            :language: python
+            :start-after: # example-begin
+
+    Args:
+        d_in_keys: Device array or DoubleBuffer containing the input keys to be sorted
+        d_out_keys: Device array to store the sorted keys (optional)
+        d_in_values: Device array or DoubleBuffer containing the input values to be sorted (optional)
+        d_out_values: Device array to store the sorted values (optional)
+        num_items: Total number of items to sort
+        num_segments: Number of segments to sort
+        start_offsets_in: Device array or iterator containing the sequence of beginning offsets
+        end_offsets_in: Device array or iterator containing the sequence of ending offsets
+        order: Sort order (ascending or descending)
+        stream: CUDA stream for the operation (optional)
+    """
+    sorter = make_segmented_sort(
+        d_in_keys,
+        d_out_keys,
+        d_in_values,
+        d_out_values,
+        start_offsets_in,
+        end_offsets_in,
+        order,
+    )
+    tmp_storage_bytes = sorter(
+        None,
+        d_in_keys,
+        d_out_keys,
+        d_in_values,
+        d_out_values,
+        num_items,
+        num_segments,
+        start_offsets_in,
+        end_offsets_in,
+        stream,
+    )
+    tmp_storage = TempStorageBuffer(tmp_storage_bytes, stream)
+    sorter(
+        tmp_storage,
+        d_in_keys,
+        d_out_keys,
+        d_in_values,
+        d_out_values,
+        num_items,
+        num_segments,
+        start_offsets_in,
+        end_offsets_in,
+        stream,
+    )

cuda/compute/algorithms/_sort/_sort_common.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from enum import Enum
+from typing import Tuple
+
+from ...typing import DeviceArrayLike
+
+
+class SortOrder(Enum):
+    ASCENDING = 0
+    DESCENDING = 1
+
+
+class DoubleBuffer:
+    def __init__(self, d_current: DeviceArrayLike, d_alternate: DeviceArrayLike):
+        self.d_buffers = [d_current, d_alternate]
+        self.selector = 0
+
+    def current(self):
+        return self.d_buffers[self.selector]
+
+    def alternate(self):
+        return self.d_buffers[1 - self.selector]
+
+
+def _get_arrays(
+    d_in_keys: DeviceArrayLike | DoubleBuffer,
+    d_out_keys: DeviceArrayLike | None,
+    d_in_values: DeviceArrayLike | DoubleBuffer | None,
+    d_out_values: DeviceArrayLike | None,
+) -> Tuple[DeviceArrayLike, DeviceArrayLike, DeviceArrayLike, DeviceArrayLike]:
+    if isinstance(d_in_keys, DoubleBuffer):
+        d_in_keys_array = d_in_keys.current()
+        d_out_keys_array = d_in_keys.alternate()
+
+        if d_in_values is not None:
+            assert isinstance(d_in_values, DoubleBuffer)
+            d_in_values_array = d_in_values.current()
+            d_out_values_array = d_in_values.alternate()
+        else:
+            d_in_values_array = None
+            d_out_values_array = None
+    else:
+        d_in_keys_array = d_in_keys
+        d_in_values_array = d_in_values
+        d_out_keys_array = d_out_keys
+        d_out_values_array = d_out_values
+
+    return d_in_keys_array, d_out_keys_array, d_in_values_array, d_out_values_array