perf(hgraph): fuse RaBitQ split storage and search for x=1..4 - #2596
perf(hgraph): fuse RaBitQ split storage and search for x=1..4#2596LightWant wants to merge 1 commit into
Conversation
|
/label status/waiting-for-review |
|
Automated pull request review completed. Review effort: Submitted 2 inline comments. |
Merge Protections🟢 All 2 merge protections satisfied — ready to merge. Show 2 satisfied protections🟢 Require kind label
🟢 Require version label
|
LHT129
left a comment
There was a problem hiding this comment.
PR Review Summary
This PR adds an opt-in rabitq_fused_datacell feature that co-locates bottom-layer HGraph neighbors, cluster id, label, x-bit filter code, and y-bit supplement code into a single cache-line-aligned node record. It supports x=1..4 filter bits with x+y <= 8, uses 16 reproducibly trained residual clusters, and adds dedicated SIMD kernels (AVX2, AVX512, AVX512VPOPCNTDQ, generic).
Overall Assessment
The implementation is well-structured and thorough. The fused layout eliminates an indirection from the search hot path, and the dedicated HGraphRaBitQSearcher with AffineScorer templates provides a clean separation from the existing search path. Test coverage is good with functional tests for round-trip serialization, duplicate label handling, and alias expansion.
Key Strengths
- Clean opt-in design: The fused path is gated behind a single boolean parameter and does not affect the existing code path
- Good test coverage: Integration tests cover serialization round-trips, duplicate labels, alias expansion, and iterator search for the fused path
- SIMD specialization: Properly dispatches to AVX2/AVX512/AVX512VPOPCNTDQ kernels with a generic fallback
- Safe feature gating: Correctly disables incompatible features (merge, tune, concurrent mutations) at the
InitFeatureslevel - Fast-math safety:
IsFiniteRaBitQValueuses bitwise checks to avoid-ffast-mathissues withstd::isfinite
Issues Found (see inline comments for details)
- Thread safety:
VisitedList::TestAndSetuses non-atomic operations but is called under shared lock in the fused searcher - Numerical robustness:
filter_rescalecomputation inEncodeFusedAffineMetadatacan overflow for edge-case inputs - Code organization:
search_one_graphis ~700 lines and would benefit from being split into smaller functions - Serialization versioning: The new format lacks an embedded version tag for forward compatibility
- Performance:
BoundedResults::Insertis O(n) per insert; acceptable for typical ef but worth documenting - Encoding performance:
EncodeFusedAffineMetadatadecodes bit-planes per dimension, which could be slow for large datasets
Recommendation
Approve with suggestions. The issues identified are non-blocking for correctness — the code handles edge cases gracefully (NaN fallbacks, feature gating). The thread safety concern in TestAndSet is pre-existing in the Set/Get methods and not introduced by this PR, but worth addressing in a follow-up.
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (8779 changed lines across 55 files).
Submitted 3 inline comments.
Reviewed commit 2c8430b.
2c8430b to
346efc7
Compare
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (9139 changed lines across 60 files).
Submitted 2 inline comments.
Reviewed commit 346efc7.
Signed-off-by: zhuangye.yxw <2510035537@qq.com> Assisted-by: Codex:gpt-5
346efc7 to
3305f72
Compare
| if (exact_filter_ip_hint and IsFiniteRaBitQValue(filter_inner_products[i])) { | ||
| ++full_count; | ||
| used_hint = | ||
| split_codes->ComputeFusedFullWithFilterIP(computer, |
There was a problem hiding this comment.
Bug: full_count double-counted when filter-IP hint path falls back to full compute
When exact_filter_ip_hint is true and the hint-based ComputeFusedFullWithFilterIP returns false, full_count is incremented both at line 79 (before the hint attempt) and at line 93 (in the fallback else branch). This double-counts every candidate that enters the hint path but falls back.
Suggested fix: remove the first ++full_count on line 79 and only increment in the success/fallback branches, or move it into the if (used_hint) branch.
| Vector<uint8_t> replacement( | ||
| static_cast<uint64_t>(new_capacity) * record_size_ + K_CACHE_LINE_SIZE - 1, 0, allocator_); | ||
| const auto replacement_address = reinterpret_cast<uintptr_t>(replacement.data()); | ||
| const uint64_t replacement_offset = |
There was a problem hiding this comment.
Thread safety: GetNodeRecord reads storage_ and aligned_offset_ without synchronization
GetNodeRecord (called from the search hot path via GetNodeView, GetCodeView, GetFusedCodeView, Prefetch, etc.) reads storage_.data() + aligned_offset_ without holding storage_mutex_. Meanwhile, Reallocate (called from Resize, ShrinkToFit) replaces storage_ and updates aligned_offset_ under storage_mutex_, but the search path never acquires even a shared lock.
This is a data race: a concurrent resize could free the old storage_ buffer while a search thread is still reading from it. The InitFeatures disables concurrent add+search for fused graphs, but Resize can also be triggered by InsertNeighborsById during build and SetFusedCodes during sync_fused_node_codes.
Consider either:
- Holding
storage_mutex_shared in the search path (may hurt performance), or - Ensuring
Reallocatenever frees the old buffer while readers exist (e.g., RCU-style deferred reclamation), or - Documenting that resize and search are externally serialized for fused graphs.
| return true; | ||
| } | ||
| } | ||
| ++rabitq_reorder_fallback_full_count; |
There was a problem hiding this comment.
Stats: rabitq_reorder_fallback_full_count over-counts for x=1 (LegacyOneBitScorer)
When K_EXACT_FILTER_HINT is false (i.e., for LegacyOneBitScorer used with x=1), the if constexpr block at line 633 is skipped entirely, and every call to refine_node unconditionally increments rabitq_reorder_fallback_full_count. This means the metric counts all rerank computations as "fallbacks" even though there is no filter-IP hint to fall back from.
Consider guarding the increment with if constexpr (Scorer::K_EXACT_FILTER_HINT) so the fallback count is only incremented when a hint was actually available but not usable.
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (9605 changed lines across 61 files).
Submitted 2 inline comments.
Reviewed commit 3305f72.
| query_sum = 0.0F; | ||
| for (uint64_t i = 0; i < this->dim_; ++i) { | ||
| auto magnitude = | ||
| static_cast<uint32_t>(query_scale * std::fabs(transformed_query[i]) + 1e-5F); |
There was a problem hiding this comment.
[P2] [P2] Reject non-finite values before integer query quantization
A fused 1+y query containing NaN or infinity reaches this cast because public search validation checks only dimension and element count. The computed scale is zero for an infinite norm or the false branch for a NaN norm, but 0 * fabs(NaN/Inf) is NaN, and converting that value to uint32_t is undefined behavior. Reject non-finite query coordinates or take a defined fallback before performing the floating-to-integer conversion.
| layout.remove_flag_bit == expected_layout.remove_flag_bit, | ||
| "fused graph remove parameters do not match construction parameters"); | ||
|
|
||
| auto codec_model = StreamReader::ReadString(reader); |
There was a problem hiding this comment.
[P2] [P2] Validate the serialized codec length before allocating it
ReadString reads the wire-supplied uint64_t length and constructs std::vector<char>(length) before its subsequent read can detect that the input is truncated. A corrupt fused index can therefore request an arbitrarily large allocation and cause bad_alloc, severe memory pressure, or process termination instead of being rejected by the surrounding layout checks. Read and validate the length against the reader's remaining bytes and a valid codec-model bound before allocating.
Change Type
Linked Issue
What Changed
rabitq_fused_datacellHGraph layout for RaBitQ split1+y,2+y,3+y, and4+y(x+y <= 8); the default remains disabled.1+y, and exact x-bit filter-inner-product reuse for2+ythrough4+yso reorder scans only the supplement planes.Test Evidence
make fmtmake lintmake testmake cov, run tests, and collect coverageTest details:
GIST1M benchmark: 960-d L2, HGraph M=32, efConstruction=200, 16 clusters, MCI disabled, topK=10, 32 query threads, efSearch=10,20,40,80,160,320,640,1280. Each point used 20 searches per process and 3 paired alternating runs.
All 32 measured efSearch/split points outperform the existing split implementation.
Compatibility Impact
rabitq_fused_datacell=true; unsupported fused-v1 combinations fail validation explicitly.Performance and Concurrency Impact
Documentation Impact
README.mdDEVELOPMENT.mdCONTRIBUTING.mddocs/docs/en/src/quantization/rabitq_split.md,docs/docs/zh/src/quantization/rabitq_split.mdRisk and Rollback
rabitq_fused_datacellfor immediate configuration rollback, or revert commit3305f72bd.Checklist
kind/bugandkind/feature; see "Linked Issue" above)[skip ci]prefix)