Deyu/perf

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<LabelT>& filter_labels, LabelT unv_label);
+
         simple_bitmask_buf& _bitmask_filters;
         std::vector<std::uint64_t> _owned_query_bitmask_buf;  // populated only by the 3-arg ctor
         std::vector<std::uint64_t>& _query_bitmask_buf;       // refs either external or _owned

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)

src/filter_match_proxy.cpp

@@ -3,6 +3,35 @@
 namespace diskann
 {
 
+template <typename LabelT>
+void bitmask_filter_match<LabelT>::build_query_mask(const std::vector<LabelT>& 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 <typename LabelT>
 bitmask_filter_match<LabelT>::bitmask_filter_match(
     simple_bitmask_buf& bitmask_filters,
@@ -15,20 +44,7 @@ bitmask_filter_match<LabelT>::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<LabelT>::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);
     }
 }
 

src/index.cpp

@@ -557,7 +557,12 @@ void Index<T, TagT, LabelT>::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<uint64_t>(_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 <typename T, typename TagT, typename LabelT>
 void Index<T, TagT, LabelT>::convert_pts_label_to_bitmask(std::vector<std::vector<LabelT>>& 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<T, TagT, LabelT>::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();

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<size_t> labels;
@@ -124,15 +124,20 @@ 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);
     }
 
     std::uint32_t bitmask_size = static_cast<std::uint32_t>(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<size_t>(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;

src/unified_label_data.cpp

@@ -134,10 +134,14 @@ void unified_label_data_base::collect_label_medoids(std::vector<uint32_t> &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<filter_match_proxy> 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<integer_proxy_owner>(_label_vector, filter_label_ints, _universal_label);
 }

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<std::uint64_t> qbuf;
+    std::vector<std::uint32_t> filter_labels = {9};
+    bitmask_filter_match<std::uint32_t> 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<std::uint64_t> 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<std::uint32_t> f1 = {5};
+        bitmask_filter_match<std::uint32_t> 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<std::uint32_t> f2 = {7}; // point 2 does NOT have label 7
+        bitmask_filter_match<std::uint32_t> 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<std::uint64_t> qbuf;
+    std::vector<std::uint32_t> filter_labels = {64};
+    bitmask_filter_match<std::uint32_t> 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<std::uint64_t> qbuf;
+    std::vector<std::uint32_t> filter_labels = {3};
+    bitmask_filter_match<std::uint32_t> 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<std::uint64_t> qbuf;
+    std::vector<std::uint32_t> filter_labels = {300}; // bit 300 lives in word 4 (tail)
+    bitmask_filter_match<std::uint32_t> 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()