[BUG]: `upper_bound`/`lower_bound` rebuild on every call (device pointers in the build cache key)
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
## Bug
`make_upper_bound` / `make_lower_bound` pass the raw device pointers of `d_data` and `d_out` into `_make_binary_search`, which is decorated with `@cache_with_registered_key_functions`. The pointers therefore become part of the build cache key:
```python
# python/cuda_cccl/cuda/compute/algorithms/_binary_search.py:241
return _make_binary_search(
d_data, d_values, d_out, comp_adapter,
_bindings.BinarySearchMode.UPPER_BOUND,
_data_pointer_or_none(d_data), # <-- cache-key discriminator
_data_pointer_or_none(d_out), # <-- cache-key discriminator
compute_capability=compute_capability,
)
```
Any caller that allocates its input or output fresh — most library code — misses the cache and pays a full NVRTC rebuild (~1.25 s) on every call. `histogram_even`, `radix_sort` and `inclusive_scan` key on dtype/shape via `_make_hashable` and build once.
Only the pointers discriminate; shapes do not. So one allocation sliced from the front hits the cache at every size, which is the basis of the workaround below.
## Reproducer
```python
import time
import cupy as cp
import cuda.compute as cc
keys = cp.arange(1_000_000, dtype=cp.int64)
values = cp.arange(16, dtype=cp.int64)
def timed(d_out):
cp.cuda.Device(0).synchronize(); t0 = time.perf_counter()
cc.upper_bound(d_data=keys, num_items=keys.size, d_values=values,
num_values=16, d_out=d_out)
cp.cuda.Device(0).synchronize(); return round((time.perf_counter() - t0) * 1e3)
reused = cp.empty(16, dtype=cp.uint64)
print("reused d_out :", [timed(reused) for _ in range(3)])
print("fresh d_out :", [timed(cp.empty(16, dtype=cp.uint64)) for _ in range(3)])
```
Output:
```
reused d_out : [1311, 0, 0] # builds once, then cached
fresh d_out : [1254, 1253, 1252] # rebuilds every call
```
Varying `d_data` has the same effect. Varying only the *shapes* does not:
```
d_out = prefix slice of one buffer : [1269, 0, 0]
varying d_values size, d_out fixed : [0, 0, 0]
```
## Impact
In a `cuda.compute` backend for `cuda-histogram`, `upper_bound` replaced the existing `cupy.searchsorted` call for non-uniform axes. Fills of varying length — what chunked input produces — ran at **1300 ms against 9 ms for the implementation it replaced**, because each new input length came with a newly allocated output and therefore a rebuild.
Two workarounds were needed, neither obvious: keep one growable output allocation and hand out prefix slices so the base pointer never changes, and avoid the algorithm entirely on the hot path. With the first alone the benchmark drops to 8.1 ms, on par with the implementation it replaced — the kernel is fine, only the key is wrong.
## Suggested fix
Drop the pointers from the key. #9596 made the equivalent change for `histogram_even` (extract the compile-relevant inputs rather than key on the arrays); `binary_search` looks like it was missed in that pass.
## Related
#9626 — different mechanism (`id()` of Python scalars in closures), same family: the build cache key is stricter than the compiled code requires.
## Environment
- `cuda-cccl` 1.1.1
- CUDA 13.3, NVIDIA RTX 6000 Ada, driver 580.167.08
- Python 3.12, cupy-cuda13x 14.1.1, numba-cuda 0.30.4
Contributor guide
Research direction
Start in python/cuda_cccl/cuda/compute/algorithms/_binary_search.py around the _make_binary_search call and compare its cache-key inputs with the equivalent change referenced in #9596. Run the supplied reproducer with reused and freshly allocated inputs. Done means repeated calls with fresh device allocations no longer trigger an NVRTC rebuild while upper_bound and lower_bound remain correct.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- build-system, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100