fix(hnsw): keep per-ISA search builds from sharing generic code

.github/scripts/ci_changes.py

@@ -234,6 +234,9 @@ def classify(paths):
                 selected.add("python_quality")
             elif path == "scripts/check-tidy.sh":
                 selected.add("cpp_tidy")
+            elif path == "scripts/check-isa-symbols.sh":
+                # Run by the static-analysis and installed-consumer jobs.
+                selected.update(("cpp", "cpp_tidy"))
             elif path != "scripts/check-includes.sh":
                 selected.add("cpp_format")
         else:

.github/scripts/test_ci_changes.py

@@ -222,6 +222,7 @@ def test_language_specific_changes(self):
             "scripts/check-tidy.sh": {"shell", "cpp_tidy", "python_quality"},
             "scripts/check-format.sh": {"shell", "cpp_format", "python_quality"},
             "scripts/check-includes.sh": {"shell"},
+            "scripts/check-isa-symbols.sh": {"shell", "cpp", "cpp_tidy"},
             ".clang-format": {"cpp_format"},
             ".clang-tidy": {"cpp_tidy"},
             ".github/workflows/test.yaml": {

.github/workflows/test.yaml

@@ -57,6 +57,8 @@ jobs:
           -Dpybind11_DIR="$pybind11_DIR"
       - name: Run clang-tidy
         run: ./scripts/check-tidy.sh build-tidy
+      - name: Check ISA-specific symbols with Clang
+        run: ./scripts/check-isa-symbols.sh build-tidy
 
   shellcheck:
     needs: changes
@@ -348,6 +350,8 @@ jobs:
           cmake --build build-install
       - name: Test portable Release build
         run: ctest --test-dir build-install --output-on-failure --no-tests=error
+      - name: Check ISA-specific symbols with GCC
+        run: ./scripts/check-isa-symbols.sh build-install
       - name: Download pinned Intel SDE
         run: |
           curl --fail --location --retry 3 --max-time 120 \

AGENTS.md

@@ -38,7 +38,7 @@ Tool versions and installation instructions are maintained in [docs/development.
 | Python sources | `./scripts/check-python.sh` and affected Python tests via `python -m pytest` |
 | Bindings or Python-visible C++ behavior | C++ and Python checks above; rebuild/install the package in the existing environment before testing, then verify the imported package/extension paths and that the extension includes the current changes |
 | Python API signatures, type declarations or accepted-input contracts | Keep `python_bindings/__init__.pyi` and runtime docstrings synchronized; run the [editor/type checks](tests/README.md#python-editor-support), including the explicitly selected mypy and Pyright tests outside `tests/python` |
-| SIMD, metrics, packing, layouts, or index formats | Before building, audit every backend and consumer sharing the contract; build and run relevant reference/compatibility tests with `RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF` |
+| SIMD, metrics, packing, layouts, or index formats | Before building, audit every backend and consumer sharing the contract; build and run relevant reference/compatibility tests with `RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF`; run `./scripts/check-isa-symbols.sh` on that build when ISA-flagged sources or headers they include change |
 | Hot-path algorithm/allocation changes or performance claims | Correctness tests first, then a representative benchmark reporting dataset, CPU, compiler, ISA, threads, latency/throughput, and quality; a correctness fix may omit benchmarks with a documented justification, without making an unmeasured performance claim |
 | Examples | Build or run the affected example |
 | Shell scripts | ShellCheck on affected scripts; see contributor instructions |

CMakeLists.txt

@@ -252,6 +252,7 @@ if(RABITQ_BUILD_INDEX)
         src/index/pipnn.cpp
         src/index/hnsw_dispatch.cpp
         src/index/hnsw_search_generic.cpp
+        src/index/visited_set_hash.cpp
         ${RABITQ_INDEX_ARCH_SOURCES}
     )
     add_library(rabitqlib::index ALIAS rabitq_index)

docs/development.md

@@ -152,6 +152,25 @@ shellcheck scripts/*.sh sample/run_cpp.sh sample/run_python.sh
 This matches CI's script set. Also pass any changed shell scripts outside
 `scripts/` to `shellcheck` explicitly.
 
+### ISA symbol check
+
+Sources compiled with ISA flags must not emit a weak function that a generic source also
+emits with different code: the linker keeps one copy for every caller, so generic code
+could execute AVX instructions. CI runs the check on portable Release builds with GCC and
+Clang:
+
+```bash
+./scripts/check-isa-symbols.sh build-tidy      # any CMake build directory
+```
+
+It recompiles the library sources from `compile_commands.json` with that build's flags
+(extra arguments such as `-O1` are appended) and fails when copies of a weak function
+differ between ISA flag sets. Only the Release (`-O3`) configuration is clean and checked.
+Lower optimization levels emit helpers the compiler otherwise inlines: Debug (`-O0`)
+emits many standard-library and Eigen functions per source, and GCC at `-O1`/`-O2`
+still shares two Eigen assignment loops between the FastScan estimator sources. The
+script only sees `src/`, needs GNU binutils, and does not cover MSVC or ARM-only builds.
+
 ### Static analysis
 
 clang-tidy performs semantic checks and is kept separate from clang-format.
@@ -753,8 +772,8 @@ index records.
 
 IVF/HNSW/SymphonyQG public index headers contain declarations and small accessors;
 they must not import encoding, scoring, rotator implementation, search scratch, Eigen
-or hnswlib. Operations compile in `src/index/{ivf,hnsw,qg}.cpp`; HNSW's per-ISA templates
-live in `src/index/hnsw_search_impl.hpp`. Preserve the advanced `hnsw_quant.hpp`
+or hnswlib. Operations compile in `src/index/{ivf,hnsw,qg}.cpp`; HNSW's per-ISA graph traversal
+templates live in `src/index/hnsw_search_impl.hpp`. Preserve the advanced `hnsw_quant.hpp`
 query-cache header. Installed-consumer tests poison implementation/vendor headers to
 check index header independence; `CoreModuleBoundary` also rejects batch declarations
 copied into core without an upward include.
@@ -922,8 +941,17 @@ a short-vector improvement does not justify applying it to wider vectors.
 
 Put calculations and scratch-storage helpers outside the dispatcher. Shared implementation
 headers use internal linkage so independently compiled backends retain their own bodies.
-This includes public headers reached from ISA-flagged sources, such as the per-ISA HNSW
-search files: keep their non-template helpers `static inline` rather than plain `inline`.
+This includes public headers reached from ISA-flagged sources: keep their non-template
+helpers `static inline` rather than plain `inline`. Template code instantiated in an
+ISA-flagged source must depend on a type local to that source, as the HNSW kernels do:
+`src/index/hnsw_search_*.cpp` keep their `Kernel` types in an anonymous namespace and
+only traverse the graph. Query preparation, centroid distances, acquiring the visited set
+and search buffer, and result formatting run in generic code
+(`HierarchicalNSW::search_knn_sorted_with`). The visited set's hash-table spill path is
+compiled once in `src/index/visited_set_hash.cpp`. Other calls inside a `Kernel` template
+that do not depend on `Kernel`, such as search-buffer operations, stay safe only while
+the compiler inlines them. `scripts/check-isa-symbols.sh` checks Release builds; see
+[ISA symbol check](#isa-symbol-check).
 The AVX FHT intrinsics in `src/simd/fht_kernels.hpp` use the private kernel namespace
 for the same reason. Rotation uses this implementation on GCC, Clang, and MSVC;
 no assembler is required. Keep the FFHT attribution and MIT license in the

include/rabitqlib/core/detail/rabitq_impl.hpp

@@ -608,13 +608,13 @@ inline T ex_bits_code_with_norm(
         }
     }
 
-    // Revert codes for negative or zero dims
+    // Revert codes for negative or zero dims. A select instead of a branch avoids a
+    // misprediction per random-sign coordinate and lets baseline builds vectorize.
     int32_t mask = (1 << ex_bits) - 1;
     for (size_t j = 0; j < dim; ++j) {
-        if (residual_arr.data()[j] <= 0) {
-            TP tmp = ex_code[j];
-            ex_code[j] = (~tmp) & mask;
-        }
+        const TP code = ex_code[j];
+        const auto reverted = static_cast<TP>((~code) & mask);
+        ex_code[j] = residual_arr.data()[j] <= 0 ? reverted : code;
     }
 
     return ipnorm_inv;

include/rabitqlib/index/detail/hnsw_dispatch.hpp

@@ -4,8 +4,18 @@
 #include <utility>
 #include <vector>
 
+#include "rabitqlib/index/detail/visited_set_fwd.hpp"
 #include "rabitqlib/index/types.hpp"
 
+namespace rabitqlib {
+template <typename T>
+class SplitSingleQuery;
+namespace buffer {
+template <typename T>
+class SearchBuffer;
+}  // namespace buffer
+}  // namespace rabitqlib
+
 namespace rabitqlib::hnsw {
 class HierarchicalNSW;
 namespace detail {
@@ -29,5 +39,64 @@ std::vector<std::pair<float, PID>> search_knn_sorted_avx512_popcnt(
     HierarchicalNSW& index, const float* rotated_query, size_t k, size_t ef_search = 0
 );
 
+// Graph traversal for one prepared query, compiled once per ISA. Query preparation,
+// the visited set, the ef_search-sized search_pool and result formatting stay in
+// generic code so that only these functions are built with ISA flags. Fills
+// `nearest` (capacity k) in ascending distance order and returns the number of results.
+using SearchLayersFn = size_t (*)(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+size_t search_layers_generic(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+size_t search_layers_neon(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+size_t search_layers_avx2(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+size_t search_layers_avx512_core(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+size_t search_layers_avx512_popcnt(
+    const HierarchicalNSW& index,
+    const SplitSingleQuery<float>& query_wrapper,
+    const float* q_to_centroids,
+    VisitedSet& visited,
+    buffer::SearchBuffer<float>& search_pool,
+    size_t k,
+    AnnCandidate<float>* nearest
+);
+
 }  // namespace detail
 }  // namespace rabitqlib::hnsw

include/rabitqlib/index/detail/visited_set_hash.hpp

@@ -40,6 +40,11 @@ class HashBasedVisitedSet {
 
     [[nodiscard]] auto hash1(const PID value) const { return value & mask_; }
 
+    // The spill-over lookups are compiled once in src/index/visited_set_hash.cpp, so
+    // ISA-specific search sources never instantiate the hash table themselves.
+    [[nodiscard]] bool spilled(PID internal_id) const;
+    void spill(PID internal_id);
+
     void initialize_table(const size_t table_size) {
         // Allocate before changing the mask so failure preserves the previous table.
         table_ = std::vector<PID, memory::AlignedAllocator<PID>>(table_size);
@@ -94,7 +99,7 @@ class HashBasedVisitedSet {
         if (val == internal_id) {
             return true;
         }
-        return (val != kPidMax && stl_hash_.find(internal_id) != stl_hash_.end());
+        return val != kPidMax && spilled(internal_id);
     }
 
     void set(PID internal_id) {
@@ -105,7 +110,7 @@ class HashBasedVisitedSet {
         if (val == kPidMax) {
             val = internal_id;
         } else {
-            stl_hash_.emplace(internal_id);
+            spill(internal_id);
         }
     }
 };

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -1,7 +1,6 @@
 // HNSW is developed from the [HNSW library](https://github.com/nmslib/hnswlib)
 #pragma once
 
-#include <algorithm>
 #include <array>
 #include <atomic>
 #include <cstddef>
@@ -36,7 +35,8 @@ struct DistanceEstimate;
 namespace rabitqlib::hnsw {
 namespace detail {
 struct QuantQuery;
-}
+struct HnswSearchAccess;
+}  // namespace detail
 
 template <typename T>
 using maxheap = std::priority_queue<T>;
@@ -169,52 +169,8 @@ class HierarchicalNSW {
     // Borrows construct() input only while its workers run; cleared on success or failure.
     const float* construction_data_{nullptr};
 
-    // Bounded priority queue implemented as a sorted vector.
-    class BoundedKNN {
-       public:
-        explicit BoundedKNN(size_t capacity) : capacity_(capacity) {}
-
-        // Insert a candidate in sorted order (ascending by distance).
-        void insert(const AnnCandidate<float>& candidate) {
-            auto it = std::upper_bound(queue_.begin(), queue_.end(), candidate);
-            queue_.insert(it, candidate);
-            // If we exceed capacity, drop the worst candidate (largest distance).
-            if (queue_.size() > capacity_) {
-                queue_.pop_back();
-            }
-        }
-
-        // Returns the worst (largest distance) candidate.
-        [[nodiscard]] const AnnCandidate<float>& worst() const { return queue_.back(); }
-
-        [[nodiscard]] size_t size() const { return queue_.size(); }
-
-        [[nodiscard]] const std::vector<AnnCandidate<float>>& candidates() const {
-            return queue_;
-        }
-
-       private:
-        size_t capacity_;
-        // Sorted in ascending order by distance so that the worst is at the back.
-        std::vector<AnnCandidate<float>> queue_;
-    };
-
     friend struct HnswTestAccess;
-    friend std::vector<std::pair<float, PID>> detail::search_knn_sorted_neon(
-        HierarchicalNSW&, const float*, size_t, size_t
-    );
-    friend std::vector<std::pair<float, PID>> detail::search_knn_sorted_generic(
-        HierarchicalNSW&, const float*, size_t, size_t
-    );
-    friend std::vector<std::pair<float, PID>> detail::search_knn_sorted_avx2(
-        HierarchicalNSW&, const float*, size_t, size_t
-    );
-    friend std::vector<std::pair<float, PID>> detail::search_knn_sorted_avx512_core(
-        HierarchicalNSW&, const float*, size_t, size_t
-    );
-    friend std::vector<std::pair<float, PID>> detail::search_knn_sorted_avx512_popcnt(
-        HierarchicalNSW&, const float*, size_t, size_t
-    );
+    friend struct detail::HnswSearchAccess;
 
     // `remove` sets the high bit of a point's stored cluster ID. construct and load
     // keep num_clusters_ at or below 2^31, so a real cluster ID never has it.
@@ -407,36 +363,50 @@ class HierarchicalNSW {
         *(reinterpret_cast<uint16_t*>(links)) = count;
     }
 
-    // ANN Search
+    // ANN Search. search_knn_sorted_with prepares the query in generic code and runs
+    // one ISA's search_layers_direct (instantiated in src/index/hnsw_search_*.cpp).
+    std::vector<std::pair<float, PID>> search_knn_sorted_with(
+        detail::SearchLayersFn search_layers,
+        const float* rotated_query,
+        size_t k,
+        size_t ef_search
+    ) const;
+
     template <class Kernel>
     void estimate_bin_direct(
-        const std::vector<float>& q_to_centroids,
+        const float* q_to_centroids,
         const SplitSingleQuery<float>& query_wrapper,
         PID candidate_id,
         DistanceEstimate& estimate
     ) const;
 
     template <class Kernel>
     void estimate_full_direct(
-        const std::vector<float>& q_to_centroids,
+        const float* q_to_centroids,
         const SplitSingleQuery<float>& query_wrapper,
         PID candidate_id,
         DistanceEstimate& estimate
     ) const;
 
     template <class Kernel>
-    std::vector<std::pair<float, PID>> search_knn_sorted_direct(
-        const float* rotated_query, size_t k, size_t ef_search
+    size_t search_layers_direct(
+        const SplitSingleQuery<float>& query_wrapper,
+        const float* q_to_centroids,
+        VisitedSet& visited,
+        buffer::SearchBuffer<float>& search_pool,
+        size_t k,
+        AnnCandidate<float>* nearest
     ) const;
 
-    template <class Kernel>
+    template <class Kernel, class Results>
     void search_base_layer_direct(
         PID entry_point_id,
-        size_t ef_search,
         size_t k,
         const SplitSingleQuery<float>& query_wrapper,
-        const std::vector<float>& q_to_centroids,
-        BoundedKNN& nearest_candidates
+        const float* q_to_centroids,
+        VisitedSet& visited,
+        buffer::SearchBuffer<float>& search_pool,
+        Results& nearest_candidates
     ) const;
 
     // Raw-vector distance used by construct; add scores links with quantized_distance.