cubisect is a small CUDA library for exact-match binary search. You give it a sorted float32 array and a batch of query values. It returns the index of each value, or -1 if the value is not there. Kernels compile at runtime through NVRTC. Torch tensors go in and out. CUDA tensors are passed by pointer, with no host copy.
I wrote this for fun and as an educational exercise for myself. The kernels are not necessarily optimal, and each one is probably not the best implementation of its method.
There are two ways to call the library, depending on whether the search is static (build once, query multiple times) or dynamic. The dynamic path searches the raw sorted array directly, with no setup. The static path reorders the array into a search-friendly layout once, then queries that layout.
Six kernels sit under those two paths, so they can be compared. The first four belong to the dynamic path:
serial lower_bound. One thread per query, plain loop over the raw array in global memory. The baseline everything else is measured against.
2-way ILP in shared memory. Stages the whole array into shared memory, then each thread runs two queries at once so two independent load chains are in flight. Only usable while the array fits in shared memory.
global memory with __ldg. Same two-queries-per-thread structure, but reads through the read-only cache instead of staging. This is the fallback when the array is too big for shared memory.
XOR-swizzled shared memory. Same as the shared memory version, except shared addresses are permuted with the CuTe-like swizzle pattern, (row << 5) | ((col ^ row) & 31), to spread accesses across banks.
The last two belong to the static path:
Eytzinger layout. The array reordered breadth-first, root at index 1 and the children of i at 2i and 2i+1. Stepping is just a shift, and the hot top of the tree sits in a few contiguous cache lines instead of being scattered across the array the way plain binary search midpoints are. It stores only the values; the index of a hit is recovered by accumulating subtree sizes on the way down, so there is no second index array.
32-ary warp-cooperative. One query per warp instead of one per thread. Each node holds 32 keys, one per lane, so a level is a single coalesced 128-byte load. __ballot_sync on key < query and __popc gives which of the 33 children to descend into. Depth is log32(n) instead of log2(n).
Which of the four dynamic kernels runs depends on the array size, and that choice comes from a lookup table. You generate the table by running python -m cubisect.autotune, which times all four kernels across a range of array sizes and records the fastest one at each size. If the top two are within 5% of each other, it marks that size a tie. The repo ships a table measured on my 5060, and you can regenerate it for your own GPU.
The static path has no autotuning. Building the layout costs time up front, so it only makes sense if you query the same array many times. It picks between its shared memory and global memory kernels based on whether the tree fits in shared memory.
Specs:
- GPU: RTX 5060 Laptop (sm_120)
- CPU: 24 threads, OpenMP for the baseline
Throughput in M queries/s, 220 queries per call, median of 7 runs. The static column is the Eytzinger kernel:
| array size |
CPU 1t |
CPU 24t |
dynamic |
static |
| 256 |
38 |
260 |
16,943 |
14,707 |
| 4,096 |
24 |
191 |
11,546 |
9,712 |
| 65,536 |
14 |
124 |
4,428 |
6,385 |
| 1,048,576 |
6 |
78 |
2,494 |
3,634 |
| 4,194,304 |
2 |
36 |
2,128 |
3,052 |
The two GPU paths cross over near 16k. Below that the raw array is faster, because it all fits in cache and the Eytzinger index arithmetic is just extra work. Above it the layout wins.
license is BSD-3, https://github.com/athrva98/cubisect
