Repository navigation
Fix the IVF-PQ list codepacking loops to stride over the whole grid - #2750
Conversation
write_list, write_list_flat and run_on_list started at the global row index but advanced by one block's worth of rows, so block b processed every row from its first one to the end of the list. Each row was encoded (or packed/unpacked) up to 16 times, sequentially in the first blocks; with pq_dim 3072 that made an encode launch take ~130 ms. Advance by the whole grid instead, so every row is processed exactly once. The codes are unchanged.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe IVF-PQ codepacking loops now advance vector indices using grid-wide strides. The copyright notice now includes affiliates. ChangesIVF-PQ codepacking
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to The grid-stride changes address repeated row processing, and the inspected packed-byte paths do not show a remaining merge-blocking issue. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/merge |
write_list,write_list_flatandrun_on_listinivf_pq_codepacking.cuhstarted at the global row index but advanced by one block's worth of rows:Their launchers already size the grid to cover every row (
encode_list_datalaunches one block per 8 rows, or 16 for 4-bit codes; pack/unpack/reconstruct launch one block per 256 rows), so blockbprocessed every row from its first one to the end of the list. Rows were encoded or packed up ton_rows / rows_per_blocktimes, and the subwarps of block 0 walked the entire list, each row a serial loop over allpq_dimsubspaces. For the ~128-row lists inNEIGHBORS_ANN_IVF_PQ_TESTthat is up to 16 passes per row; atpq_dim = 3072oneencode_list_data_interleaved_kernellaunch took ~127–146 ms.Affected paths:
encode_list_data, used byivf_pq::helpers::codepacker::extend_list, for any list longer than one block (8–16 rows).pack_list_data,unpack_list_data,{un}pack_contiguous_list_dataandreconstruct_list_datafor lists longer than 256 rows. The C API (cuvsIvfPqIndexUnpackContiguousListData) and Pythonindex.lists()go through these; unpacking a 100k-row list did ~2·10⁷ row-unpacks instead of 10⁵.The main
build/extendencoder (process_and_fill_codes_kernel) does not use these loops and is unaffected.This PR makes the three loops grid-stride loops by multiplying the stride by
gridDim.x. Launch configurations are unchanged.encode_vectors, or the pack/unpack/reconstruct action) with the same lane layout, codebook scan order and reduction order, and is written to the same location. The only difference is that each row is now processed exactly once instead of up to 16 times.pq_bits < 8, the FLAT-layout encoder andunpack_contiguousdo bitfield read-modify-writes on shared bytes, and two blocks processing the same row could collide.The commit also updates the file's SPDX copyright header to the canonical form required by the copyright pre-commit hook.
Testing
All IVF-PQ test executables and the full
ctestsuite pass. The existing tests already check that packing round-trips byte-exactly and that reconstruct → extend reproduces the vectors.Measurements
Single-process wall time on an RTX 6000 Ada (48 GB), otherwise idle (no ctest parallelism). The old and new
libcuvs.sowere run alternately, 2 or more repetitions each, with the order reversed between repetitions; mean (min–max). The two builds differ only by this change.NEIGHBORS_ANN_IVF_PQ_TESTThe other executables measured the same way (IVF-Flat, IVF-SQ, CAGRA, multi-GPU, tiered index, dynamic batching, all-neighbors) were unchanged within noise. All tests passed in every run.