diff --git a/include/filter_match_proxy.h b/include/filter_match_proxy.h index 2b58f8a296..469ff49031 100644 --- a/include/filter_match_proxy.h +++ b/include/filter_match_proxy.h @@ -32,6 +32,12 @@ namespace diskann virtual void prefetch_bitmask(uint32_t id) override; private: + // Single place that sizes _query_bitmask_buf (padded to >=4 words so the + // AVX2 256-bit load in test_full_mask_val never reads past the buffer) + // and folds the filter + universal labels into the query mask. Both ctors + // delegate here after binding _query_bitmask_buf to their chosen storage. + void build_query_mask(const std::vector& filter_labels, LabelT unv_label); + simple_bitmask_buf& _bitmask_filters; std::vector _owned_query_bitmask_buf; // populated only by the 3-arg ctor std::vector& _query_bitmask_buf; // refs either external or _owned diff --git a/include/label_bitmask.h b/include/label_bitmask.h index b601e145e4..c10df9e08a 100644 --- a/include/label_bitmask.h +++ b/include/label_bitmask.h @@ -29,12 +29,27 @@ struct simple_bitmask_full_val struct simple_bitmask_buf { + // Extra trailing uint64 words kept past the logical per-point extent so the + // AVX2 fast path in simple_bitmask::test_full_mask_val can issue a 256-bit + // (4-word) load on the last node without reading past the allocation. Never + // participates in addressing or iteration (those use _bitmask_size), so it + // is inert on the scalar path. + static constexpr std::uint64_t AVX2_TAIL_PADDING = 4; + simple_bitmask_buf() = default; simple_bitmask_buf(std::uint64_t capacity, std::uint64_t bitmask_size) { - _buf.resize(capacity); _bitmask_size = bitmask_size; + _buf.resize(capacity + AVX2_TAIL_PADDING, 0); + } + + // Size the buffer to hold num_points per-point bitmasks plus AVX2 padding. + // Callers pass the logical extent; the padding is added here so no call site + // has to remember it. + void resize_for_points(std::uint64_t num_points) + { + _buf.resize(num_points * _bitmask_size + AVX2_TAIL_PADDING, 0); } std::uint64_t* get_bitmask(std::uint64_t index) diff --git a/src/filter_match_proxy.cpp b/src/filter_match_proxy.cpp index 9d9b609bf6..ce33435402 100644 --- a/src/filter_match_proxy.cpp +++ b/src/filter_match_proxy.cpp @@ -3,6 +3,35 @@ namespace diskann { +template +void bitmask_filter_match::build_query_mask(const std::vector& filter_labels, LabelT unv_label) +{ + // _query_bitmask_buf is already bound (by the ctor init list) to either the + // caller-supplied scratch buffer or the owned buffer. Size it to at least + // AVX2_TAIL_PADDING words so the 256-bit load in test_full_mask_val never + // reads past the end. + // + // clear() first so the resize zero-fills every word: a reused scratch buffer + // may already be at padded_size, in which case a bare resize is a no-op that + // leaves stale high words from a previous query -- test_full_mask_val's + // unconditional 256-bit AND would then treat them as real filter bits. + const std::uint64_t padded_size = + std::max(_bitmask_filters._bitmask_size, simple_bitmask_buf::AVX2_TAIL_PADDING); + _query_bitmask_buf.clear(); + _query_bitmask_buf.resize(padded_size, 0); + _bitmask_full_val._mask = _query_bitmask_buf.data(); + + for (const auto& filter_label : filter_labels) + { + auto bitmask_val = simple_bitmask::get_bitmask_val(filter_label); + _bitmask_full_val.merge_bitmask_val(bitmask_val); + } + + // if unv isn't set, it will be default value 0 + auto bitmask_val = simple_bitmask::get_bitmask_val(unv_label); + _bitmask_full_val.merge_bitmask_val(bitmask_val); +} + template bitmask_filter_match::bitmask_filter_match( simple_bitmask_buf& bitmask_filters, @@ -15,20 +44,7 @@ bitmask_filter_match::bitmask_filter_match( // _bitmask_size == 0 means no filter is set if (_bitmask_filters._bitmask_size > 0) { - // Pad to at least 4 words (32 bytes) for safe AVX2 256-bit loads - size_t padded_size = std::max(_bitmask_filters._bitmask_size, (std::uint64_t)4); - query_bitmask_buf.resize(padded_size, 0); - _bitmask_full_val._mask = query_bitmask_buf.data(); - - for (const auto& filter_label : filter_labels) - { - auto bitmask_val = simple_bitmask::get_bitmask_val(filter_label); - _bitmask_full_val.merge_bitmask_val(bitmask_val); - } - - // if unv isn't set, it will be default value 0 - auto bitmask_val = simple_bitmask::get_bitmask_val(unv_label); - _bitmask_full_val.merge_bitmask_val(bitmask_val); + build_query_mask(filter_labels, unv_label); } } @@ -42,17 +58,7 @@ bitmask_filter_match::bitmask_filter_match( { if (_bitmask_filters._bitmask_size > 0) { - _query_bitmask_buf.resize(_bitmask_filters._bitmask_size, 0); - _bitmask_full_val._mask = _query_bitmask_buf.data(); - - for (const auto& filter_label : filter_labels) - { - auto bitmask_val = simple_bitmask::get_bitmask_val(filter_label); - _bitmask_full_val.merge_bitmask_val(bitmask_val); - } - - auto bitmask_val = simple_bitmask::get_bitmask_val(unv_label); - _bitmask_full_val.merge_bitmask_val(bitmask_val); + build_query_mask(filter_labels, unv_label); } } diff --git a/src/index.cpp b/src/index.cpp index 3a9c2d9b3a..76d7422598 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -557,7 +557,12 @@ void Index::save_unified(const char *filename, const std::vecto } else if (_bitmask_buf._buf.size() > 0) { - const uint64_t bitmap_bytes = _bitmask_buf._buf.size() * sizeof(uint64_t); + // _buf is over-allocated by up to 4 words for AVX2 padding + // (convert_pts_label_to_bitmask). The unified format stores exactly + // npts * bitmask_size words, matching the reader's expectation in + // unified_label_data_bitmask::load_encoding. + const uint64_t exact_words = static_cast(_nd) * _bitmask_buf._bitmask_size; + const uint64_t bitmap_bytes = exact_words * sizeof(uint64_t); writer.write_labels_bitmask(total_labels, universal, dict_bytes.data(), dict_bytes.size(), _bitmask_buf._buf.data(), bitmap_bytes); } @@ -2475,8 +2480,7 @@ template void Index::convert_pts_label_to_bitmask(std::vector>& pts_to_labels, simple_bitmask_buf& bitmask_buf, size_t num_labels) { _bitmask_buf._bitmask_size = simple_bitmask::get_bitmask_size(num_labels + 1); - // Add 4 extra uint64 words at end for safe AVX2 256-bit loads on last node - _bitmask_buf._buf.resize(pts_to_labels.size() * _bitmask_buf._bitmask_size + 4, 0); + _bitmask_buf.resize_for_points(pts_to_labels.size()); for (size_t i = 0; i < pts_to_labels.size(); i++) { @@ -2573,8 +2577,12 @@ void Index::aggregate_points_by_bitmask_label( std::advance(itr, lbl); auto& x = *itr; + // Pad to >=AVX2_TAIL_PADDING words: test_full_mask_val issues an + // unconditional 256-bit load over this query mask, so it must be at least + // 4 words long. clear() first so resize zero-fills every word (no stale + // bits carried across label iterations). label_bitmask.clear(); - label_bitmask.resize(_bitmask_buf._bitmask_size, 0); + label_bitmask.resize(std::max(_bitmask_buf._bitmask_size, simple_bitmask_buf::AVX2_TAIL_PADDING), 0); simple_bitmask_full_val bitmask_full_val; bitmask_full_val._mask = label_bitmask.data(); diff --git a/src/label_helper.cpp b/src/label_helper.cpp index df41fe0f4f..eee7fe5a5a 100644 --- a/src/label_helper.cpp +++ b/src/label_helper.cpp @@ -50,11 +50,11 @@ bool label_helper::parse_label_file_in_bitset( bitmask_buf._bitmask_size = simple_bitmask::get_bitmask_size(num_labels + 1); if (num_points > line_cnt) { - bitmask_buf._buf.resize(num_points * bitmask_buf._bitmask_size, 0); + bitmask_buf.resize_for_points(num_points); } else { - bitmask_buf._buf.resize(line_cnt * bitmask_buf._bitmask_size, 0); + bitmask_buf.resize_for_points(line_cnt); } tsl::robin_set labels; @@ -124,7 +124,12 @@ bool label_helper::write_bitmask_to_file(const std::string& bitmask_label_file, throw diskann::ANNException(std::string("Failed to open file ") + bitmask_label_file, -1); } - if (bitmask_buf._buf.size() != bitmask_buf._bitmask_size * num_points) + // convert_pts_label_to_bitmask over-allocates by up to 4 words for AVX2 + // padding, so the buffer may be larger than the exact per-point extent. + // Require at least the exact extent; only the unpadded words are written so + // the file format (num_points * bitmask_size words) stays canonical. + const size_t exact_words = bitmask_buf._bitmask_size * num_points; + if (bitmask_buf._buf.size() < exact_words) { throw diskann::ANNException(std::string("Bitmask buffer is empty"), -1); } @@ -132,7 +137,7 @@ bool label_helper::write_bitmask_to_file(const std::string& bitmask_label_file, std::uint32_t bitmask_size = static_cast(bitmask_buf._bitmask_size); outfile.write((char*)(&num_points), sizeof(std::uint32_t)); outfile.write((char*)(&bitmask_size), sizeof(std::uint32_t)); - outfile.write((char *)bitmask_buf._buf.data(), bitmask_buf._buf.size() * sizeof(std::uint64_t)); + outfile.write((char *)bitmask_buf._buf.data(), exact_words * sizeof(std::uint64_t)); outfile.close(); return true; @@ -164,14 +169,14 @@ bool label_helper::read_bitmask_from_file(const std::string& bitmask_label_file, } bitmask_buf._bitmask_size = bitmask_size; - // num_points > num_points_in_file in dynamic index + // num_points > num_points_in_file in dynamic index if (num_points > static_cast(num_points_in_file)) { - bitmask_buf._buf.resize(num_points * bitmask_buf._bitmask_size, 0); + bitmask_buf.resize_for_points(num_points); } else { - bitmask_buf._buf.resize(num_points_in_file * bitmask_buf._bitmask_size, 0); + bitmask_buf.resize_for_points(num_points_in_file); } num_points = num_points_in_file; diff --git a/src/unified_label_data.cpp b/src/unified_label_data.cpp index 07f7ab7767..6566e75003 100644 --- a/src/unified_label_data.cpp +++ b/src/unified_label_data.cpp @@ -134,10 +134,14 @@ void unified_label_data_base::collect_label_medoids(std::vector &out) void unified_label_data_bitmask::load_encoding(UnifiedIndexReader &r, const UnifiedIndexHeader &h, uint64_t npts) { const uint64_t bitmask_words = simple_bitmask::get_bitmask_size(h.total_labels); - _bitmask_buf = simple_bitmask_buf(npts * bitmask_words, bitmask_words); + const uint64_t exact_words = npts * bitmask_words; + // Ctor sizes _buf to exact_words + AVX2 tail padding; the padding stays zero + // and is excluded from the region-size check and zero-copy read below (both + // use exact_words), keeping the on-disk format canonical. + _bitmask_buf = simple_bitmask_buf(exact_words, bitmask_words); if (h.per_point_labels_len > 0) { - const uint64_t expected_bytes = _bitmask_buf._buf.size() * sizeof(std::uint64_t); + const uint64_t expected_bytes = exact_words * sizeof(std::uint64_t); if (h.per_point_labels_len != expected_bytes) { throw ANNException("unified_label_data_bitmask: bitmask region size mismatch", -1, __FUNCSIG__, __FILE__, @@ -217,6 +221,10 @@ std::unique_ptr unified_label_data_integer::make_match_proxy { return inner.contain_filtered_label(id); } + void prefetch_bitmask(uint32_t id) override + { + inner.prefetch_bitmask(id); + } }; return std::make_unique(_label_vector, filter_label_ints, _universal_label); } diff --git a/tests/filter_match_proxy_tests.cpp b/tests/filter_match_proxy_tests.cpp index 18317faf4a..76b7513f74 100644 --- a/tests/filter_match_proxy_tests.cpp +++ b/tests/filter_match_proxy_tests.cpp @@ -329,4 +329,132 @@ BOOST_AUTO_TEST_CASE(integer_contain_matches_on_filter_or_unv) BOOST_TEST(m.contain_filtered_label(3) == true); // has 4 } +// ---- AVX2 correctness corner cases ----------------------------------------- +// +// test_full_mask_val's AVX2 fast path (_bitmask_size <= 4) issues an +// UNCONDITIONAL 256-bit load over both the query and the node bitmask, then a +// single 256-bit AND. This only agrees with the scalar semantics ("intersect +// the first _bitmask_size words") if both operands are zero in the padding +// lanes past _bitmask_size. These tests pin the invariants that make that true; +// each one previously failed when the query buffer was not zero-padded to 4 +// words or was reused without re-zeroing. + +// A single-word bitmask (size 1) means the AVX2 load reads 3 words of the NEXT +// three points as the query's high lanes. If a match on point p leaked in from +// point p+1/p+2/p+3's labels, this would report a false positive. Only the +// point that actually carries the filter label may match. +BOOST_AUTO_TEST_CASE(bitmask_size1_no_cross_point_leak) +{ + // 64 bits -> size 1. Give ONLY point 5 the filter label; its neighbours + // 6/7/8 carry other labels that must not leak into point 5's test. + auto filters = make_bitmask_buf(/*num_points*/ 16, /*total_bits*/ 64, + /*point_id*/ 5, /*set_labels*/ {9}); + // Set unrelated labels on the following points (whose words the AVX2 load + // over point 5 will also touch as high lanes). + { + simple_bitmask bm6(filters.get_bitmask(6), filters._bitmask_size); + bm6.set(9); // same label as the filter, but on a DIFFERENT point + simple_bitmask bm7(filters.get_bitmask(7), filters._bitmask_size); + bm7.set(9); + simple_bitmask bm8(filters.get_bitmask(8), filters._bitmask_size); + bm8.set(9); + } + BOOST_TEST(filters._bitmask_size == (std::uint64_t)1); + + std::vector qbuf; + std::vector filter_labels = {9}; + bitmask_filter_match m(filters, qbuf, filter_labels, /*unv*/ 0); + + BOOST_TEST(m.contain_filtered_label(5) == true); // really has label 9 + BOOST_TEST(m.contain_filtered_label(4) == false); // no labels + // Points 6/7/8 DO carry label 9, so they legitimately match -- the point is + // that testing point 5 must not have leaked THEIR bits into point 5's result + // (already asserted by 5==true / 4==false above). Confirm the boundary point + // just past the filter label owner behaves per its own bits. + BOOST_TEST(m.contain_filtered_label(6) == true); +} + +// Reusing one query buffer across proxies is the exact scenario that broke: +// the second ctor sees a buffer already at size 4 and must still zero every +// word (clear+resize), or stale high/low words from the first proxy produce +// false matches. +BOOST_AUTO_TEST_CASE(bitmask_reused_query_buffer_is_rezeroed) +{ + auto filters = make_bitmask_buf(8, 64, /*point*/ 2, /*labels*/ {5}); + std::vector qbuf; + + // First proxy: filter label 5, and deliberately a label whose bit lands in + // a high word would require size>1; here size==1 so all bits are word 0. + { + std::vector f1 = {5}; + bitmask_filter_match m1(filters, qbuf, f1, /*unv*/ 0); + BOOST_TEST(m1.contain_filtered_label(2) == true); + BOOST_TEST(qbuf.size() == (size_t)4); // padded + } + // qbuf now retains proxy #1's bits (filter label 5 in word 0). A second + // proxy filtering a DISJOINT label must not inherit label 5. + { + std::vector f2 = {7}; // point 2 does NOT have label 7 + bitmask_filter_match m2(filters, qbuf, f2, /*unv*/ 0); + // If qbuf's word 0 still held label 5's bit, point 2 (which has 5) would + // falsely match. It must not. + BOOST_TEST(m2.contain_filtered_label(2) == false); + } +} + +// Labels on 64-bit word boundaries: bit 63 is the last bit of word 0, bit 64 +// the first of word 1. With size 2 the AVX2 fast path still applies. A filter +// on bit 64 must match a point carrying bit 64 and not one carrying bit 63. +BOOST_AUTO_TEST_CASE(bitmask_word_boundary_labels_63_and_64) +{ + auto filters = make_bitmask_buf(4, /*total_bits*/ 128, /*point*/ 0, /*labels*/ {63}); + { + simple_bitmask bm1(filters.get_bitmask(1), filters._bitmask_size); + bm1.set(64); + } + BOOST_TEST(filters._bitmask_size == (std::uint64_t)2); + + std::vector qbuf; + std::vector filter_labels = {64}; + bitmask_filter_match m(filters, qbuf, filter_labels, /*unv*/ 0); + + BOOST_TEST(m.contain_filtered_label(1) == true); // has bit 64 + BOOST_TEST(m.contain_filtered_label(0) == false); // has bit 63, not 64 +} + +// Universal label participates in the query mask via merge. A point carrying +// only the universal label must match even when it lacks every filter label. +BOOST_AUTO_TEST_CASE(bitmask_universal_label_matches) +{ + auto filters = make_bitmask_buf(4, 64, /*point*/ 0, /*labels*/ {7}); // unv label + { + simple_bitmask bm1(filters.get_bitmask(1), filters._bitmask_size); + bm1.set(3); // a real filter label + } + std::vector qbuf; + std::vector filter_labels = {3}; + bitmask_filter_match m(filters, qbuf, filter_labels, /*unv*/ 7); + + BOOST_TEST(m.contain_filtered_label(0) == true); // has universal label 7 + BOOST_TEST(m.contain_filtered_label(1) == true); // has filter label 3 + BOOST_TEST(m.contain_filtered_label(2) == false); // has neither +} + +// A bitmask larger than 4 words exercises the AVX2 blocked loop plus the +// scalar tail (size not a multiple of 4). A match whose only shared bit lives +// in the tail words must still be found. +BOOST_AUTO_TEST_CASE(bitmask_large_size_tail_match) +{ + // 320 bits -> size 5: one 4-word AVX2 block + a 1-word scalar tail. + auto filters = make_bitmask_buf(4, /*total_bits*/ 320, /*point*/ 2, /*labels*/ {300}); + BOOST_TEST(filters._bitmask_size == (std::uint64_t)5); + + std::vector qbuf; + std::vector filter_labels = {300}; // bit 300 lives in word 4 (tail) + bitmask_filter_match m(filters, qbuf, filter_labels, /*unv*/ 0); + + BOOST_TEST(m.contain_filtered_label(2) == true); // shared bit in the tail + BOOST_TEST(m.contain_filtered_label(0) == false); +} + BOOST_AUTO_TEST_SUITE_END()