Speed up IndexFlatPanorama search with query blocking - #5447
Conversation
|
@mnorris11, could you take a look at this when you get a chance? Thanks! |
We are a bit behind in reviewing things. @AlSchlo @aknayar , how does this look? |
aknayar
left a comment
aknayar
left a comment
There was a problem hiding this comment.
@mnorris11 Did a high-level pass—this definitely seems like a viable optimization. I would make sure it is thoroughly benchmarked for different configurations (e.g., 8-64 queries + 8-32 threads, for which I noticed some regressions locally). Also I'm not sure how we'd extend it for IVFFlatPanorama (given that we can no longer trivially exploit the property of each query accessing the exact same base vectors) but perhaps there's a clean solution I'm missing. @mulugetam Thank you for the PR :)
| size_t n_batches = (index.ntotal + bs - 1) / bs; | ||
|
|
||
| size_t n_blocks = (size_t(n) + qbs - 1) / qbs; | ||
| [[maybe_unused]] int nt = std::min(int(n_blocks), omp_get_max_threads()); |
There was a problem hiding this comment.
I don't fully understand this. Personally I would try to keep the same number of threads and assign each one the same n/nt slice as before (rounded to qbs queries).
There was a problem hiding this comment.
Good catch. Fixed by treating qbs as a maximum: the effective block size shrinks to ceil(n / max_threads), so the thread count matches the unblocked path and blocking degrades gracefully to the original schedule when queries are scarce.
| const uint8_t* storage_base = codes_base + batch_offset; | ||
|
|
||
| // Per-query init of the active set (identical to the scalar path). | ||
| std::vector<uint8_t> first_level_full(block_size); |
There was a problem hiding this comment.
nit: This might not change anything but should we hoist this allocation out into the caller?
There was a problem hiding this comment.
Done. first_level_full is now caller-allocated alongside the other per-thread scratch and passed in as a pointer.
5d1c5ed to
8f88464
Compare
|
Thanks for the review @aknayar @mnorris11 Below is summary of the new result and the gains I see on my machine. Raw result: https://gist.github.com/mulugetam/bda9537802ed63f891cbe3b6370deb86#file-pr5447-stdout The Pano baseline QPS = 0 block size And below is the result comparing the existing implementation with this PR for query-block-size = 0. Raw result: https://gist.github.com/mulugetam/bda9537802ed63f891cbe3b6370deb86#file-current-stdout |
|
Thanks @mulugetam. One more thing I would recommend is benching this optimization with a |
Thanks @aknayar. Below is measured on top of #5441, which changes the default batch size to 1024. Raw data: https://gist.github.com/mulugetam/bda9537802ed63f891cbe3b6370deb86#file-pr5447_and_5441-stdout |
|
@mnorris11 has imported this pull request. If you are a Meta employee, you can view this in D114380042. |
| /// 1 selects the original query-at-a-time path. Results are identical for any | ||
| /// value; this only trades off cache behavior. Set before searching; changing | ||
| /// it concurrently with in-flight searches is not thread-safe. | ||
| FAISS_API extern size_t panorama_query_block_size; |
There was a problem hiding this comment.
Can this be moved out of IndexFlat.h/cpp to Panorama.h/cpp?
There was a problem hiding this comment.
Thanks, @mnorris11! I just did that and rebased it.
Could you also take a look at #5341? We plan to apply it to speed up OpenSearch-KNN, and I think it would be a big benefit.
Process queries in blocks so each DB level-storage block is streamed from
memory once and reused across the block while resident in a faster cache
level, raising cache-bandwidth efficiency on the bandwidth-bound Panorama dot
kernel. The DB-batch loop stays outside the block (preserving each query's
threshold evolution) and within a batch the loop runs level-outer /
query-inner.
Controlled by the global panorama_query_block_size, following the existing
FAISS tunable idiom (cf. distance_compute_blas_query_bs). It defaults to an
enabled value of 32; 0 or 1 selects the original query-at-a-time path. The
value is a maximum: the effective block size shrinks to ceil(n / max_threads)
when queries are scarce, so the thread count always matches the unblocked
path (min(n, max_threads)) and blocking degrades gracefully to the original
schedule instead of trading threads for full blocks. The filter kernel is
non-allocating; all per-thread scratch is allocated once by the caller.
Pure loop-order transform: each query keeps its own active set, threshold,
and pruning decisions, so results are bit-identical to the original path
(verified 432/432 equality cases + unchanged SIFT1M recall, and across
metrics {L2, IP} x threads {1, 4, 32, 240} x block sizes {2, 7, 32, 64}).
Top-k only; range search falls back to the original path since
RangeSearchPartialResult requires per-query-contiguous appends.
SIFT1M speedups vs the unblocked path: ~1.8-2.1x single-thread at any batch
size, 1.5-2.5x multithreaded once each thread holds a full block, up to
~3.3x at nq=1000 with 16-32 threads, and no regression (1.02-1.07x) in the
small-batch corner where nq is close to the thread count and the effective
block size collapses by design.
Also extend bench_flat_l2_panorama so the speedup is measurable directly:
--query-block-size, --nq, and --threads accept comma-separated sweeps (the
indexes are built once; only nq, the thread count, and the block-size global
change between timed runs), and --repeat times each configuration N times
reporting the fastest, needed at small nq where single runs are noisy. The
summary prints per-(nq, threads) speedups vs plain Flat and vs the Panorama
baseline. E.g.:
python benchs/bench_flat_l2_panorama.py --dataset sift1m \
--nq 8,16,32,64 --threads 8,16,32 --query-block-size 0,32 \
--repeat 20
Signed-off-by: Mulugeta Mammo <mulugeta.mammo@intel.com>
Move panorama_query_block_size to faiss/impl/Panorama.{h,cpp}
The global belongs with the Panorama implementation, not IndexFlat.
IndexFlat.h already includes Panorama.h so all existing users still
compile without changes.
Signed-off-by: Mulugeta Mammo <mulugeta.mammo@intel.com>
b73411f to
c4ac96d
Compare
|
@mnorris11 fixed the merge conflict after #5441 got merged. |
3 participants
The Panorama flat-search dot-product kernel is memory-bandwidth bound: every query repeatedly streams the same level-major database layout from memory. This change improves cache reuse by processing queries in blocks (configured via
panorama_query_block_size). The search loop is reordered so that the database-batch loop wraps a level-outer/query-inner traversal, allowing each level block to be loaded once and reused across all queries in the block while it remains in cache.This optimization delivers approximately 2.2× speedup for
n_levels=8and up to 2.8× forn_levels=4.The change is purely a loop-order transformation. Each query maintains its own active set, pruning threshold, and traversal state, producing bit-identical results to the original query-at-a-time implementation.
Results
SIFT1M,
n_levels=8, single core, mean of 3 runsThe optimization consistently achieves 2.2–2.3× speedups for
n_levels=8and up to 2.8× forn_levels=4. It also maintains or improves multithreaded performance by reducing shared-cache contention. Recall and the percentage of dimensions scanned remain unchanged.Implementation Notes
panorama_query_block_size(FAISS_API extern size_t), following the same tuning pattern asdistance_compute_blas_query_bs. The default is 32. Setting it to0or1restores the original implementation. Configure it before issuing searches; changing it while searches are in flight is not thread-safe.SingleResultHandleris maintained for each query slot, sincebegin()reinitializes the heap and must be called once per query, not once per batch.qbs > 1 && !use_radius) becauseRangeSearchPartialResultrequires per-query contiguous appends, which are incompatible with the interleaved execution schedule.Verification
A dedicated test (
test_query_blocking_bit_identicalintests/test_flat_panorama.py) verifies bit-identical distances and IDs across a broad configuration sweep and all supported SIMD implementations. End-to-end SIFT1M recall is unchanged.Reproducing the Results
OPENBLAS_NUM_THREADS=1 taskset -c 2 \ python benchs/bench_flat_l2_panorama.py \ --dataset sift1m --nq 1000 --query-block-size 0,16,32,64