Add cuda.compute APIs for upper_bound and lower_bound (#7250)

* Python bindings to binary search

* Add Python API for binary search

* Add tests, examples and benchmarks for binary search.

* Docs clarification

* Update binary_search to use new registered caching mechanism

* Update binary_search to use new registered caching mechanism

* Change to private method

* Actually, we don't need to work with the same pointers between __init__ and __call__

* Remove our hack to return indices

---------

Co-authored-by: Ashwin Srinath <shwina@users.noreply.github.com>

Author: shwina

python/cuda_cccl/benchmarks/compute/bench_binary_search.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+import cupy as cp
+import numpy as np
+import pytest
+
+import cuda.compute
+
+
+def lower_bound_run(d_data, d_values, d_out, build_only):
+    searcher = cuda.compute.make_lower_bound(d_data, d_values, d_out)
+    if not build_only:
+        searcher(d_data, d_values, d_out, len(d_data), len(d_values))
+    cp.cuda.runtime.deviceSynchronize()
+
+
+def upper_bound_run(d_data, d_values, d_out, build_only):
+    searcher = cuda.compute.make_upper_bound(d_data, d_values, d_out)
+    if not build_only:
+        searcher(d_data, d_values, d_out, len(d_data), len(d_values))
+    cp.cuda.runtime.deviceSynchronize()
+
+
+@pytest.mark.parametrize("bench_fixture", ["compile_benchmark", "benchmark"])
+def bench_lower_bound(bench_fixture, request, size):
+    actual_size = 100 if bench_fixture == "compile_benchmark" else size
+    d_data = cp.sort(cp.random.randint(0, 1000, actual_size, dtype=np.int32))
+    d_values = cp.random.randint(0, 1000, actual_size, dtype=np.int32)
+    d_out = cp.empty_like(d_values, dtype=np.uintp)
+
+    def run():
+        lower_bound_run(
+            d_data, d_values, d_out, build_only=(bench_fixture == "compile_benchmark")
+        )
+
+    fixture = request.getfixturevalue(bench_fixture)
+    fixture(run)
+
+
+@pytest.mark.parametrize("bench_fixture", ["compile_benchmark", "benchmark"])
+def bench_upper_bound(bench_fixture, request, size):
+    actual_size = 100 if bench_fixture == "compile_benchmark" else size
+    d_data = cp.sort(cp.random.randint(0, 1000, actual_size, dtype=np.int32))
+    d_values = cp.random.randint(0, 1000, actual_size, dtype=np.int32)
+    d_out = cp.empty_like(d_values, dtype=np.uintp)
+
+    def run():
+        upper_bound_run(
+            d_data, d_values, d_out, build_only=(bench_fixture == "compile_benchmark")
+        )
+
+    fixture = request.getfixturevalue(bench_fixture)
+    fixture(run)

python/cuda_cccl/cuda/compute/__init__.py

@@ -21,10 +21,12 @@ def __getattr__(name):
         exclusive_scan,
         histogram_even,
         inclusive_scan,
+        lower_bound,
         make_binary_transform,
         make_exclusive_scan,
         make_histogram_even,
         make_inclusive_scan,
+        make_lower_bound,
         make_merge_sort,
         make_radix_sort,
         make_reduce_into,
@@ -34,6 +36,7 @@ def __getattr__(name):
         make_three_way_partition,
         make_unary_transform,
         make_unique_by_key,
+        make_upper_bound,
         merge_sort,
         radix_sort,
         reduce_into,
@@ -43,6 +46,7 @@ def __getattr__(name):
         three_way_partition,
         unary_transform,
         unique_by_key,
+        upper_bound,
     )
     from .determinism import Determinism
     from .iterators import (
@@ -72,11 +76,13 @@ def __getattr__(name):
         "gpu_struct",
         "histogram_even",
         "inclusive_scan",
+        "lower_bound",
         "make_binary_transform",
         "make_exclusive_scan",
         "make_select",
         "make_histogram_even",
         "make_inclusive_scan",
+        "make_lower_bound",
         "make_merge_sort",
         "make_radix_sort",
         "make_reduce_into",
@@ -85,6 +91,7 @@ def __getattr__(name):
         "make_three_way_partition",
         "make_unary_transform",
         "make_unique_by_key",
+        "make_upper_bound",
         "merge_sort",
         "OpKind",
         "Determinism",
@@ -101,5 +108,6 @@ def __getattr__(name):
         "three_way_partition",
         "unary_transform",
         "unique_by_key",
+        "upper_bound",
         "ZipIterator",
     ]

python/cuda_cccl/cuda/compute/_bindings.pyi

@@ -69,6 +69,11 @@ class Determinism(IntEnum):
     RUN_TO_RUN = ...
     GPU_TO_GPU = ...
 
+class BinarySearchMode(IntEnum):
+    _value_: int
+    LOWER_BOUND = ...
+    UPPER_BOUND = ...
+
 class Op:
     def __init__(
         self,
@@ -465,6 +470,31 @@ class DeviceHistogramBuildResult:
         stream,
     ) -> None: ...
 
+# -------------------
+# DeviceBinarySearch
+# -------------------
+
+class DeviceBinarySearchBuildResult:
+    def __init__(
+        self,
+        mode: BinarySearchMode,
+        d_data: Iterator,
+        d_values: Iterator,
+        d_out: Iterator,
+        comparison_op: Op,
+        info: CommonData,
+    ): ...
+    def compute(
+        self,
+        d_data: Iterator,
+        num_items: int,
+        d_values: Iterator,
+        num_values: int,
+        d_out: Iterator,
+        comparison_op: Op,
+        stream,
+    ) -> None: ...
+
 # -----------------
 # DeviceSegmentedSort
 # -----------------

python/cuda_cccl/cuda/compute/_bindings_impl.pyx

@@ -130,6 +130,10 @@ cdef extern from "cccl/c/types.h":
         RUN_TO_RUN "CCCL_RUN_TO_RUN"
         GPU_TO_GPU "CCCL_GPU_TO_GPU"
 
+    cpdef enum cccl_binary_search_mode_t:
+        LOWER_BOUND "CCCL_BINARY_SEARCH_LOWER_BOUND"
+        UPPER_BOUND "CCCL_BINARY_SEARCH_UPPER_BOUND"
+
 cdef void arg_type_check(
     str arg_name,
     object expected_type,
@@ -147,6 +151,7 @@ IteratorKind = cccl_iterator_kind_t
 SortOrder = cccl_sort_order_t
 InitKind = cccl_init_kind_t
 Determinism = cccl_determinism_t
+BinarySearchMode = cccl_binary_search_mode_t
 
 cdef void _validate_alignment(int alignment) except *:
     """
@@ -2188,6 +2193,127 @@ cdef class DeviceHistogramBuildResult:
         )
 
 
+# -------------------
+#   DeviceBinarySearch
+# -------------------
+cdef extern from "cccl/c/binary_search.h":
+    cdef struct cccl_device_binary_search_build_result_t 'cccl_device_binary_search_build_result_t':
+        int cc
+        void* cubin
+        size_t cubin_size
+        CUlibrary library
+        CUkernel kernel
+
+    cdef CUresult cccl_device_binary_search_build(
+        cccl_device_binary_search_build_result_t*,
+        cccl_binary_search_mode_t,
+        cccl_iterator_t,
+        cccl_iterator_t,
+        cccl_iterator_t,
+        cccl_op_t,
+        int, int, const char*, const char*, const char*, const char*
+    ) nogil
+
+    cdef CUresult cccl_device_binary_search(
+        cccl_device_binary_search_build_result_t,
+        cccl_iterator_t,
+        uint64_t,
+        cccl_iterator_t,
+        uint64_t,
+        cccl_iterator_t,
+        cccl_op_t,
+        CUstream
+    ) nogil
+
+    cdef CUresult cccl_device_binary_search_cleanup(
+        cccl_device_binary_search_build_result_t *build_ptr
+    ) nogil
+
+
+cdef class DeviceBinarySearchBuildResult:
+    cdef cccl_device_binary_search_build_result_t build_data
+
+    def __dealloc__(DeviceBinarySearchBuildResult self):
+        cdef CUresult status = -1
+        with nogil:
+            status = cccl_device_binary_search_cleanup(&self.build_data)
+        if (status != 0):
+            print(f"Return code {status} encountered during binary_search result cleanup")
+
+    def __cinit__(
+        DeviceBinarySearchBuildResult self,
+        cccl_binary_search_mode_t mode,
+        Iterator d_data,
+        Iterator d_values,
+        Iterator d_out,
+        Op op,
+        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_binary_search_build_result_t))
+        with nogil:
+            status = cccl_device_binary_search_build(
+                &self.build_data,
+                mode,
+                d_data.iter_data,
+                d_values.iter_data,
+                d_out.iter_data,
+                op.op_data,
+                cc_major,
+                cc_minor,
+                cub_path,
+                thrust_path,
+                libcudacxx_path,
+                ctk_path,
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed building binary_search, error code: {status}"
+            )
+
+    cpdef void compute(
+        DeviceBinarySearchBuildResult self,
+        Iterator d_data,
+        size_t num_items,
+        Iterator d_values,
+        size_t num_values,
+        Iterator d_out,
+        Op op,
+        stream
+    ):
+        cdef CUresult status = -1
+        cdef CUstream c_stream = <CUstream><uintptr_t>(stream) if stream else NULL
+
+        with nogil:
+            status = cccl_device_binary_search(
+                self.build_data,
+                d_data.iter_data,
+                <uint64_t>num_items,
+                d_values.iter_data,
+                <uint64_t>num_values,
+                d_out.iter_data,
+                op.op_data,
+                c_stream
+            )
+        if status != 0:
+            raise RuntimeError(
+                f"Failed executing binary_search, error code: {status}"
+            )
+
+    def _get_cubin(self):
+        return PyBytes_FromStringAndSize(
+            <const char*>self.build_data.cubin,
+            self.build_data.cubin_size
+        )
+
+
 # ----------------------------------
 # DeviceThreeWayPartitionBuildResult
 # ----------------------------------

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

@@ -3,6 +3,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+from ._binary_search import lower_bound as lower_bound
+from ._binary_search import make_lower_bound as make_lower_bound
+from ._binary_search import make_upper_bound as make_upper_bound
+from ._binary_search import upper_bound as upper_bound
 from ._histogram import histogram_even as histogram_even
 from ._histogram import make_histogram_even as make_histogram_even
 from ._reduce import make_reduce_into as make_reduce_into
@@ -33,6 +37,10 @@
 __all__ = [
     "reduce_into",
     "make_reduce_into",
+    "lower_bound",
+    "make_lower_bound",
+    "upper_bound",
+    "make_upper_bound",
     "inclusive_scan",
     "make_inclusive_scan",
     "exclusive_scan",