NVIDIA / NVIDIA/cccl

[BUG]: `upper_bound`/`lower_bound` rebuild on every call (device pointers in the build cache key)

Open
#10,689 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.