From f26257ae12d9b213356f3fd3e7f7d50ead74b0e5 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:14:58 +0000 Subject: [PATCH 1/7] feat: add HNSW RaBitQ with rotated storage and staged search Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- include/knowhere/comp/index_param.h | 1 + .../index/index_node_data_mock_wrapper.h | 5 + .../index/index_node_thread_pool_wrapper.h | 5 + include/knowhere/index/index_table.h | 4 + src/common/prometheus_client.cc | 1 + src/index/hnsw/faiss_hnsw.cc | 291 ++++- src/index/hnsw/faiss_hnsw_config.h | 51 + src/index/hnsw/impl/IndexBruteForceWrapper.cc | 9 +- src/index/hnsw/impl/IndexHNSWWrapper.cc | 42 +- src/index/hnsw/impl/IndexHNSWWrapper.h | 8 + src/index/hnsw/impl/RaBitQSearchParameters.h | 34 + src/io/memory_io.cc | 6 + tests/ut/test_hnsw_rabitq.cc | 480 +++++++++ tests/ut/test_hnsw_rabitq_acceptance.cc | 992 ++++++++++++++++++ .../cppcontrib/knowhere/IndexHNSWRaBitQ.cpp | 250 +++++ .../cppcontrib/knowhere/IndexHNSWRaBitQ.h | 86 ++ .../cppcontrib/knowhere/impl/HnswSearcher.h | 39 +- .../knowhere/impl/RaBitQBuildUtils.h | 59 ++ .../cppcontrib/knowhere/impl/RaBitQSearch.h | 160 +++ .../knowhere/impl/StagedDistanceComputer.h | 12 + .../cppcontrib/knowhere/impl/index_read.cpp | 185 +++- .../cppcontrib/knowhere/impl/index_write.cpp | 151 ++- thirdparty/faiss/faiss/utils/rabitq_simd.h | 31 +- .../faiss/utils/simd_impl/rabitq_avx2.cpp | 44 +- .../faiss/utils/simd_impl/rabitq_avx512.cpp | 153 ++- .../utils/simd_impl/rabitq_avx512_spr.cpp | 17 +- 26 files changed, 2943 insertions(+), 173 deletions(-) create mode 100644 src/index/hnsw/impl/RaBitQSearchParameters.h create mode 100644 tests/ut/test_hnsw_rabitq.cc create mode 100644 tests/ut/test_hnsw_rabitq_acceptance.cc create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQBuildUtils.h create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/StagedDistanceComputer.h diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 58f10d3f2..58f53f4f5 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -58,6 +58,7 @@ constexpr const char* INDEX_HNSW = "HNSW"; constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ"; constexpr const char* INDEX_HNSW_PQ = "HNSW_PQ"; constexpr const char* INDEX_HNSW_PRQ = "HNSW_PRQ"; +constexpr const char* INDEX_HNSW_RABITQ = "HNSW_RABITQ"; constexpr const char* INDEX_DISKANN = "DISKANN"; constexpr const char* INDEX_AISAQ = "AISAQ"; diff --git a/include/knowhere/index/index_node_data_mock_wrapper.h b/include/knowhere/index/index_node_data_mock_wrapper.h index be6ba5885..9dd05e0c1 100644 --- a/include/knowhere/index/index_node_data_mock_wrapper.h +++ b/include/knowhere/index/index_node_data_mock_wrapper.h @@ -76,6 +76,11 @@ class IndexNodeDataMockWrapper : public IndexNode { return index_node_->HasRawData(metric_type); } + bool + IsIndexRefineEnabled() const override { + return index_node_->IsIndexRefineEnabled(); + } + expected GetIndexMeta(std::unique_ptr cfg) const override { return index_node_->GetIndexMeta(std::move(cfg)); diff --git a/include/knowhere/index/index_node_thread_pool_wrapper.h b/include/knowhere/index/index_node_thread_pool_wrapper.h index 92d9af7c8..ea401b16a 100644 --- a/include/knowhere/index/index_node_thread_pool_wrapper.h +++ b/include/knowhere/index/index_node_thread_pool_wrapper.h @@ -57,6 +57,11 @@ class IndexNodeThreadPoolWrapper : public IndexNode { return index_node_->HasRawData(metric_type); } + bool + IsIndexRefineEnabled() const override { + return index_node_->IsIndexRefineEnabled(); + } + expected GetIndexMeta(std::unique_ptr cfg) const override { return index_node_->GetIndexMeta(std::move(cfg)); diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index b5b4863bb..4d2d97c69 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -105,6 +105,10 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_HNSW_PRQ, VecType::VECTOR_BFLOAT16}, {IndexEnum::INDEX_HNSW_PRQ, VecType::VECTOR_INT8}, + {IndexEnum::INDEX_HNSW_RABITQ, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_HNSW_RABITQ, VecType::VECTOR_FLOAT16}, + {IndexEnum::INDEX_HNSW_RABITQ, VecType::VECTOR_BFLOAT16}, + // diskann {IndexEnum::INDEX_DISKANN, VecType::VECTOR_FLOAT}, {IndexEnum::INDEX_DISKANN, VecType::VECTOR_FLOAT16}, diff --git a/src/common/prometheus_client.cc b/src/common/prometheus_client.cc index e1e4b31ac..7458d383e 100644 --- a/src/common/prometheus_client.cc +++ b/src/common/prometheus_client.cc @@ -57,6 +57,7 @@ KnownIndexTypes() { IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW_PQ, IndexEnum::INDEX_HNSW_PRQ, + IndexEnum::INDEX_HNSW_RABITQ, IndexEnum::INDEX_DISKANN, IndexEnum::INDEX_AISAQ, IndexEnum::INDEX_MINHASH_LSH, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 089b503ab..a033fcf45 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -9,14 +9,18 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include #include #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -46,6 +50,7 @@ #include "index/hnsw/impl/IndexConditionalWrapper.h" #include "index/hnsw/impl/IndexHNSWWrapper.h" #include "index/hnsw/impl/IndexWrapperCosine.h" +#include "index/hnsw/impl/RaBitQSearchParameters.h" #include "index/refine/refine_utils.h" #include "io/memory_io.h" #include "knowhere/bitsetview_idselector.h" @@ -784,16 +789,6 @@ get_index_data_format(const faiss::Index* index) { return std::nullopt; } -// cloned from IndexHNSW.cpp -faiss::DistanceComputer* -storage_distance_computer(const faiss::Index* storage) { - if (faiss::cppcontrib::knowhere::is_similarity_metric(storage->metric_type)) { - return new faiss::NegativeDistanceComputer(storage->get_distance_computer()); - } else { - return storage->get_distance_computer(); - } -} - // there are chances that each partition split by scalar distribution is too small that we could not even train pq on it // bcz 256 points are needed for a 8-bit pq training in faiss // combine some small partitions to get a bigger one @@ -886,7 +881,8 @@ class FaissHnswIterator : public IndexIterator { const std::shared_ptr>& labels_in, std::unique_ptr&& query_in, const BitsetView& bitset_in, const int32_t ef_in, bool larger_is_closer, const float refine_ratio = 0.5f, const std::vector& label_to_internal_offset_in = {}, - const uint32_t mv_base_offset_in = 0, bool use_knowhere_search_pool = true) + const uint32_t mv_base_offset_in = 0, bool use_knowhere_search_pool = true, + const SearchParametersHNSWWrapper* storage_params = nullptr) : IndexIterator(larger_is_closer, use_knowhere_search_pool, refine_ratio), index{index_in}, labels{labels_in}, @@ -920,7 +916,11 @@ class FaissHnswIterator : public IndexIterator { workspace.hnsw = &index_hnsw->hnsw; // wrap a sign, if needed - workspace.qdis = std::unique_ptr(storage_distance_computer(index_hnsw)); + workspace.qdis.reset(storage_params ? storage_params->storage_distance_computer(index_hnsw) + : index_hnsw->get_distance_computer()); + if (larger_is_closer) { + workspace.qdis.reset(new faiss::NegativeDistanceComputer(workspace.qdis.release())); + } if (refine_ratio != 0) { // the refine is needed @@ -960,7 +960,11 @@ class FaissHnswIterator : public IndexIterator { workspace.hnsw = &index_hnsw->hnsw; // wrap a sign, if needed - workspace.qdis = std::unique_ptr(storage_distance_computer(index_hnsw)); + workspace.qdis.reset(storage_params ? storage_params->storage_distance_computer(index_hnsw) + : index_hnsw->get_distance_computer()); + if (larger_is_closer) { + workspace.qdis.reset(new faiss::NegativeDistanceComputer(workspace.qdis.release())); + } } // set query @@ -1155,6 +1159,11 @@ class FaissHnswIterator : public IndexIterator { // class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { public: + virtual std::unique_ptr + CreateSearchParameters(const FaissHnswConfig&) const { + return std::make_unique(); + } + BaseFaissRegularIndexHNSWNode(const int32_t& version, const Object& object, DataFormatEnum data_format_in) : BaseFaissRegularIndexNode(version, object), data_format{data_format_in} { } @@ -1359,7 +1368,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { const auto rows = dataset->GetRows(); const auto* data = dataset->GetTensor(); - const auto hnsw_cfg = static_cast(*cfg); + const auto& hnsw_cfg = static_cast(*cfg); const auto k = hnsw_cfg.k.value(); BitsetView bitset(bitset_); @@ -1412,7 +1421,8 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - knowhere::SearchParametersHNSWWrapper hnsw_search_params; + auto search_parameters = CreateSearchParameters(hnsw_cfg); + auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { hnsw_search_params.efSearch = hnsw_cfg.ef.value(); } @@ -1660,7 +1670,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { const auto rows = dataset->GetRows(); const auto* data = dataset->GetTensor(); - const auto hnsw_cfg = static_cast(*cfg); + const auto& hnsw_cfg = static_cast(*cfg); BitsetView bitset(bitset_); auto index_id = getIndexToSearchByScalarInfo(bitset); if (index_id < 0) { @@ -1705,7 +1715,8 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - knowhere::SearchParametersHNSWWrapper hnsw_search_params; + auto search_parameters = CreateSearchParameters(hnsw_cfg); + auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { hnsw_search_params.efSearch = hnsw_cfg.ef.value(); @@ -1993,6 +2004,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { const bool larger_is_closer = (IsMetricType(hnsw_cfg.metric_type.value(), knowhere::metric::IP) || is_cosine); const auto ef = hnsw_cfg.ef.value_or(kIteratorSeedEf); + const auto storage_search_params = CreateSearchParameters(hnsw_cfg); const auto& id_map = GetIdMap(); const auto* result_id_map = SearchResultIdMap(id_map); @@ -2033,7 +2045,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { auto it = std::make_shared( indexes[index_id], labels.empty() ? nullptr : labels[index_id], std::move(cur_query), bitset, ef, larger_is_closer, iterator_refine_ratio, label_to_internal_offset, mv_base_offset, - use_knowhere_search_pool); + use_knowhere_search_pool, storage_search_params.get()); it->SetResultIdMap(result_id_map); // store vec[i] = it; @@ -3032,6 +3044,247 @@ class BaseFaissRegularIndexHNSWPQNodeTemplate : public BaseFaissRegularIndexHNSW } }; +// Build an exact HNSW graph with Flat storage, train RaBitQ independently, +// then replace the Flat storage after both indexes have received the data. +// RaBitQ does not currently provide the symmetric code-to-code distance that +// HNSW graph construction requires, so the finalized index is intentionally +// immutable. +class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode { + public: + std::unique_ptr + CreateSearchParameters(const FaissHnswConfig& config) const override { + auto params = std::make_unique(); + params->storage_params.qb = dynamic_cast(config).rbq_bits_query.value_or(4); + return params; + } + + BaseFaissRegularIndexHNSWRaBitQNode(const int32_t& version, const Object& object, DataFormatEnum data_format) + : BaseFaissRegularIndexHNSWNode(version, object, data_format) { + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + std::string + Type() const override { + return knowhere::IndexEnum::INDEX_HNSW_RABITQ; + } + + bool + IsAdditionalScalarSupported(bool) const override { + return false; + } + + bool + IsIndexRefineEnabled() const override { + return !indexes.empty() && std::all_of(indexes.begin(), indexes.end(), [](const auto& index) { + return index != nullptr && + dynamic_cast(index.get()) != nullptr; + }); + } + + protected: + std::vector> tmp_index_rabitq; + + Status + TrainInternal(const DataSetPtr dataset, const Config& cfg) override { + const auto rows = dataset->GetRows(); + const auto dim = dataset->GetDim(); + const auto& hnsw_cfg = static_cast(cfg); + + auto metric = Str2FaissMetricType(hnsw_cfg.metric_type.value()); + if (!metric.has_value() || + (metric.value() != faiss::METRIC_L2 && metric.value() != faiss::METRIC_INNER_PRODUCT)) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ only supports L2, IP and COSINE metrics"; + return Status::invalid_metric_type; + } + const bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); + const auto& scalar_info_map = + dataset->Get>>>(meta::SCALAR_INFO); + if (!scalar_info_map.empty()) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ does not support building with scalar info"; + return Status::invalid_args; + } + + auto float_ds_ptr = convert_ds_to_float(dataset, data_format); + if (float_ds_ptr == nullptr) { + LOG_KNOWHERE_ERROR_ << "Unsupported data format"; + return Status::invalid_args; + } + const auto* data = static_cast(float_ds_ptr->GetTensor()); + + try { + std::unique_ptr hnsw_index; + if (is_cosine) { + hnsw_index = + std::make_unique(dim, hnsw_cfg.M.value()); + } else { + hnsw_index = std::make_unique(dim, hnsw_cfg.M.value(), + metric.value()); + } + hnsw_index->hnsw.efConstruction = hnsw_cfg.efConstruction.value(); + + const auto rbq_bits = static_cast(hnsw_cfg.rbq_bits.value()); + auto rabitq_index = std::make_unique(dim, metric.value(), rbq_bits); + // Query quantization accelerates the coarse estimate; full scoring uses FP32. + rabitq_index->qb = 4; + rabitq_index->centered = false; + auto rotation = std::make_unique(dim, dim); + std::unique_ptr transformed_rabitq; + if (is_cosine) { + transformed_rabitq = std::make_unique( + rotation.get(), rabitq_index.get()); + } else { + transformed_rabitq = std::make_unique(rotation.get(), rabitq_index.get()); + } + transformed_rabitq->own_fields = true; + rotation.release(); + rabitq_index.release(); + + std::unique_ptr final_index; + if (hnsw_cfg.refine.value_or(false) && hnsw_cfg.refine_type.has_value()) { + const auto hnsw_d = hnsw_index->storage->d; + const auto hnsw_metric_type = hnsw_index->storage->metric_type; + auto final_index_cnd = pick_refine_index(data_format, hnsw_cfg.refine_type, std::move(hnsw_index), + hnsw_d, hnsw_metric_type); + if (!final_index_cnd.has_value()) { + return Status::invalid_args; + } + final_index = std::move(final_index_cnd.value()); + } else { + final_index = std::move(hnsw_index); + } + + LOG_KNOWHERE_INFO_ << "Training exact HNSW graph storage"; + final_index->train(rows, data); + LOG_KNOWHERE_INFO_ << "Training RaBitQ storage"; + transformed_rabitq->train(rows, data); + + indexes[0] = std::move(final_index); + tmp_index_rabitq.clear(); + tmp_index_rabitq.emplace_back(std::move(transformed_rabitq)); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what(); + return Status::faiss_inner_error; + } + + return Status::success; + } + + Status + AddInternal(const DataSetPtr dataset, const Config&) override { + if (isIndexEmpty()) { + LOG_KNOWHERE_ERROR_ << "Can not add data to an empty index."; + return Status::empty_index; + } + if (tmp_index_rabitq.size() != indexes.size() || tmp_index_rabitq.empty() || tmp_index_rabitq[0] == nullptr) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ is immutable after its initial Add"; + return Status::not_implemented; + } + + const auto& scalar_info_map = + dataset->Get>>>(meta::SCALAR_INFO); + if (!scalar_info_map.empty()) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ does not support building with scalar info"; + return Status::invalid_args; + } + + try { + LOG_KNOWHERE_INFO_ << "Adding " << dataset->GetRows() << " rows to exact HNSW graph"; + auto status = add_to_index(indexes[0].get(), dataset, data_format); + if (status != Status::success) { + return status; + } + + LOG_KNOWHERE_INFO_ << "Adding " << dataset->GetRows() << " rows to RaBitQ storage"; + // Bound the rotation buffer only while populating RBQ storage; + // leave graph/refine construction and the common add API unchanged. + if (data_format == DataFormatEnum::fp32) { + faiss::cppcontrib::knowhere::rabitq_build::add_in_blocks( + *tmp_index_rabitq[0], dataset->GetRows(), static_cast(dataset->GetTensor())); + } else { + // Non-FP32 conversion already feeds storage in 4096-row blocks. + status = add_to_index(tmp_index_rabitq[0].get(), dataset, data_format); + } + if (status != Status::success) { + return status; + } + + faiss::cppcontrib::knowhere::IndexRefine* index_refine = + dynamic_cast(indexes[0].get()); + auto* index_hnsw = index_refine != nullptr + ? dynamic_cast(index_refine->base_index) + : dynamic_cast(indexes[0].get()); + if (index_hnsw == nullptr) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ build produced an unexpected base index"; + return Status::invalid_index_error; + } + + const bool is_cosine = faiss::cppcontrib::knowhere::is_cosine_index(index_hnsw->storage); + std::unique_ptr index_hnsw_rabitq; + if (is_cosine) { + index_hnsw_rabitq = std::make_unique(); + } else { + index_hnsw_rabitq = std::make_unique(); + } + // C++ slicing is intentional: preserve the exact graph while + // changing only the runtime HNSW type and its vector storage. + static_cast(*index_hnsw_rabitq) = + static_cast(*index_hnsw); + + // Validate the replacement before relinquishing either owner so a + // malformed storage cannot leave the exact graph half-finalized. + auto* flat_storage = index_hnsw->storage; + index_hnsw_rabitq->storage = tmp_index_rabitq[0].get(); + index_hnsw_rabitq->own_fields = false; + if (is_cosine) { + dynamic_cast(index_hnsw_rabitq.get()) + ->validate_cosine_storage(); + } else { + index_hnsw_rabitq->validate_storage(); + } + index_hnsw_rabitq->own_fields = true; + tmp_index_rabitq[0].release(); + index_hnsw->storage = nullptr; + delete flat_storage; + + if (index_refine != nullptr) { + delete index_refine->base_index; + index_refine->base_index = index_hnsw_rabitq.release(); + } else { + indexes[0] = std::move(index_hnsw_rabitq); + } + tmp_index_rabitq.clear(); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what(); + return Status::faiss_inner_error; + } + + return Status::success; + } +}; + +template +class BaseFaissRegularIndexHNSWRaBitQNodeTemplate : public BaseFaissRegularIndexHNSWRaBitQNode { + public: + BaseFaissRegularIndexHNSWRaBitQNodeTemplate(const int32_t& version, const Object& object) + : BaseFaissRegularIndexHNSWRaBitQNode(version, object, datatype_v) { + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + const auto& hnsw_cfg = static_cast(config); + return has_lossless_refine_index(hnsw_cfg.refine, hnsw_cfg.refine_type, datatype_v); + } +}; + // this index trains PRQ and HNSW+FLAT separately, then constructs HNSW+PRQ class BaseFaissRegularIndexHNSWPRQNode : public BaseFaissRegularIndexHNSWNode { public: @@ -3354,5 +3607,7 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH knowhere::feature::EMB_LIST) KNOWHERE_SIMPLE_REGISTER_DENSE_INT_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexHNSWPRQNodeTemplate, knowhere::feature::MMAP | knowhere::feature::MV | knowhere::feature::EMB_LIST) +KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_RABITQ, BaseFaissRegularIndexHNSWRaBitQNodeTemplate, + knowhere::feature::NONE) } // namespace knowhere diff --git a/src/index/hnsw/faiss_hnsw_config.h b/src/index/hnsw/faiss_hnsw_config.h index e35766be7..6e4b7eef2 100644 --- a/src/index/hnsw/faiss_hnsw_config.h +++ b/src/index/hnsw/faiss_hnsw_config.h @@ -195,6 +195,57 @@ class FaissHnswPqConfig : public FaissHnswConfig { } }; +class FaissHnswRaBitQConfig : public FaissHnswConfig { + public: + // Number of bits per database vector dimension. + CFG_INT rbq_bits; + // Request-local coarse estimator query precision. + CFG_INT rbq_bits_query; + + KNOWHERE_DECLARE_CONFIG(FaissHnswRaBitQConfig) { + KNOWHERE_CONFIG_DECLARE_FIELD(rbq_bits) + .description("number of RaBitQ bits per database vector dimension") + .set_default(1) + .set_range(1, 9) + .for_train() + .for_static(); + KNOWHERE_CONFIG_DECLARE_FIELD(rbq_bits_query) + .description("query bits for the RaBitQ coarse estimator; 0 uses FP32") + .set_default(4) + .set_range(0, 8) + .for_search() + .for_range_search() + .for_iterator(); + } + + Status + CheckAndAdjust(PARAM_TYPE param_type, std::string* err_msg) override { + const auto base_status = FaissHnswConfig::CheckAndAdjust(param_type, err_msg); + if (base_status != Status::success) { + return base_status; + } + + const auto metric = str_to_lower(metric_type.value_or(knowhere::metric::L2)); + if (metric != "l2" && metric != "ip" && metric != "cosine") { + return HandleError(err_msg, "HNSW_RABITQ only supports L2, IP and COSINE metrics", + Status::invalid_metric_type); + } + + if ((param_type == PARAM_TYPE::DESERIALIZE || param_type == PARAM_TYPE::DESERIALIZE_FROM_FILE) && + enable_mmap.value_or(false)) { + return HandleError(err_msg, "HNSW_RABITQ does not support mmap loading", Status::invalid_args); + } + if (param_type == PARAM_TYPE::TRAIN && refine_type.has_value() && + !WhetherAcceptableRefineType(refine_type.value())) { + return HandleError(err_msg, + "invalid refine type : " + refine_type.value() + + ", optional types are [sq4u, sq6, sq8, fp16, bf16, fp32, flat]", + Status::invalid_args); + } + return Status::success; + } +}; + class FaissHnswPrqConfig : public FaissHnswConfig { public: // number of subquantizer splits diff --git a/src/index/hnsw/impl/IndexBruteForceWrapper.cc b/src/index/hnsw/impl/IndexBruteForceWrapper.cc index ece7aa7b5..37ba1a4be 100644 --- a/src/index/hnsw/impl/IndexBruteForceWrapper.cc +++ b/src/index/hnsw/impl/IndexBruteForceWrapper.cc @@ -23,6 +23,7 @@ #include #include +#include "index/hnsw/impl/IndexHNSWWrapper.h" #include "knowhere/bitsetview.h" #include "knowhere/bitsetview_idselector.h" @@ -56,7 +57,9 @@ IndexBruteForceWrapper::search(faiss::idx_t n, const float* __restrict x, faiss: const faiss::SearchParameters* __restrict params) const { FAISS_THROW_IF_NOT(k > 0); - std::unique_ptr dis(index->get_distance_computer()); + const auto* hnsw_params = dynamic_cast(params); + std::unique_ptr dis(hnsw_params ? hnsw_params->storage_distance_computer(index) + : index->get_distance_computer()); // no parallelism by design for (idx_t i = 0; i < n; i++) { @@ -116,7 +119,9 @@ IndexBruteForceWrapper::range_search(faiss::idx_t n, const float* x, float radiu RH_min bres_min(result, radius); RH_max bres_max(result, radius); - std::unique_ptr dis(index->get_distance_computer()); + const auto* hnsw_params = dynamic_cast(params); + std::unique_ptr dis(hnsw_params ? hnsw_params->storage_distance_computer(index) + : index->get_distance_computer()); // no parallelism by design for (idx_t i = 0; i < n; i++) { diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.cc b/src/index/hnsw/impl/IndexHNSWWrapper.cc index 3638fc595..9135f7a62 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWWrapper.cc @@ -13,10 +13,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -30,6 +32,7 @@ #include "index/hnsw/impl/DummyVisitor.h" #include "index/hnsw/impl/FederVisitor.h" +#include "index/hnsw/impl/RaBitQSearchParameters.h" #include "knowhere/bitsetview.h" #include "knowhere/bitsetview_idselector.h" @@ -38,6 +41,7 @@ #endif namespace knowhere { +namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; /************************************************************** * Utilities @@ -103,6 +107,34 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r kAlpha = params->kAlpha; } + const auto* rbq_params = dynamic_cast(params); + // Use the optimized multi-bit path only when its selector/visitor contract + // is satisfied. RBQ1, filtering and feder use the compatible searcher below. + const auto* rabitq_index = dynamic_cast(index_hnsw); + const auto* bitset_sel = params ? dynamic_cast(params->sel) : nullptr; + const bool unfiltered = !params || !params->sel || (bitset_sel && bitset_sel->bitset_view.empty()); + if (rabitq_index && rabitq_index->rabitq_index()->rabitq.nb_bits > 1 && unfiltered && (!params || !params->feder)) { + rabitq_search::search(*rabitq_index, n, x, k, distances, labels, params ? params->efSearch : hnsw.efSearch, + params ? params->check_relative_distance : hnsw.check_relative_distance, + rbq_params ? &rbq_params->storage_params : nullptr, + [&](const rabitq_search::SearchStats& counts) { + const size_t hops = counts.expanded + counts.upper_expanded; +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + knowhere::knowhere_hnsw_search_hops.Observe(hops); +#endif + if (params && params->hnsw_stats) { + params->hnsw_stats->combine({.n1 = 1, + .n2 = size_t(counts.exhausted), + .ndis = counts.estimate + counts.upper_full, + .nhops = hops}); + } + }); + if (faiss::cppcontrib::knowhere::is_similarity_metric(index->metric_type)) { + for (idx_t i = 0; i < k * n; ++i) distances[i] = -distances[i]; + } + return; + } + // set up hnsw_stats faiss::cppcontrib::knowhere::HNSWStats* __restrict const hnsw_stats = (params == nullptr) ? nullptr : params->hnsw_stats; @@ -118,7 +150,9 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r faiss::cppcontrib::knowhere::Bitset::create_uninitialized(index->ntotal); // create a distance computer - std::unique_ptr dis(storage_distance_computer(index_hnsw->storage)); + std::unique_ptr dis( + rabitq_index ? rabitq_index->get_staged_distance_computer(rbq_params ? &rbq_params->storage_params : nullptr) + : storage_distance_computer(index_hnsw->storage)); // no parallelism by design for (idx_t i = 0; i < n; i++) { @@ -271,7 +305,11 @@ IndexHNSWWrapper::range_search(idx_t n, const float* __restrict x, float radius_ faiss::cppcontrib::knowhere::Bitset::create_uninitialized(index->ntotal); // create a distance computer - std::unique_ptr dis(storage_distance_computer(index_hnsw->storage)); + std::unique_ptr dis(params ? params->storage_distance_computer(index_hnsw) + : index_hnsw->get_distance_computer()); + if (faiss::cppcontrib::knowhere::is_similarity_metric(index_hnsw->metric_type)) { + dis.reset(new faiss::NegativeDistanceComputer(dis.release())); + } // radius float radius = radius_in; diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.h b/src/index/hnsw/impl/IndexHNSWWrapper.h index e92e8c9a6..a2bef8cac 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWWrapper.h @@ -13,6 +13,8 @@ #include #include +#include +#include #include #include @@ -30,6 +32,12 @@ struct SearchParametersHNSWWrapper : public faiss::cppcontrib::knowhere::SearchP // filtering parameter float kAlpha = 1.0f; + // Request-local storage factory, also used by brute-force fallback. + virtual faiss::DistanceComputer* + storage_distance_computer(const faiss::Index* index) const { + return index->get_distance_computer(); + } + inline ~SearchParametersHNSWWrapper() { } }; diff --git a/src/index/hnsw/impl/RaBitQSearchParameters.h b/src/index/hnsw/impl/RaBitQSearchParameters.h new file mode 100644 index 000000000..2ce4eed17 --- /dev/null +++ b/src/index/hnsw/impl/RaBitQSearchParameters.h @@ -0,0 +1,34 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +#pragma once + +#include + +#include "index/hnsw/impl/IndexHNSWWrapper.h" + +namespace knowhere { + +// Owns the per-request storage parameters; no shared index state is changed. +struct SearchParametersHNSWRaBitQWrapper : SearchParametersHNSWWrapper { + faiss::RaBitQSearchParameters storage_params; + + faiss::DistanceComputer* + storage_distance_computer(const faiss::Index* index) const override { + const auto* rbq = dynamic_cast(index); + FAISS_THROW_IF_NOT_MSG(rbq, "RaBitQ search parameters require RaBitQ storage"); + auto* dc = rbq->get_staged_distance_computer(&storage_params); + // The staged adapter is smaller-is-better; BF expects public metric units. + return index->metric_type == faiss::METRIC_INNER_PRODUCT ? new faiss::NegativeDistanceComputer(dc) : dc; + } +}; + +} // namespace knowhere diff --git a/src/io/memory_io.cc b/src/io/memory_io.cc index d84acc922..2a8382c3c 100644 --- a/src/io/memory_io.cc +++ b/src/io/memory_io.cc @@ -20,6 +20,9 @@ static size_t magic_num = 2; size_t MemoryIOWriter::operator()(const void* ptr, size_t size, size_t nitems) { + if (size == 0 || nitems == 0) { + return 0; + } auto total_need = size * nitems + rp_; if (!data_) { // data == nullptr @@ -49,6 +52,9 @@ MemoryIOWriter::operator()(const void* ptr, size_t size, size_t nitems) { size_t MemoryIOReader::operator()(void* ptr, size_t size, size_t nitems) { + if (size == 0 || nitems == 0) { + return 0; + } if (rp_ >= total_) { return 0; } diff --git a/tests/ut/test_hnsw_rabitq.cc b/tests/ut/test_hnsw_rabitq.cc new file mode 100644 index 000000000..0a1611a62 --- /dev/null +++ b/tests/ut/test_hnsw_rabitq.cc @@ -0,0 +1,480 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "catch2/catch_approx.hpp" +#include "catch2/catch_test_macros.hpp" +#include "index/hnsw/impl/IndexHNSWWrapper.h" +#include "knowhere/bitsetview.h" +#include "knowhere/comp/knowhere_config.h" +#include "knowhere/index/index_factory.h" +#include "knowhere/utils.h" +#include "utils.h" + +namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; + +TEST_CASE("RaBitQ qb4 SIMD matches scalar including masked tails", "[hnsw_rabitq_core]") { +#if defined(__GNUC__) && defined(__x86_64__) + if (!__builtin_cpu_supports("avx512f") || !__builtin_cpu_supports("avx512bw") || + !__builtin_cpu_supports("avx512dq") || !__builtin_cpu_supports("avx512vl")) + return; + for (size_t bytes : {1, 7, 8, 15, 16, 31, 32, 63, 64, 65, 96, 192, 193}) { + for (int seed = 0; seed < 8; ++seed) { + std::vector data(bytes + 1), query(bytes * 4 + 1); + for (size_t i = 0; i < data.size(); ++i) data[i] = (i * 31 + seed * 73) % 256; + for (size_t i = 0; i < query.size(); ++i) query[i] = (i * 17 + seed * 47) % 256; + const auto expected = faiss::rabitq::bitwise_and_dot_product_with_popcount( + query.data() + 1, data.data() + 1, bytes, 4); + const auto actual = faiss::rabitq::bitwise_and_dot_product_with_popcount( + query.data() + 1, data.data() + 1, bytes, 4); + REQUIRE(actual.dot_product == expected.dot_product); + REQUIRE(actual.popcount == expected.popcount); + } + } +#endif +} + +TEST_CASE("RaBitQ SIMD full scorers independently match scalar for multi-bit tails", "[hnsw_rabitq_core]") { +#if defined(__GNUC__) && defined(__x86_64__) + const bool avx512 = __builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512bw") && + __builtin_cpu_supports("avx512dq") && __builtin_cpu_supports("avx512vl"); + const bool avx2_supported = __builtin_cpu_supports("avx2"); + for (size_t d : {1, 7, 8, 15, 16, 23, 24, 31, 65, 200, 768, 1536}) { + for (size_t ex : {1, 2, 3, 4, 5, 6, 7, 8}) { + for (int seed = 1; seed <= 3; ++seed) { + CAPTURE(d, ex, seed); + std::vector signs((d + 7) / 8), extra((d * ex + 7) / 8 + (ex == 8 ? 0 : 32)); + std::vector query(d); + for (size_t i = 0; i < signs.size(); ++i) signs[i] = (i * 73 + seed * 19) % 256; + for (size_t i = 0; i < extra.size(); ++i) extra[i] = (i * 131 + seed * 37) % 256; + for (size_t i = 0; i < d; ++i) query[i] = std::sin(float(i) * .37f + seed); + const float cb = -float(1u << ex) + .5f; + const float ref = faiss::rabitq::multibit::compute_inner_product( + signs.data(), extra.data(), query.data(), d, ex, cb); + if (avx512) { + const float actual = faiss::rabitq::multibit::compute_inner_product( + signs.data(), extra.data(), query.data(), d, ex, cb); + REQUIRE(actual == Catch::Approx(ref).epsilon(1e-5).margin(1e-3)); + } + if (avx2_supported) { + const float avx2 = faiss::rabitq::multibit::compute_inner_product( + signs.data(), extra.data(), query.data(), d, ex, cb); + REQUIRE(avx2 == Catch::Approx(ref).epsilon(1e-5).margin(1e-3)); + } + } + } + } +#endif +} + +TEST_CASE("RaBitQ traversal retains all results when k covers the graph", "[hnsw_rabitq_core]") { + namespace fk = faiss::cppcontrib::knowhere; + for (const std::string metric : {"L2", "IP", "COSINE"}) { + CAPTURE(metric); + const bool similarity = metric != "L2", cosine = metric == "COSINE"; + const auto metric_type = similarity ? faiss::METRIC_INNER_PRODUCT : faiss::METRIC_L2; + constexpr int n = 128, dim = 65; + auto base = GenDataSet(n, dim, 121); + auto queries = GenDataSet(4, dim, 122); + const auto* x = static_cast(base->GetTensor()); + // Use one connected topology for this full-coverage distance-order test. + fk::IndexHNSWFlat fp32(dim, 16, faiss::METRIC_L2); + fp32.add(n, x); + auto* rq = new faiss::IndexRaBitQ(dim, metric_type, 8); + rq->qb = 4; + auto* rr = new faiss::RandomRotationMatrix(dim, dim); + std::unique_ptr storage_owner(cosine ? new fk::IndexPreTransformRaBitQCosine(rr, rq) + : new faiss::IndexPreTransform(rr, rq)); + auto& storage = *storage_owner; + storage.own_fields = true; + storage.train(n, x); + storage.add(n, x); + std::unique_ptr graph_owner(cosine ? new fk::IndexHNSWRaBitQCosine() + : new fk::IndexHNSWRaBitQ()); + auto& graph = *graph_owner; + graph.d = dim; + graph.ntotal = n; + graph.metric_type = metric_type; + graph.storage = &storage; + graph.own_fields = false; + graph.hnsw = std::move(fp32.hnsw); + std::unique_ptr full(storage.get_distance_computer()); + std::vector distances(4 * n); + std::vector labels(4 * n); + rabitq_search::search(graph, 4, static_cast(queries->GetTensor()), n, distances.data(), + labels.data(), n, true); + std::vector api_distances(4 * n); + std::vector api_labels(4 * n); + knowhere::IndexHNSWWrapper api(&graph); + knowhere::SearchParametersHNSWWrapper params; + params.efSearch = n; + api.search(4, static_cast(queries->GetTensor()), n, api_distances.data(), api_labels.data(), + ¶ms); + for (int i = 0; i < 4 * n; ++i) { + REQUIRE(api_labels[i] == labels[i]); + REQUIRE(api_distances[i] == Catch::Approx((similarity ? -1.f : 1.f) * distances[i]).margin(1e-5)); + } + for (int q = 0; q < 4; ++q) { + full->set_query(static_cast(queries->GetTensor()) + q * dim); + std::vector> expected; + for (int i = 0; i < n; ++i) expected.emplace_back((similarity ? -1.f : 1.f) * (*full)(i), i); + std::sort(expected.begin(), expected.end()); + for (int i = 0; i < n; ++i) { + CAPTURE(q, i, distances[q * n + i], expected[i].first); + REQUIRE(labels[q * n + i] == expected[i].second); + REQUIRE(distances[q * n + i] == Catch::Approx(expected[i].first).margin(1e-5)); + } + } + } +} + +TEST_CASE("RaBitQ staged distances preserve metric and cosine threshold semantics", "[hnsw_rabitq]") { + namespace fk = faiss::cppcontrib::knowhere; + for (const auto* metric : {"L2", "IP", "COSINE"}) { + const bool cosine = std::string(metric) == "COSINE"; + const bool similarity = std::string(metric) != "L2"; + for (int qb : {0, 4}) { + auto base = GenDataSet(128, 65, 4201); + auto query = GenDataSet(4, 65, 4202); + const auto* x = static_cast(base->GetTensor()); + auto* rr = new faiss::RandomRotationMatrix(65, 65); + auto* rq = new faiss::IndexRaBitQ(65, similarity ? faiss::METRIC_INNER_PRODUCT : faiss::METRIC_L2, 8); + rq->qb = qb; + std::unique_ptr storage(cosine ? new fk::IndexPreTransformRaBitQCosine(rr, rq) + : new faiss::IndexPreTransform(rr, rq)); + storage->own_fields = true; + storage->train(128, x); + storage->add(128, x); + std::unique_ptr graph(cosine ? new fk::IndexHNSWRaBitQCosine() + : new fk::IndexHNSWRaBitQ()); + graph->d = 65; + graph->metric_type = rq->metric_type; + graph->storage = storage.get(); + graph->own_fields = false; + std::unique_ptr staged_owner(graph->get_staged_distance_computer()); + auto* staged = dynamic_cast(staged_owner.get()); + REQUIRE(staged != nullptr); + std::unique_ptr full(storage->get_distance_computer()); + std::unique_ptr raw_owner(rq->get_FlatCodesDistanceComputer()); + auto* raw = dynamic_cast(raw_owner.get()); + REQUIRE(raw != nullptr); + std::vector rotated(65); + for (int q = 0; q < 4; ++q) { + const auto* v = static_cast(query->GetTensor()) + q * 65; + staged->set_query(v); + full->set_query(v); + rr->apply_noalloc(1, v, rotated.data()); + raw->set_query(rotated.data()); + for (int i = 0; i < 128; ++i) { + const float expected = (similarity ? -1 : 1) * (*full)(i); + REQUIRE((*staged)(i) == Catch::Approx(expected).margin(1e-5)); + REQUIRE(staged->evaluate(i, std::numeric_limits::infinity()) == + Catch::Approx(expected).margin(1e-5)); + const float scale = cosine ? dynamic_cast(storage.get()) + ->get_inverse_l2_norms()[i] / + std::sqrt(faiss::fvec_norm_L2sqr(v, 65)) + : 1; + const float estimate = + raw->distance_to_code_1bit(raw->codes + i * raw->code_size) * scale * (similarity ? -1 : 1); + REQUIRE(staged->evaluate(i, -std::numeric_limits::infinity()) == + Catch::Approx(estimate).margin(1e-5)); + } + REQUIRE(staged->estimate_count == 256); + REQUIRE(staged->refine_count == 128); + // Finite thresholds: compare the scaled inequality against the + // previous raw-threshold division, away from rounding ties. + for (int i = 0; i < 16; ++i) { + const auto* code = raw->codes + i * raw->code_size; + const auto* factors = + reinterpret_cast(code + (raw->d + 7) / 8); + const float scale = cosine ? dynamic_cast(storage.get()) + ->get_inverse_l2_norms()[i] / + std::sqrt(faiss::fvec_norm_L2sqr(v, 65)) + : 1.f; + const float estimate = raw->distance_to_code_1bit(code); + for (float offset : {-10.f, -1.f, .125f, 1.f, 10.f}) { + const float threshold = (similarity ? -estimate : estimate) * scale + offset; + const bool refine = faiss::rabitq_utils::should_refine_candidate( + estimate, factors->f_error, raw->g_error, (similarity ? -threshold : threshold) / scale, + similarity); + const float expected = + (refine ? raw->distance_to_code_full(code) : estimate) * scale * (similarity ? -1.f : 1.f); + const auto before = staged->refine_count; + REQUIRE(staged->evaluate(i, threshold) == Catch::Approx(expected).margin(1e-5)); + REQUIRE(staged->refine_count - before == size_t(refine)); + } + } + } + } + } +} + +TEST_CASE("HNSW RaBitQ metrics and serialized search", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (int qb : {0, 4}) { + CAPTURE(metric, qb); + auto base = GenDataSet(1024, 65, 3101); + auto query = GenDataSet(16, 65, 3102); + const auto* data = static_cast(base->GetTensor()); + std::vector original(data, data + 1024 * 65); + knowhere::Json config = {{"dim", 65}, {"metric_type", metric}, {"k", 20}, + {"M", 16}, {"efConstruction", 100}, {"ef", 200}, + {"rbq_bits", 8}, {"rbq_bits_query", qb}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(base, config) == knowhere::Status::success); + REQUIRE(std::equal(original.begin(), original.end(), data)); + auto before = index.Search(query, config, nullptr); + REQUIRE(before.has_value()); + for (int excluded : {128, 960}) { + std::vector mask(128, 0); + for (int i = 0; i < excluded; ++i) mask[i / 8] |= uint8_t(1u << (i % 8)); + knowhere::BitsetView filter(mask.data(), 1024); + auto filtered = index.Search(query, config, filter); + REQUIRE(filtered.has_value()); + for (int i = 0; i < 320; ++i) { + REQUIRE(filtered.value()->GetIds()[i] >= excluded); + REQUIRE(filtered.value()->GetIds()[i] < 1024); + } + } + for (int q = 0; q < 16; ++q) + for (int j = 0; j < 20; ++j) { + int i = q * 20 + j; + REQUIRE(before.value()->GetIds()[i] >= 0); + REQUIRE(before.value()->GetIds()[i] < 1024); + REQUIRE(std::isfinite(before.value()->GetDistance()[i])); + if (j) { + if (std::string(metric) == "L2") + REQUIRE(before.value()->GetDistance()[i - 1] <= before.value()->GetDistance()[i]); + else + REQUIRE(before.value()->GetDistance()[i - 1] >= before.value()->GetDistance()[i]); + } + } + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + auto loaded = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(loaded.Deserialize(binary, config) == knowhere::Status::success); + auto after = loaded.Search(query, config, nullptr); + REQUIRE(after.has_value()); + for (int i = 0; i < 320; ++i) { + REQUIRE(before.value()->GetIds()[i] == after.value()->GetIds()[i]); + REQUIRE(before.value()->GetDistance()[i] == after.value()->GetDistance()[i]); + } + } + } +} + +TEST_CASE("RaBitQ public search supports all database and query bit widths", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto base = GenDataSet(128, 33, 731); + auto query = GenDataSet(3, 33, 732); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (int bits = 1; bits <= 9; ++bits) { + CAPTURE(metric, bits); + knowhere::Json config = {{"dim", 33}, {"metric_type", metric}, {"k", 10}, + {"M", 8}, {"efConstruction", 64}, {"ef", 64}}; + // Also exercise the public default (RBQ1). + if (bits != 1) + config["rbq_bits"] = bits; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(base, config) == knowhere::Status::success); + knowhere::BinarySet original; + REQUIRE(index.Serialize(original) == knowhere::Status::success); + knowhere::DataSetPtr unquantized; + for (int qb = 0; qb <= 8; ++qb) { + CAPTURE(qb); + config["rbq_bits_query"] = qb; + auto result = index.Search(query, config, nullptr); + REQUIRE(result.has_value()); + if (qb == 0) + unquantized = result.value(); + if (qb == 1 && bits == 1) { + // Detect accidental parameter slicing/default-only routing. + bool different = false; + for (int i = 0; i < 30; ++i) { + different |= result.value()->GetIds()[i] != unquantized->GetIds()[i] || + result.value()->GetDistance()[i] != unquantized->GetDistance()[i]; + } + REQUIRE(different); + } + for (int i = 0; i < 30; ++i) { + REQUIRE(result.value()->GetIds()[i] >= 0); + REQUIRE(std::isfinite(result.value()->GetDistance()[i])); + } + for (int excluded : {16, 124, 128}) { + CAPTURE(excluded); + std::vector mask(16, 0); + for (int i = 0; i < excluded; ++i) mask[i / 8] |= uint8_t(1u << (i % 8)); + auto filtered = index.Search(query, config, knowhere::BitsetView(mask.data(), 128)); + REQUIRE(filtered.has_value()); + for (int i = 0; i < 30; ++i) { + const auto id = filtered.value()->GetIds()[i]; + REQUIRE((id == -1 || (id >= excluded && id < 128))); + } + } + } + // Request qb never mutates serialized storage defaults or codes. + knowhere::BinarySet after; + REQUIRE(index.Serialize(after) == knowhere::Status::success); + REQUIRE(original.binary_map_.size() == after.binary_map_.size()); + for (const auto& [name, bin] : original.binary_map_) { + const auto copy = after.GetByName(name); + REQUIRE(bin->size == copy->size); + REQUIRE(std::memcmp(bin->data.get(), copy->data.get(), bin->size) == 0); + } + for (int invalid : {-1, 9}) { + config["rbq_bits_query"] = invalid; + REQUIRE_FALSE(index.Search(query, config, nullptr).has_value()); + } + } + } +} + +TEST_CASE("RaBitQ request qb survives refine range iterator and concurrent search", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto base = GenDataSet(256, 33, 833); + auto query = GenDataSet(1, 33, 834); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (bool refine : {false, true}) { + CAPTURE(metric, refine); + knowhere::Json config = {{"dim", 33}, {"metric_type", metric}, {"k", 10}, {"M", 8}, {"efConstruction", 64}, + {"ef", 128}, {"rbq_bits", 9}}; + if (refine) { + config["refine"] = true; + config["refine_type"] = "fp16"; + config["refine_k"] = 1.5; + } + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(base, config) == knowhere::Status::success); + std::vector references; + std::vector> jobs; + for (int qb : {0, 4, 8}) { + auto request = config; + request["rbq_bits_query"] = qb; + auto result = index.Search(query, request, nullptr); + REQUIRE(result.has_value()); + references.push_back(result.value()); + jobs.push_back(std::async(std::launch::async, + [&, request] { return index.Search(query, request, nullptr).value(); })); + auto iterators = index.AnnIterator(query, request, nullptr); + REQUIRE(iterators.has_value()); + auto& it = iterators.value()[0]; + for (int n = 0; n < 10; ++n) { + REQUIRE(it->HasNext().value()); + const auto [id, distance] = it->Next().value(); + REQUIRE(id >= 0); + REQUIRE(id < 256); + REQUIRE(std::isfinite(distance)); + } + request["radius"] = std::string(metric) == "L2" ? 1e6 : -1e6; + auto range = index.RangeSearch(query, request, nullptr); + REQUIRE(range.has_value()); + REQUIRE(range.value()->GetLims()[1] > 0); + request["trace_visit"] = true; + REQUIRE(index.Search(query, request, nullptr).has_value()); + } + for (size_t j = 0; j < jobs.size(); ++j) { + const auto result = jobs[j].get(); + for (int i = 0; i < 10; ++i) { + REQUIRE(result->GetIds()[i] == references[j]->GetIds()[i]); + REQUIRE(result->GetDistance()[i] == references[j]->GetDistance()[i]); + } + } + } + } +} + +TEST_CASE("RaBitQ 9-bit storage supports floating input formats and loaded request parameters", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto base = GenDataSet(128, 33, 913); + auto query = GenDataSet(2, 33, 914); + auto* base_values = const_cast(static_cast(base->GetTensor())); + auto* query_values = const_cast(static_cast(query->GetTensor())); + for (int i = 0; i < 128 * 33; ++i) base_values[i] -= 50.f; + for (int i = 0; i < 2 * 33; ++i) query_values[i] -= 50.f; + std::fill_n(base_values, 33, 0.f); + std::fill_n(query_values, 33, 0.f); + auto exercise = [&](auto tag) { + using T = decltype(tag); + auto typed_base = knowhere::data_type_conversion(*base); + auto typed_query = knowhere::data_type_conversion(*query); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + CAPTURE(metric, sizeof(T)); + knowhere::Json config = {{"dim", 33}, {"metric_type", metric}, {"k", 10}, {"M", 8}, {"efConstruction", 64}, + {"ef", 100}, {"rbq_bits", 9}}; + auto index = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version).value(); + REQUIRE(index.Build(typed_base, config) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + auto loaded = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version).value(); + REQUIRE(loaded.Deserialize(binary, config) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + config["rbq_bits_query"] = qb; + auto a = index.Search(typed_query, config, nullptr); + auto b = loaded.Search(typed_query, config, nullptr); + REQUIRE(a.has_value()); + REQUIRE(b.has_value()); + for (int i = 0; i < 20; ++i) { + REQUIRE(a.value()->GetIds()[i] == b.value()->GetIds()[i]); + REQUIRE(std::isfinite(a.value()->GetDistance()[i])); + REQUIRE(a.value()->GetDistance()[i] == b.value()->GetDistance()[i]); + } + } + } + }; + exercise(knowhere::fp32{}); + exercise(knowhere::fp16{}); + exercise(knowhere::bf16{}); +} + +TEST_CASE("Generic HNSW parameter factory preserves SQ and PQ searches", "[hnsw_rabitq_regression]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto base = GenDataSet(1024, 32, 931); + auto query = GenDataSet(2, 32, 932); + for (const auto* name : {"HNSW_SQ", "HNSW_PQ"}) { + for (const auto* metric : {"L2", "IP", "COSINE"}) { + CAPTURE(name, metric); + knowhere::Json config = { + {"dim", 32}, {"metric_type", metric}, {"k", 10}, {"M", 8}, {"efConstruction", 64}, + {"ef", 128}, {"sq_type", "SQ8"}, {"m", 4}, {"nbits", 4}}; + auto index = knowhere::IndexFactory::Instance().Create(name, version).value(); + REQUIRE(index.Build(base, config) == knowhere::Status::success); + auto result = index.Search(query, config, nullptr); + REQUIRE(result.has_value()); + auto iterators = index.AnnIterator(query, config, nullptr); + REQUIRE(iterators.has_value()); + REQUIRE(iterators.value()[0]->HasNext().value()); + REQUIRE(std::isfinite(iterators.value()[0]->Next().value().second)); + std::vector mask(128, 255); + mask.back() = 0; + auto filtered = index.Search(query, config, knowhere::BitsetView(mask.data(), 1024)); + REQUIRE(filtered.has_value()); + for (int i = 0; i < 20; ++i) { + auto id = filtered.value()->GetIds()[i]; + REQUIRE((id == -1 || (id >= 1016 && id < 1024))); + } + } + } +} diff --git a/tests/ut/test_hnsw_rabitq_acceptance.cc b/tests/ut/test_hnsw_rabitq_acceptance.cc new file mode 100644 index 000000000..e6686cc80 --- /dev/null +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -0,0 +1,992 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "catch2/catch_approx.hpp" +#include "catch2/catch_test_macros.hpp" +#include "index/hnsw/impl/IndexHNSWWrapper.h" +#include "index/hnsw/impl/RaBitQSearchParameters.h" +#include "io/memory_io.h" +#include "knowhere/bitsetview.h" +#include "knowhere/bitsetview_idselector.h" +#include "knowhere/comp/knowhere_config.h" +#include "knowhere/feder/HNSW.h" +#include "knowhere/index/index_factory.h" +#include "knowhere/utils.h" +#include "utils.h" + +namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; + +TEST_CASE("RBQ bounded add preserves input slices and rejects invalid sizes", "[hnsw_rabitq_acceptance][rbq_build]") { + using faiss::cppcontrib::knowhere::rabitq_build::add_in_blocks; + struct RecordingIndex : faiss::IndexFlatL2 { + std::vector> calls; + RecordingIndex() : faiss::IndexFlatL2(3) { + } + void + add(faiss::idx_t n, const float* x) override { + calls.emplace_back(n, x); + ntotal += n; + } + }; + std::vector data(8193 * 3); + for (const auto n : {0, 1, 4095, 4096, 4097, 8193}) { + RecordingIndex index; + add_in_blocks(index, n, n ? data.data() : nullptr); + REQUIRE(index.ntotal == n); + REQUIRE(index.calls.size() == size_t((n + 4095) / 4096)); + for (size_t i = 0; i < index.calls.size(); ++i) { + REQUIRE(index.calls[i].first == std::min(4096, n - int(i) * 4096)); + REQUIRE(index.calls[i].second == data.data() + i * 4096 * 3); + } + } + RecordingIndex index; + add_in_blocks(index, 17, data.data(), 7); + REQUIRE(index.calls.size() == 3); + REQUIRE(index.calls.back().first == 3); + add_in_blocks(index, 2, data.data()); + REQUIRE(index.ntotal == 19); + const auto calls = index.calls.size(); + REQUIRE_THROWS(add_in_blocks(index, -1, data.data())); + REQUIRE_THROWS(add_in_blocks(index, 1, data.data(), 0)); + REQUIRE_THROWS(add_in_blocks(index, 1, data.data(), -1)); + REQUIRE_THROWS(add_in_blocks(index, 1, nullptr)); + REQUIRE_THROWS(add_in_blocks(index, std::numeric_limits::max(), data.data())); + index.d = 0; + REQUIRE_THROWS(add_in_blocks(index, 1, data.data())); + REQUIRE(index.calls.size() == calls); +} + +TEST_CASE("RBQ bounded storage encoding preserves codes norms and serialization", + "[hnsw_rabitq_acceptance][rbq_build]") { + namespace fk = faiss::cppcontrib::knowhere; + constexpr int n = 4101, d = 33; + auto dataset = GenDataSet(n, d, 29091); + const auto* x = static_cast(dataset->GetTensor()); + for (const bool cosine : {false, true}) { + for (const auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + if (cosine && metric != faiss::METRIC_INNER_PRODUCT) + continue; + for (const uint8_t bits : {1, 4, 8, 9}) { + CAPTURE(cosine, metric, int(bits)); + auto make_storage = [&]() -> std::unique_ptr { + auto rotation = std::make_unique(d, d); + auto leaf = std::make_unique(d, metric, bits); + std::unique_ptr storage; + if (cosine) { + storage = std::make_unique(rotation.get(), leaf.get()); + } else { + storage = std::make_unique(rotation.get(), leaf.get()); + } + storage->own_fields = true; + rotation.release(); + leaf.release(); + storage->train(n, x); + return storage; + }; + auto reference = make_storage(); + auto actual = make_storage(); + // Independent reproduction of the pre-refactor add_to_index loop. + for (int offset = 0; offset < n; offset += 4096) { + reference->add(std::min(4096, n - offset), x + offset * d); + } + fk::rabitq_build::add_in_blocks(*actual, n, x); + REQUIRE(actual->ntotal == n); + auto* rbq = dynamic_cast(actual->index); + auto* expected = dynamic_cast(reference->index); + REQUIRE(rbq != nullptr); + REQUIRE(expected != nullptr); + REQUIRE(rbq->center == expected->center); + REQUIRE(rbq->codes == expected->codes); + if (cosine) { + auto* norms = dynamic_cast(actual.get()); + REQUIRE(norms != nullptr); + norms->validate_norms(); + for (const int row : {0, 4095, 4096, 4100}) { + const auto norm2 = faiss::fvec_norm_L2sqr(x + row * d, d); + REQUIRE(norms->get_inverse_l2_norms()[row] == Catch::Approx(1 / std::sqrt(norm2))); + } + } + faiss::VectorIOWriter before, after; + fk::write_index(reference.get(), &before); + fk::write_index(actual.get(), &after); + REQUIRE(before.data == after.data); + faiss::VectorIOReader reader; + reader.data = after.data; + std::unique_ptr loaded(fk::read_index(&reader)); + REQUIRE(loaded->ntotal == n); + faiss::VectorIOWriter roundtrip; + fk::write_index(loaded.get(), &roundtrip); + REQUIRE(roundtrip.data == after.data); + std::unique_ptr a(actual->get_distance_computer()); + std::unique_ptr b(loaded->get_distance_computer()); + a->set_query(x); + b->set_query(x); + for (const int row : {0, 4095, 4096, 4100}) REQUIRE((*a)(row) == (*b)(row)); + } + } + } +} + +TEST_CASE("RBQ bounded add keeps outer refinement in the original input space", "[hnsw_rabitq_acceptance][rbq_build]") { + namespace fk = faiss::cppcontrib::knowhere; + constexpr int n = 4101, d = 16; + auto dataset = GenDataSet(n, d, 29092); + const auto* x = static_cast(dataset->GetTensor()); + faiss::RandomRotationMatrix rotation(d, d); + faiss::IndexRaBitQ leaf(d, faiss::METRIC_L2, 4); + faiss::IndexPreTransform storage(&rotation, &leaf); + fk::IndexRefineFlat refine(&storage); + refine.train(n, x); + fk::rabitq_build::add_in_blocks(refine, n, x); + REQUIRE(refine.ntotal == n); + REQUIRE(storage.ntotal == n); + REQUIRE(leaf.ntotal == n); + REQUIRE(refine.refine_index->ntotal == n); + std::vector reconstructed(d); + for (const int row : {0, 4095, 4096, 4100}) { + refine.reconstruct(row, reconstructed.data()); + REQUIRE(std::equal(reconstructed.begin(), reconstructed.end(), x + row * d)); + } +} + +TEST_CASE("RaBitQ memory serialization permits empty transfers without touching null pointers", + "[hnsw_rabitq_acceptance]") { + knowhere::MemoryIOWriter writer; + REQUIRE(writer(nullptr, sizeof(float), 0) == 0); + REQUIRE(writer(nullptr, 0, 17) == 0); + REQUIRE(writer.tellg() == 0); + const uint32_t value = 0x12345678; + REQUIRE(writer(&value, sizeof(value), 1) == 1); + std::unique_ptr bytes(writer.data()); + REQUIRE(writer(nullptr, sizeof(float), 0) == 0); + REQUIRE(writer.tellg() == sizeof(value)); + knowhere::MemoryIOReader reader(bytes.get(), writer.tellg()); + REQUIRE(reader(nullptr, sizeof(float), 0) == 0); + REQUIRE(reader(nullptr, 0, 17) == 0); + REQUIRE(reader.tellg() == 0); + uint32_t copy = 0; + REQUIRE(reader(©, sizeof(copy), 1) == 1); + REQUIRE(copy == value); + REQUIRE(reader(nullptr, sizeof(float), 0) == 0); +} + +TEST_CASE("RaBitQ byte-aligned bitwise kernels match independent byte oracle", "[hnsw_rabitq_core]") { + auto exercise = [&](auto level_tag) { + constexpr auto level = decltype(level_tag)::value; + for (size_t size : {1, 7, 8, 9, 15, 31, 32, 63, 64, 65, 192}) + for (size_t qb = 1; qb <= 8; ++qb) { + std::vector data(size + 1), query(size * qb + 1); + for (size_t i = 0; i < data.size(); ++i) data[i] = i * 37 + 11; + for (size_t i = 0; i < query.size(); ++i) query[i] = i * 71 + 19; + const auto* x = data.data() + 1; + const auto* q = query.data() + 1; + uint64_t dot = 0, xor_dot = 0, pop = 0; + for (size_t i = 0; i < size; ++i) { + pop += __builtin_popcount(unsigned(x[i])); + for (size_t bit = 0; bit < qb; ++bit) { + dot += uint64_t(__builtin_popcount(unsigned(x[i] & q[bit * size + i]))) << bit; + xor_dot += uint64_t(__builtin_popcount(unsigned(x[i] ^ q[bit * size + i]))) << bit; + } + } + REQUIRE(faiss::rabitq::bitwise_and_dot_product(q, x, size, qb) == dot); + REQUIRE(faiss::rabitq::bitwise_xor_dot_product(q, x, size, qb) == xor_dot); + REQUIRE(faiss::rabitq::popcount(x, size) == pop); + auto fused = faiss::rabitq::bitwise_and_dot_product_with_popcount(q, x, size, qb); + REQUIRE(fused.dot_product == dot); + REQUIRE(fused.popcount == pop); + } + }; + exercise(std::integral_constant{}); +#if defined(__GNUC__) && defined(__x86_64__) + if (__builtin_cpu_supports("avx2")) + exercise(std::integral_constant{}); + if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512bw") && __builtin_cpu_supports("avx512dq") && + __builtin_cpu_supports("avx512vl")) + exercise(std::integral_constant{}); +#endif +} + +TEST_CASE("RaBitQ iterators own query and parameter state while the parent index is retained", + "[hnsw_rabitq_acceptance]") { + for (int qb : {0, 4, 8}) { + // Common IndexIterator borrows the node-owned result IdMap. The parent + // index must outlive iteration; only request-local objects are released. + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + auto create_iterator = [&] { + auto base = GenDataSet(128, 33, 2001); + auto query = GenDataSet(1, 33, 2002); + knowhere::Json cfg = { + {"dim", 33}, {"metric_type", "COSINE"}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, {"k", 128}, + {"rbq_bits", 9}, {"rbq_bits_query", qb}}; + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + auto all = index.Search(query, cfg, nullptr); + REQUIRE(all.has_value()); + std::vector distance(128); + for (int i = 0; i < 128; ++i) distance[all.value()->GetIds()[i]] = all.value()->GetDistance()[i]; + auto iterators = index.AnnIterator(query, cfg, nullptr); + REQUIRE(iterators.has_value()); + cfg["rbq_bits_query"] = 9; + cfg["radius"] = 0.5; + REQUIRE_FALSE(index.Search(query, cfg, nullptr).has_value()); + REQUIRE_FALSE(index.AnnIterator(query, cfg, nullptr).has_value()); + REQUIRE_FALSE(index.RangeSearch(query, cfg, nullptr).has_value()); + return std::make_pair(iterators.value()[0], distance); + }; + auto [iterator, distance] = create_iterator(); + std::set ids; + while (iterator->HasNext().value()) { + const auto [id, value] = iterator->Next().value(); + REQUIRE(id >= 0); + REQUIRE(id < 128); + REQUIRE(ids.insert(id).second); + REQUIRE(value == Catch::Approx(distance[id]).margin(1e-5)); + } + REQUIRE_FALSE(ids.empty()); + } +} + +// Opt-in data-backed diagnostic; ordinary CI does not require benchmark files. +TEST_CASE("RaBitQ original COSINE and normalized IP real-data metric diagnostic", "[.hnsw_rabitq_realdata]") { + const char* data_root = std::getenv("KNOWHERE_RBQ_ACCEPTANCE_DATA"); + REQUIRE(data_root != nullptr); + constexpr int n = 4096, nq = 16, k = 100; + auto read_prefix = [&](const std::string& path, int rows, int& dim) { + std::ifstream input(path, std::ios::binary); + REQUIRE(input.good()); + uint32_t header[2]; + input.read(reinterpret_cast(header), sizeof(header)); + REQUIRE(input.good()); + REQUIRE(header[0] >= uint32_t(rows)); + REQUIRE(header[1] > 0); + REQUIRE(header[1] <= 8192); + dim = header[1]; + std::vector values(size_t(rows) * dim); + input.read(reinterpret_cast(values.data()), values.size() * sizeof(float)); + REQUIRE(input.good()); + return values; + }; + for (const auto* dataset : {"cohere", "openai"}) { + int d, qd; + const auto prefix = std::string(data_root) + "/" + dataset + "/" + dataset; + auto raw = read_prefix(prefix + ".fbin", n, d); + auto raw_q = read_prefix(prefix + "_query.fbin", nq, qd); + REQUIRE(d == qd); + auto normalized = raw, normalized_q = raw_q; + faiss::fvec_renorm_L2(d, n, normalized.data()); + faiss::fvec_renorm_L2(d, nq, normalized_q.data()); + faiss::IndexFlatIP exact(d); + exact.add(n, normalized.data()); + std::vector gt_distance(nq * k); + std::vector gt(nq * k); + exact.search(nq, normalized_q.data(), k, gt_distance.data(), gt.data()); + for (int bits : {8, 9}) + for (const auto* metric : {"COSINE", "IP"}) { + CAPTURE(dataset, bits, metric); + const bool cosine = std::string(metric) == "COSINE"; + auto base = knowhere::GenDataSet(n, d, cosine ? raw.data() : normalized.data()); + auto queries = knowhere::GenDataSet(nq, d, cosine ? raw_q.data() : normalized_q.data()); + base->SetIsOwner(false); + queries->SetIsOwner(false); + knowhere::Json cfg = {{"dim", d}, {"metric_type", metric}, {"rbq_bits", bits}, + {"M", 30}, {"efConstruction", 360}, {"ef", 500}, + {"k", k}, {"rbq_bits_query", 4}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + auto result = index.Search(queries, cfg, nullptr); + REQUIRE(result.has_value()); + int hits = 0; + double error = 0; + for (int q = 0; q < nq; ++q) + for (int i = 0; i < k; ++i) { + auto id = result.value()->GetIds()[q * k + i]; + REQUIRE(id >= 0); + REQUIRE(id < n); + hits += std::find(gt.begin() + q * k, gt.begin() + (q + 1) * k, id) != gt.begin() + (q + 1) * k; + const float expected = + faiss::fvec_inner_product(normalized_q.data() + q * d, normalized.data() + id * d, d); + const float delta = std::abs(result.value()->GetDistance()[q * k + i] - expected); + REQUIRE(delta < 0.03f); + error += delta; + } + std::cout << "RBQ_METRIC_DATA dataset=" << dataset << " bits=" << bits << " metric=" << metric + << " n=" << n << " nq=" << nq << " recall100=" << double(hits) / (nq * k) + << " mean_abs_error=" << error / (nq * k) << '\n'; + } + } +} + +TEST_CASE("RaBitQ same-index full distances agree across available SIMD levels", "[hnsw_rabitq_acceptance]") { + struct RestoreLevel { + faiss::SIMDLevel level = faiss::SIMDConfig::get_level(); + ~RestoreLevel() { + faiss::SIMDConfig::set_level(level); + } + } restore; + auto base = GenDataSet(128, 65, 1981); + auto query = GenDataSet(1, 65, 1982); + for (const auto* metric : {"L2", "IP", "COSINE"}) + for (int bits : {1, 4, 8, 9}) { + knowhere::Json cfg = {{"dim", 65}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, + {"k", 128}, {"rbq_bits", bits}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + faiss::SIMDConfig::set_level(faiss::SIMDLevel::NONE); + auto reference = index.Search(query, cfg, nullptr); + REQUIRE(reference.has_value()); + std::vector distances(128); + for (int i = 0; i < 128; ++i) + distances[reference.value()->GetIds()[i]] = reference.value()->GetDistance()[i]; + for (auto level : {faiss::SIMDLevel::NONE, faiss::SIMDLevel::AVX2, faiss::SIMDLevel::AVX512}) { + if (!faiss::SIMDConfig::is_simd_level_available(level)) + continue; + CAPTURE(metric, bits, qb, static_cast(level)); + faiss::SIMDConfig::set_level(level); + auto result = index.Search(query, cfg, nullptr); + REQUIRE(result.has_value()); + std::set ids; + for (int i = 0; i < 128; ++i) { + auto id = result.value()->GetIds()[i]; + REQUIRE(id >= 0); + REQUIRE(id < 128); + REQUIRE(ids.insert(id).second); + REQUIRE(result.value()->GetDistance()[i] == + Catch::Approx(distances[id]).epsilon(1e-5).margin(1e-4)); + } + } + } + } +} + +TEST_CASE("RaBitQ advertised refiners rerank the requested expanded candidate set", "[hnsw_rabitq_acceptance]") { + namespace fk = faiss::cppcontrib::knowhere; + auto base = GenDataSet(256, 33, 1991); + auto query = GenDataSet(1, 33, 1992); + const auto* q = static_cast(query->GetTensor()); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (const auto* refine : {"SQ4U", "SQ6", "SQ8", "FP16", "BF16", "FP32", "FLAT"}) { + CAPTURE(metric, refine); + knowhere::Json cfg = { + {"dim", 33}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, + {"k", 10}, {"rbq_bits", 4}, {"refine", true}, {"refine_type", refine}, {"refine_k", 1.3}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + auto blob = binary.binary_map_.begin()->second; + faiss::VectorIOReader reader; + reader.data.assign(blob->data.get(), blob->data.get() + blob->size); + std::unique_ptr decoded(fk::read_index(&reader)); + auto* refiner = dynamic_cast(decoded.get()); + REQUIRE(refiner != nullptr); + auto* graph = dynamic_cast(refiner->base_index); + REQUIRE(graph != nullptr); + knowhere::IndexHNSWWrapper wrapper(graph); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + auto result = index.Search(query, cfg, nullptr); + REQUIRE(result.has_value()); + knowhere::SearchParametersHNSWRaBitQWrapper params; + params.efSearch = 128; + params.storage_params.qb = qb; + std::vector d(13); + std::vector ids(13); + wrapper.search(1, q, 13, d.data(), ids.data(), ¶ms); + std::unique_ptr dc(refiner->refine_index->get_distance_computer()); + dc->set_query(q); + std::vector> expected; + const float sign = std::string(metric) == "L2" ? 1.f : -1.f; + for (auto id : ids) { + REQUIRE(id >= 0); + float distance = (*dc)(id); + if (std::string(metric) == "COSINE") { + const auto* norms = dynamic_cast(graph->storage); + REQUIRE(norms != nullptr); + distance *= norms->get_inverse_l2_norms()[id] / std::sqrt(faiss::fvec_norm_L2sqr(q, 33)); + } + expected.emplace_back(sign * distance, id); + } + std::sort(expected.begin(), expected.end()); + for (int i = 0; i < 10; ++i) { + REQUIRE(result.value()->GetIds()[i] == expected[i].second); + REQUIRE(result.value()->GetDistance()[i] == Catch::Approx(sign * expected[i].first).margin(1e-5)); + } + } + } + } +} + +TEST_CASE("RaBitQ feder trace contains valid edges and retains exhaustive results", "[hnsw_rabitq_acceptance]") { + auto base = GenDataSet(128, 33, 1951); + auto query = GenDataSet(1, 33, 1952); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (int bits : {1, 4, 9}) { + CAPTURE(metric, bits); + knowhere::Json cfg = {{"dim", 33}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, + {"k", 10}, {"rbq_bits", bits}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + cfg["trace_visit"] = false; + auto plain = index.Search(query, cfg, nullptr); + cfg["trace_visit"] = true; + auto traced = index.Search(query, cfg, nullptr); + REQUIRE(plain.has_value()); + REQUIRE(traced.has_value()); + auto info = + knowhere::Json::parse(traced.value()->GetJsonInfo()).get(); + size_t edges = 0; + for (auto& level : info.GetInfos()) + for (const auto& edge : level.GetRecords()) { + const auto [from, to, distance] = edge; + REQUIRE(from >= 0); + REQUIRE(from < 128); + REQUIRE(to >= 0); + REQUIRE(to < 128); + REQUIRE(std::isfinite(distance)); + ++edges; + } + REQUIRE(edges > 0); + const auto ids = knowhere::Json::parse(traced.value()->GetJsonIdSet()).get>(); + REQUIRE_FALSE(ids.empty()); + for (auto id : ids) { + REQUIRE(id >= 0); + REQUIRE(id < 128); + } + for (int i = 0; i < 10; ++i) { + REQUIRE(plain.value()->GetIds()[i] == traced.value()->GetIds()[i]); + REQUIRE(plain.value()->GetDistance()[i] == traced.value()->GetDistance()[i]); + } + } + } + } +} + +TEST_CASE("RaBitQ file serialization rejects truncated and incompatible indexes", "[hnsw_rabitq_acceptance]") { + namespace fk = faiss::cppcontrib::knowhere; + auto base = GenDataSet(128, 33, 1931); + auto query = GenDataSet(2, 33, 1932); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + CAPTURE(metric); + knowhere::Json cfg = {{"dim", 33}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", 100}, + {"k", 10}, {"rbq_bits", 9}}; + auto create = [&] { + return knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + }; + auto index = create(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + REQUIRE(binary.binary_map_.size() == 1); + const auto [name, blob] = *binary.binary_map_.begin(); + auto rejected = [&](const std::vector& bytes) { + knowhere::BinarySet bad; + std::shared_ptr copy(new uint8_t[bytes.size()]); + std::copy(bytes.begin(), bytes.end(), copy.get()); + bad.Append(name, copy, bytes.size()); + auto loaded = create(); + REQUIRE(loaded.Deserialize(bad, cfg) != knowhere::Status::success); + }; + for (size_t length : {size_t(1), size_t(4), size_t(blob->size / 2), size_t(blob->size - 1)}) + rejected(std::vector(blob->data.get(), blob->data.get() + length)); + for (const auto* fourcc : {"IHNr", "IHRK", "IHRC", "BAD!"}) { + std::vector bytes(blob->data.get(), blob->data.get() + blob->size); + std::copy_n(fourcc, 4, bytes.begin()); + rejected(bytes); + } + faiss::VectorIOReader reader; + reader.data.assign(blob->data.get(), blob->data.get() + blob->size); + std::unique_ptr decoded(fk::read_index(&reader)); + auto* graph = dynamic_cast(decoded.get()); + REQUIRE(graph != nullptr); + auto* storage = const_cast(graph->pretransform_index()); + auto* rq = const_cast(graph->rabitq_index()); + auto* rr = dynamic_cast(storage->chain[0]); + REQUIRE_NOTHROW(graph->validate_storage()); + ++rq->code_size; + REQUIRE_THROWS(graph->validate_storage()); + --rq->code_size; + auto last = rr->A.back(); + rr->A.pop_back(); + REQUIRE_THROWS(graph->validate_storage()); + rr->A.push_back(last); + rr->have_bias = true; + REQUIRE_THROWS(graph->validate_storage()); + rr->have_bias = false; + const auto old_metric = rq->rabitq.metric_type; + rq->rabitq.metric_type = faiss::METRIC_L1; + REQUIRE_THROWS(graph->validate_storage()); + rq->rabitq.metric_type = old_metric; + if (auto* cosine = dynamic_cast(graph)) { + auto* cs = dynamic_cast(storage); + auto norm = cs->inverse_norms_storage.inverse_l2_norms.back(); + cs->inverse_norms_storage.inverse_l2_norms.pop_back(); + REQUIRE_THROWS(cosine->validate_cosine_storage()); + cs->inverse_norms_storage.inverse_l2_norms.push_back(norm); + cs->inverse_norms_storage.inverse_l2_norms[0] = std::numeric_limits::quiet_NaN(); + REQUIRE_THROWS(cosine->validate_cosine_storage()); + } + // A real file, independently loaded through the public Knowhere API. + auto pattern = (std::filesystem::temp_directory_path() / "knowhere-rabitq-XXXXXX").string(); + std::vector filename(pattern.begin(), pattern.end()); + filename.push_back('\0'); + const int fd = mkstemp(filename.data()); + REQUIRE(fd >= 0); + struct RemoveFile { + std::string path; + ~RemoveFile() { + std::error_code ignored; + std::filesystem::remove(path, ignored); + } + } cleanup{filename.data()}; + std::unique_ptr file(fdopen(fd, "w+b"), &std::fclose); + if (!file) + close(fd); + REQUIRE(file != nullptr); + REQUIRE(std::fwrite(blob->data.get(), 1, blob->size, file.get()) == blob->size); + REQUIRE(std::fflush(file.get()) == 0); + const std::string path = filename.data(); + auto loaded = create(); + REQUIRE(loaded.DeserializeFromFile(path, cfg) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + auto a = index.Search(query, cfg, nullptr); + auto b = loaded.Search(query, cfg, nullptr); + REQUIRE(a.has_value()); + REQUIRE(b.has_value()); + for (int i = 0; i < 20; ++i) { + REQUIRE(a.value()->GetIds()[i] == b.value()->GetIds()[i]); + REQUIRE(a.value()->GetDistance()[i] == b.value()->GetDistance()[i]); + } + } + REQUIRE(index.Add(base, cfg) != knowhere::Status::success); + REQUIRE(index.Count() == 128); + cfg["enable_mmap"] = true; + REQUIRE(create().DeserializeFromFile(path, cfg) != knowhere::Status::success); + } +} + +TEST_CASE("Shared Faiss IVF RaBitQ bits and query parameters survive serialization", "[hnsw_rabitq_regression]") { + auto base = GenDataSet(512, 33, 1961); + auto query = GenDataSet(2, 33, 1962); + auto* x = static_cast(base->GetTensor()); + auto* q = static_cast(query->GetTensor()); + for (auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + for (int bits : {1, 4, 8, 9}) { + CAPTURE(metric, bits); + faiss::IndexFlat quantizer(33, metric); + faiss::IndexIVFRaBitQ index(&quantizer, 33, 4, metric, true, bits); + index.train(512, x); + index.add(512, x); + faiss::VectorIOWriter writer; + faiss::write_index(&index, &writer); + faiss::VectorIOReader reader; + reader.data = writer.data; + std::unique_ptr loaded(faiss::read_index(&reader)); + for (int qb : {0, 4, 8}) { + CAPTURE(qb); + faiss::IVFRaBitQSearchParameters params; + params.qb = qb; + params.nprobe = 4; + std::vector a(20), b(20); + std::vector ia(20), ib(20); + index.search(2, q, 10, a.data(), ia.data(), ¶ms); + loaded->search(2, q, 10, b.data(), ib.data(), ¶ms); + REQUIRE(ia == ib); + for (int i = 0; i < 20; ++i) { + REQUIRE(ia[i] >= 0); + REQUIRE(ia[i] < 512); + REQUIRE(std::isfinite(a[i])); + REQUIRE(a[i] == b[i]); + } + } + } + } +} + +TEST_CASE("RaBitQ invalid build parameters are explicitly rejected", "[hnsw_rabitq_acceptance]") { + auto base = GenDataSet(64, 33, 1971); + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + for (auto feature : {knowhere::feature::MMAP, knowhere::feature::MV, knowhere::feature::EMB_LIST}) { + REQUIRE_FALSE(knowhere::IndexFactory::Instance().FeatureCheck(knowhere::IndexEnum::INDEX_HNSW_RABITQ, feature)); + } + knowhere::Json cfg = {{"dim", 33}, {"metric_type", "L2"}, {"M", 8}, {"efConstruction", 64}, {"rbq_bits", 4}}; + for (const auto& change : + {knowhere::Json{{"rbq_bits", 0}}, knowhere::Json{{"rbq_bits", 10}}, knowhere::Json{{"metric_type", "HAMMING"}}, + knowhere::Json{{"refine", true}, {"refine_type", "unknown"}}}) { + auto request = cfg; + request.update(change); + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE_FALSE(index.IsAdditionalScalarSupported(false)); + REQUIRE_FALSE(index.IsAdditionalScalarSupported(true)); + REQUIRE(index.Build(base, request) != knowhere::Status::success); + } +} + +TEST_CASE("RaBitQ request state remains stable over repeated concurrent load lifecycles", "[hnsw_rabitq_acceptance]") { + auto base = GenDataSet(256, 33, 1941); + auto query = GenDataSet(1, 33, 1942); + for (int cycle = 0; cycle < 4; ++cycle) { + CAPTURE(cycle); + knowhere::Json cfg = { + {"dim", 33}, {"metric_type", "COSINE"}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, + {"k", 10}, {"rbq_bits", cycle % 2 ? 1 : 9}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet original; + REQUIRE(index.Serialize(original) == knowhere::Status::success); + auto loaded = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(loaded.Deserialize(original, cfg) == knowhere::Status::success); + std::vector refs; + for (int qb : {0, 4, 8}) { + auto request = cfg; + request["rbq_bits_query"] = qb; + refs.push_back(index.Search(query, request, nullptr).value()); + } + std::vector> jobs; + for (int task = 0; task < 12; ++task) { + jobs.push_back(std::async(std::launch::async, [&, task] { + const int slot = task % 3; + const int qbs[] = {0, 4, 8}; + auto request = cfg; + request["rbq_bits_query"] = qbs[slot]; + for (int repeat = 0; repeat < 20; ++repeat) { + auto result = (task % 2 ? loaded : index).Search(query, request, nullptr); + if (!result.has_value()) + return false; + for (int i = 0; i < 10; ++i) + if (result.value()->GetIds()[i] != refs[slot]->GetIds()[i] || + result.value()->GetDistance()[i] != refs[slot]->GetDistance()[i]) + return false; + } + return true; + })); + } + for (auto& job : jobs) REQUIRE(job.get()); + knowhere::BinarySet after; + REQUIRE(index.Serialize(after) == knowhere::Status::success); + for (const auto& [name, blob] : original.binary_map_) { + const auto copy = after.GetByName(name); + REQUIRE(copy->size == blob->size); + REQUIRE(std::memcmp(copy->data.get(), blob->data.get(), blob->size) == 0); + } + } +} + +TEST_CASE("RaBitQ filtered results and exhausted iterators match full-code references", "[hnsw_rabitq_acceptance]") { + constexpr int n = 256, d = 33, k = 10; + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto base = GenDataSet(n, d, 1901); + auto query = GenDataSet(1, d, 1902); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (int bits : {1, 4, 9}) { + CAPTURE(metric, bits); + knowhere::Json cfg = {{"dim", d}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", 128}, + {"k", k}, {"rbq_bits", bits}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + REQUIRE(binary.binary_map_.size() == 1); + const auto blob = binary.binary_map_.begin()->second; + faiss::VectorIOReader reader; + reader.data.assign(blob->data.get(), blob->data.get() + blob->size); + std::unique_ptr decoded(faiss::cppcontrib::knowhere::read_index(&reader)); + auto* graph = dynamic_cast(decoded.get()); + REQUIRE(graph != nullptr); + for (int qb : {0, 4, 8}) { + CAPTURE(qb); + cfg["rbq_bits_query"] = qb; + auto all_cfg = cfg; + all_cfg["k"] = n; + all_cfg["ef"] = n; + auto all = index.Search(query, all_cfg, nullptr); + REQUIRE(all.has_value()); + std::vector reference(n); + for (int i = 0; i < n; ++i) { + REQUIRE(all.value()->GetIds()[i] >= 0); + reference[all.value()->GetIds()[i]] = all.value()->GetDistance()[i]; + } + // An IP graph can contain unreachable vertices. Check iterator + // completeness against its actual directed graph, not an + // assumption that every constructed HNSW is strongly connected. + auto score = [&](int id) { return std::string(metric) == "L2" ? reference[id] : -reference[id]; }; + int nearest = graph->hnsw.entry_point; + for (int level = graph->hnsw.max_level; level > 0; --level) { + bool improved = true; + while (improved) { + improved = false; + size_t begin, end; + graph->hnsw.neighbor_range(nearest, level, &begin, &end); + for (size_t j = begin; j < end; ++j) { + int id = graph->hnsw.neighbors[j]; + if (id < 0) + break; + if (score(id) < score(nearest)) { + nearest = id; + improved = true; + } + } + } + } + std::set reachable{nearest}; + std::vector pending{nearest}; + for (size_t i = 0; i < pending.size(); ++i) { + size_t begin, end; + graph->hnsw.neighbor_range(pending[i], 0, &begin, &end); + for (size_t j = begin; j < end; ++j) { + int id = graph->hnsw.neighbors[j]; + if (id < 0) + break; + if (reachable.insert(id).second) + pending.push_back(id); + } + } + for (int excluded : {0, 64, 200, 235, 236, 238, 239, 248, 249, 252, 256}) { + CAPTURE(excluded); + std::vector mask(n / 8, 0); + for (int i = 0; i < excluded; ++i) mask[i / 8] |= uint8_t(1u << (i % 8)); + knowhere::BitsetView filter(mask.data(), n); + auto result = index.Search(query, cfg, filter); + REQUIRE(result.has_value()); + std::set ids; + std::vector expected; + for (int i = 0; i < n; ++i) { + if (all.value()->GetIds()[i] >= excluded) + expected.push_back(all.value()->GetIds()[i]); + } + int hits = 0; + for (int i = 0; i < k; ++i) { + const auto id = result.value()->GetIds()[i]; + if (i >= n - excluded) { + REQUIRE(id == -1); + continue; + } + REQUIRE(id >= excluded); + REQUIRE(id < n); + REQUIRE(ids.insert(id).second); + REQUIRE(result.value()->GetDistance()[i] == Catch::Approx(reference[id]).margin(1e-4)); + hits += std::find(expected.begin(), expected.begin() + std::min(k, n - excluded), id) != + expected.begin() + std::min(k, n - excluded); + if (excluded >= 236) + REQUIRE(id == expected[i]); + } + REQUIRE(ids.size() == size_t(std::min(k, n - excluded))); + int reachable_topk = 0; + for (int i = 0; i < std::min(k, n - excluded); ++i) reachable_topk += reachable.count(expected[i]); + // Finite-ef filtered graph search is approximate. Record its + // recall instead of inventing a universal minimum for this + // random graph fixture. BF routing above must be exact. + std::cout << "RBQ_FILTER_RECALL metric=" << metric << " bits=" << bits << " qb=" << qb + << " excluded=" << excluded << " hits=" << hits << " reachable_topk=" << reachable_topk + << '\n'; + if (excluded == 235) { + filter.set_filter_count(excluded); + // Diagnostic control: same graph and codes, but ordinary + // HNSW with full distances (no staged probability window). + namespace fk = faiss::cppcontrib::knowhere; + auto* rq = const_cast(graph->rabitq_index()); + rq->qb = qb; + fk::IndexHNSW plain; + plain.d = d; + plain.ntotal = n; + plain.is_trained = true; + plain.metric_type = graph->metric_type; + plain.hnsw = graph->hnsw; + plain.storage = graph->storage; + plain.own_fields = false; + knowhere::IndexHNSWWrapper wrapper(&plain); + knowhere::BitsetViewIDSelector selector(filter); + knowhere::SearchParametersHNSWWrapper params; + params.efSearch = 128; + params.sel = &selector; + params.kAlpha = filter.filter_ratio() * 0.7f; + std::vector full_dist(k); + std::vector full_ids(k); + wrapper.search(1, static_cast(query->GetTensor()), k, full_dist.data(), + full_ids.data(), ¶ms); + int full_hits = 0; + for (auto id : full_ids) + full_hits += std::find(expected.begin(), expected.begin() + k, id) != expected.begin() + k; + std::cout << "RBQ_FILTER_CONTROL metric=" << metric << " bits=" << bits << " qb=" << qb + << " staged_hits=" << hits << " full_hits=" << full_hits << '\n'; + if (bits == 1) + for (int i = 0; i < k; ++i) { + REQUIRE(full_ids[i] == result.value()->GetIds()[i]); + REQUIRE(full_dist[i] == Catch::Approx(result.value()->GetDistance()[i]).margin(1e-5)); + } + } + + auto iterators = index.AnnIterator(query, cfg, filter); + REQUIRE(iterators.has_value()); + std::set iter_ids; + auto it = iterators.value()[0]; + while (it->HasNext().value()) { + const auto [id, distance] = it->Next().value(); + REQUIRE(id >= excluded); + REQUIRE(id < n); + REQUIRE(iter_ids.insert(id).second); + REQUIRE(distance == Catch::Approx(reference[id]).margin(1e-4)); + REQUIRE(iter_ids.size() <= size_t(n - excluded)); + } + auto expected_reachable = reachable; + for (int i = 0; i < excluded; ++i) expected_reachable.erase(i); + REQUIRE(iter_ids == expected_reachable); + REQUIRE_FALSE(it->HasNext().value()); + } + } + } + } +} + +TEST_CASE("RaBitQ range boundaries and range_filter match full-code reference", "[hnsw_rabitq_acceptance]") { + constexpr int n = 128, d = 33; + auto base = GenDataSet(n, d, 1911); + auto query = GenDataSet(1, d, 1912); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + const bool l2 = std::string(metric) == "L2"; + knowhere::Json cfg = {{"dim", d}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", n}, + {"k", n}, {"rbq_bits", 4}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + auto all = index.Search(query, cfg, nullptr); + REQUIRE(all.has_value()); + for (bool empty : {false, true}) { + for (bool band : {false, true}) { + CAPTURE(metric, qb, empty, band); + auto request = cfg; + float radius = empty ? all.value()->GetDistance()[0] + (l2 ? -1.f : 1.f) + : (all.value()->GetDistance()[63] + all.value()->GetDistance()[64]) / 2; + float lower = (all.value()->GetDistance()[15] + all.value()->GetDistance()[16]) / 2; + request["radius"] = radius; + if (band && !empty) + request["range_filter"] = lower; + auto range = index.RangeSearch(query, request, nullptr); + REQUIRE(range.has_value()); + std::set expected, actual; + for (int i = 0; i < n; ++i) { + float distance = all.value()->GetDistance()[i]; + if ((l2 ? distance < radius : distance > radius) && + (!(band && !empty) || (l2 ? distance >= lower : distance <= lower))) + expected.insert(all.value()->GetIds()[i]); + } + for (size_t i = 0; i < range.value()->GetLims()[1]; ++i) { + REQUIRE(actual.insert(range.value()->GetIds()[i]).second); + } + REQUIRE(actual == expected); + } + } + } + } +} + +TEST_CASE("RaBitQ FP16 BF16 and FP32 refine return refiner distances", "[hnsw_rabitq_acceptance]") { + constexpr int n = 128, d = 33, k = 20; + auto base = GenDataSet(n, d, 1921); + auto query = GenDataSet(1, d, 1922); + const auto* x = static_cast(base->GetTensor()); + const auto* q = static_cast(query->GetTensor()); + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (const auto* refine : {"FP16", "BF16", "FP32"}) { + CAPTURE(metric, refine); + knowhere::Json cfg = { + {"dim", d}, {"metric_type", metric}, {"M", 16}, {"efConstruction", 100}, {"ef", n}, + {"k", k}, {"rbq_bits", 4}, {"refine", true}, {"refine_type", refine}, {"refine_k", 1.3}}; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + for (int qb : {0, 4, 8}) { + cfg["rbq_bits_query"] = qb; + auto result = index.Search(query, cfg, nullptr); + REQUIRE(result.has_value()); + for (int i = 0; i < k; ++i) { + auto id = result.value()->GetIds()[i]; + REQUIRE(id >= 0); + std::vector decoded(d); + for (int j = 0; j < d; ++j) { + const float value = x[id * d + j]; + decoded[j] = std::string(refine) == "FP16" ? float(knowhere::fp16(value)) + : std::string(refine) == "BF16" ? faiss::decode_bf16(faiss::encode_bf16(value)) + : value; + } + float expected = std::string(metric) == "L2" ? faiss::fvec_L2sqr(q, decoded.data(), d) + : faiss::fvec_inner_product(q, decoded.data(), d); + if (std::string(metric) == "COSINE") + expected /= std::sqrt(faiss::fvec_norm_L2sqr(q, d) * faiss::fvec_norm_L2sqr(x + id * d, d)); + REQUIRE(result.value()->GetDistance()[i] == Catch::Approx(expected).epsilon(1e-5).margin(1e-4)); + if (i) { + if (std::string(metric) == "L2") + REQUIRE(result.value()->GetDistance()[i - 1] <= result.value()->GetDistance()[i]); + else + REQUIRE(result.value()->GetDistance()[i - 1] >= result.value()->GetDistance()[i]); + } + } + } + } + } +} diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp new file mode 100644 index 000000000..1eeab6555 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -0,0 +1,250 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include + +#include +#include + +namespace faiss::cppcontrib::knowhere { + +namespace { +struct RaBitQStagedDistanceComputer final : StagedDistanceComputer { + const faiss::VectorTransform& rotation; + std::unique_ptr dc; + const float* norms; + bool similarity; + float query_inverse_norm = 1; + std::vector rotated; + + explicit RaBitQStagedDistanceComputer(const IndexHNSWRaBitQ& index, + const faiss::RaBitQSearchParameters* params) + : rotation(*index.pretransform_index()->chain[0]), + norms(nullptr), similarity(index.metric_type == METRIC_INNER_PRODUCT), + rotated(index.d) { + auto* raw = params ? index.rabitq_index()->get_quantized_distance_computer(params->qb, params->centered) + : index.rabitq_index()->get_FlatCodesDistanceComputer(); + auto* typed = dynamic_cast(raw); + if (!typed) { delete raw; FAISS_THROW_MSG("RaBitQ distance computer required"); } + dc.reset(typed); + if (auto* cosine = dynamic_cast(&index)) { + norms = cosine->get_inverse_l2_norms(); + } + } + void set_query(const float* q) override { + estimate_count = refine_count = 0; + rotation.apply_noalloc(1, q, rotated.data()); + dc->set_query(rotated.data()); + const float norm2 = norms ? faiss::fvec_norm_L2sqr(q, rotation.d_in) : 1; + query_inverse_norm = norm2 > 0 ? 1 / std::sqrt(norm2) : 1; + } + float scale(idx_t id) const { return norms ? norms[id] * query_inverse_norm : 1; } + float operator()(idx_t id) override { + float d = (*dc)(id) * scale(id); + return similarity ? -d : d; + } + float symmetric_dis(idx_t, idx_t) override { + FAISS_THROW_MSG("staged storage is search-only; construct graph with FP32"); + } + float evaluate(idx_t id, float threshold) override { + if (dc->nb_bits == 1) return (*this)(id); + const auto* code = dc->codes + id * dc->code_size; + const float estimate = dc->distance_to_code_1bit(code); + ++estimate_count; + const auto* factors = reinterpret_cast( + code + (dc->d + 7) / 8); + const float s = scale(id); + float d = estimate; + // Compare in output-distance units: positive cosine scale preserves + // ordering, avoiding a division for every visited candidate. + const float error = factors->f_error * dc->g_error; + const bool refine = similarity ? (estimate + error) * s > -threshold + : std::max(0.0f, estimate - error) < threshold; + if (refine) { + d = dc->distance_to_code_full(code); + ++refine_count; + } + return (similarity ? -d : d) * s; + } +}; +} // namespace + +faiss::DistanceComputer* IndexHNSWRaBitQ::get_staged_distance_computer( + const faiss::RaBitQSearchParameters* params) const { + return new RaBitQStagedDistanceComputer(*this, params); +} + +IndexPreTransformRaBitQCosine::IndexPreTransformRaBitQCosine() = default; + +IndexPreTransformRaBitQCosine::IndexPreTransformRaBitQCosine( + faiss::VectorTransform* transform, + faiss::Index* index_in) + : faiss::IndexPreTransform(transform, index_in) {} + +void IndexPreTransformRaBitQCosine::add(idx_t n, const float* x) { + faiss::IndexPreTransform::add(n, x); + inverse_norms_storage.add(x, n, d); +} + +void IndexPreTransformRaBitQCosine::reset() { + faiss::IndexPreTransform::reset(); + inverse_norms_storage.reset(); +} + +faiss::DistanceComputer* IndexPreTransformRaBitQCosine::get_distance_computer() + const { + FAISS_THROW_IF_NOT_MSG( + inverse_norms_storage.inverse_l2_norms.size() == + static_cast(ntotal), + "cosine RaBitQ inverse norm count must match ntotal"); + return new WithCosineNormDistanceComputer( + get_inverse_l2_norms(), + d, + std::unique_ptr( + faiss::IndexPreTransform::get_distance_computer())); +} + +const float* IndexPreTransformRaBitQCosine::get_inverse_l2_norms() const { + return inverse_norms_storage.inverse_l2_norms.data(); +} + +void IndexPreTransformRaBitQCosine::validate_norms() const { + FAISS_THROW_IF_NOT_MSG( + inverse_norms_storage.inverse_l2_norms.size() == + static_cast(ntotal), + "cosine RaBitQ inverse norm count must match ntotal"); + for (const float inverse_norm : inverse_norms_storage.inverse_l2_norms) { + FAISS_THROW_IF_NOT_MSG( + std::isfinite(inverse_norm) && inverse_norm > 0.0f, + "cosine RaBitQ inverse norms must be finite and positive"); + } +} + +IndexHNSWRaBitQ::IndexHNSWRaBitQ() = default; + +void IndexHNSWRaBitQ::add(idx_t, const float*) { + FAISS_THROW_MSG( + "IndexHNSWRaBitQ does not support incremental add: build the " + "HNSW graph with exact storage before attaching RaBitQ storage"); +} + +const faiss::IndexPreTransform* IndexHNSWRaBitQ::pretransform_index() const { + return dynamic_cast(storage); +} + +const faiss::IndexRaBitQ* IndexHNSWRaBitQ::rabitq_index() const { + const auto* pretransform = pretransform_index(); + return pretransform + ? dynamic_cast(pretransform->index) + : nullptr; +} + +void IndexHNSWRaBitQ::validate_storage() const { + FAISS_THROW_IF_NOT_MSG( + metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT, + "IndexHNSWRaBitQ only supports L2 and inner product metrics"); + FAISS_THROW_IF_NOT_MSG( + storage != nullptr, "IndexHNSWRaBitQ requires non-null storage"); + + const auto* pretransform = pretransform_index(); + FAISS_THROW_IF_NOT_MSG( + pretransform != nullptr, + "IndexHNSWRaBitQ storage must be IndexPreTransform"); + FAISS_THROW_IF_NOT_MSG( + pretransform->chain.size() == 1, + "IndexHNSWRaBitQ storage must contain exactly one transform"); + + const auto* rotation = dynamic_cast( + pretransform->chain[0]); + FAISS_THROW_IF_NOT_MSG( + rotation != nullptr, + "IndexHNSWRaBitQ transform must be RandomRotationMatrix"); + + const auto* rabitq = rabitq_index(); + FAISS_THROW_IF_NOT_MSG( + rabitq != nullptr, + "IndexHNSWRaBitQ pretransform leaf must be IndexRaBitQ"); + + FAISS_THROW_IF_NOT_MSG( + d == pretransform->d && metric_type == pretransform->metric_type && + ntotal == pretransform->ntotal, + "IndexHNSWRaBitQ outer index and pretransform metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + is_trained && pretransform->is_trained && rotation->is_trained && + rabitq->is_trained, + "IndexHNSWRaBitQ requires fully trained storage"); + FAISS_THROW_IF_NOT_MSG( + pretransform->index != nullptr && + pretransform->ntotal == rabitq->ntotal && + pretransform->metric_type == rabitq->metric_type, + "IndexHNSWRaBitQ pretransform and RaBitQ metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + rotation->d_in == d && rotation->d_out == rabitq->d && + rotation->d_in == rotation->d_out, + "IndexHNSWRaBitQ requires a square rotation matching index dimensions"); + FAISS_THROW_IF_NOT_MSG( + rotation->is_orthonormal && !rotation->have_bias && + rotation->b.empty() && + rotation->A.size() == + static_cast(rotation->d_in) * + rotation->d_out, + "IndexHNSWRaBitQ rotation matrix has invalid storage"); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.d == static_cast(rabitq->d) && + rabitq->rabitq.metric_type == rabitq->metric_type, + "IndexHNSWRaBitQ RaBitQ quantizer metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.nb_bits >= 1 && rabitq->rabitq.nb_bits <= 9, + "IndexHNSWRaBitQ RaBitQ nb_bits must be in [1, 9]"); + + const size_t expected_code_size = + rabitq->rabitq.compute_code_size(rabitq->d, rabitq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.code_size == expected_code_size && + rabitq->code_size == expected_code_size, + "IndexHNSWRaBitQ RaBitQ code size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->codes.size() == + static_cast(rabitq->ntotal) * expected_code_size, + "IndexHNSWRaBitQ RaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->center.size() == static_cast(rabitq->d), + "IndexHNSWRaBitQ RaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->qb <= 8, "IndexHNSWRaBitQ RaBitQ qb must be in [0, 8]"); + FAISS_THROW_IF_NOT_MSG( + !rabitq->centered, "IndexHNSWRaBitQ V1 requires centered=false"); +} + +IndexHNSWRaBitQCosine::IndexHNSWRaBitQCosine() = default; + +const float* IndexHNSWRaBitQCosine::get_inverse_l2_norms() const { + const auto* cosine_storage = + dynamic_cast(storage); + return cosine_storage ? cosine_storage->get_inverse_l2_norms() : nullptr; +} + +void IndexHNSWRaBitQCosine::validate_cosine_storage() const { + validate_storage(); + const auto* cosine_storage = + dynamic_cast(storage); + FAISS_THROW_IF_NOT_MSG( + cosine_storage != nullptr, + "IndexHNSWRaBitQCosine requires cosine-aware pretransform storage"); + FAISS_THROW_IF_NOT_MSG( + metric_type == METRIC_INNER_PRODUCT, + "IndexHNSWRaBitQCosine requires inner product storage"); + cosine_storage->validate_norms(); +} + +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h new file mode 100644 index 000000000..7c38ab1a0 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h @@ -0,0 +1,86 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace faiss::cppcontrib::knowhere { + +// Private Knowhere serialization tag. Upstream Faiss reserves "IHNr" for +// its incompatible direct-build/staged-search IndexHNSWRaBitQ format. +inline constexpr char kHnswRaBitQFourcc[] = "IHRS"; +inline constexpr char kHnswRaBitQCosineFourcc[] = "IHCS"; +inline constexpr char kRaBitQPreTransformCosineFourcc[] = "IRKC"; + +/** Random-rotation + RaBitQ storage with Knowhere cosine semantics. + * + * Original vectors are quantized without permanently normalizing them. The + * underlying RaBitQ distance computer estimates inner product; this wrapper + * applies the stored database inverse norm and the query inverse norm. + */ +struct IndexPreTransformRaBitQCosine : faiss::IndexPreTransform, + HasInverseL2Norms { + L2NormsStorage inverse_norms_storage; + + IndexPreTransformRaBitQCosine(); + IndexPreTransformRaBitQCosine( + faiss::VectorTransform* transform, + faiss::Index* index); + + void add(idx_t n, const float* x) override; + void reset() override; + faiss::DistanceComputer* get_distance_computer() const override; + const float* get_inverse_l2_norms() const override; + + void validate_norms() const; +}; + +/** HNSW graph backed by a randomly-rotated standalone RaBitQ index. + * + * The storage layout is deliberately strict: + * + * IndexPreTransform + * -> RandomRotationMatrix + * -> faiss::IndexRaBitQ + * + * RaBitQ does not implement code-to-code symmetric distances, so this index + * is immutable after its graph and storage have been assembled. Build the + * graph with exact storage first, then attach the trained/populated RaBitQ + * storage to this runtime type. + */ +struct IndexHNSWRaBitQ : IndexHNSW { + IndexHNSWRaBitQ(); + + void add(idx_t n, const float* x) override; + + const faiss::IndexPreTransform* pretransform_index() const; + + const faiss::IndexRaBitQ* rabitq_index() const; + + faiss::DistanceComputer* get_staged_distance_computer( + const faiss::RaBitQSearchParameters* params = nullptr) const; + + /** Validate the complete runtime/storage shape and serialized invariants. + * Throws FaissException on malformed state. */ + void validate_storage() const; +}; + +/** Cosine runtime marker for HNSW backed by cosine-aware RaBitQ storage. */ +struct IndexHNSWRaBitQCosine : IndexHNSWRaBitQ, HasInverseL2Norms { + IndexHNSWRaBitQCosine(); + + const float* get_inverse_l2_norms() const override; + void validate_cosine_storage() const; +}; + +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h index 70c06d6c3..d316daab6 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h @@ -34,6 +34,7 @@ // Knowhere-specific headers #include +#include namespace faiss { namespace cppcontrib { @@ -89,6 +90,22 @@ struct v2_hnsw_searcher { // the pointer is not owned. const faiss::cppcontrib::knowhere::SearchParametersHNSW* params; + StagedDistanceComputer* staged = nullptr; + size_t staged_k = 0; + std::priority_queue staged_results; + + float staged_threshold() const { + return staged_results.size() < staged_k + ? std::numeric_limits::infinity() : staged_results.top(); + } + void record_staged_result(float distance, int status) { + if (!staged || !staged_k || status == knowhere::Neighbor::kInvalid) return; + if (staged_results.size() < staged_k) staged_results.push(distance); + else if (distance < staged_results.top()) { + staged_results.pop(); staged_results.push(distance); + } + } + // v2_hnsw_searcher( const faiss::cppcontrib::knowhere::HNSW& hnsw_, @@ -231,6 +248,17 @@ struct v2_hnsw_searcher { ndis += 1; if (counter == 4) { + // Staged evaluation preserves per-candidate threshold updates. + if (staged && staged_k && level == 0) { + for (size_t i = 0; i < 4; ++i) { + const float d = staged->evaluate(saved_indices[i], staged_threshold()); + graph_visitor.visit_edge(level, node_id, saved_indices[i], d); + record_staged_result(d, saved_statuses[i]); + func_add_candidate(knowhere::Neighbor(saved_indices[i], d, saved_statuses[i])); + } + counter = 0; + continue; + } // evaluate 4x distances at once float dis[4] = {0, 0, 0, 0}; qdis.distances_batch_4( @@ -266,7 +294,10 @@ struct v2_hnsw_searcher { // process leftovers for (size_t id4 = 0; id4 < counter; id4++) { // evaluate a single distance - const float dis = qdis(saved_indices[id4]); + const float dis = staged && staged_k && level == 0 + ? staged->evaluate(saved_indices[id4], staged_threshold()) + : qdis(saved_indices[id4]); + record_staged_result(dis, saved_statuses[id4]); // record a traversed edge graph_visitor.visit_edge(level, node_id, saved_indices[id4], dis); @@ -369,6 +400,10 @@ struct v2_hnsw_searcher { // grab some needed parameters const int efSearch = params ? params->efSearch : hnsw.efSearch; + staged = dynamic_cast(&qdis); + staged_k = static_cast(k); + staged_results = {}; + // yes. // greedy search on upper levels. @@ -404,6 +439,8 @@ struct v2_hnsw_searcher { } visited_nodes[nearest] = true; + record_staged_result(d_nearest, filter.is_member(nearest) + ? knowhere::Neighbor::kValid : knowhere::Neighbor::kInvalid); } // perform the search of the level 0. diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQBuildUtils.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQBuildUtils.h new file mode 100644 index 000000000..840520c80 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQBuildUtils.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +namespace faiss::cppcontrib::knowhere::rabitq_build { + +/** Populate a trained RaBitQ storage pipeline in bounded input batches. + * + * The caller supplies the complete add pipeline (including any pretransform, + * cosine norms or raw-space refiner). Each original input slice is passed to + * its existing virtual add(), preserving its encoding and ID semantics. Do + * not pass a graph-building index: changing graph add boundaries can change + * topology. IVF and file-backed callers retain their own training and I/O. + * + * This is an execution policy, not an Index subtype or serialized property. + * It bounds rows per add call, not total RSS or training allocations. The + * default preserves HNSW RaBitQ's original 4096-row encoding boundaries. + * No input ownership, transformation or shared mutable state is introduced. + * On failure, completed batches remain added; no transactional rollback is + * promised beyond the underlying Index::add contract. + */ +inline void add_in_blocks( + faiss::Index& storage, + idx_t n, + const float* x, + idx_t block_rows = 4096) { + FAISS_THROW_IF_NOT_MSG(n >= 0, "negative RaBitQ input count"); + FAISS_THROW_IF_NOT_MSG(block_rows > 0, "RaBitQ block size must be positive"); + FAISS_THROW_IF_NOT_MSG(storage.d > 0, "invalid RaBitQ input dimension"); + if (n == 0) { + return; + } + FAISS_THROW_IF_NOT_MSG(x != nullptr, "null RaBitQ input"); + const size_t dim = static_cast(storage.d); + FAISS_THROW_IF_NOT_MSG( + static_cast(n) <= + std::numeric_limits::max() / sizeof(float) / dim, + "RaBitQ input size overflow"); + for (idx_t offset = 0; offset < n;) { + const idx_t count = std::min(block_rows, n - offset); + storage.add(count, x + static_cast(offset) * dim); + offset += count; + } +} + +} // namespace faiss::cppcontrib::knowhere::rabitq_build diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h new file mode 100644 index 000000000..7df8bce14 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h @@ -0,0 +1,160 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * Licensed under the MIT license in thirdparty/faiss/LICENSE. + * + * Port of Faiss #5526 (d8a85956) bounded traversal. + * Graph adjacency is read from Knowhere without copying or changing the graph. + * Extended to L2/IP/COSINE, deliberately limited to unfiltered KNN. + * Callers must dispatch filtered/visitor requests to a compatible searcher. + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace faiss::cppcontrib::knowhere::rabitq_search { +struct SearchStats { + size_t estimate = 0, refine = 0, expanded = 0, upper_full = 0, upper_expanded = 0; + bool exhausted = false; +}; + +template +SearchStats search_one(const faiss::cppcontrib::knowhere::HNSW& graph, + faiss::RaBitQDistanceComputer& rq, VT& vt, + faiss::ResultHandler& res, int ef, bool relative, + bool similarity = false, const float* norms = nullptr, + float query_inverse_norm = 1) { + SearchStats stats; + using HC = faiss::CMax; + int32_t nearest = graph.entry_point; + auto scale = [&](int32_t id) { return norms ? norms[id] * query_inverse_norm : 1.f; }; + auto convert = [&](int32_t id, float raw) { return (similarity ? -raw : raw) * scale(id); }; + float nearest_distance = convert(nearest, rq(nearest)); + ++stats.upper_full; + for (int level = graph.max_level; level >= 1; --level) { + for (;;) { + ++stats.upper_expanded; + const int32_t previous = nearest; + size_t begin, end; + graph.neighbor_range(nearest, level, &begin, &end); + int32_t ids[4]; + int count = 0; + auto update = [&](int32_t id, float d) { + d = convert(id, d); + if (d < nearest_distance) { nearest = id; nearest_distance = d; } + }; + for (size_t j = begin; j < end && graph.neighbors[j] >= 0; ++j) { + ids[count++] = graph.neighbors[j]; + ++stats.upper_full; + if (count == 4) { + float d[4]; + rq.distances_batch_4(ids[0], ids[1], ids[2], ids[3], d[0], d[1], d[2], d[3]); + for (int i = 0; i < 4; ++i) update(ids[i], d[i]); + count = 0; + } + } + for (int i = 0; i < count; ++i) update(ids[i], rq(ids[i])); + if (previous == nearest) break; + } + } + + faiss::MinimaxHeapT candidates(ef); + candidates.push(nearest, nearest_distance); + vt.reserve(ef); + if (nearest_distance < res.threshold) res.add_result(nearest_distance, nearest); + vt.set(nearest); + while (candidates.size() > 0) { + float d0; + const int32_t node = candidates.pop_min(&d0); + if (relative && candidates.count_below(d0) >= ef) break; + size_t begin, end; + graph.neighbor_range(node, 0, &begin, &end); + size_t limit = begin; + for (size_t j = begin; j < end; ++j) { + if (graph.neighbors[j] < 0) break; + vt.prefetch(graph.neighbors[j]); + ++limit; + } + int32_t ids[4]; + int count = 0; + float threshold = res.threshold; + auto evaluate = [&] { + for (int i = 0; i < count; ++i) { + const auto* code = rq.codes + static_cast(ids[i]) * rq.code_size; + const float estimate = rq.distance_to_code_1bit(code); + ++stats.estimate; + const auto* factors = reinterpret_cast( + code + (rq.d + 7) / 8); + const float s = scale(ids[i]); + const float error = factors->f_error * rq.g_error; + float distance = estimate; + const bool refine = similarity ? (estimate + error) * s > -threshold + : std::max(0.f, estimate - error) < threshold; + if (refine) { + distance = rq.distance_to_code_full(code); + ++stats.refine; + } + distance = (similarity ? -distance : distance) * s; + if (distance < threshold && res.add_result(distance, ids[i])) threshold = res.threshold; + candidates.push(ids[i], distance); + } + }; + for (size_t j = begin; j < limit; ++j) { + ids[count] = graph.neighbors[j]; + count += vt.set(ids[count]) ? 1 : 0; + if (count == 4) { evaluate(); count = 0; } + } + if (count) evaluate(); + ++stats.expanded; + if (!relative && stats.expanded > static_cast(ef)) break; + } + stats.exhausted = candidates.size() == 0; + return stats; +} + +inline void search(const faiss::cppcontrib::knowhere::IndexHNSWRaBitQ& index, + faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, + faiss::idx_t* labels, int ef, bool relative, + const faiss::RaBitQSearchParameters* params = nullptr, + const std::function& on_query = {}) { + FAISS_THROW_IF_NOT(index.metric_type == faiss::METRIC_L2 || + index.metric_type == faiss::METRIC_INNER_PRODUCT); + const bool similarity = index.metric_type == faiss::METRIC_INNER_PRODUCT; + const auto* cosine = dynamic_cast(&index); + const float* norms = cosine ? cosine->get_inverse_l2_norms() : nullptr; + FAISS_THROW_IF_NOT(index.rabitq_index()->rabitq.nb_bits > 1); + auto raw = std::unique_ptr( + params ? index.rabitq_index()->get_quantized_distance_computer(params->qb, params->centered) + : index.rabitq_index()->get_FlatCodesDistanceComputer()); + auto& rq = dynamic_cast(*raw); + // Reuse Faiss's visited table across queries without clearing the full + // table on each search; advance its generation after processing the query. + auto& vt = faiss::VisitedTable::get_reusable(index.ntotal); + faiss::HeapBlockResultHandler> block(n, distances, labels, k); + decltype(block)::SingleResultHandler result(block); + std::vector rotated(index.d); + for (faiss::idx_t i = 0; i < n; ++i) { + result.begin(i); + index.pretransform_index()->chain[0]->apply_noalloc(1, x + i * index.d, rotated.data()); + rq.set_query(rotated.data()); + const float norm2 = norms ? faiss::fvec_norm_L2sqr(x + i * index.d, index.d) : 1.f; + const float query_inverse_norm = norm2 > 0 ? 1.f / std::sqrt(norm2) : 1.f; + SearchStats stats; + if (auto* vector = dynamic_cast(&vt)) + stats = search_one(index.hnsw, rq, *vector, result, std::max(ef, k), relative, + similarity, norms, query_inverse_norm); + else + stats = search_one(index.hnsw, rq, dynamic_cast(vt), result, + std::max(ef, k), relative, similarity, norms, query_inverse_norm); + result.end(); + vt.advance(); + if (on_query) on_query(stats); + } +} +} // namespace faiss::cppcontrib::knowhere::rabitq_search diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/StagedDistanceComputer.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/StagedDistanceComputer.h new file mode 100644 index 000000000..1a45aa1ac --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/StagedDistanceComputer.h @@ -0,0 +1,12 @@ +#pragma once + +#include + +namespace faiss::cppcontrib::knowhere { +// Distances and thresholds follow the graph searcher's smaller-is-better convention. +struct StagedDistanceComputer : faiss::DistanceComputer { + size_t estimate_count = 0; + size_t refine_count = 0; + virtual float evaluate(idx_t id, float threshold) = 0; +}; +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp index ceaf05193..0166cadb5 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp @@ -15,36 +15,38 @@ #include #include +#include #include #include #include -#include #include #include #include -#include #include -#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include +#include #include #include #include #include #include -#include -#include #include #include #include #include -#include -#include +#include #include #include @@ -58,11 +60,8 @@ #include #include - - namespace faiss::cppcontrib::knowhere { - uint32_t read_value(IOReader* f) { uint32_t h; READ1(h) @@ -216,7 +215,10 @@ void read_xb_vector(VectorT& target, IOReader* f) { * Read **************************************************************/ -static void read_index_header(Index* idx, IOReader* f, bool* is_cosine_out = nullptr) { +static void read_index_header( + Index* idx, + IOReader* f, + bool* is_cosine_out = nullptr) { READ1(idx->d); READ1(idx->ntotal); @@ -362,24 +364,27 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { "read_InvertedLists:" " WARN! inverted lists not stored with IVF object\n"); return nullptr; - } else if (h == fourcc ("iloa") && !(io_flags & IO_FLAG_MMAP)) { + } else if (h == fourcc("iloa") && !(io_flags & IO_FLAG_MMAP)) { size_t nlist; size_t code_size; - std::vector list_length; + std::vector list_length; READ1(nlist); READ1(code_size); READVECTOR(list_length); - auto ails = new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); + auto ails = + new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); size_t n; READ1(n); #ifdef USE_GPU - ails->pin_readonly_ids = std::make_shared(n * sizeof(idx_t)); - ails->pin_readonly_codes = std::make_shared(n * code_size * sizeof(uint8_t)); - READANDCHECK((idx_t *) ails->pin_readonly_ids->data, n); - READANDCHECK((uint8_t *) ails->pin_readonly_codes->data, n * code_size); + ails->pin_readonly_ids = + std::make_shared(n * sizeof(idx_t)); + ails->pin_readonly_codes = std::make_shared( + n * code_size * sizeof(uint8_t)); + READANDCHECK((idx_t*)ails->pin_readonly_ids->data, n); + READANDCHECK((uint8_t*)ails->pin_readonly_codes->data, n * code_size); #else ails->readonly_ids.resize(n); - ails->readonly_codes.resize(n*code_size); + ails->readonly_codes.resize(n * code_size); READANDCHECK(ails->readonly_ids.data(), n); READANDCHECK(ails->readonly_codes.data(), n * code_size); #endif @@ -391,7 +396,8 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { READ1(segment_size); bool save_norm = io_flags & IO_FLAG_WITH_NORM; - auto lca = new ConcurrentArrayInvertedLists(nlist, code_size, segment_size, save_norm); + auto lca = new ConcurrentArrayInvertedLists( + nlist, code_size, segment_size, save_norm); std::vector sizes(nlist); read_ArrayInvertedLists_sizes(f, sizes); for (size_t i = 0; i < lca->nlist; i++) { @@ -402,12 +408,15 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { if (n > 0) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { - size_t seg_size = lca->get_segment_size(i , j); + size_t seg_size = lca->get_segment_size(i, j); size_t seg_off = lca->get_segment_offset(i, j); - READANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); + READANDCHECK( + lca->codes[i][j].data_.data(), + seg_size * lca->code_size); READANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (save_norm) { - READANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); + READANDCHECK( + lca->code_norms[i][j].data_.data(), seg_size); } } } @@ -585,9 +594,7 @@ static void read_ProductLocalSearchQuantizer( } } -static void read_ScalarQuantizer( - ::faiss::ScalarQuantizer* ivsc, - IOReader* f) { +static void read_ScalarQuantizer(::faiss::ScalarQuantizer* ivsc, IOReader* f) { READ1(ivsc->qtype); READ1(ivsc->rangestat); READ1(ivsc->rangestat_arg); @@ -685,6 +692,41 @@ static void read_RaBitQuantizer( } } +static void finalize_and_validate_RaBitQ_index(::faiss::IndexRaBitQ* idxq) { + FAISS_THROW_IF_NOT_MSG( + idxq->metric_type == METRIC_L2 || + idxq->metric_type == METRIC_INNER_PRODUCT, + "IndexRaBitQ only supports L2 and inner product metrics"); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.d == static_cast(idxq->d) && + idxq->rabitq.metric_type == idxq->metric_type, + "IndexRaBitQ quantizer metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.nb_bits >= 1 && idxq->rabitq.nb_bits <= 9, + "IndexRaBitQ nb_bits must be in [1, 9]"); + + const size_t expected_code_size = + idxq->rabitq.compute_code_size(idxq->d, idxq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.code_size == expected_code_size, + "IndexRaBitQ quantizer code size mismatch"); + idxq->code_size = expected_code_size; + FAISS_THROW_IF_NOT_MSG( + idxq->codes.size() == + static_cast(idxq->ntotal) * expected_code_size, + "IndexRaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->center.empty() || + idxq->center.size() == static_cast(idxq->d), + "IndexRaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_FMT( + idxq->qb <= 8, + "invalid RaBitQ qb=%d (must be in [0, 8])", + idxq->qb); + // The V1 cppcontrib wire format intentionally has no centered field. + idxq->centered = false; +} + static void read_direct_map(DirectMap* dm, IOReader* f) { char maintain_direct_map; READ1(maintain_direct_map); @@ -699,11 +741,11 @@ static void read_direct_map(DirectMap* dm, IOReader* f) { map[it.first] = it.second; } } - // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` - // read branch is gone — see the symmetric comment in index_write.cpp. - // Old files (if any) with `type == 3` would fail to round-trip here - // since the enum value no longer exists; in practice CC indexes - // were never written through this path. + // Path-D step 10.9: the former `if (dm->type == + // DirectMap::ConcurrentArray)` read branch is gone — see the symmetric + // comment in index_write.cpp. Old files (if any) with `type == 3` would + // fail to round-trip here since the enum value no longer exists; in + // practice CC indexes were never written through this path. } static void read_ivf_header( @@ -791,7 +833,8 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(wire_l2_norms); // reconstruct inverse norms from wire L2 norms - idxf->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_l2_norms); + idxf->inverse_norms_storage = + L2NormsStorage::from_l2_norms(wire_l2_norms); FAISS_THROW_IF_NOT( idxf->codes.size() == idxf->ntotal * idxf->code_size); @@ -828,7 +871,8 @@ Index* read_index(IOReader* f, int io_flags) { idxfc->code_size = idxf->code_size; idxfc->codes = std::move(idxf->codes); // reconstruct inverse norms from wire L2 norms - idxfc->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_code_norms); + idxfc->inverse_norms_storage = + L2NormsStorage::from_l2_norms(wire_code_norms); delete idxf; idxf = idxfc; } @@ -849,7 +893,7 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(idxp->inverse_norms_storage.inverse_l2_norms); if (!(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { - idxp->pq.compute_sdc_table (); + idxp->pq.compute_sdc_table(); } idx = idxp; @@ -875,7 +919,7 @@ Index* read_index(IOReader* f, int io_flags) { // the following "if" block is Knowhere-specific if (h == fourcc("IxPq")) { - idxp->pq.compute_sdc_table (); + idxp->pq.compute_sdc_table(); } idx = idxp; @@ -1070,8 +1114,7 @@ Index* read_index(IOReader* f, int io_flags) { // either enum name, and route legacy data to // IndexBinaryScalarQuantizer. const int legacy_qt_1bit_direct_marker = 9; - if (static_cast(idxs->sq.qtype) == - legacy_qt_1bit_direct_marker) { + if (static_cast(idxs->sq.qtype) == legacy_qt_1bit_direct_marker) { IndexBinaryScalarQuantizer* bsq = new IndexBinaryScalarQuantizer( static_cast(idxs->d), idxs->metric_type); bsq->ntotal = idxs->ntotal; @@ -1082,6 +1125,17 @@ Index* read_index(IOReader* f, int io_flags) { } else { idx = idxs; } + } else if ( + h == fourcc("Ixrq") || h == fourcc("Ixrr")) { + auto idxq = std::make_unique<::faiss::IndexRaBitQ>(); + read_index_header(idxq.get(), f); + read_RaBitQuantizer( + &idxq->rabitq, f, /*multi_bit=*/h != fourcc("Ixrq")); + READVECTOR(idxq->codes); + READVECTOR(idxq->center); + READ1(idxq->qb); + finalize_and_validate_RaBitQ_index(idxq.get()); + idx = idxq.release(); } else if (h == fourcc("IvSQ")) { // legacy IndexIVFScalarQuantizer* ivsc = new IndexIVFScalarQuantizer(); std::vector> ids; @@ -1108,8 +1162,25 @@ Index* read_index(IOReader* f, int io_flags) { h == fourcc("IvPQ") || h == fourcc("IvQR") || h == fourcc("IwPQ") || h == fourcc("IwQR")) { idx = read_ivfpq(f, h, io_flags); + } else if (h == fourcc(kRaBitQPreTransformCosineFourcc)) { + auto owner = std::make_unique(); + auto* ixpt = owner.get(); + ixpt->own_fields = true; + read_index_header(ixpt, f); + int nt; + READ1(nt); + FAISS_THROW_IF_NOT_MSG( + nt >= 0, "negative transform count in cosine RaBitQ storage"); + for (int i = 0; i < nt; i++) { + ixpt->chain.push_back(read_VectorTransform(f)); + } + ixpt->index = read_index(f, io_flags); + READVECTOR(ixpt->inverse_norms_storage.inverse_l2_norms); + ixpt->validate_norms(); + idx = owner.release(); } else if (h == fourcc("IxPT")) { - IndexPreTransform* ixpt = new IndexPreTransform(); + auto owner = std::make_unique(); + auto* ixpt = owner.get(); ixpt->own_fields = true; read_index_header(ixpt, f); int nt; @@ -1122,7 +1193,7 @@ Index* read_index(IOReader* f, int io_flags) { ixpt->chain.push_back(read_VectorTransform(f)); } ixpt->index = read_index(f, io_flags); - idx = ixpt; + idx = owner.release(); } else if (h == fourcc("Imiq")) { MultiIndexQuantizer* imiq = new MultiIndexQuantizer(); read_index_header(imiq, f); @@ -1151,7 +1222,8 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxrf->k_factor); if (dynamic_cast<::faiss::IndexFlat*>(idxrf->refine_index)) { // then make a RefineFlat with it. Refine index may be a baseline - // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware subclass. + // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware + // subclass. IndexRefine* idxrf_old = idxrf; idxrf = new IndexRefineFlat(); *idxrf = *idxrf_old; @@ -1161,11 +1233,22 @@ Index* read_index(IOReader* f, int io_flags) { idxrf->own_refine_index = true; idx = idxrf; } else if ( - h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") || - h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHN9") || - h == fourcc("IHN8") || h == fourcc("IHNa") || h == fourcc("IHNb") || - h == fourcc("IHN7") || h == fourcc("IHN6") || h == fourcc("IHN5")) { + h == fourcc(kHnswRaBitQFourcc) || + h == fourcc(kHnswRaBitQCosineFourcc) || h == fourcc("IHNf") || + h == fourcc("IHNp") || h == fourcc("IHNs") || h == fourcc("IHN2") || + h == fourcc("IHNc") || h == fourcc("IHN9") || h == fourcc("IHN8") || + h == fourcc("IHNa") || h == fourcc("IHNb") || h == fourcc("IHN7") || + h == fourcc("IHN6") || h == fourcc("IHN5")) { IndexHNSW* idxhnsw = nullptr; + std::unique_ptr idxhnsw_rabitq_owner; + if (h == fourcc(kHnswRaBitQFourcc)) { + idxhnsw_rabitq_owner = std::make_unique(); + idxhnsw = idxhnsw_rabitq_owner.get(); + } + if (h == fourcc(kHnswRaBitQCosineFourcc)) { + idxhnsw_rabitq_owner = std::make_unique(); + idxhnsw = idxhnsw_rabitq_owner.get(); + } if (h == fourcc("IHNf")) idxhnsw = new IndexHNSWFlat(); if (h == fourcc("IHNp")) @@ -1198,6 +1281,13 @@ Index* read_index(IOReader* f, int io_flags) { read_HNSW(&idxhnsw->hnsw, f); idxhnsw->storage = read_index(f, io_flags); idxhnsw->own_fields = idxhnsw->storage != nullptr; + if (h == fourcc(kHnswRaBitQFourcc)) { + dynamic_cast(idxhnsw)->validate_storage(); + } + if (h == fourcc(kHnswRaBitQCosineFourcc)) { + dynamic_cast(idxhnsw) + ->validate_cosine_storage(); + } if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { dynamic_cast(idxhnsw->storage)->pq.compute_sdc_table(); } @@ -1222,7 +1312,7 @@ Index* read_index(IOReader* f, int io_flags) { delete idxhnsw; idxhnsw = newh; } - idx = idxhnsw; + idx = idxhnsw_rabitq_owner ? idxhnsw_rabitq_owner.release() : idxhnsw; } else if (h == fourcc("IwPf")) { ::faiss::IndexIVFPQFastScan* ivpq = new ::faiss::IndexIVFPQFastScan(); read_ivf_header(ivpq, f); @@ -1333,7 +1423,8 @@ Index* read_index(IOReader* f, int io_flags) { // field); Iwrr is baseline multi-bit and does serialize nb_bits. auto ivrq = new IndexIVFRaBitQ(); read_ivf_header(ivrq, f); - read_RaBitQuantizer(&ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); + read_RaBitQuantizer( + &ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); READ1(ivrq->code_size); READ1(ivrq->by_residual); READ1(ivrq->qb); @@ -1484,4 +1575,4 @@ IndexBinary* read_index_binary(const char* fname, int io_flags) { } } -} +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp index cf6301302..3e84d5da2 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp @@ -18,30 +18,32 @@ #include #include +#include #include #include -#include #include -#include #include -#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include +#include #include #include #include -#include -#include #include #include #include #include -#include -#include +#include #include #include @@ -69,11 +71,8 @@ * leak memory. **************************************************************/ - - namespace faiss::cppcontrib::knowhere { - /************************************************************* * Write **************************************************************/ @@ -94,7 +93,7 @@ static void write_index_header(const Index* idx, IOWriter* f) { WRITE1(dummy32); idx_t dummy = 0; WRITE1(dummy); - + WRITE1(idx->is_trained); WRITE1(idx->metric_type); if (idx->metric_type > 1) { @@ -305,8 +304,9 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { } } } - } else if (const auto & lca = - dynamic_cast(ils)) { + } else if ( + const auto& lca = + dynamic_cast(ils)) { uint32_t h = fourcc("ilca"); WRITE1(h); WRITE1(lca->nlist); @@ -350,16 +350,20 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { size_t seg_size = lca->get_segment_size(i, j); - WRITEANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); + WRITEANDCHECK( + lca->codes[i][j].data_.data(), + seg_size * lca->code_size); WRITEANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (lca->save_norm) { - WRITEANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); + WRITEANDCHECK( + lca->code_norms[i][j].data_.data(), seg_size); } } } } - } else if (const auto & oa = - dynamic_cast(ils)) { + } else if ( + const auto& oa = + dynamic_cast(ils)) { uint32_t h = fourcc("iloa"); WRITE1(h); WRITE1(oa->nlist); @@ -369,16 +373,16 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t n = oa->pin_readonly_ids->size() / sizeof(InvertedLists::idx_t); WRITE1(n); WRITEANDCHECK((InvertedLists::idx_t*)oa->pin_readonly_ids->data, n); - WRITEANDCHECK((uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); + WRITEANDCHECK( + (uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); #else size_t n = oa->readonly_ids.size(); WRITE1(n); WRITEANDCHECK(oa->readonly_ids.data(), n); WRITEANDCHECK(oa->readonly_codes.data(), n * oa->code_size); #endif - } else if (const auto & od = - dynamic_cast(ils)) { - uint32_t h = fourcc ("ilod"); + } else if (const auto& od = dynamic_cast(ils)) { + uint32_t h = fourcc("ilod"); WRITE1(h); WRITE1(ils->nlist); WRITE1(ils->code_size); @@ -387,7 +391,7 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { { std::vector v( - od->slots.begin(), od->slots.end()); + od->slots.begin(), od->slots.end()); WRITEVECTOR(v); } { @@ -502,6 +506,38 @@ static void write_RaBitQuantizer( } } +static void validate_RaBitQ_index_for_write(const ::faiss::IndexRaBitQ* idxq) { + FAISS_THROW_IF_NOT_MSG( + idxq->metric_type == METRIC_L2 || + idxq->metric_type == METRIC_INNER_PRODUCT, + "IndexRaBitQ only supports L2 and inner product metrics"); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.d == static_cast(idxq->d) && + idxq->rabitq.metric_type == idxq->metric_type, + "IndexRaBitQ quantizer metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.nb_bits >= 1 && idxq->rabitq.nb_bits <= 9, + "IndexRaBitQ nb_bits must be in [1, 9]"); + const size_t expected_code_size = + idxq->rabitq.compute_code_size(idxq->d, idxq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.code_size == expected_code_size && + idxq->code_size == expected_code_size, + "IndexRaBitQ code size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->codes.size() == + static_cast(idxq->ntotal) * expected_code_size, + "IndexRaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->center.empty() || + idxq->center.size() == static_cast(idxq->d), + "IndexRaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_MSG(idxq->qb <= 8, "IndexRaBitQ qb must be in [0, 8]"); + FAISS_THROW_IF_NOT_MSG( + !idxq->centered, + "cppcontrib IndexRaBitQ V1 serialization requires centered=false"); +} + static void write_direct_map(const DirectMap* dm, IOWriter* f) { char maintain_direct_map = (char)dm->type; // for backwards compatibility with bool @@ -514,11 +550,12 @@ static void write_direct_map(const DirectMap* dm, IOWriter* f) { std::copy(map.begin(), map.end(), v.begin()); WRITEVECTOR(v); } - // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` - // write branch is gone — fork DirectMap no longer supports that - // variant. CC indexes now carry their own `cc_direct_map` member - // (ConcurrentDirectMap) which is not serialized through this path - // (CC indexes have no serialize stage; see ivf.cc:619 comment). + // Path-D step 10.9: the former `if (dm->type == + // DirectMap::ConcurrentArray)` write branch is gone — fork DirectMap no + // longer supports that variant. CC indexes now carry their own + // `cc_direct_map` member (ConcurrentDirectMap) which is not serialized + // through this path (CC indexes have no serialize stage; see ivf.cc:619 + // comment). } static void write_ivf_header(const IndexIVF* ivf, IOWriter* f) { @@ -536,13 +573,15 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { // eg. for a storage component of HNSW that is set to nullptr uint32_t h = fourcc("null"); WRITE1(h); - } else if (const IndexFlatCosine* idxf = dynamic_cast(idx)) { + } else if ( + const IndexFlatCosine* idxf = + dynamic_cast(idx)) { uint32_t h = fourcc("IxF9"); WRITE1(h); write_index_header(idx, f); WRITEXBVECTOR(idxf->codes); // we're storing real l2 norms, because of - // backward compatibility issues. + // backward compatibility issues. WRITEVECTOR(idxf->inverse_norms_storage.as_l2_norms()); } else if ( const ::faiss::IndexFlat* idxf = @@ -556,7 +595,9 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(h); write_index_header(idx, f); WRITEXBVECTOR(idxf->codes); - } else if (const IndexPQCosine* idxp = dynamic_cast(idx)) { + } else if ( + const IndexPQCosine* idxp = + dynamic_cast(idx)) { uint32_t h = fourcc("IxP7"); WRITE1(h); write_index_header(idx, f); @@ -598,7 +639,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITEVECTOR(idxr_2->codes); } else if ( const IndexProductResidualQuantizerCosine* idxpr = - dynamic_cast(idx)) { + dynamic_cast( + idx)) { uint32_t h = fourcc("IxP5"); WRITE1(h); write_index_header(idx, f); @@ -710,6 +752,18 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index_header(idx, f); write_ScalarQuantizer(&idxs->sq, f); WRITEVECTOR(idxs->codes); + } else if ( + const ::faiss::IndexRaBitQ* idxq = + dynamic_cast(idx)) { + validate_RaBitQ_index_for_write(idxq); + const bool multi_bit = idxq->rabitq.nb_bits > 1; + uint32_t h = multi_bit ? fourcc("Ixrr") : fourcc("Ixrq"); + WRITE1(h); + write_index_header(idxq, f); + write_RaBitQuantizer(&idxq->rabitq, f, multi_bit); + WRITEVECTOR(idxq->codes); + WRITEVECTOR(idxq->center); + WRITE1(idxq->qb); } else if ( const IndexIVFFlat* ivfl = dynamic_cast(idx)) { @@ -742,6 +796,19 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(ivpq->code_size); write_ProductQuantizer(&ivpq->pq, f); write_InvertedLists(ivpq->invlists, f); + } else if ( + const auto* cosine_rabitq = + dynamic_cast(idx)) { + cosine_rabitq->validate_norms(); + uint32_t h = fourcc(kRaBitQPreTransformCosineFourcc); + WRITE1(h); + write_index_header(cosine_rabitq, f); + int nt = cosine_rabitq->chain.size(); + WRITE1(nt); + for (int i = 0; i < nt; i++) + write_VectorTransform(cosine_rabitq->chain[i], f); + write_index(cosine_rabitq->index, f); + WRITEVECTOR(cosine_rabitq->inverse_norms_storage.inverse_l2_norms); } else if ( const IndexPreTransform* ixpt = dynamic_cast(idx)) { @@ -780,8 +847,23 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index(idxrf->refine_index, f); WRITE1(idxrf->k_factor); } else if (const IndexHNSW* idxhnsw = dynamic_cast(idx)) { - uint32_t h = dynamic_cast(idx) ? fourcc("IHNf") - : dynamic_cast(idx) ? fourcc("IHNp") + const auto* hnsw_rabitq_cosine = + dynamic_cast(idx); + const auto* hnsw_rabitq = dynamic_cast(idx); + if (hnsw_rabitq) { + FAISS_THROW_IF_NOT_MSG( + !(io_flags & IO_FLAG_SKIP_STORAGE), + "IndexHNSWRaBitQ cannot be serialized without its RaBitQ storage"); + if (hnsw_rabitq_cosine) { + hnsw_rabitq_cosine->validate_cosine_storage(); + } else { + hnsw_rabitq->validate_storage(); + } + } + uint32_t h = hnsw_rabitq_cosine ? fourcc(kHnswRaBitQCosineFourcc) + : hnsw_rabitq ? fourcc(kHnswRaBitQFourcc) + : dynamic_cast(idx) ? fourcc("IHNf") + : dynamic_cast(idx) ? fourcc("IHNp") // IndexHNSWBinary reuses the legacy IHNs fourcc so // on-disk bytes match what IndexHNSWSQ(QT_1bit_direct, // metric) used to produce. Readers dispatch to @@ -888,7 +970,6 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { } } - void write_index(const Index* idx, FILE* f, int io_flags) { FileIOWriter writer(f); write_index(idx, &writer, io_flags); @@ -967,4 +1048,4 @@ void write_index_binary(const IndexBinary* idx, const char* fname) { write_index_binary(idx, &writer); } -} +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/utils/rabitq_simd.h b/thirdparty/faiss/faiss/utils/rabitq_simd.h index 5dac1d60a..9f49700f4 100644 --- a/thirdparty/faiss/faiss/utils/rabitq_simd.h +++ b/thirdparty/faiss/faiss/utils/rabitq_simd.h @@ -131,6 +131,14 @@ void quantize_query_values( // NONE specializations — scalar fallbacks +// RaBitQ codes and query bit planes are byte-aligned, including their tails. +// memcpy preserves the unaligned load contract without pointer-alignment UB. +inline uint64_t load_u64_unaligned(const uint8_t* ptr) { + uint64_t value; + std::memcpy(&value, ptr, sizeof(value)); + return value; +} + template <> inline uint64_t bitwise_and_dot_product( const uint8_t* query, @@ -140,9 +148,9 @@ inline uint64_t bitwise_and_dot_product( uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv & yv) << j; } } @@ -167,10 +175,10 @@ inline BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< uint64_t popcount_sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); dot_product += popcount64(qv & yv) << j; } } @@ -194,9 +202,9 @@ inline uint64_t bitwise_xor_dot_product( uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv ^ yv) << j; } } @@ -215,7 +223,7 @@ inline uint64_t popcount(const uint8_t* data, size_t size) { uint64_t sum = 0; size_t offset = 0; for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -357,6 +365,15 @@ inline float ip_scalar( size_t ex_bits, float cb) { float result = 0.0f; + if (ex_bits == 8) { + // RBQ9 is byte-aligned, including the last dimension. Do not require + // trailing factor bytes for the scalar reference or SIMD tail. + for (size_t i = start; i < d; ++i) { + const int sb = (sign_bits[i / 8] >> (i % 8)) & 1; + result += rotated_q[i] * (static_cast((sb << 8) + ex_code[i]) + cb); + } + return result; + } const int sign_shift = static_cast(ex_bits); const uint64_t code_mask = (1ULL << ex_bits) - 1; for (size_t i = start; i < d; i++) { diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp index 0d63504dc..0fb319e9f 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp @@ -275,9 +275,9 @@ uint64_t bitwise_and_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const uint64_t yv = *(const uint64_t*)(data + offset); + const uint64_t yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const uint64_t qv = *(const uint64_t*)(query + j * size + offset); + const uint64_t qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv & yv) << j; } } @@ -336,10 +336,10 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< dot_product += reduce_add_128(dot_128); popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const uint64_t yv = *(const uint64_t*)(data + offset); + const uint64_t yv = load_u64_unaligned(data + offset); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const uint64_t qv = *(const uint64_t*)(query + j * size + offset); + const uint64_t qv = load_u64_unaligned(query + j * size + offset); dot_product += popcount64(qv & yv) << j; } } @@ -391,9 +391,9 @@ uint64_t bitwise_xor_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv ^ yv) << j; } } @@ -427,7 +427,7 @@ uint64_t popcount(const uint8_t* data, size_t size) { } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -517,8 +517,8 @@ inline float ip_1exbit_avx2( return result; } -#ifdef __BMI2__ -inline float ip_bitplane_avx2( +#if defined(__GNUC__) && defined(__x86_64__) +__attribute__((target("bmi2"), noinline)) float ip_bitplane_avx2( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, @@ -583,12 +583,34 @@ float compute_inner_product( size_t d, size_t ex_bits, float cb) { + if (ex_bits == 8) { + // Eight byte-aligned extra codes, plus their independent sign bits. + __m256 acc = _mm256_setzero_ps(); + const __m256 weight = _mm256_set1_ps(256.f); + const __m256 offset = _mm256_set1_ps(cb); + const __m256i positions = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m128i bytes = _mm_loadl_epi64( + reinterpret_cast(ex_code + i)); + const __m256 extra = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(bytes)); + const __m256i mask = _mm256_cmpgt_epi32( + _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), positions), + _mm256_setzero_si256()); + const __m256 recon = _mm256_add_ps( + extra, _mm256_and_ps(_mm256_castsi256_ps(mask), weight)); + acc = _mm256_fmadd_ps(_mm256_loadu_ps(rotated_q + i), + _mm256_add_ps(recon, offset), acc); + } + return hsum_avx2(acc) + + ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + } if (ex_bits == 1) { return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ - if (ex_bits <= 7) { +#if defined(__GNUC__) && defined(__x86_64__) + if (ex_bits <= 7 && __builtin_cpu_supports("bmi2")) { return ip_bitplane_avx2(sign_bits, ex_code, rotated_q, d, ex_bits, cb); } #endif diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp index 991255f77..60afe80e7 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp @@ -378,9 +378,9 @@ uint64_t bitwise_and_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv & yv) << j; } } @@ -394,6 +394,32 @@ uint64_t bitwise_and_dot_product( return sum; } +#if defined(__GNUC__) && defined(__x86_64__) +namespace { +// Ice Lake already has VPOPCNTDQ; requiring the full SPR feature set here +// unnecessarily selects the shuffle-based fallback. Isolate the optional ISA. +__attribute__((target("avx512vpopcntdq"), noinline)) +BitwiseAndDotProductResult bitwise_q4_vpopcnt( + const uint8_t* query, const uint8_t* data, size_t size) { + __m512i dots = _mm512_setzero_si512(); + __m512i pops = _mm512_setzero_si512(); + for (size_t off = 0; off < size; off += 64) { + const size_t count = std::min(size - off, size_t(64)); + const __mmask64 mask = count == 64 ? ~__mmask64(0) : (__mmask64(1) << count) - 1; + const __m512i x = _mm512_maskz_loadu_epi8(mask, data + off); + pops = _mm512_add_epi64(pops, _mm512_popcnt_epi64(x)); + for (int bit = 0; bit < 4; ++bit) { + const __m512i q = _mm512_maskz_loadu_epi8(mask, query + bit * size + off); + const __m512i p = _mm512_popcnt_epi64(_mm512_and_si512(q, x)); + dots = _mm512_add_epi64(dots, _mm512_slli_epi64(p, bit)); + } + } + return {static_cast(_mm512_reduce_add_epi64(dots)), + static_cast(_mm512_reduce_add_epi64(pops))}; +} +} // namespace +#endif + template <> BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< SIMDLevel::AVX512>( @@ -401,6 +427,11 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< const uint8_t* data, size_t size, size_t qb) { +#if defined(__GNUC__) && defined(__x86_64__) + if (qb == 4 && __builtin_cpu_supports("avx512vpopcntdq")) { + return bitwise_q4_vpopcnt(query, data, size); + } +#endif uint64_t dot_product = 0; uint64_t popcount_sum = 0; size_t offset = 0; @@ -457,10 +488,10 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< dot_product += reduce_add_128(dot_128); popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); popcount_sum += popcount64(yv); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); dot_product += popcount64(qv & yv) << j; } } @@ -527,9 +558,9 @@ uint64_t bitwise_xor_dot_product( } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (int j = 0; j < qb; j++) { - const auto qv = *(const uint64_t*)(query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += popcount64(qv ^ yv) << j; } } @@ -572,7 +603,7 @@ uint64_t popcount(const uint8_t* data, size_t size) { } sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *(const uint64_t*)(data + offset); + const auto yv = load_u64_unaligned(data + offset); sum += popcount64(yv); } for (; offset < size; ++offset) { @@ -664,60 +695,82 @@ inline float ip_1exbit_avx512( return result; } -// AVX2+BMI2 bitplane kernel used as fallback for ex_bits >= 2. -// AVX512 TU has AVX2 available. BMI2 guarded separately since -// VIA Eden X4 has AVX2 without BMI2. -#ifdef __BMI2__ -inline float ip_bitplane_avx2( +// Faiss #5526 AVX512 bitplane kernel: 16 dimensions per iteration. +// BMI2 is isolated in this function and checked by the caller at runtime. +#if defined(__GNUC__) && defined(__x86_64__) +__attribute__((target("bmi2"), noinline)) float ip_bitplane_avx512( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, size_t d, size_t ex_bits, float cb) { - __m256 acc = _mm256_setzero_ps(); - const __m256 v_one = _mm256_set1_ps(1.0f); - const __m256i bit_pos = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); - const __m256i zero = _mm256_setzero_si256(); - const __m256 v_cb = _mm256_set1_ps(cb); + __m512 acc = _mm512_setzero_ps(); + const __m512 v_cb = _mm512_set1_ps(cb); uint64_t pext_masks[7]; - __m256 v_weights[8]; + __m512 v_weights[8]; for (size_t b = 0; b < ex_bits; b++) { uint64_t m = 0; for (int j = 0; j < 8; j++) { m |= (1ULL << (b + j * ex_bits)); } pext_masks[b] = m; - v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); + v_weights[b] = _mm512_set1_ps(static_cast(1u << b)); } - v_weights[ex_bits] = _mm256_set1_ps(static_cast(1u << ex_bits)); + v_weights[ex_bits] = _mm512_set1_ps(static_cast(1u << ex_bits)); size_t i = 0; - for (; i + 8 <= d; i += 8) { - __m256i sb_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(sign_bits[i / 8]), bit_pos), - zero); - __m256 recon = _mm256_mul_ps( - _mm256_and_ps(_mm256_castsi256_ps(sb_cmp), v_one), - v_weights[ex_bits]); + for (; i + 16 <= d; i += 16) { + uint16_t sb = 0; + memcpy(&sb, sign_bits + (i / 8), sizeof(uint16_t)); + __m512 recon = _mm512_maskz_mov_ps( + static_cast<__mmask16>(sb), v_weights[ex_bits]); - uint64_t ex64 = 0; - memcpy(&ex64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); + uint64_t lo64 = 0; + uint64_t hi64 = 0; + memcpy(&lo64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); + memcpy(&hi64, ex_code + ((i / 8) + 1) * ex_bits, sizeof(uint64_t)); for (size_t b = 0; b < ex_bits; b++) { - auto plane = static_cast(_pext_u64(ex64, pext_masks[b])); - __m256i p_cmp = _mm256_cmpgt_epi32( - _mm256_and_si256(_mm256_set1_epi32(plane), bit_pos), zero); - __m256 p_f = _mm256_and_ps(_mm256_castsi256_ps(p_cmp), v_one); - recon = _mm256_fmadd_ps(p_f, v_weights[b], recon); + const uint32_t plane = + static_cast(_pext_u64(lo64, pext_masks[b])) | + (static_cast(_pext_u64(hi64, pext_masks[b])) + << 8); + recon = _mm512_mask_add_ps( + recon, static_cast<__mmask16>(plane), recon, v_weights[b]); } - __m256 rq = _mm256_loadu_ps(rotated_q + i); - acc = _mm256_fmadd_ps(rq, _mm256_add_ps(recon, v_cb), acc); + __m512 rq = _mm512_loadu_ps(rotated_q + i); + acc = _mm512_fmadd_ps(rq, _mm512_add_ps(recon, v_cb), acc); + } + + // Half-width step: keeps the scalar tail under 8 dims when d is a multiple + // of 8 but not of 16 (e.g. 200, 1000). The upper 8 lanes are masked off + // throughout, and rotated_q is loaded masked so nothing is read past the + // end. + if (i + 8 <= d) { + const __mmask16 low8 = static_cast<__mmask16>(0x00ff); + __m512 recon = _mm512_maskz_mov_ps( + static_cast<__mmask16>(sign_bits[i / 8]), v_weights[ex_bits]); + + uint64_t lo64 = 0; + memcpy(&lo64, ex_code + (i / 8) * ex_bits, sizeof(uint64_t)); + + for (size_t b = 0; b < ex_bits; b++) { + const uint32_t plane = + static_cast(_pext_u64(lo64, pext_masks[b])); + recon = _mm512_mask_add_ps( + recon, static_cast<__mmask16>(plane), recon, v_weights[b]); + } + + __m512 rq = _mm512_maskz_loadu_ps(low8, rotated_q + i); + acc = _mm512_fmadd_ps( + rq, _mm512_mask_add_ps(recon, low8, recon, v_cb), acc); + i += 8; } - float result = hsum_avx2(acc); + float result = _mm512_reduce_add_ps(acc); result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } @@ -733,13 +786,33 @@ float compute_inner_product( size_t d, size_t ex_bits, float cb) { + if (ex_bits == 8) { + // RBQ9 has one byte per extra code: no bit-plane extraction or BMI2. + __m512 acc = _mm512_setzero_ps(); + const __m512 weight = _mm512_set1_ps(256.f); + const __m512 offset = _mm512_set1_ps(cb); + size_t i = 0; + for (; i + 16 <= d; i += 16) { + uint16_t signs; + memcpy(&signs, sign_bits + i / 8, sizeof(signs)); + const __m128i bytes = _mm_loadu_si128( + reinterpret_cast(ex_code + i)); + __m512 recon = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(bytes)); + recon = _mm512_mask_add_ps(recon, signs, recon, weight); + acc = _mm512_fmadd_ps( + _mm512_loadu_ps(rotated_q + i), + _mm512_add_ps(recon, offset), acc); + } + return _mm512_reduce_add_ps(acc) + + ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + } if (ex_bits == 1) { return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ - if (ex_bits <= 7) { - return ip_bitplane_avx2(sign_bits, ex_code, rotated_q, d, ex_bits, cb); +#if defined(__GNUC__) && defined(__x86_64__) + if (ex_bits <= 7 && __builtin_cpu_supports("bmi2")) { + return ip_bitplane_avx512(sign_bits, ex_code, rotated_q, d, ex_bits, cb); } #endif return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb); diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp index 0f8951f80..1ee004db2 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp @@ -167,10 +167,9 @@ uint64_t bitwise_and_dot_product( // 64-bit scalar tail. for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += static_cast(popcount64(qv & yv)) << j; } } @@ -258,11 +257,10 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< popcount_sum += reduce_add_128(pop_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + const auto yv = load_u64_unaligned(data + offset); popcount_sum += popcount64(yv); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); dot_product += static_cast(popcount64(qv & yv)) << j; } } @@ -339,10 +337,9 @@ uint64_t bitwise_xor_dot_product( sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + const auto yv = load_u64_unaligned(data + offset); for (size_t j = 0; j < qb; j++) { - const auto qv = *reinterpret_cast( - query + j * size + offset); + const auto qv = load_u64_unaligned(query + j * size + offset); sum += static_cast(popcount64(qv ^ yv)) << j; } } @@ -392,7 +389,7 @@ uint64_t popcount(const uint8_t* data, size_t size) { sum += reduce_add_128(sum_128); for (size_t step = 64 / 8; offset + step <= size; offset += step) { - const auto yv = *reinterpret_cast(data + offset); + const auto yv = load_u64_unaligned(data + offset); sum += popcount64(yv); } for (; offset < size; ++offset) { From 75018a2393140c81085f315f8103e6fdee88a727 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:25:46 +0000 Subject: [PATCH 2/7] fix: isolate HNSW RaBitQ search and harden load and iterator contracts Move RaBitQ dispatch and staged evaluation out of generic HNSW code. Reuse the refinement predicate and count full refinements in search statistics. Reject non-RaBitQ payloads before replacing a live index, retain filtered iterator bridges, and add targeted and SQ/PQ regression coverage. Remove unrelated serialization formatting changes. Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 71 ++++++++-- src/index/hnsw/impl/HnswSearchDispatch.h | 44 ++++++ .../hnsw/impl/IndexConditionalWrapper.cc | 8 +- src/index/hnsw/impl/IndexConditionalWrapper.h | 5 +- src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc | 78 ++++++++++ src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h | 23 +++ src/index/hnsw/impl/IndexHNSWWrapper.cc | 133 +++--------------- src/index/hnsw/impl/IndexHNSWWrapper.h | 14 ++ src/index/hnsw/impl/RaBitQSearchParameters.h | 7 +- tests/ut/test_hnsw_pending.cc | 84 +++++++++++ tests/ut/test_hnsw_rabitq.cc | 74 +++++++++- tests/ut/test_hnsw_rabitq_acceptance.cc | 50 ++++++- .../cppcontrib/knowhere/IndexHNSWRaBitQ.cpp | 6 +- .../cppcontrib/knowhere/impl/HnswSearcher.h | 124 ++++++---------- .../faiss/cppcontrib/knowhere/impl/Neighbor.h | 9 ++ .../knowhere/impl/RaBitQDistanceEvaluation.h | 57 ++++++++ .../cppcontrib/knowhere/impl/RaBitQSearch.h | 6 +- .../cppcontrib/knowhere/impl/index_read.cpp | 69 ++++----- .../cppcontrib/knowhere/impl/index_write.cpp | 49 +++---- 19 files changed, 628 insertions(+), 283 deletions(-) create mode 100644 src/index/hnsw/impl/HnswSearchDispatch.h create mode 100644 src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc create mode 100644 src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h create mode 100644 tests/ut/test_hnsw_pending.cc create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index a033fcf45..a9e8626a2 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -1398,9 +1398,11 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { // whether a user wants a refine const bool whether_to_enable_refine = hnsw_cfg.refine_k.has_value(); + auto search_parameters = CreateSearchParameters(hnsw_cfg); // set up an index wrapper - auto [index_wrapper, is_refined] = create_conditional_hnsw_wrapper( - indexes[index_id].get(), hnsw_cfg, whether_bf_search.value_or(false), whether_to_enable_refine); + auto [index_wrapper, is_refined] = + create_conditional_hnsw_wrapper(indexes[index_id].get(), hnsw_cfg, whether_bf_search.value_or(false), + whether_to_enable_refine, search_parameters.get()); if (index_wrapper == nullptr) { return expected::Err(Status::invalid_args, "an input index seems to be unrelated to HNSW"); @@ -1410,8 +1412,8 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { std::unique_ptr bf_index_wrapper = nullptr; faiss::Index* bf_index_wrapper_ptr = nullptr; if (!whether_bf_search.value_or(false)) { - std::tie(bf_index_wrapper, is_refined) = - create_conditional_hnsw_wrapper(indexes[index_id].get(), hnsw_cfg, true, whether_to_enable_refine); + std::tie(bf_index_wrapper, is_refined) = create_conditional_hnsw_wrapper( + indexes[index_id].get(), hnsw_cfg, true, whether_to_enable_refine, search_parameters.get()); if (bf_index_wrapper == nullptr) { return expected::Err(Status::invalid_args, "an input index seems to be unrelated to HNSW"); } @@ -1421,7 +1423,6 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - auto search_parameters = CreateSearchParameters(hnsw_cfg); auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { hnsw_search_params.efSearch = hnsw_cfg.ef.value(); @@ -1704,9 +1705,11 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { // whether a user wants a refine const bool whether_to_enable_refine = true; + auto search_parameters = CreateSearchParameters(hnsw_cfg); // set up an index wrapper - auto [index_wrapper, is_refined] = create_conditional_hnsw_wrapper( - indexes[index_id].get(), hnsw_cfg, whether_bf_search.value_or(false), whether_to_enable_refine); + auto [index_wrapper, is_refined] = + create_conditional_hnsw_wrapper(indexes[index_id].get(), hnsw_cfg, whether_bf_search.value_or(false), + whether_to_enable_refine, search_parameters.get()); if (index_wrapper == nullptr) { return expected::Err(Status::invalid_args, "an input index seems to be unrelated to HNSW"); @@ -1715,7 +1718,6 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - auto search_parameters = CreateSearchParameters(hnsw_cfg); auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { @@ -3051,6 +3053,26 @@ class BaseFaissRegularIndexHNSWPQNodeTemplate : public BaseFaissRegularIndexHNSW // immutable. class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode { public: + Status + Deserialize(const BinarySet& binset, std::shared_ptr) override { + auto binary = binset.GetByName(Type()); + if (!binary) + return Status::invalid_binary_set; + MemoryIOReader reader(binary->data.get(), binary->size); + return LoadRaBitQ(reader); + } + + Status + DeserializeFromFile(const std::string& filename, std::shared_ptr) override { + try { + faiss::FileIOReader reader(filename.c_str()); + return LoadRaBitQ(reader); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "RaBitQ file load failed: " << e.what(); + return Status::faiss_inner_error; + } + } + std::unique_ptr CreateSearchParameters(const FaissHnswConfig& config) const override { auto params = std::make_unique(); @@ -3093,6 +3115,39 @@ class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode protected: std::vector> tmp_index_rabitq; + Status + LoadRaBitQ(faiss::IOReader& reader) { + try { + // Validate before replacing the live index. Other readable Faiss + // types are not valid HNSW_RABITQ payloads, including MV containers. + auto loaded = std::unique_ptr(faiss::cppcontrib::knowhere::read_index(&reader)); + const auto* refine = dynamic_cast(loaded.get()); + const auto* rbq = dynamic_cast( + refine ? refine->base_index : loaded.get()); + if (!rbq) + return Status::invalid_serialized_index_type; + if (const auto* cosine = dynamic_cast(rbq)) { + cosine->validate_cosine_storage(); + } else { + rbq->validate_storage(); + } + if (refine) { + const auto* storage = refine->refine_index; + if (!storage || !refine->is_trained || !storage->is_trained || refine->d != rbq->d || + storage->d != rbq->d || refine->ntotal != rbq->ntotal || storage->ntotal != rbq->ntotal || + refine->metric_type != rbq->metric_type || storage->metric_type != rbq->metric_type) { + return Status::invalid_serialized_index_type; + } + } + indexes.assign(1, std::shared_ptr(loaded.release())); + tmp_index_rabitq.clear(); + return Status::success; + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "RaBitQ load failed: " << e.what(); + return is_faiss_fourcc_error(e.what()) ? Status::invalid_serialized_index_type : Status::faiss_inner_error; + } + } + Status TrainInternal(const DataSetPtr dataset, const Config& cfg) override { const auto rows = dataset->GetRows(); diff --git a/src/index/hnsw/impl/HnswSearchDispatch.h b/src/index/hnsw/impl/HnswSearchDispatch.h new file mode 100644 index 000000000..4a0529ebc --- /dev/null +++ b/src/index/hnsw/impl/HnswSearchDispatch.h @@ -0,0 +1,44 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once +#include +#include + +#include + +#include "index/hnsw/impl/DummyVisitor.h" +#include "index/hnsw/impl/FederVisitor.h" +#include "index/hnsw/impl/IndexHNSWWrapper.h" +#include "knowhere/bitsetview_idselector.h" + +namespace knowhere { + +// Reuse selector/visitor dispatch without exposing codec types to common HNSW. +template +faiss::cppcontrib::knowhere::HNSWStats +search_hnsw_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& distance, + faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, + const SearchParametersHNSWWrapper* params) { + auto run = [&](auto& visitor, const auto& selector) { + using Visitor = std::remove_reference_t; + using Selector = std::decay_t; + faiss::cppcontrib::knowhere::v2_hnsw_searcher + searcher{graph, distance, visitor, visited, selector, params ? params->kAlpha : 0.0f, params}; + return searcher.search(k, distances, labels); + }; + auto visit = [&](const auto& selector) { + if (params && params->feder) { + FederVisitor visitor(params->feder); + return run(visitor, selector); + } + DummyVisitor visitor; + return run(visitor, selector); + }; + const auto* selector = params ? dynamic_cast(params->sel) : nullptr; + if (selector && !selector->bitset_view.empty()) + return visit(*selector); + faiss::IDSelectorAll all; + return visit(all); +} +} // namespace knowhere diff --git a/src/index/hnsw/impl/IndexConditionalWrapper.cc b/src/index/hnsw/impl/IndexConditionalWrapper.cc index e543d5dd9..97d6bb50b 100644 --- a/src/index/hnsw/impl/IndexConditionalWrapper.cc +++ b/src/index/hnsw/impl/IndexConditionalWrapper.cc @@ -101,7 +101,7 @@ WhetherPerformBruteForceRangeSearch(const faiss::Index* index, const FaissHnswCo // index was trained with the refine. std::tuple, bool> create_conditional_hnsw_wrapper(faiss::Index* index, const FaissHnswConfig& hnsw_cfg, const bool whether_bf_search, - const bool whether_to_enable_refine) { + const bool whether_to_enable_refine, const SearchParametersHNSWWrapper* search_params) { const bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), knowhere::metric::COSINE); // check if we have a refine available. @@ -129,7 +129,8 @@ create_conditional_hnsw_wrapper(faiss::Index* index, const FaissHnswConfig& hnsw base_wrapper = std::make_unique(index_hnsw); } else { // use hnsw-search wrapper - base_wrapper = std::make_unique(index_hnsw); + base_wrapper = search_params ? search_params->create_hnsw_wrapper(index_hnsw) + : std::make_unique(index_hnsw); } // check if a user wants a refined result @@ -201,7 +202,8 @@ create_conditional_hnsw_wrapper(faiss::Index* index, const FaissHnswConfig& hnsw base_wrapper = std::make_unique(index_hnsw); } else { // use hnsw-search wrapper - base_wrapper = std::make_unique(index_hnsw); + base_wrapper = search_params ? search_params->create_hnsw_wrapper(index_hnsw) + : std::make_unique(index_hnsw); } return {std::move(base_wrapper), false}; diff --git a/src/index/hnsw/impl/IndexConditionalWrapper.h b/src/index/hnsw/impl/IndexConditionalWrapper.h index 84d86ee13..dd6b000ee 100644 --- a/src/index/hnsw/impl/IndexConditionalWrapper.h +++ b/src/index/hnsw/impl/IndexConditionalWrapper.h @@ -23,6 +23,8 @@ namespace knowhere { +struct SearchParametersHNSWWrapper; + struct HnswSearchThresholds { static constexpr float kHnswSearchKnnBFFilterThreshold = 0.93f; static constexpr float kHnswSearchRangeBFFilterThreshold = 0.97f; @@ -48,6 +50,7 @@ WhetherPerformBruteForceRangeSearch(const faiss::Index* index, const FaissHnswCo // index was trained with the refine. std::tuple, bool> create_conditional_hnsw_wrapper(faiss::Index* index, const FaissHnswConfig& hnsw_cfg, const bool whether_bf_search, - const bool whether_to_enable_refine); + const bool whether_to_enable_refine, + const SearchParametersHNSWWrapper* search_params = nullptr); } // namespace knowhere diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc new file mode 100644 index 000000000..074462141 --- /dev/null +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc @@ -0,0 +1,78 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" + +#include +#include +#include + +#include "index/hnsw/impl/HnswSearchDispatch.h" +#include "index/hnsw/impl/RaBitQSearchParameters.h" +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) +#include "knowhere/prometheus_client.h" +#endif + +namespace knowhere { +using idx_t = faiss::idx_t; +namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; + +void +IndexHNSWRaBitQWrapper::search(faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, faiss::idx_t* labels, + const faiss::SearchParameters* params_in) const { + FAISS_THROW_IF_NOT(k > 0); + const auto* index_hnsw = dynamic_cast(index); + FAISS_THROW_IF_NOT(index_hnsw && index_hnsw->storage); + const auto* params = dynamic_cast(params_in); + FAISS_THROW_IF_NOT_MSG(!params_in || params, "params type invalid"); + if (index_hnsw->hnsw.entry_point == -1) { + IndexHNSWWrapper::search(n, x, k, distances, labels, params_in); + return; + } + const auto& hnsw = index_hnsw->hnsw; + const auto* rbq_params = dynamic_cast(params); + // Use the optimized multi-bit path only when its selector/visitor contract + // is satisfied. RBQ1, filtering and feder use the compatible searcher below. + const auto* bitset_sel = params ? dynamic_cast(params->sel) : nullptr; + const bool unfiltered = !params || !params->sel || (bitset_sel && bitset_sel->bitset_view.empty()); + if (index_hnsw->rabitq_index()->rabitq.nb_bits > 1 && unfiltered && (!params || !params->feder)) { + rabitq_search::search( + *index_hnsw, n, x, k, distances, labels, params ? params->efSearch : hnsw.efSearch, + params ? params->check_relative_distance : hnsw.check_relative_distance, + rbq_params ? &rbq_params->storage_params : nullptr, [&](const rabitq_search::SearchStats& counts) { + const size_t hops = counts.expanded + counts.upper_expanded; +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + knowhere::knowhere_hnsw_search_hops.Observe(hops); +#endif + if (params && params->hnsw_stats) { + params->hnsw_stats->combine({.n1 = 1, + .n2 = size_t(counts.exhausted), + .ndis = counts.estimate + counts.refine + counts.upper_full, + .nhops = hops}); + } + }); + if (faiss::cppcontrib::knowhere::is_similarity_metric(index->metric_type)) { + for (idx_t i = 0; i < k * n; ++i) distances[i] = -distances[i]; + } + return; + } + + IndexHNSWWrapper::search(n, x, k, distances, labels, params_in); +} + +std::unique_ptr +IndexHNSWRaBitQWrapper::graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const { + const auto* rbq = dynamic_cast(index); + FAISS_THROW_IF_NOT(rbq); + const auto* rbq_params = dynamic_cast(params); + return std::unique_ptr( + rbq->get_staged_distance_computer(rbq_params ? &rbq_params->storage_params : nullptr)); +} + +faiss::cppcontrib::knowhere::HNSWStats +IndexHNSWRaBitQWrapper::search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, + faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, + faiss::idx_t* labels, const SearchParametersHNSWWrapper* params) const { + return search_hnsw_query(graph, dc, visited, k, distances, labels, params); +} +} // namespace knowhere diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h new file mode 100644 index 000000000..9fc6eb39f --- /dev/null +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h @@ -0,0 +1,23 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "index/hnsw/impl/IndexHNSWWrapper.h" + +namespace knowhere { +struct IndexHNSWRaBitQWrapper : IndexHNSWWrapper { + using IndexHNSWWrapper::IndexHNSWWrapper; + void + search(faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, faiss::idx_t* labels, + const faiss::SearchParameters* params) const override; + + protected: + std::unique_ptr + graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const override; + faiss::cppcontrib::knowhere::HNSWStats + search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, + faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, + const SearchParametersHNSWWrapper* params) const override; +}; +} // namespace knowhere diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.cc b/src/index/hnsw/impl/IndexHNSWWrapper.cc index 9135f7a62..3b57bde11 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWWrapper.cc @@ -13,12 +13,10 @@ #include #include -#include #include #include #include #include -#include #include #include #include @@ -32,7 +30,7 @@ #include "index/hnsw/impl/DummyVisitor.h" #include "index/hnsw/impl/FederVisitor.h" -#include "index/hnsw/impl/RaBitQSearchParameters.h" +#include "index/hnsw/impl/HnswSearchDispatch.h" #include "knowhere/bitsetview.h" #include "knowhere/bitsetview_idselector.h" @@ -41,7 +39,6 @@ #endif namespace knowhere { -namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; /************************************************************** * Utilities @@ -71,6 +68,24 @@ IndexHNSWWrapper::IndexHNSWWrapper(faiss::cppcontrib::knowhere::IndexHNSW* under : faiss::cppcontrib::knowhere::IndexWrapper(underlying_index) { } +std::unique_ptr +SearchParametersHNSWWrapper::create_hnsw_wrapper(faiss::cppcontrib::knowhere::IndexHNSW* index) const { + return std::make_unique(index); +} + +std::unique_ptr +IndexHNSWWrapper::graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper*) const { + return std::unique_ptr(storage_distance_computer(index->storage)); +} + +faiss::cppcontrib::knowhere::HNSWStats +IndexHNSWWrapper::search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, + faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, + faiss::idx_t* labels, const SearchParametersHNSWWrapper* params) const { + return search_hnsw_query(graph, dc, visited, k, distances, labels, params); +} + void IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __restrict distances, idx_t* __restrict labels, const faiss::SearchParameters* __restrict params_in) const { @@ -99,40 +114,9 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r const SearchParametersHNSWWrapper* params = nullptr; const faiss::cppcontrib::knowhere::HNSW& hnsw = index_hnsw->hnsw; - float kAlpha = 0.0f; if (params_in) { params = dynamic_cast(params_in); FAISS_THROW_IF_NOT_MSG(params, "params type invalid"); - - kAlpha = params->kAlpha; - } - - const auto* rbq_params = dynamic_cast(params); - // Use the optimized multi-bit path only when its selector/visitor contract - // is satisfied. RBQ1, filtering and feder use the compatible searcher below. - const auto* rabitq_index = dynamic_cast(index_hnsw); - const auto* bitset_sel = params ? dynamic_cast(params->sel) : nullptr; - const bool unfiltered = !params || !params->sel || (bitset_sel && bitset_sel->bitset_view.empty()); - if (rabitq_index && rabitq_index->rabitq_index()->rabitq.nb_bits > 1 && unfiltered && (!params || !params->feder)) { - rabitq_search::search(*rabitq_index, n, x, k, distances, labels, params ? params->efSearch : hnsw.efSearch, - params ? params->check_relative_distance : hnsw.check_relative_distance, - rbq_params ? &rbq_params->storage_params : nullptr, - [&](const rabitq_search::SearchStats& counts) { - const size_t hops = counts.expanded + counts.upper_expanded; -#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) - knowhere::knowhere_hnsw_search_hops.Observe(hops); -#endif - if (params && params->hnsw_stats) { - params->hnsw_stats->combine({.n1 = 1, - .n2 = size_t(counts.exhausted), - .ndis = counts.estimate + counts.upper_full, - .nhops = hops}); - } - }); - if (faiss::cppcontrib::knowhere::is_similarity_metric(index->metric_type)) { - for (idx_t i = 0; i < k * n; ++i) distances[i] = -distances[i]; - } - return; } // set up hnsw_stats @@ -149,10 +133,7 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r faiss::cppcontrib::knowhere::Bitset bitset_visited_nodes = faiss::cppcontrib::knowhere::Bitset::create_uninitialized(index->ntotal); - // create a distance computer - std::unique_ptr dis( - rabitq_index ? rabitq_index->get_staged_distance_computer(rbq_params ? &rbq_params->storage_params : nullptr) - : storage_distance_computer(index_hnsw->storage)); + auto dis = graph_distance_computer(index_hnsw, params); // no parallelism by design for (idx_t i = 0; i < n; i++) { @@ -162,78 +143,8 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r // prepare the table of visited elements bitset_visited_nodes.clear(); - // a visitor - knowhere::feder::hnsw::FederResult* feder = (params == nullptr) ? nullptr : params->feder; - - // future results - faiss::cppcontrib::knowhere::HNSWStats local_stats; - - // set up a filter - faiss::IDSelector* sel = (params == nullptr) ? nullptr : params->sel; - - // try knowhere-specific filter - if (const knowhere::BitsetViewIDSelector* __restrict bw_idselector = - dynamic_cast(sel); - bw_idselector && !bw_idselector->bitset_view.empty()) { - // with filter, no mapping - - // feder templating is important, bcz it removes an unneeded 'CALL' instruction. - if (feder == nullptr) { - // no feder - DummyVisitor graph_visitor; - - using searcher_type = - faiss::cppcontrib::knowhere::v2_hnsw_searcher; - - searcher_type searcher{hnsw, *(dis.get()), graph_visitor, bitset_visited_nodes, - *bw_idselector, kAlpha, params}; - - local_stats = searcher.search(k, distances + i * k, labels + i * k); - } else { - // use feder - FederVisitor graph_visitor(feder); - - using searcher_type = - faiss::cppcontrib::knowhere::v2_hnsw_searcher; - - searcher_type searcher{hnsw, *(dis.get()), graph_visitor, bitset_visited_nodes, - *bw_idselector, kAlpha, params}; - - local_stats = searcher.search(k, distances + i * k, labels + i * k); - } - } else { - // no filter - faiss::IDSelectorAll sel_all; - - // feder templating is important, bcz it removes an unneeded 'CALL' instruction. - if (feder == nullptr) { - // no feder - DummyVisitor graph_visitor; - - using searcher_type = faiss::cppcontrib::knowhere::v2_hnsw_searcher< - faiss::DistanceComputer, DummyVisitor, faiss::cppcontrib::knowhere::Bitset, faiss::IDSelectorAll>; - - searcher_type searcher{hnsw, *(dis.get()), graph_visitor, bitset_visited_nodes, - sel_all, kAlpha, params}; - - local_stats = searcher.search(k, distances + i * k, labels + i * k); - } else { - // use feder - FederVisitor graph_visitor(feder); - - using searcher_type = faiss::cppcontrib::knowhere::v2_hnsw_searcher< - faiss::DistanceComputer, FederVisitor, faiss::cppcontrib::knowhere::Bitset, faiss::IDSelectorAll>; - - searcher_type searcher{hnsw, *(dis.get()), graph_visitor, bitset_visited_nodes, - sel_all, kAlpha, params}; - - local_stats = searcher.search(k, distances + i * k, labels + i * k); - } - } + const auto local_stats = + search_query(hnsw, *dis, bitset_visited_nodes, k, distances + i * k, labels + i * k, params); // record some statistics #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.h b/src/index/hnsw/impl/IndexHNSWWrapper.h index a2bef8cac..6f213aa2a 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWWrapper.h @@ -13,11 +13,13 @@ #include #include +#include #include #include #include #include +#include #include "knowhere/feder/HNSW.h" @@ -38,6 +40,9 @@ struct SearchParametersHNSWWrapper : public faiss::cppcontrib::knowhere::SearchP return index->get_distance_computer(); } + virtual std::unique_ptr + create_hnsw_wrapper(faiss::cppcontrib::knowhere::IndexHNSW* index) const; + inline ~SearchParametersHNSWWrapper() { } }; @@ -59,6 +64,15 @@ struct IndexHNSWWrapper : public faiss::cppcontrib::knowhere::IndexWrapper { void range_search(faiss::idx_t n, const float* x, float radius, faiss::RangeSearchResult* result, const faiss::SearchParameters* params) const override; + + protected: + virtual std::unique_ptr + graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const; + virtual faiss::cppcontrib::knowhere::HNSWStats + search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, + faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, + const SearchParametersHNSWWrapper* params) const; }; } // namespace knowhere diff --git a/src/index/hnsw/impl/RaBitQSearchParameters.h b/src/index/hnsw/impl/RaBitQSearchParameters.h index 2ce4eed17..3b62f09b2 100644 --- a/src/index/hnsw/impl/RaBitQSearchParameters.h +++ b/src/index/hnsw/impl/RaBitQSearchParameters.h @@ -13,7 +13,7 @@ #include -#include "index/hnsw/impl/IndexHNSWWrapper.h" +#include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" namespace knowhere { @@ -21,6 +21,11 @@ namespace knowhere { struct SearchParametersHNSWRaBitQWrapper : SearchParametersHNSWWrapper { faiss::RaBitQSearchParameters storage_params; + std::unique_ptr + create_hnsw_wrapper(faiss::cppcontrib::knowhere::IndexHNSW* index) const override { + return std::make_unique(index); + } + faiss::DistanceComputer* storage_distance_computer(const faiss::Index* index) const override { const auto* rbq = dynamic_cast(index); diff --git a/tests/ut/test_hnsw_pending.cc b/tests/ut/test_hnsw_pending.cc new file mode 100644 index 000000000..549a3cf4a --- /dev/null +++ b/tests/ut/test_hnsw_pending.cc @@ -0,0 +1,84 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include + +#include + +#include "catch2/catch_test_macros.hpp" +#include "index/hnsw/impl/DummyVisitor.h" + +TEST_CASE("HNSW retains filtered bridges hidden by a tighter result threshold", "[hnsw_pending]") { + namespace fk = faiss::cppcontrib::knowhere; + fk::NeighborSetDoublePopList candidates(1); + fk::IteratorMinHeap pending; + candidates.insert({0, 10, fk::Neighbor::kValid}, &pending); + REQUIRE(candidates.pop().id == 0); + candidates.insert({1, 5, fk::Neighbor::kInvalid}, &pending); + candidates.insert({2, 1, fk::Neighbor::kValid}, &pending); + REQUIRE(candidates.pop().id == 2); + REQUIRE_FALSE(candidates.has_next()); + candidates.save_pending(pending); + REQUIRE(pending.top().id == 1); +} + +TEST_CASE("HNSW pending traversal reaches a vertex behind a filtered bridge", "[hnsw_pending]") { + namespace fk = faiss::cppcontrib::knowhere; + // 0 -> [1(filtered), 2], 1 -> 3. After accepting 2 the threshold drops + // below 1's distance. Vertex 3 is reachable only through the saved bridge. + fk::HNSW graph(2); + graph.entry_point = 0; + graph.max_level = 0; + graph.levels.assign(4, 1); + graph.offsets = {0, 4, 8, 12, 16}; + graph.neighbors.resize(16); + std::fill(graph.neighbors.data(), graph.neighbors.data() + 16, -1); + graph.neighbors[0] = 1; + graph.neighbors[1] = 2; + graph.neighbors[4] = 3; + struct Distances : faiss::DistanceComputer { + float + operator()(faiss::idx_t id) override { + const float d[] = {10, 5, 1, 0.5f}; + return d[id]; + } + void + set_query(const float*) override { + } + float + symmetric_dis(faiss::idx_t, faiss::idx_t) override { + return 0; + } + } dc; + struct Filter { + bool + is_member(faiss::idx_t id) const { + return id != 1; + } + } filter; + knowhere::DummyVisitor visitor; + auto visited = fk::Bitset::create_cleared(4); + visited.set(0); + fk::v2_hnsw_searcher searcher( + graph, dc, visitor, visited, filter, 1.0f, nullptr); + fk::NeighborSetDoublePopList candidates(1); + candidates.insert({0, 10, fk::Neighbor::kValid}); + fk::IteratorMinHeap pending; + searcher.search_on_a_level(candidates, 0, &pending); + REQUIRE_FALSE(visited.get(3)); + std::set emitted{candidates[0].id}; + float alpha = 1; + while (!pending.empty()) { + const auto next = pending.top(); + pending.pop(); + searcher.evaluate_single_node(next.id, 0, alpha, [&](fk::Neighbor item) { + pending.push(item); + return true; + }); + if (filter.is_member(next.id)) + emitted.insert(next.id); + } + REQUIRE(visited.get(3)); + REQUIRE(emitted == std::set{0, 2, 3}); +} diff --git a/tests/ut/test_hnsw_rabitq.cc b/tests/ut/test_hnsw_rabitq.cc index 0a1611a62..e5c0c9972 100644 --- a/tests/ut/test_hnsw_rabitq.cc +++ b/tests/ut/test_hnsw_rabitq.cc @@ -5,7 +5,9 @@ #include #include #include +#include #include +#include #include #include @@ -16,8 +18,11 @@ #include "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" +#include "index/hnsw/impl/HnswSearchDispatch.h" +#include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" #include "index/hnsw/impl/IndexHNSWWrapper.h" #include "knowhere/bitsetview.h" +#include "knowhere/bitsetview_idselector.h" #include "knowhere/comp/knowhere_config.h" #include "knowhere/index/index_factory.h" #include "knowhere/utils.h" @@ -25,6 +30,34 @@ namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; +namespace { +// Independent reproduction of the original full-distance candidate batch: +// four-at-a-time, with scalar tails and no threshold/result state. +struct OriginalFullEvaluation { + void + begin(size_t) { + } + void + record(float, int) { + } + template + size_t + compute(DC& dc, const size_t* ids, const int*, size_t count, int, Emit&& emit) { + size_t i = 0; + for (; i + 4 <= count; i += 4) { + float a, b, c, d; + dc.distances_batch_4(ids[i], ids[i + 1], ids[i + 2], ids[i + 3], a, b, c, d); + emit(i, a); + emit(i + 1, b); + emit(i + 2, c); + emit(i + 3, d); + } + for (; i < count; ++i) emit(i, dc(ids[i])); + return 0; + } +}; +} // namespace + TEST_CASE("RaBitQ qb4 SIMD matches scalar including masked tails", "[hnsw_rabitq_core]") { #if defined(__GNUC__) && defined(__x86_64__) if (!__builtin_cpu_supports("avx512f") || !__builtin_cpu_supports("avx512bw") || @@ -117,7 +150,7 @@ TEST_CASE("RaBitQ traversal retains all results when k covers the graph", "[hnsw labels.data(), n, true); std::vector api_distances(4 * n); std::vector api_labels(4 * n); - knowhere::IndexHNSWWrapper api(&graph); + knowhere::IndexHNSWRaBitQWrapper api(&graph); knowhere::SearchParametersHNSWWrapper params; params.efSearch = n; api.search(4, static_cast(queries->GetTensor()), n, api_distances.data(), api_labels.data(), @@ -463,6 +496,45 @@ TEST_CASE("Generic HNSW parameter factory preserves SQ and PQ searches", "[hnsw_ REQUIRE(index.Build(base, config) == knowhere::Status::success); auto result = index.Search(query, config, nullptr); REQUIRE(result.has_value()); + // Same graph and codes: compare the new default evaluator with + // the pre-refactor full batch rule, including moderate filtering. + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + auto blob = binary.binary_map_.begin()->second; + faiss::VectorIOReader reader; + reader.data.assign(blob->data.get(), blob->data.get() + blob->size); + std::unique_ptr decoded(faiss::cppcontrib::knowhere::read_index(&reader)); + auto* graph = dynamic_cast(decoded.get()); + REQUIRE(graph != nullptr); + for (int excluded : {0, 128}) { + std::vector bits(128, 0); + std::fill_n(bits.begin(), excluded / 8, uint8_t(255)); + knowhere::BitsetView filter(bits.data(), 1024, excluded); + knowhere::BitsetViewIDSelector selector(filter); + knowhere::SearchParametersHNSWWrapper parameters; + parameters.efSearch = 128; + parameters.sel = excluded ? &selector : nullptr; + parameters.kAlpha = filter.filter_ratio() * 0.7f; + auto actual = index.Search(query, config, excluded ? filter : knowhere::BitsetView{}); + REQUIRE(actual.has_value()); + for (int q = 0; q < 2; ++q) { + std::unique_ptr dc(graph->storage->get_distance_computer()); + const bool similarity = std::string(metric) != "L2"; + if (similarity) + dc.reset(new faiss::NegativeDistanceComputer(dc.release())); + dc->set_query(static_cast(query->GetTensor()) + q * 32); + auto visited = faiss::cppcontrib::knowhere::Bitset::create_cleared(1024); + float distances[10]; + faiss::idx_t ids[10]; + knowhere::search_hnsw_query(graph->hnsw, *dc, visited, 10, distances, ids, + ¶meters); + for (int j = 0; j < 10; ++j) { + REQUIRE(actual.value()->GetIds()[q * 10 + j] == ids[j]); + REQUIRE(actual.value()->GetDistance()[q * 10 + j] == + (similarity ? -distances[j] : distances[j])); + } + } + } auto iterators = index.AnnIterator(query, config, nullptr); REQUIRE(iterators.has_value()); REQUIRE(iterators.value()[0]->HasNext().value()); diff --git a/tests/ut/test_hnsw_rabitq_acceptance.cc b/tests/ut/test_hnsw_rabitq_acceptance.cc index e6686cc80..8db99e690 100644 --- a/tests/ut/test_hnsw_rabitq_acceptance.cc +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -422,7 +422,7 @@ TEST_CASE("RaBitQ advertised refiners rerank the requested expanded candidate se REQUIRE(refiner != nullptr); auto* graph = dynamic_cast(refiner->base_index); REQUIRE(graph != nullptr); - knowhere::IndexHNSWWrapper wrapper(graph); + knowhere::IndexHNSWRaBitQWrapper wrapper(graph); for (int qb : {0, 4, 8}) { cfg["rbq_bits_query"] = qb; auto result = index.Search(query, cfg, nullptr); @@ -614,6 +614,54 @@ TEST_CASE("RaBitQ file serialization rejects truncated and incompatible indexes" } } +TEST_CASE("RaBitQ rejects valid non-RaBitQ payloads without replacing the live index", "[hnsw_rabitq_acceptance]") { + auto base = GenDataSet(128, 33, 2031); + auto query = GenDataSet(2, 33, 2032); + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + knowhere::Json cfg = {{"dim", 33}, {"metric_type", "L2"}, {"M", 8}, {"efConstruction", 64}, {"ef", 80}, + {"k", 10}, {"rbq_bits", 4}, {"sq_type", "SQ8"}}; + auto index = knowhere::IndexFactory::Instance().Create("HNSW_RABITQ", version).value(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + auto before = index.Search(query, cfg, nullptr); + REQUIRE(before.has_value()); + for (const auto* wrong_type : {"HNSW", "HNSW_SQ"}) { + CAPTURE(wrong_type); + auto wrong = knowhere::IndexFactory::Instance().Create(wrong_type, version).value(); + REQUIRE(wrong.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(wrong.Serialize(binary) == knowhere::Status::success); + const auto blob = binary.binary_map_.begin()->second; + knowhere::BinarySet mislabelled; + mislabelled.Append("HNSW_RABITQ", blob->data, blob->size); + REQUIRE(index.Deserialize(mislabelled, cfg) == knowhere::Status::invalid_serialized_index_type); + auto pattern = (std::filesystem::temp_directory_path() / "knowhere-rabitq-type-XXXXXX").string(); + std::vector filename(pattern.begin(), pattern.end()); + filename.push_back('\0'); + const int fd = mkstemp(filename.data()); + REQUIRE(fd >= 0); + struct Cleanup { + std::string path; + ~Cleanup() { + std::error_code error; + std::filesystem::remove(path, error); + } + } cleanup{filename.data()}; + std::unique_ptr file(fdopen(fd, "wb"), &std::fclose); + if (!file) + close(fd); + REQUIRE(file != nullptr); + REQUIRE(std::fwrite(blob->data.get(), 1, blob->size, file.get()) == blob->size); + REQUIRE(std::fflush(file.get()) == 0); + REQUIRE(index.DeserializeFromFile(filename.data(), cfg) == knowhere::Status::invalid_serialized_index_type); + auto after = index.Search(query, cfg, nullptr); + REQUIRE(after.has_value()); + for (int i = 0; i < 20; ++i) { + REQUIRE(after.value()->GetIds()[i] == before.value()->GetIds()[i]); + REQUIRE(after.value()->GetDistance()[i] == before.value()->GetDistance()[i]); + } + } +} + TEST_CASE("Shared Faiss IVF RaBitQ bits and query parameters survive serialization", "[hnsw_rabitq_regression]") { auto base = GenDataSet(512, 33, 1961); auto query = GenDataSet(2, 33, 1962); diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp index 1eeab6555..727bdde65 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -6,6 +6,7 @@ */ #include +#include #include #include @@ -67,9 +68,8 @@ struct RaBitQStagedDistanceComputer final : StagedDistanceComputer { float d = estimate; // Compare in output-distance units: positive cosine scale preserves // ordering, avoiding a division for every visited candidate. - const float error = factors->f_error * dc->g_error; - const bool refine = similarity ? (estimate + error) * s > -threshold - : std::max(0.0f, estimate - error) < threshold; + const bool refine = rabitq_search::should_refine( + estimate, factors->f_error, dc->g_error, threshold, similarity, s); if (refine) { d = dc->distance_to_code_full(code); ++refine_count; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h index d316daab6..eaf5f513a 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h @@ -34,7 +34,6 @@ // Knowhere-specific headers #include -#include namespace faiss { namespace cppcontrib { @@ -47,6 +46,26 @@ constexpr bool track_hnsw_stats = true; } // namespace +// Default evaluator preserves ordinary full-distance, batch-four execution. +struct FullDistanceEvaluation { + void begin(size_t) {} + void record(float, int) {} + + template + size_t compute(DC& dc, const size_t* ids, const int*, size_t count, + int, Emit&& emit) { + if (count == 4) { + float d[4]; + dc.distances_batch_4(ids[0], ids[1], ids[2], ids[3], + d[0], d[1], d[2], d[3]); + for (size_t i = 0; i < count; ++i) emit(i, d[i]); + } else { + for (size_t i = 0; i < count; ++i) emit(i, dc(ids[i])); + } + return 0; + } +}; + // Accomodates all the search logic and variables. /// * DistanceComputerT is responsible for computing distances /// * GraphVisitorT records visited edges @@ -58,7 +77,8 @@ template < typename DistanceComputerT, typename GraphVisitorT, typename VisitedT, - typename FilterT> + typename FilterT, + typename EvaluationT = FullDistanceEvaluation> struct v2_hnsw_searcher { using storage_idx_t = faiss::cppcontrib::knowhere::HNSW::storage_idx_t; using idx_t = faiss::idx_t; @@ -90,21 +110,7 @@ struct v2_hnsw_searcher { // the pointer is not owned. const faiss::cppcontrib::knowhere::SearchParametersHNSW* params; - StagedDistanceComputer* staged = nullptr; - size_t staged_k = 0; - std::priority_queue staged_results; - - float staged_threshold() const { - return staged_results.size() < staged_k - ? std::numeric_limits::infinity() : staged_results.top(); - } - void record_staged_result(float distance, int status) { - if (!staged || !staged_k || status == knowhere::Neighbor::kInvalid) return; - if (staged_results.size() < staged_k) staged_results.push(distance); - else if (distance < staged_results.top()) { - staged_results.pop(); staged_results.push(distance); - } - } + EvaluationT evaluation; // v2_hnsw_searcher( @@ -248,76 +254,31 @@ struct v2_hnsw_searcher { ndis += 1; if (counter == 4) { - // Staged evaluation preserves per-candidate threshold updates. - if (staged && staged_k && level == 0) { - for (size_t i = 0; i < 4; ++i) { - const float d = staged->evaluate(saved_indices[i], staged_threshold()); - graph_visitor.visit_edge(level, node_id, saved_indices[i], d); - record_staged_result(d, saved_statuses[i]); - func_add_candidate(knowhere::Neighbor(saved_indices[i], d, saved_statuses[i])); - } - counter = 0; - continue; - } - // evaluate 4x distances at once - float dis[4] = {0, 0, 0, 0}; - qdis.distances_batch_4( - saved_indices[0], - saved_indices[1], - saved_indices[2], - saved_indices[3], - dis[0], - dis[1], - dis[2], - dis[3]); - - for (size_t id4 = 0; id4 < 4; id4++) { - // record a traversed edge - graph_visitor.visit_edge( - level, node_id, saved_indices[id4], dis[id4]); - - // add a record of visited nodes - knowhere::Neighbor nn( - saved_indices[id4], dis[id4], saved_statuses[id4]); - if (func_add_candidate(nn)) { -#if defined(USE_PREFETCH) - // TODO - // _mm_prefetch(get_linklist0(v), _MM_HINT_T0); -#endif - } - } + ndis += evaluation.compute( + qdis, saved_indices, saved_statuses, counter, level, + [&](size_t i, float distance) { + graph_visitor.visit_edge(level, node_id, saved_indices[i], distance); + func_add_candidate(knowhere::Neighbor( + saved_indices[i], distance, saved_statuses[i])); + }); counter = 0; } } - // process leftovers - for (size_t id4 = 0; id4 < counter; id4++) { - // evaluate a single distance - const float dis = staged && staged_k && level == 0 - ? staged->evaluate(saved_indices[id4], staged_threshold()) - : qdis(saved_indices[id4]); - record_staged_result(dis, saved_statuses[id4]); - - // record a traversed edge - graph_visitor.visit_edge(level, node_id, saved_indices[id4], dis); - - // add a record of visited - knowhere::Neighbor nn(saved_indices[id4], dis, saved_statuses[id4]); - if (func_add_candidate(nn)) { -#if defined(USE_PREFETCH) - // TODO - // _mm_prefetch(get_linklist0(v), _MM_HINT_T0); -#endif - } - } + // Evaluate the remaining candidates with the same policy. + ndis += evaluation.compute( + qdis, saved_indices, saved_statuses, counter, level, + [&](size_t i, float distance) { + graph_visitor.visit_edge(level, node_id, saved_indices[i], distance); + func_add_candidate(knowhere::Neighbor( + saved_indices[i], distance, saved_statuses[i])); + }); - // update stats if (track_hnsw_stats) { stats.ndis = ndis; stats.nhops = 1; } - // done return stats; } @@ -357,6 +318,9 @@ struct v2_hnsw_searcher { } } + if (disqualified) { + retset.save_pending(*disqualified); + } // done return stats; } @@ -400,9 +364,7 @@ struct v2_hnsw_searcher { // grab some needed parameters const int efSearch = params ? params->efSearch : hnsw.efSearch; - staged = dynamic_cast(&qdis); - staged_k = static_cast(k); - staged_results = {}; + evaluation.begin(static_cast(k)); // yes. // greedy search on upper levels. @@ -439,7 +401,7 @@ struct v2_hnsw_searcher { } visited_nodes[nearest] = true; - record_staged_result(d_nearest, filter.is_member(nearest) + evaluation.record(d_nearest, filter.is_member(nearest) ? knowhere::Neighbor::kValid : knowhere::Neighbor::kInvalid); } diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/Neighbor.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/Neighbor.h index b2ec3cae5..b6be9f4f1 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/Neighbor.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/Neighbor.h @@ -187,6 +187,15 @@ class NeighborSetDoublePopList { (invalid_ns_->has_next() && invalid_ns_->cur().distance < valid_ns_->at_search_back_dist()); } + // A tightened valid-result threshold can hide already queued filtered + // bridges. Iterators must retain them even though this search is finished. + void + save_pending(IteratorMinHeap& pending) { + while (invalid_ns_->has_next()) { + pending.push(invalid_ns_->pop()); + } + } + inline const Neighbor& operator[](size_t i) { return (*valid_ns_)[i]; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h new file mode 100644 index 000000000..fcbd85cb1 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h @@ -0,0 +1,57 @@ +/* Copyright (c) Meta Platforms, Inc. and affiliates. + * Licensed under the MIT license in thirdparty/faiss/LICENSE. */ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace faiss::cppcontrib::knowhere::rabitq_search { + +// Both graph traversal implementations compare in smaller-is-better units. +// This is a probability window, not a deterministic lower bound. +inline bool should_refine(float estimate, float f_error, float g_error, + float threshold, bool similarity, float scale) { + const float error = f_error * g_error; + return similarity ? (estimate + error) * scale > -threshold + : std::max(0.0f, estimate - error) < threshold; +} + +// Used only by the RaBitQ wrapper's filtered/feder/RBQ1 kNN specialization. +struct DistanceEvaluation { + size_t k = 0; + std::priority_queue results; + + void begin(size_t count) { k = count; results = {}; } + float threshold() const { + return results.size() < k ? std::numeric_limits::infinity() : results.top(); + } + void record(float distance, int status) { + if (!k || status == Neighbor::kInvalid) return; + if (results.size() < k) results.push(distance); + else if (distance < results.top()) { + results.pop(); + results.push(distance); + } + } + template + size_t compute(DC& dc, const size_t* ids, const int* statuses, + size_t count, int level, Emit&& emit) { + if (k && level == 0) { + auto& staged = static_cast(dc); + const auto before = staged.refine_count; + for (size_t i = 0; i < count; ++i) { + const float distance = staged.evaluate(ids[i], threshold()); + record(distance, statuses[i]); + emit(i, distance); + } + return staged.refine_count - before; + } + FullDistanceEvaluation full; + return full.compute(dc, ids, statuses, count, level, std::forward(emit)); + } +}; +} // namespace faiss::cppcontrib::knowhere::rabitq_search diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h index 7df8bce14..537ef0d5d 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -92,10 +93,9 @@ SearchStats search_one(const faiss::cppcontrib::knowhere::HNSW& graph, const auto* factors = reinterpret_cast( code + (rq.d + 7) / 8); const float s = scale(ids[i]); - const float error = factors->f_error * rq.g_error; float distance = estimate; - const bool refine = similarity ? (estimate + error) * s > -threshold - : std::max(0.f, estimate - error) < threshold; + const bool refine = should_refine( + estimate, factors->f_error, rq.g_error, threshold, similarity, s); if (refine) { distance = rq.distance_to_code_full(code); ++stats.refine; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp index 0166cadb5..1be477c09 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp @@ -215,10 +215,7 @@ void read_xb_vector(VectorT& target, IOReader* f) { * Read **************************************************************/ -static void read_index_header( - Index* idx, - IOReader* f, - bool* is_cosine_out = nullptr) { +static void read_index_header(Index* idx, IOReader* f, bool* is_cosine_out = nullptr) { READ1(idx->d); READ1(idx->ntotal); @@ -364,27 +361,24 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { "read_InvertedLists:" " WARN! inverted lists not stored with IVF object\n"); return nullptr; - } else if (h == fourcc("iloa") && !(io_flags & IO_FLAG_MMAP)) { + } else if (h == fourcc ("iloa") && !(io_flags & IO_FLAG_MMAP)) { size_t nlist; size_t code_size; - std::vector list_length; + std::vector list_length; READ1(nlist); READ1(code_size); READVECTOR(list_length); - auto ails = - new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); + auto ails = new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); size_t n; READ1(n); #ifdef USE_GPU - ails->pin_readonly_ids = - std::make_shared(n * sizeof(idx_t)); - ails->pin_readonly_codes = std::make_shared( - n * code_size * sizeof(uint8_t)); - READANDCHECK((idx_t*)ails->pin_readonly_ids->data, n); - READANDCHECK((uint8_t*)ails->pin_readonly_codes->data, n * code_size); + ails->pin_readonly_ids = std::make_shared(n * sizeof(idx_t)); + ails->pin_readonly_codes = std::make_shared(n * code_size * sizeof(uint8_t)); + READANDCHECK((idx_t *) ails->pin_readonly_ids->data, n); + READANDCHECK((uint8_t *) ails->pin_readonly_codes->data, n * code_size); #else ails->readonly_ids.resize(n); - ails->readonly_codes.resize(n * code_size); + ails->readonly_codes.resize(n*code_size); READANDCHECK(ails->readonly_ids.data(), n); READANDCHECK(ails->readonly_codes.data(), n * code_size); #endif @@ -396,8 +390,7 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { READ1(segment_size); bool save_norm = io_flags & IO_FLAG_WITH_NORM; - auto lca = new ConcurrentArrayInvertedLists( - nlist, code_size, segment_size, save_norm); + auto lca = new ConcurrentArrayInvertedLists(nlist, code_size, segment_size, save_norm); std::vector sizes(nlist); read_ArrayInvertedLists_sizes(f, sizes); for (size_t i = 0; i < lca->nlist; i++) { @@ -408,15 +401,12 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { if (n > 0) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { - size_t seg_size = lca->get_segment_size(i, j); + size_t seg_size = lca->get_segment_size(i , j); size_t seg_off = lca->get_segment_offset(i, j); - READANDCHECK( - lca->codes[i][j].data_.data(), - seg_size * lca->code_size); + READANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); READANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (save_norm) { - READANDCHECK( - lca->code_norms[i][j].data_.data(), seg_size); + READANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); } } } @@ -594,7 +584,9 @@ static void read_ProductLocalSearchQuantizer( } } -static void read_ScalarQuantizer(::faiss::ScalarQuantizer* ivsc, IOReader* f) { +static void read_ScalarQuantizer( + ::faiss::ScalarQuantizer* ivsc, + IOReader* f) { READ1(ivsc->qtype); READ1(ivsc->rangestat); READ1(ivsc->rangestat_arg); @@ -741,11 +733,11 @@ static void read_direct_map(DirectMap* dm, IOReader* f) { map[it.first] = it.second; } } - // Path-D step 10.9: the former `if (dm->type == - // DirectMap::ConcurrentArray)` read branch is gone — see the symmetric - // comment in index_write.cpp. Old files (if any) with `type == 3` would - // fail to round-trip here since the enum value no longer exists; in - // practice CC indexes were never written through this path. + // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` + // read branch is gone — see the symmetric comment in index_write.cpp. + // Old files (if any) with `type == 3` would fail to round-trip here + // since the enum value no longer exists; in practice CC indexes + // were never written through this path. } static void read_ivf_header( @@ -833,8 +825,7 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(wire_l2_norms); // reconstruct inverse norms from wire L2 norms - idxf->inverse_norms_storage = - L2NormsStorage::from_l2_norms(wire_l2_norms); + idxf->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_l2_norms); FAISS_THROW_IF_NOT( idxf->codes.size() == idxf->ntotal * idxf->code_size); @@ -871,8 +862,7 @@ Index* read_index(IOReader* f, int io_flags) { idxfc->code_size = idxf->code_size; idxfc->codes = std::move(idxf->codes); // reconstruct inverse norms from wire L2 norms - idxfc->inverse_norms_storage = - L2NormsStorage::from_l2_norms(wire_code_norms); + idxfc->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_code_norms); delete idxf; idxf = idxfc; } @@ -893,7 +883,7 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(idxp->inverse_norms_storage.inverse_l2_norms); if (!(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { - idxp->pq.compute_sdc_table(); + idxp->pq.compute_sdc_table (); } idx = idxp; @@ -919,7 +909,7 @@ Index* read_index(IOReader* f, int io_flags) { // the following "if" block is Knowhere-specific if (h == fourcc("IxPq")) { - idxp->pq.compute_sdc_table(); + idxp->pq.compute_sdc_table (); } idx = idxp; @@ -1114,7 +1104,8 @@ Index* read_index(IOReader* f, int io_flags) { // either enum name, and route legacy data to // IndexBinaryScalarQuantizer. const int legacy_qt_1bit_direct_marker = 9; - if (static_cast(idxs->sq.qtype) == legacy_qt_1bit_direct_marker) { + if (static_cast(idxs->sq.qtype) == + legacy_qt_1bit_direct_marker) { IndexBinaryScalarQuantizer* bsq = new IndexBinaryScalarQuantizer( static_cast(idxs->d), idxs->metric_type); bsq->ntotal = idxs->ntotal; @@ -1222,8 +1213,7 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxrf->k_factor); if (dynamic_cast<::faiss::IndexFlat*>(idxrf->refine_index)) { // then make a RefineFlat with it. Refine index may be a baseline - // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware - // subclass. + // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware subclass. IndexRefine* idxrf_old = idxrf; idxrf = new IndexRefineFlat(); *idxrf = *idxrf_old; @@ -1423,8 +1413,7 @@ Index* read_index(IOReader* f, int io_flags) { // field); Iwrr is baseline multi-bit and does serialize nb_bits. auto ivrq = new IndexIVFRaBitQ(); read_ivf_header(ivrq, f); - read_RaBitQuantizer( - &ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); + read_RaBitQuantizer(&ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); READ1(ivrq->code_size); READ1(ivrq->by_residual); READ1(ivrq->qb); diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp index 3e84d5da2..6415be569 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp @@ -304,9 +304,8 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { } } } - } else if ( - const auto& lca = - dynamic_cast(ils)) { + } else if (const auto & lca = + dynamic_cast(ils)) { uint32_t h = fourcc("ilca"); WRITE1(h); WRITE1(lca->nlist); @@ -350,20 +349,16 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { size_t seg_size = lca->get_segment_size(i, j); - WRITEANDCHECK( - lca->codes[i][j].data_.data(), - seg_size * lca->code_size); + WRITEANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); WRITEANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (lca->save_norm) { - WRITEANDCHECK( - lca->code_norms[i][j].data_.data(), seg_size); + WRITEANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); } } } } - } else if ( - const auto& oa = - dynamic_cast(ils)) { + } else if (const auto & oa = + dynamic_cast(ils)) { uint32_t h = fourcc("iloa"); WRITE1(h); WRITE1(oa->nlist); @@ -373,16 +368,16 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t n = oa->pin_readonly_ids->size() / sizeof(InvertedLists::idx_t); WRITE1(n); WRITEANDCHECK((InvertedLists::idx_t*)oa->pin_readonly_ids->data, n); - WRITEANDCHECK( - (uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); + WRITEANDCHECK((uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); #else size_t n = oa->readonly_ids.size(); WRITE1(n); WRITEANDCHECK(oa->readonly_ids.data(), n); WRITEANDCHECK(oa->readonly_codes.data(), n * oa->code_size); #endif - } else if (const auto& od = dynamic_cast(ils)) { - uint32_t h = fourcc("ilod"); + } else if (const auto & od = + dynamic_cast(ils)) { + uint32_t h = fourcc ("ilod"); WRITE1(h); WRITE1(ils->nlist); WRITE1(ils->code_size); @@ -391,7 +386,7 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { { std::vector v( - od->slots.begin(), od->slots.end()); + od->slots.begin(), od->slots.end()); WRITEVECTOR(v); } { @@ -550,12 +545,11 @@ static void write_direct_map(const DirectMap* dm, IOWriter* f) { std::copy(map.begin(), map.end(), v.begin()); WRITEVECTOR(v); } - // Path-D step 10.9: the former `if (dm->type == - // DirectMap::ConcurrentArray)` write branch is gone — fork DirectMap no - // longer supports that variant. CC indexes now carry their own - // `cc_direct_map` member (ConcurrentDirectMap) which is not serialized - // through this path (CC indexes have no serialize stage; see ivf.cc:619 - // comment). + // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` + // write branch is gone — fork DirectMap no longer supports that + // variant. CC indexes now carry their own `cc_direct_map` member + // (ConcurrentDirectMap) which is not serialized through this path + // (CC indexes have no serialize stage; see ivf.cc:619 comment). } static void write_ivf_header(const IndexIVF* ivf, IOWriter* f) { @@ -573,9 +567,7 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { // eg. for a storage component of HNSW that is set to nullptr uint32_t h = fourcc("null"); WRITE1(h); - } else if ( - const IndexFlatCosine* idxf = - dynamic_cast(idx)) { + } else if (const IndexFlatCosine* idxf = dynamic_cast(idx)) { uint32_t h = fourcc("IxF9"); WRITE1(h); write_index_header(idx, f); @@ -595,9 +587,7 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(h); write_index_header(idx, f); WRITEXBVECTOR(idxf->codes); - } else if ( - const IndexPQCosine* idxp = - dynamic_cast(idx)) { + } else if (const IndexPQCosine* idxp = dynamic_cast(idx)) { uint32_t h = fourcc("IxP7"); WRITE1(h); write_index_header(idx, f); @@ -639,8 +629,7 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITEVECTOR(idxr_2->codes); } else if ( const IndexProductResidualQuantizerCosine* idxpr = - dynamic_cast( - idx)) { + dynamic_cast(idx)) { uint32_t h = fourcc("IxP5"); WRITE1(h); write_index_header(idx, f); From 533a5836560f6da43457e34389efcaea4e1b997a Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 07:49:45 +0000 Subject: [PATCH 3/7] perf: reuse Knowhere HNSW traversal with batched RaBitQ estimates Remove the separate native RaBitQ traversal. Keep search state, queues and stopping rules in the existing Knowhere searcher, with RaBitQ-specific staged distance evaluation and qb4 SIMD/prefetch batching. Preserve candidate-order refinement and request-local query parameters. Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc | 49 ------ src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h | 3 - tests/ut/test_hnsw_rabitq.cc | 111 ++++++++++-- tests/ut/test_hnsw_rabitq_acceptance.cc | 3 - .../cppcontrib/knowhere/IndexHNSWRaBitQ.cpp | 63 +------ .../knowhere/impl/RaBitQDistanceEvaluation.h | 58 ++++--- .../cppcontrib/knowhere/impl/RaBitQSearch.h | 160 ------------------ .../impl/RaBitQStagedDistanceComputer.h | 85 ++++++++++ .../faiss/faiss/impl/RaBitQuantizer.cpp | 29 ++++ thirdparty/faiss/faiss/impl/RaBitQuantizer.h | 8 + thirdparty/faiss/faiss/utils/rabitq_simd.h | 14 ++ .../faiss/utils/simd_impl/rabitq_avx512.cpp | 49 ++++++ 12 files changed, 324 insertions(+), 308 deletions(-) delete mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQStagedDistanceComputer.h diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc index 074462141..6daba7d42 100644 --- a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc @@ -2,63 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 #include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" -#include #include -#include #include "index/hnsw/impl/HnswSearchDispatch.h" #include "index/hnsw/impl/RaBitQSearchParameters.h" -#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) -#include "knowhere/prometheus_client.h" -#endif namespace knowhere { -using idx_t = faiss::idx_t; namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; -void -IndexHNSWRaBitQWrapper::search(faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, faiss::idx_t* labels, - const faiss::SearchParameters* params_in) const { - FAISS_THROW_IF_NOT(k > 0); - const auto* index_hnsw = dynamic_cast(index); - FAISS_THROW_IF_NOT(index_hnsw && index_hnsw->storage); - const auto* params = dynamic_cast(params_in); - FAISS_THROW_IF_NOT_MSG(!params_in || params, "params type invalid"); - if (index_hnsw->hnsw.entry_point == -1) { - IndexHNSWWrapper::search(n, x, k, distances, labels, params_in); - return; - } - const auto& hnsw = index_hnsw->hnsw; - const auto* rbq_params = dynamic_cast(params); - // Use the optimized multi-bit path only when its selector/visitor contract - // is satisfied. RBQ1, filtering and feder use the compatible searcher below. - const auto* bitset_sel = params ? dynamic_cast(params->sel) : nullptr; - const bool unfiltered = !params || !params->sel || (bitset_sel && bitset_sel->bitset_view.empty()); - if (index_hnsw->rabitq_index()->rabitq.nb_bits > 1 && unfiltered && (!params || !params->feder)) { - rabitq_search::search( - *index_hnsw, n, x, k, distances, labels, params ? params->efSearch : hnsw.efSearch, - params ? params->check_relative_distance : hnsw.check_relative_distance, - rbq_params ? &rbq_params->storage_params : nullptr, [&](const rabitq_search::SearchStats& counts) { - const size_t hops = counts.expanded + counts.upper_expanded; -#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) - knowhere::knowhere_hnsw_search_hops.Observe(hops); -#endif - if (params && params->hnsw_stats) { - params->hnsw_stats->combine({.n1 = 1, - .n2 = size_t(counts.exhausted), - .ndis = counts.estimate + counts.refine + counts.upper_full, - .nhops = hops}); - } - }); - if (faiss::cppcontrib::knowhere::is_similarity_metric(index->metric_type)) { - for (idx_t i = 0; i < k * n; ++i) distances[i] = -distances[i]; - } - return; - } - - IndexHNSWWrapper::search(n, x, k, distances, labels, params_in); -} - std::unique_ptr IndexHNSWRaBitQWrapper::graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, const SearchParametersHNSWWrapper* params) const { diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h index 9fc6eb39f..80b85f00b 100644 --- a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h @@ -7,9 +7,6 @@ namespace knowhere { struct IndexHNSWRaBitQWrapper : IndexHNSWWrapper { using IndexHNSWWrapper::IndexHNSWWrapper; - void - search(faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, faiss::idx_t* labels, - const faiss::SearchParameters* params) const override; protected: std::unique_ptr diff --git a/tests/ut/test_hnsw_rabitq.cc b/tests/ut/test_hnsw_rabitq.cc index e5c0c9972..e1f3662ba 100644 --- a/tests/ut/test_hnsw_rabitq.cc +++ b/tests/ut/test_hnsw_rabitq.cc @@ -3,7 +3,7 @@ #include #include -#include +#include #include #include #include @@ -14,6 +14,8 @@ #include #include #include +#include +#include #include #include "catch2/catch_approx.hpp" @@ -58,6 +60,99 @@ struct OriginalFullEvaluation { }; } // namespace +TEST_CASE("RaBitQ threshold heap matches priority queue after every update", "[hnsw_rabitq_core]") { + using faiss::cppcontrib::knowhere::Neighbor; + rabitq_search::DistanceEvaluation evaluation; + std::mt19937 rng(12345); + for (size_t k : {1, 2, 3, 10, 100, 511}) { + for (int mode = 0; mode < 3; ++mode) { + evaluation.begin(k); + std::priority_queue reference; + for (int i = 0; i < 10000; ++i) { + float distance = mode == 0 ? static_cast(static_cast(rng() % 2001) - 1000) + : static_cast(mode == 1 ? i : -i); + if (i % 89 == 0) + distance = std::numeric_limits::infinity(); + if (i % 97 == 0) + distance = -std::numeric_limits::infinity(); + const int status = i % 7 == 0 ? Neighbor::kInvalid : Neighbor::kValid; + evaluation.record(distance, status); + if (status != Neighbor::kInvalid) { + if (reference.size() < k) + reference.push(distance); + else if (distance < reference.top()) { + reference.pop(); + reference.push(distance); + } + } + const float expected = reference.size() < k ? std::numeric_limits::infinity() : reference.top(); + REQUIRE(evaluation.threshold() == expected); + REQUIRE(std::is_heap(evaluation.results.begin(), evaluation.results.end())); + } + } + } + evaluation.begin(0); + evaluation.record(1.0f, Neighbor::kValid); + REQUIRE(evaluation.results.empty()); +} + +TEST_CASE("RaBitQ batch4 integer kernel matches independent scalar candidates", "[hnsw_rabitq_core]") { +#if defined(__GNUC__) && defined(__x86_64__) + if (!__builtin_cpu_supports("avx512f") || !__builtin_cpu_supports("avx512bw") || + !__builtin_cpu_supports("avx512dq") || !__builtin_cpu_supports("avx512vl")) + return; + for (size_t bytes : {1, 7, 8, 15, 16, 31, 32, 63, 64, 65, 96, 120, 128, 192, 193}) { + std::vector query(bytes * 4 + 1); + std::vector data[4]; + const uint8_t* codes[4]; + for (size_t i = 0; i < query.size(); ++i) query[i] = (i * 17 + 47) % 256; + for (int lane = 0; lane < 4; ++lane) { + data[lane].resize(bytes + 1); + for (size_t i = 0; i < data[lane].size(); ++i) data[lane][i] = (i * 31 + lane * 73) % 256; + codes[lane] = data[lane].data() + 1; + } + faiss::rabitq::BitwiseAndDotProductResult actual[4]; + faiss::rabitq::bitwise_q4_batch_4(query.data() + 1, codes, bytes, actual); + for (int lane = 0; lane < 4; ++lane) { + const auto expected = faiss::rabitq::bitwise_and_dot_product_with_popcount( + query.data() + 1, codes[lane], bytes, 4); + REQUIRE(actual[lane].dot_product == expected.dot_product); + REQUIRE(actual[lane].popcount == expected.popcount); + } + } +#endif +} + +TEST_CASE("RaBitQ batch4 estimates exactly match scalar calls across query modes", "[hnsw_rabitq_core]") { + for (int dim : {65, 960, 1024, 1536}) { + auto base = GenDataSet(32, dim, 9201); + auto query = GenDataSet(2, dim, 9202); + for (auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + for (int bits : {1, 4, 8, 9}) { + faiss::IndexRaBitQ index(dim, metric, bits); + index.train(32, static_cast(base->GetTensor())); + index.add(32, static_cast(base->GetTensor())); + for (int qb : {0, 4, 8}) { + for (bool centered : {false, true}) { + std::unique_ptr owner( + index.get_quantized_distance_computer(qb, centered)); + auto* dc = dynamic_cast(owner.get()); + REQUIRE(dc != nullptr); + for (int q = 0; q < 2; ++q) { + dc->set_query(static_cast(query->GetTensor()) + q * dim); + const uint8_t* codes[4]; + float actual[4]; + for (int i = 0; i < 4; ++i) codes[i] = dc->codes + (i * 7 + q) * dc->code_size; + dc->distance_to_code_1bit_batch_4(codes, actual); + for (int i = 0; i < 4; ++i) REQUIRE(actual[i] == dc->distance_to_code_1bit(codes[i])); + } + } + } + } + } + } +} + TEST_CASE("RaBitQ qb4 SIMD matches scalar including masked tails", "[hnsw_rabitq_core]") { #if defined(__GNUC__) && defined(__x86_64__) if (!__builtin_cpu_supports("avx512f") || !__builtin_cpu_supports("avx512bw") || @@ -146,19 +241,10 @@ TEST_CASE("RaBitQ traversal retains all results when k covers the graph", "[hnsw std::unique_ptr full(storage.get_distance_computer()); std::vector distances(4 * n); std::vector labels(4 * n); - rabitq_search::search(graph, 4, static_cast(queries->GetTensor()), n, distances.data(), - labels.data(), n, true); - std::vector api_distances(4 * n); - std::vector api_labels(4 * n); knowhere::IndexHNSWRaBitQWrapper api(&graph); knowhere::SearchParametersHNSWWrapper params; params.efSearch = n; - api.search(4, static_cast(queries->GetTensor()), n, api_distances.data(), api_labels.data(), - ¶ms); - for (int i = 0; i < 4 * n; ++i) { - REQUIRE(api_labels[i] == labels[i]); - REQUIRE(api_distances[i] == Catch::Approx((similarity ? -1.f : 1.f) * distances[i]).margin(1e-5)); - } + api.search(4, static_cast(queries->GetTensor()), n, distances.data(), labels.data(), ¶ms); for (int q = 0; q < 4; ++q) { full->set_query(static_cast(queries->GetTensor()) + q * dim); std::vector> expected; @@ -167,7 +253,8 @@ TEST_CASE("RaBitQ traversal retains all results when k covers the graph", "[hnsw for (int i = 0; i < n; ++i) { CAPTURE(q, i, distances[q * n + i], expected[i].first); REQUIRE(labels[q * n + i] == expected[i].second); - REQUIRE(distances[q * n + i] == Catch::Approx(expected[i].first).margin(1e-5)); + REQUIRE(distances[q * n + i] == + Catch::Approx((similarity ? -1.f : 1.f) * expected[i].first).margin(1e-5)); } } } diff --git a/tests/ut/test_hnsw_rabitq_acceptance.cc b/tests/ut/test_hnsw_rabitq_acceptance.cc index 8db99e690..f7b93e5ce 100644 --- a/tests/ut/test_hnsw_rabitq_acceptance.cc +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -7,7 +7,6 @@ #include #include #include -#include #include #include #include @@ -42,8 +41,6 @@ #include "knowhere/utils.h" #include "utils.h" -namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; - TEST_CASE("RBQ bounded add preserves input slices and rejects invalid sizes", "[hnsw_rabitq_acceptance][rbq_build]") { using faiss::cppcontrib::knowhere::rabitq_build::add_in_blocks; struct RecordingIndex : faiss::IndexFlatL2 { diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp index 727bdde65..5f2a98875 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -6,7 +6,7 @@ */ #include -#include +#include #include #include @@ -19,69 +19,10 @@ namespace faiss::cppcontrib::knowhere { -namespace { -struct RaBitQStagedDistanceComputer final : StagedDistanceComputer { - const faiss::VectorTransform& rotation; - std::unique_ptr dc; - const float* norms; - bool similarity; - float query_inverse_norm = 1; - std::vector rotated; - - explicit RaBitQStagedDistanceComputer(const IndexHNSWRaBitQ& index, - const faiss::RaBitQSearchParameters* params) - : rotation(*index.pretransform_index()->chain[0]), - norms(nullptr), similarity(index.metric_type == METRIC_INNER_PRODUCT), - rotated(index.d) { - auto* raw = params ? index.rabitq_index()->get_quantized_distance_computer(params->qb, params->centered) - : index.rabitq_index()->get_FlatCodesDistanceComputer(); - auto* typed = dynamic_cast(raw); - if (!typed) { delete raw; FAISS_THROW_MSG("RaBitQ distance computer required"); } - dc.reset(typed); - if (auto* cosine = dynamic_cast(&index)) { - norms = cosine->get_inverse_l2_norms(); - } - } - void set_query(const float* q) override { - estimate_count = refine_count = 0; - rotation.apply_noalloc(1, q, rotated.data()); - dc->set_query(rotated.data()); - const float norm2 = norms ? faiss::fvec_norm_L2sqr(q, rotation.d_in) : 1; - query_inverse_norm = norm2 > 0 ? 1 / std::sqrt(norm2) : 1; - } - float scale(idx_t id) const { return norms ? norms[id] * query_inverse_norm : 1; } - float operator()(idx_t id) override { - float d = (*dc)(id) * scale(id); - return similarity ? -d : d; - } - float symmetric_dis(idx_t, idx_t) override { - FAISS_THROW_MSG("staged storage is search-only; construct graph with FP32"); - } - float evaluate(idx_t id, float threshold) override { - if (dc->nb_bits == 1) return (*this)(id); - const auto* code = dc->codes + id * dc->code_size; - const float estimate = dc->distance_to_code_1bit(code); - ++estimate_count; - const auto* factors = reinterpret_cast( - code + (dc->d + 7) / 8); - const float s = scale(id); - float d = estimate; - // Compare in output-distance units: positive cosine scale preserves - // ordering, avoiding a division for every visited candidate. - const bool refine = rabitq_search::should_refine( - estimate, factors->f_error, dc->g_error, threshold, similarity, s); - if (refine) { - d = dc->distance_to_code_full(code); - ++refine_count; - } - return (similarity ? -d : d) * s; - } -}; -} // namespace faiss::DistanceComputer* IndexHNSWRaBitQ::get_staged_distance_computer( const faiss::RaBitQSearchParameters* params) const { - return new RaBitQStagedDistanceComputer(*this, params); + return new rabitq_search::RaBitQStagedDistanceComputer(*this, params); } IndexPreTransformRaBitQCosine::IndexPreTransformRaBitQCosine() = default; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h index fcbd85cb1..9dd997bbd 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h @@ -3,46 +3,64 @@ #pragma once #include -#include +#include #include #include -#include +#include #include namespace faiss::cppcontrib::knowhere::rabitq_search { -// Both graph traversal implementations compare in smaller-is-better units. -// This is a probability window, not a deterministic lower bound. -inline bool should_refine(float estimate, float f_error, float g_error, - float threshold, bool similarity, float scale) { - const float error = f_error * g_error; - return similarity ? (estimate + error) * scale > -threshold - : std::max(0.0f, estimate - error) < threshold; -} - -// Used only by the RaBitQ wrapper's filtered/feder/RBQ1 kNN specialization. +// Used only by the RaBitQ wrapper's Knowhere-traversal specialization. struct DistanceEvaluation { size_t k = 0; - std::priority_queue results; + std::vector results; - void begin(size_t count) { k = count; results = {}; } + void begin(size_t count) { k = count; results.clear(); results.reserve(k); } float threshold() const { - return results.size() < k ? std::numeric_limits::infinity() : results.top(); + return results.size() < k ? std::numeric_limits::infinity() : results.front(); } void record(float distance, int status) { if (!k || status == Neighbor::kInvalid) return; - if (results.size() < k) results.push(distance); - else if (distance < results.top()) { - results.pop(); - results.push(distance); + if (results.size() < k) { + results.push_back(distance); + std::push_heap(results.begin(), results.end()); + } else if (distance < results.front()) { + // Replace the worst retained distance with one sift-down, instead + // of repairing the threshold heap separately for pop and push. + size_t parent = 0; + for (size_t child = 1; child < k; child = 2 * parent + 1) { + if (child + 1 < k && results[child] < results[child + 1]) ++child; + if (!(distance < results[child])) break; + results[parent] = results[child]; + parent = child; + } + results[parent] = distance; } } template size_t compute(DC& dc, const size_t* ids, const int* statuses, size_t count, int level, Emit&& emit) { if (k && level == 0) { - auto& staged = static_cast(dc); + auto& staged = static_cast(dc); const auto before = staged.refine_count; + if (count == 4 && staged.dc->nb_bits > 1) { + const uint8_t* codes[4]; + for (size_t i = 0; i < 4; ++i) { + codes[i] = staged.dc->codes + ids[i] * staged.dc->code_size; + } + float estimates[4]; + staged.dc->distance_to_code_1bit_batch_4(codes, estimates); + staged.estimate_count += 4; + for (size_t i = 0; i < 4; ++i) { + // Only estimates are batched. Refine and update the threshold + // in exactly the original candidate order. + const float distance = staged.evaluate_estimate(ids[i], codes[i], estimates[i], threshold()); + record(distance, statuses[i]); + emit(i, distance); + } + return staged.refine_count - before; + } for (size_t i = 0; i < count; ++i) { const float distance = staged.evaluate(ids[i], threshold()); record(distance, statuses[i]); diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h deleted file mode 100644 index 537ef0d5d..000000000 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQSearch.h +++ /dev/null @@ -1,160 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * Licensed under the MIT license in thirdparty/faiss/LICENSE. - * - * Port of Faiss #5526 (d8a85956) bounded traversal. - * Graph adjacency is read from Knowhere without copying or changing the graph. - * Extended to L2/IP/COSINE, deliberately limited to unfiltered KNN. - * Callers must dispatch filtered/visitor requests to a compatible searcher. - */ -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace faiss::cppcontrib::knowhere::rabitq_search { -struct SearchStats { - size_t estimate = 0, refine = 0, expanded = 0, upper_full = 0, upper_expanded = 0; - bool exhausted = false; -}; - -template -SearchStats search_one(const faiss::cppcontrib::knowhere::HNSW& graph, - faiss::RaBitQDistanceComputer& rq, VT& vt, - faiss::ResultHandler& res, int ef, bool relative, - bool similarity = false, const float* norms = nullptr, - float query_inverse_norm = 1) { - SearchStats stats; - using HC = faiss::CMax; - int32_t nearest = graph.entry_point; - auto scale = [&](int32_t id) { return norms ? norms[id] * query_inverse_norm : 1.f; }; - auto convert = [&](int32_t id, float raw) { return (similarity ? -raw : raw) * scale(id); }; - float nearest_distance = convert(nearest, rq(nearest)); - ++stats.upper_full; - for (int level = graph.max_level; level >= 1; --level) { - for (;;) { - ++stats.upper_expanded; - const int32_t previous = nearest; - size_t begin, end; - graph.neighbor_range(nearest, level, &begin, &end); - int32_t ids[4]; - int count = 0; - auto update = [&](int32_t id, float d) { - d = convert(id, d); - if (d < nearest_distance) { nearest = id; nearest_distance = d; } - }; - for (size_t j = begin; j < end && graph.neighbors[j] >= 0; ++j) { - ids[count++] = graph.neighbors[j]; - ++stats.upper_full; - if (count == 4) { - float d[4]; - rq.distances_batch_4(ids[0], ids[1], ids[2], ids[3], d[0], d[1], d[2], d[3]); - for (int i = 0; i < 4; ++i) update(ids[i], d[i]); - count = 0; - } - } - for (int i = 0; i < count; ++i) update(ids[i], rq(ids[i])); - if (previous == nearest) break; - } - } - - faiss::MinimaxHeapT candidates(ef); - candidates.push(nearest, nearest_distance); - vt.reserve(ef); - if (nearest_distance < res.threshold) res.add_result(nearest_distance, nearest); - vt.set(nearest); - while (candidates.size() > 0) { - float d0; - const int32_t node = candidates.pop_min(&d0); - if (relative && candidates.count_below(d0) >= ef) break; - size_t begin, end; - graph.neighbor_range(node, 0, &begin, &end); - size_t limit = begin; - for (size_t j = begin; j < end; ++j) { - if (graph.neighbors[j] < 0) break; - vt.prefetch(graph.neighbors[j]); - ++limit; - } - int32_t ids[4]; - int count = 0; - float threshold = res.threshold; - auto evaluate = [&] { - for (int i = 0; i < count; ++i) { - const auto* code = rq.codes + static_cast(ids[i]) * rq.code_size; - const float estimate = rq.distance_to_code_1bit(code); - ++stats.estimate; - const auto* factors = reinterpret_cast( - code + (rq.d + 7) / 8); - const float s = scale(ids[i]); - float distance = estimate; - const bool refine = should_refine( - estimate, factors->f_error, rq.g_error, threshold, similarity, s); - if (refine) { - distance = rq.distance_to_code_full(code); - ++stats.refine; - } - distance = (similarity ? -distance : distance) * s; - if (distance < threshold && res.add_result(distance, ids[i])) threshold = res.threshold; - candidates.push(ids[i], distance); - } - }; - for (size_t j = begin; j < limit; ++j) { - ids[count] = graph.neighbors[j]; - count += vt.set(ids[count]) ? 1 : 0; - if (count == 4) { evaluate(); count = 0; } - } - if (count) evaluate(); - ++stats.expanded; - if (!relative && stats.expanded > static_cast(ef)) break; - } - stats.exhausted = candidates.size() == 0; - return stats; -} - -inline void search(const faiss::cppcontrib::knowhere::IndexHNSWRaBitQ& index, - faiss::idx_t n, const float* x, faiss::idx_t k, float* distances, - faiss::idx_t* labels, int ef, bool relative, - const faiss::RaBitQSearchParameters* params = nullptr, - const std::function& on_query = {}) { - FAISS_THROW_IF_NOT(index.metric_type == faiss::METRIC_L2 || - index.metric_type == faiss::METRIC_INNER_PRODUCT); - const bool similarity = index.metric_type == faiss::METRIC_INNER_PRODUCT; - const auto* cosine = dynamic_cast(&index); - const float* norms = cosine ? cosine->get_inverse_l2_norms() : nullptr; - FAISS_THROW_IF_NOT(index.rabitq_index()->rabitq.nb_bits > 1); - auto raw = std::unique_ptr( - params ? index.rabitq_index()->get_quantized_distance_computer(params->qb, params->centered) - : index.rabitq_index()->get_FlatCodesDistanceComputer()); - auto& rq = dynamic_cast(*raw); - // Reuse Faiss's visited table across queries without clearing the full - // table on each search; advance its generation after processing the query. - auto& vt = faiss::VisitedTable::get_reusable(index.ntotal); - faiss::HeapBlockResultHandler> block(n, distances, labels, k); - decltype(block)::SingleResultHandler result(block); - std::vector rotated(index.d); - for (faiss::idx_t i = 0; i < n; ++i) { - result.begin(i); - index.pretransform_index()->chain[0]->apply_noalloc(1, x + i * index.d, rotated.data()); - rq.set_query(rotated.data()); - const float norm2 = norms ? faiss::fvec_norm_L2sqr(x + i * index.d, index.d) : 1.f; - const float query_inverse_norm = norm2 > 0 ? 1.f / std::sqrt(norm2) : 1.f; - SearchStats stats; - if (auto* vector = dynamic_cast(&vt)) - stats = search_one(index.hnsw, rq, *vector, result, std::max(ef, k), relative, - similarity, norms, query_inverse_norm); - else - stats = search_one(index.hnsw, rq, dynamic_cast(vt), result, - std::max(ef, k), relative, similarity, norms, query_inverse_norm); - result.end(); - vt.advance(); - if (on_query) on_query(stats); - } -} -} // namespace faiss::cppcontrib::knowhere::rabitq_search diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQStagedDistanceComputer.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQStagedDistanceComputer.h new file mode 100644 index 000000000..f27df0add --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQStagedDistanceComputer.h @@ -0,0 +1,85 @@ +/* Copyright (c) Meta Platforms, Inc. and affiliates. + * Licensed under the MIT license in thirdparty/faiss/LICENSE. */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace faiss::cppcontrib::knowhere::rabitq_search { +// Both traversal implementations compare in smaller-is-better units. +// This probability window is not a deterministic lower bound. +inline bool should_refine(float estimate, float f_error, float g_error, + float threshold, bool similarity, float scale) { + const float error = f_error * g_error; + return similarity ? (estimate + error) * scale > -threshold + : std::max(0.0f, estimate - error) < threshold; +} + +// Concrete RBQ adapter exposed only to the RBQ evaluator so its per-candidate +// threshold logic can be inlined, without specializing the common searcher. +struct RaBitQStagedDistanceComputer final : StagedDistanceComputer { + const faiss::VectorTransform& rotation; + std::unique_ptr dc; + const float* norms; + bool similarity; + float query_inverse_norm = 1; + std::vector rotated; + + explicit RaBitQStagedDistanceComputer(const IndexHNSWRaBitQ& index, + const faiss::RaBitQSearchParameters* params) + : rotation(*index.pretransform_index()->chain[0]), + norms(nullptr), similarity(index.metric_type == METRIC_INNER_PRODUCT), + rotated(index.d) { + auto* raw = params ? index.rabitq_index()->get_quantized_distance_computer(params->qb, params->centered) + : index.rabitq_index()->get_FlatCodesDistanceComputer(); + auto* typed = dynamic_cast(raw); + if (!typed) { delete raw; FAISS_THROW_MSG("RaBitQ distance computer required"); } + dc.reset(typed); + if (auto* cosine = dynamic_cast(&index)) { + norms = cosine->get_inverse_l2_norms(); + } + } + void set_query(const float* q) override { + estimate_count = refine_count = 0; + rotation.apply_noalloc(1, q, rotated.data()); + dc->set_query(rotated.data()); + const float norm2 = norms ? faiss::fvec_norm_L2sqr(q, rotation.d_in) : 1; + query_inverse_norm = norm2 > 0 ? 1 / std::sqrt(norm2) : 1; + } + float scale(idx_t id) const { return norms ? norms[id] * query_inverse_norm : 1; } + float operator()(idx_t id) override { + float d = (*dc)(id) * scale(id); + return similarity ? -d : d; + } + float symmetric_dis(idx_t, idx_t) override { + FAISS_THROW_MSG("staged storage is search-only; construct graph with FP32"); + } + float evaluate(idx_t id, float threshold) override { + if (dc->nb_bits == 1) return (*this)(id); + const auto* code = dc->codes + id * dc->code_size; + const float estimate = dc->distance_to_code_1bit(code); + ++estimate_count; + return evaluate_estimate(id, code, estimate, threshold); + } + float evaluate_estimate(idx_t id, const uint8_t* code, float estimate, float threshold) { + const auto* factors = reinterpret_cast( + code + (dc->d + 7) / 8); + const float s = scale(id); + float d = estimate; + // Compare in output-distance units: positive cosine scale preserves + // ordering, avoiding a division for every visited candidate. + const bool refine = rabitq_search::should_refine( + estimate, factors->f_error, dc->g_error, threshold, similarity, s); + if (refine) { + d = dc->distance_to_code_full(code); + ++refine_count; + } + return (similarity ? -d : d) * s; + } +}; +} // namespace faiss::cppcontrib::knowhere::rabitq_search diff --git a/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp b/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp index d9103484a..c7bcc8fc8 100644 --- a/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp +++ b/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp @@ -6,6 +6,7 @@ */ #include +#include #include #include @@ -510,6 +511,34 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer { return distance_to_code_1bit_impl(code, base_fac, size); } + void distance_to_code_1bit_batch_4( + const uint8_t* const* codes_in, float* distances) final { + if (qb != 4 || centered) { + RaBitQDistanceComputer::distance_to_code_1bit_batch_4(codes_in, distances); + return; + } + const size_t size = (d + 7) / 8; + const size_t prefix = size + (nb_bits == 1 ? sizeof(SignBitFactors) : sizeof(SignBitFactorsWithError)); + for (int i = 0; i < 4; ++i) { + for (size_t offset = 0; offset < prefix; offset += 64) { + prefetch_L1(codes_in[i] + offset); + } + } + rabitq::BitwiseAndDotProductResult results[4]; + rabitq::bitwise_q4_batch_4(rearranged_rotated_qq.data(), codes_in, size, results); + for (int i = 0; i < 4; ++i) { + const auto* factors = reinterpret_cast(codes_in[i] + size); + float final_dot = 0; + final_dot += query_fac.c1 * results[i].dot_product; + final_dot += query_fac.c2 * results[i].popcount; + final_dot -= query_fac.c34; + const float pre_dist = factors->or_minus_c_l2sqr + query_fac.qr_to_c_L2sqr - + 2 * factors->dp_multiplier * final_dot; + distances[i] = metric_type == METRIC_L2 ? std::max(0.0f, pre_dist) + : -0.5f * (pre_dist - query_fac.qr_norm_L2sqr); + } + } + // Compute full distance using 1-bit + ex-bits (accurate) float distance_to_code_full(const uint8_t* code) final { FAISS_ASSERT(code != nullptr); diff --git a/thirdparty/faiss/faiss/impl/RaBitQuantizer.h b/thirdparty/faiss/faiss/impl/RaBitQuantizer.h index 7766df35c..c8e31bf73 100644 --- a/thirdparty/faiss/faiss/impl/RaBitQuantizer.h +++ b/thirdparty/faiss/faiss/impl/RaBitQuantizer.h @@ -130,6 +130,14 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer { // Compute 1-bit distance estimate (fast) virtual float distance_to_code_1bit(const uint8_t* code) = 0; + // Independent estimates only; callers retain sequential refinement decisions. + virtual void distance_to_code_1bit_batch_4( + const uint8_t* const* codes_in, float* distances) { + for (size_t i = 0; i < 4; ++i) { + distances[i] = distance_to_code_1bit(codes_in[i]); + } + } + // Compute full multi-bit distance (accurate) virtual float distance_to_code_full(const uint8_t* code) = 0; diff --git a/thirdparty/faiss/faiss/utils/rabitq_simd.h b/thirdparty/faiss/faiss/utils/rabitq_simd.h index 9f49700f4..74946bb30 100644 --- a/thirdparty/faiss/faiss/utils/rabitq_simd.h +++ b/thirdparty/faiss/faiss/utils/rabitq_simd.h @@ -53,6 +53,20 @@ BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount( size_t size, size_t qb); +template +inline void bitwise_q4_batch_4( + const uint8_t* query, const uint8_t* const* data, size_t size, + BitwiseAndDotProductResult* results) { + for (size_t i = 0; i < 4; ++i) { + results[i] = bitwise_and_dot_product_with_popcount(query, data[i], size, 4); + } +} + +template <> +void bitwise_q4_batch_4( + const uint8_t* query, const uint8_t* const* data, size_t size, + BitwiseAndDotProductResult* results); + /** * Compute dot product between query and binary data using popcount on XOR. * diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp index 60afe80e7..c80c2b4b3 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp @@ -417,9 +417,58 @@ BitwiseAndDotProductResult bitwise_q4_vpopcnt( return {static_cast(_mm512_reduce_add_epi64(dots)), static_cast(_mm512_reduce_add_epi64(pops))}; } +__attribute__((target("avx512vpopcntdq"), noinline)) +void bitwise_q4_vpopcnt_batch_4( + const uint8_t* query, const uint8_t* const* data, size_t size, + BitwiseAndDotProductResult* results) { + __m512i dots[4], pops[4]; + for (int i = 0; i < 4; ++i) { + dots[i] = pops[i] = _mm512_setzero_si512(); + } + for (size_t off = 0; off < size; off += 64) { + const size_t count = std::min(size - off, size_t(64)); + const __mmask64 mask = count == 64 ? ~__mmask64(0) : (__mmask64(1) << count) - 1; + __m512i x[4]; + for (int i = 0; i < 4; ++i) { + x[i] = _mm512_maskz_loadu_epi8(mask, data[i] + off); + pops[i] = _mm512_add_epi64(pops[i], _mm512_popcnt_epi64(x[i])); + } + for (int bit = 0; bit < 4; ++bit) { + // Load each query plane once for four independent database codes. + const __m512i q = _mm512_maskz_loadu_epi8(mask, query + bit * size + off); + for (int i = 0; i < 4; ++i) { + const __m512i p = _mm512_popcnt_epi64(_mm512_and_si512(q, x[i])); + dots[i] = _mm512_add_epi64(dots[i], _mm512_slli_epi64(p, bit)); + } + } + } + for (int i = 0; i < 4; ++i) { + results[i] = {static_cast(_mm512_reduce_add_epi64(dots[i])), + static_cast(_mm512_reduce_add_epi64(pops[i]))}; + } +} } // namespace #endif +template <> +BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount( + const uint8_t* query, const uint8_t* data, size_t size, size_t qb); + +template <> +void bitwise_q4_batch_4( + const uint8_t* query, const uint8_t* const* data, size_t size, + BitwiseAndDotProductResult* results) { +#if defined(__GNUC__) && defined(__x86_64__) + if (__builtin_cpu_supports("avx512vpopcntdq")) { + bitwise_q4_vpopcnt_batch_4(query, data, size, results); + return; + } +#endif + for (int i = 0; i < 4; ++i) { + results[i] = bitwise_and_dot_product_with_popcount(query, data[i], size, 4); + } +} + template <> BitwiseAndDotProductResult bitwise_and_dot_product_with_popcount< SIMDLevel::AVX512>( From 32ec27c18e2a5146beadbdb8bf3eaba8ec8b5387 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 08:22:21 +0000 Subject: [PATCH 4/7] refactor: consolidate HNSW storage distance-computer factory Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc | 4 +-- src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h | 4 +-- src/index/hnsw/impl/IndexHNSWWrapper.cc | 31 ++++++------------- src/index/hnsw/impl/IndexHNSWWrapper.h | 7 +++-- 4 files changed, 18 insertions(+), 28 deletions(-) diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc index 6daba7d42..48bd0f11f 100644 --- a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc @@ -11,8 +11,8 @@ namespace knowhere { namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; std::unique_ptr -IndexHNSWRaBitQWrapper::graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, - const SearchParametersHNSWWrapper* params) const { +IndexHNSWRaBitQWrapper::storage_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const { const auto* rbq = dynamic_cast(index); FAISS_THROW_IF_NOT(rbq); const auto* rbq_params = dynamic_cast(params); diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h index 80b85f00b..ff4ed41ee 100644 --- a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h @@ -10,8 +10,8 @@ struct IndexHNSWRaBitQWrapper : IndexHNSWWrapper { protected: std::unique_ptr - graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, - const SearchParametersHNSWWrapper* params) const override; + storage_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const override; faiss::cppcontrib::knowhere::HNSWStats search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.cc b/src/index/hnsw/impl/IndexHNSWWrapper.cc index 3b57bde11..a6cef659c 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWWrapper.cc @@ -40,24 +40,6 @@ namespace knowhere { -/************************************************************** - * Utilities - **************************************************************/ - -namespace { - -// cloned from IndexHNSW.cpp -faiss::DistanceComputer* -storage_distance_computer(const faiss::Index* storage) { - if (faiss::cppcontrib::knowhere::is_similarity_metric(storage->metric_type)) { - return new faiss::NegativeDistanceComputer(storage->get_distance_computer()); - } else { - return storage->get_distance_computer(); - } -} - -} // namespace - /************************************************************** * IndexHNSWWrapper implementation **************************************************************/ @@ -74,9 +56,14 @@ SearchParametersHNSWWrapper::create_hnsw_wrapper(faiss::cppcontrib::knowhere::In } std::unique_ptr -IndexHNSWWrapper::graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, - const SearchParametersHNSWWrapper*) const { - return std::unique_ptr(storage_distance_computer(index->storage)); +IndexHNSWWrapper::storage_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper*) const { + const auto* storage = index->storage; + if (faiss::cppcontrib::knowhere::is_similarity_metric(storage->metric_type)) { + return std::unique_ptr( + new faiss::NegativeDistanceComputer(storage->get_distance_computer())); + } + return std::unique_ptr(storage->get_distance_computer()); } faiss::cppcontrib::knowhere::HNSWStats @@ -133,7 +120,7 @@ IndexHNSWWrapper::search(idx_t n, const float* __restrict x, idx_t k, float* __r faiss::cppcontrib::knowhere::Bitset bitset_visited_nodes = faiss::cppcontrib::knowhere::Bitset::create_uninitialized(index->ntotal); - auto dis = graph_distance_computer(index_hnsw, params); + auto dis = storage_distance_computer(index_hnsw, params); // no parallelism by design for (idx_t i = 0; i < n; i++) { diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.h b/src/index/hnsw/impl/IndexHNSWWrapper.h index 6f213aa2a..90f899c2a 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWWrapper.h @@ -66,9 +66,12 @@ struct IndexHNSWWrapper : public faiss::cppcontrib::knowhere::IndexWrapper { const faiss::SearchParameters* params) const override; protected: + // Graph-search factory: returns smaller-is-better distances. The index is + // the owning HNSW index, allowing subclasses to access storage metadata. + // Unlike the request-parameter factory, this may expose staged evaluation. virtual std::unique_ptr - graph_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, - const SearchParametersHNSWWrapper* params) const; + storage_distance_computer(const faiss::cppcontrib::knowhere::IndexHNSW* index, + const SearchParametersHNSWWrapper* params) const; virtual faiss::cppcontrib::knowhere::HNSWStats search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, From f0e4433ec6a988b53551a63dfe1fb2c50f620e80 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:10:12 +0000 Subject: [PATCH 5/7] refactor: clarify HNSW RaBitQ policies and validation boundaries Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 15 +++--- src/index/hnsw/impl/HnswSearchDispatch.h | 6 +-- src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc | 5 +- tests/ut/test_hnsw_rabitq.cc | 4 +- tests/ut/test_hnsw_rabitq_acceptance.cc | 53 ++++++++++++++++--- .../cppcontrib/knowhere/IndexHNSWRaBitQ.cpp | 35 ++---------- .../cppcontrib/knowhere/IndexHNSWRaBitQ.h | 9 ++-- .../knowhere/impl/HnswDistanceEvaluation.h | 37 +++++++++++++ .../cppcontrib/knowhere/impl/HnswSearcher.h | 27 ++-------- ...ation.h => RaBitQHnswDistanceEvaluation.h} | 14 +++-- .../cppcontrib/knowhere/impl/index_read.cpp | 6 +-- .../cppcontrib/knowhere/impl/index_write.cpp | 6 +-- 12 files changed, 126 insertions(+), 91 deletions(-) create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswDistanceEvaluation.h rename thirdparty/faiss/faiss/cppcontrib/knowhere/impl/{RaBitQDistanceEvaluation.h => RaBitQHnswDistanceEvaluation.h} (84%) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index a9e8626a2..4561d0f2f 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -1412,8 +1412,8 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { std::unique_ptr bf_index_wrapper = nullptr; faiss::Index* bf_index_wrapper_ptr = nullptr; if (!whether_bf_search.value_or(false)) { - std::tie(bf_index_wrapper, is_refined) = create_conditional_hnsw_wrapper( - indexes[index_id].get(), hnsw_cfg, true, whether_to_enable_refine, search_parameters.get()); + std::tie(bf_index_wrapper, is_refined) = + create_conditional_hnsw_wrapper(indexes[index_id].get(), hnsw_cfg, true, whether_to_enable_refine); if (bf_index_wrapper == nullptr) { return expected::Err(Status::invalid_args, "an input index seems to be unrelated to HNSW"); } @@ -3126,11 +3126,8 @@ class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode refine ? refine->base_index : loaded.get()); if (!rbq) return Status::invalid_serialized_index_type; - if (const auto* cosine = dynamic_cast(rbq)) { - cosine->validate_cosine_storage(); - } else { - rbq->validate_storage(); - } + // read_index already checked the storage and HNSW composition. + // Only the outer Knowhere type and optional refine relation remain. if (refine) { const auto* storage = refine->refine_index; if (!storage || !refine->is_trained || !storage->is_trained || refine->d != rbq->d || @@ -3301,9 +3298,9 @@ class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode index_hnsw_rabitq->own_fields = false; if (is_cosine) { dynamic_cast(index_hnsw_rabitq.get()) - ->validate_cosine_storage(); + ->check_cosine_storage_compatibility(); } else { - index_hnsw_rabitq->validate_storage(); + index_hnsw_rabitq->check_storage_compatibility(); } index_hnsw_rabitq->own_fields = true; tmp_index_rabitq[0].release(); diff --git a/src/index/hnsw/impl/HnswSearchDispatch.h b/src/index/hnsw/impl/HnswSearchDispatch.h index 4a0529ebc..4b47b5985 100644 --- a/src/index/hnsw/impl/HnswSearchDispatch.h +++ b/src/index/hnsw/impl/HnswSearchDispatch.h @@ -14,7 +14,7 @@ namespace knowhere { // Reuse selector/visitor dispatch without exposing codec types to common HNSW. -template +template faiss::cppcontrib::knowhere::HNSWStats search_hnsw_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& distance, faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, @@ -22,8 +22,8 @@ search_hnsw_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::Distanc auto run = [&](auto& visitor, const auto& selector) { using Visitor = std::remove_reference_t; using Selector = std::decay_t; - faiss::cppcontrib::knowhere::v2_hnsw_searcher + faiss::cppcontrib::knowhere::v2_hnsw_searcher< + faiss::DistanceComputer, Visitor, faiss::cppcontrib::knowhere::Bitset, Selector, DistanceEvaluationT> searcher{graph, distance, visitor, visited, selector, params ? params->kAlpha : 0.0f, params}; return searcher.search(k, distances, labels); }; diff --git a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc index 48bd0f11f..c061aa94c 100644 --- a/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 #include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" -#include +#include #include "index/hnsw/impl/HnswSearchDispatch.h" #include "index/hnsw/impl/RaBitQSearchParameters.h" @@ -24,6 +24,7 @@ faiss::cppcontrib::knowhere::HNSWStats IndexHNSWRaBitQWrapper::search_query(const faiss::cppcontrib::knowhere::HNSW& graph, faiss::DistanceComputer& dc, faiss::cppcontrib::knowhere::Bitset& visited, faiss::idx_t k, float* distances, faiss::idx_t* labels, const SearchParametersHNSWWrapper* params) const { - return search_hnsw_query(graph, dc, visited, k, distances, labels, params); + return search_hnsw_query(graph, dc, visited, k, distances, labels, + params); } } // namespace knowhere diff --git a/tests/ut/test_hnsw_rabitq.cc b/tests/ut/test_hnsw_rabitq.cc index e1f3662ba..e3d58bc0d 100644 --- a/tests/ut/test_hnsw_rabitq.cc +++ b/tests/ut/test_hnsw_rabitq.cc @@ -3,7 +3,7 @@ #include #include -#include +#include #include #include #include @@ -62,7 +62,7 @@ struct OriginalFullEvaluation { TEST_CASE("RaBitQ threshold heap matches priority queue after every update", "[hnsw_rabitq_core]") { using faiss::cppcontrib::knowhere::Neighbor; - rabitq_search::DistanceEvaluation evaluation; + rabitq_search::RaBitQHnswDistanceEvaluation evaluation; std::mt19937 rng(12345); for (size_t k : {1, 2, 3, 10, 100, 511}) { for (int mode = 0; mode < 3; ++mode) { diff --git a/tests/ut/test_hnsw_rabitq_acceptance.cc b/tests/ut/test_hnsw_rabitq_acceptance.cc index f7b93e5ce..d8636199d 100644 --- a/tests/ut/test_hnsw_rabitq_acceptance.cc +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -547,29 +547,68 @@ TEST_CASE("RaBitQ file serialization rejects truncated and incompatible indexes" auto* storage = const_cast(graph->pretransform_index()); auto* rq = const_cast(graph->rabitq_index()); auto* rr = dynamic_cast(storage->chain[0]); - REQUIRE_NOTHROW(graph->validate_storage()); + // Test the serialization boundary, not an outer validator duplicating + // checks owned by RaBitQ codes and cosine norm storage. + auto write_decoded = [&] { + faiss::VectorIOWriter writer; + fk::write_index(decoded.get(), &writer); + }; + REQUIRE_NOTHROW(graph->check_storage_compatibility()); + REQUIRE_NOTHROW(write_decoded()); ++rq->code_size; - REQUIRE_THROWS(graph->validate_storage()); + REQUIRE_NOTHROW(graph->check_storage_compatibility()); + REQUIRE_THROWS(write_decoded()); --rq->code_size; auto last = rr->A.back(); rr->A.pop_back(); - REQUIRE_THROWS(graph->validate_storage()); + REQUIRE_THROWS(graph->check_storage_compatibility()); rr->A.push_back(last); rr->have_bias = true; - REQUIRE_THROWS(graph->validate_storage()); + REQUIRE_THROWS(graph->check_storage_compatibility()); rr->have_bias = false; const auto old_metric = rq->rabitq.metric_type; rq->rabitq.metric_type = faiss::METRIC_L1; - REQUIRE_THROWS(graph->validate_storage()); + REQUIRE_THROWS(write_decoded()); rq->rabitq.metric_type = old_metric; + const auto saved_center = rq->center; + rq->center.clear(); + REQUIRE_THROWS(write_decoded()); + rq->center = saved_center; + const auto code_bytes = rq->codes.size(); + const auto last_code = rq->codes[code_bytes - 1]; + rq->codes.resize(code_bytes - 1); + REQUIRE_THROWS(write_decoded()); + rq->codes.resize(code_bytes); + rq->codes[code_bytes - 1] = last_code; + const auto saved_bits = rq->rabitq.nb_bits; + rq->rabitq.nb_bits = 10; + REQUIRE_THROWS(write_decoded()); + rq->rabitq.nb_bits = saved_bits; + const auto saved_qb = rq->qb; + rq->qb = 9; + REQUIRE_THROWS(write_decoded()); + rq->qb = saved_qb; + rq->centered = true; + REQUIRE_NOTHROW(graph->check_storage_compatibility()); + REQUIRE_THROWS(write_decoded()); + rq->centered = false; + ++storage->ntotal; + REQUIRE_THROWS(graph->check_storage_compatibility()); + REQUIRE_THROWS(write_decoded()); + --storage->ntotal; + REQUIRE_NOTHROW(write_decoded()); if (auto* cosine = dynamic_cast(graph)) { auto* cs = dynamic_cast(storage); auto norm = cs->inverse_norms_storage.inverse_l2_norms.back(); cs->inverse_norms_storage.inverse_l2_norms.pop_back(); - REQUIRE_THROWS(cosine->validate_cosine_storage()); + REQUIRE_NOTHROW(cosine->check_cosine_storage_compatibility()); + REQUIRE_THROWS(write_decoded()); cs->inverse_norms_storage.inverse_l2_norms.push_back(norm); + const auto first_norm = cs->inverse_norms_storage.inverse_l2_norms[0]; cs->inverse_norms_storage.inverse_l2_norms[0] = std::numeric_limits::quiet_NaN(); - REQUIRE_THROWS(cosine->validate_cosine_storage()); + REQUIRE_THROWS(write_decoded()); + cs->inverse_norms_storage.inverse_l2_norms[0] = first_norm; + REQUIRE_NOTHROW(write_decoded()); } // A real file, independently loaded through the public Knowhere API. auto pattern = (std::filesystem::temp_directory_path() / "knowhere-rabitq-XXXXXX").string(); diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp index 5f2a98875..75c3d2c15 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -90,7 +90,7 @@ const faiss::IndexRaBitQ* IndexHNSWRaBitQ::rabitq_index() const { : nullptr; } -void IndexHNSWRaBitQ::validate_storage() const { +void IndexHNSWRaBitQ::check_storage_compatibility() const { FAISS_THROW_IF_NOT_MSG( metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT, "IndexHNSWRaBitQ only supports L2 and inner product metrics"); @@ -125,8 +125,7 @@ void IndexHNSWRaBitQ::validate_storage() const { rabitq->is_trained, "IndexHNSWRaBitQ requires fully trained storage"); FAISS_THROW_IF_NOT_MSG( - pretransform->index != nullptr && - pretransform->ntotal == rabitq->ntotal && + pretransform->ntotal == rabitq->ntotal && pretransform->metric_type == rabitq->metric_type, "IndexHNSWRaBitQ pretransform and RaBitQ metadata mismatch"); FAISS_THROW_IF_NOT_MSG( @@ -140,31 +139,6 @@ void IndexHNSWRaBitQ::validate_storage() const { static_cast(rotation->d_in) * rotation->d_out, "IndexHNSWRaBitQ rotation matrix has invalid storage"); - FAISS_THROW_IF_NOT_MSG( - rabitq->rabitq.d == static_cast(rabitq->d) && - rabitq->rabitq.metric_type == rabitq->metric_type, - "IndexHNSWRaBitQ RaBitQ quantizer metadata mismatch"); - FAISS_THROW_IF_NOT_MSG( - rabitq->rabitq.nb_bits >= 1 && rabitq->rabitq.nb_bits <= 9, - "IndexHNSWRaBitQ RaBitQ nb_bits must be in [1, 9]"); - - const size_t expected_code_size = - rabitq->rabitq.compute_code_size(rabitq->d, rabitq->rabitq.nb_bits); - FAISS_THROW_IF_NOT_MSG( - rabitq->rabitq.code_size == expected_code_size && - rabitq->code_size == expected_code_size, - "IndexHNSWRaBitQ RaBitQ code size mismatch"); - FAISS_THROW_IF_NOT_MSG( - rabitq->codes.size() == - static_cast(rabitq->ntotal) * expected_code_size, - "IndexHNSWRaBitQ RaBitQ codes size mismatch"); - FAISS_THROW_IF_NOT_MSG( - rabitq->center.size() == static_cast(rabitq->d), - "IndexHNSWRaBitQ RaBitQ center size mismatch"); - FAISS_THROW_IF_NOT_MSG( - rabitq->qb <= 8, "IndexHNSWRaBitQ RaBitQ qb must be in [0, 8]"); - FAISS_THROW_IF_NOT_MSG( - !rabitq->centered, "IndexHNSWRaBitQ V1 requires centered=false"); } IndexHNSWRaBitQCosine::IndexHNSWRaBitQCosine() = default; @@ -175,8 +149,8 @@ const float* IndexHNSWRaBitQCosine::get_inverse_l2_norms() const { return cosine_storage ? cosine_storage->get_inverse_l2_norms() : nullptr; } -void IndexHNSWRaBitQCosine::validate_cosine_storage() const { - validate_storage(); +void IndexHNSWRaBitQCosine::check_cosine_storage_compatibility() const { + check_storage_compatibility(); const auto* cosine_storage = dynamic_cast(storage); FAISS_THROW_IF_NOT_MSG( @@ -185,7 +159,6 @@ void IndexHNSWRaBitQCosine::validate_cosine_storage() const { FAISS_THROW_IF_NOT_MSG( metric_type == METRIC_INNER_PRODUCT, "IndexHNSWRaBitQCosine requires inner product storage"); - cosine_storage->validate_norms(); } } // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h index 7c38ab1a0..006666d2f 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h @@ -70,9 +70,10 @@ struct IndexHNSWRaBitQ : IndexHNSW { faiss::DistanceComputer* get_staged_distance_computer( const faiss::RaBitQSearchParameters* params = nullptr) const; - /** Validate the complete runtime/storage shape and serialized invariants. - * Throws FaissException on malformed state. */ - void validate_storage() const; + /** Check the HNSW/pretransform/storage composition at build and IO boundaries. + * Leaf code buffers and cosine norms are checked by their own IO handlers. + * This does not validate all graph data or the serialization format. */ + void check_storage_compatibility() const; }; /** Cosine runtime marker for HNSW backed by cosine-aware RaBitQ storage. */ @@ -80,7 +81,7 @@ struct IndexHNSWRaBitQCosine : IndexHNSWRaBitQ, HasInverseL2Norms { IndexHNSWRaBitQCosine(); const float* get_inverse_l2_norms() const override; - void validate_cosine_storage() const; + void check_cosine_storage_compatibility() const; }; } // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswDistanceEvaluation.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswDistanceEvaluation.h new file mode 100644 index 000000000..2dbe590ef --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswDistanceEvaluation.h @@ -0,0 +1,37 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +namespace faiss::cppcontrib::knowhere { + +// Compile-time HNSW distance-evaluation policy contract (no common base class): +// begin(k) initializes query-local state; record(distance, status) seeds it. +// compute receives up to four candidates and emits their distances in order. +// Its return value counts additional distance evaluations for HNSW statistics; +// the searcher already counts one evaluation per candidate. +// Policies must match the supplied DC type. They do not own traversal queues. +// +// Default policy: directly evaluate the storage distance, including batch-four. +// "Default" does not imply FP32 accuracy: the DC may use SQ/PQ or another codec. +struct DefaultHnswDistanceEvaluation { + void begin(size_t) {} + void record(float, int) {} + + template + size_t compute(DC& dc, const size_t* ids, const int*, size_t count, + int, Emit&& emit) { + if (count == 4) { + float d[4]; + dc.distances_batch_4(ids[0], ids[1], ids[2], ids[3], + d[0], d[1], d[2], d[3]); + for (size_t i = 0; i < count; ++i) emit(i, d[i]); + } else { + for (size_t i = 0; i < count; ++i) emit(i, dc(ids[i])); + } + return 0; + } +}; + +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h index eaf5f513a..651d91be3 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/HnswSearcher.h @@ -34,6 +34,7 @@ // Knowhere-specific headers #include +#include namespace faiss { namespace cppcontrib { @@ -46,31 +47,13 @@ constexpr bool track_hnsw_stats = true; } // namespace -// Default evaluator preserves ordinary full-distance, batch-four execution. -struct FullDistanceEvaluation { - void begin(size_t) {} - void record(float, int) {} - - template - size_t compute(DC& dc, const size_t* ids, const int*, size_t count, - int, Emit&& emit) { - if (count == 4) { - float d[4]; - dc.distances_batch_4(ids[0], ids[1], ids[2], ids[3], - d[0], d[1], d[2], d[3]); - for (size_t i = 0; i < count; ++i) emit(i, d[i]); - } else { - for (size_t i = 0; i < count; ++i) emit(i, dc(ids[i])); - } - return 0; - } -}; - // Accomodates all the search logic and variables. /// * DistanceComputerT is responsible for computing distances /// * GraphVisitorT records visited edges /// * VisitedT is responsible for tracking visited nodes /// * FilterT is resposible for filtering unneeded nodes +/// * DistanceEvaluationT selects candidate evaluation via the compile-time +/// policy contract in HnswDistanceEvaluation.h (independent of traversal). /// Interfaces of all templates are tweaked to accept standard Faiss structures /// with dynamic dispatching. Custom Knowhere structures are also accepted. template < @@ -78,7 +61,7 @@ template < typename GraphVisitorT, typename VisitedT, typename FilterT, - typename EvaluationT = FullDistanceEvaluation> + typename DistanceEvaluationT = DefaultHnswDistanceEvaluation> struct v2_hnsw_searcher { using storage_idx_t = faiss::cppcontrib::knowhere::HNSW::storage_idx_t; using idx_t = faiss::idx_t; @@ -110,7 +93,7 @@ struct v2_hnsw_searcher { // the pointer is not owned. const faiss::cppcontrib::knowhere::SearchParametersHNSW* params; - EvaluationT evaluation; + DistanceEvaluationT evaluation; // v2_hnsw_searcher( diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQHnswDistanceEvaluation.h similarity index 84% rename from thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h rename to thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQHnswDistanceEvaluation.h index 9dd997bbd..b1494d898 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQDistanceEvaluation.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQHnswDistanceEvaluation.h @@ -2,17 +2,21 @@ * Licensed under the MIT license in thirdparty/faiss/LICENSE. */ #pragma once -#include +#include +#include #include #include #include #include #include +#include namespace faiss::cppcontrib::knowhere::rabitq_search { -// Used only by the RaBitQ wrapper's Knowhere-traversal specialization. -struct DistanceEvaluation { +// Alternative to DefaultHnswDistanceEvaluation under the same policy contract. +// Requires RaBitQStagedDistanceComputer; delegates non-staged work to the +// default policy without inheriting from it or depending on the searcher. +struct RaBitQHnswDistanceEvaluation { size_t k = 0; std::vector results; @@ -68,8 +72,8 @@ struct DistanceEvaluation { } return staged.refine_count - before; } - FullDistanceEvaluation full; - return full.compute(dc, ids, statuses, count, level, std::forward(emit)); + DefaultHnswDistanceEvaluation default_evaluation; + return default_evaluation.compute(dc, ids, statuses, count, level, std::forward(emit)); } }; } // namespace faiss::cppcontrib::knowhere::rabitq_search diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp index 1be477c09..f820b7382 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp @@ -708,7 +708,7 @@ static void finalize_and_validate_RaBitQ_index(::faiss::IndexRaBitQ* idxq) { static_cast(idxq->ntotal) * expected_code_size, "IndexRaBitQ codes size mismatch"); FAISS_THROW_IF_NOT_MSG( - idxq->center.empty() || + (!idxq->is_trained && idxq->center.empty()) || idxq->center.size() == static_cast(idxq->d), "IndexRaBitQ center size mismatch"); FAISS_THROW_IF_NOT_FMT( @@ -1272,11 +1272,11 @@ Index* read_index(IOReader* f, int io_flags) { idxhnsw->storage = read_index(f, io_flags); idxhnsw->own_fields = idxhnsw->storage != nullptr; if (h == fourcc(kHnswRaBitQFourcc)) { - dynamic_cast(idxhnsw)->validate_storage(); + dynamic_cast(idxhnsw)->check_storage_compatibility(); } if (h == fourcc(kHnswRaBitQCosineFourcc)) { dynamic_cast(idxhnsw) - ->validate_cosine_storage(); + ->check_cosine_storage_compatibility(); } if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { dynamic_cast(idxhnsw->storage)->pq.compute_sdc_table(); diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp index 6415be569..1dc2c8fe0 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp @@ -524,7 +524,7 @@ static void validate_RaBitQ_index_for_write(const ::faiss::IndexRaBitQ* idxq) { static_cast(idxq->ntotal) * expected_code_size, "IndexRaBitQ codes size mismatch"); FAISS_THROW_IF_NOT_MSG( - idxq->center.empty() || + (!idxq->is_trained && idxq->center.empty()) || idxq->center.size() == static_cast(idxq->d), "IndexRaBitQ center size mismatch"); FAISS_THROW_IF_NOT_MSG(idxq->qb <= 8, "IndexRaBitQ qb must be in [0, 8]"); @@ -844,9 +844,9 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { !(io_flags & IO_FLAG_SKIP_STORAGE), "IndexHNSWRaBitQ cannot be serialized without its RaBitQ storage"); if (hnsw_rabitq_cosine) { - hnsw_rabitq_cosine->validate_cosine_storage(); + hnsw_rabitq_cosine->check_cosine_storage_compatibility(); } else { - hnsw_rabitq->validate_storage(); + hnsw_rabitq->check_storage_compatibility(); } } uint32_t h = hnsw_rabitq_cosine ? fourcc(kHnswRaBitQCosineFourcc) From 6d0a48764e1516a86ecd8fb323983a8f13cb76d8 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:40:52 +0000 Subject: [PATCH 6/7] fix: normalize cosine refine distances in HNSW ID lookup Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 15 +++- tests/ut/test_hnsw_rabitq_acceptance.cc | 96 +++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 4561d0f2f..1634a0575 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -1601,7 +1601,20 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { const faiss::cppcontrib::knowhere::IndexRefine* index_refine = dynamic_cast(indexes[index_id].get()); if (index_refine != nullptr) { - dist_computer.reset(index_refine->refine_index->get_distance_computer()); + const auto* graph = + dynamic_cast(index_refine->base_index); + const auto* norms = + graph ? dynamic_cast(graph->storage) + : nullptr; + if (is_cosine && norms) { + // Match Search: refine storage holds unnormalized vectors, + // while the base storage owns their original inverse norms. + IndexWrapperCosine refine_wrapper(index_refine->refine_index, + norms->get_inverse_l2_norms()); + dist_computer.reset(refine_wrapper.get_distance_computer()); + } else { + dist_computer.reset(index_refine->refine_index->get_distance_computer()); + } } else { dist_computer.reset(indexes[index_id]->get_distance_computer()); } diff --git a/tests/ut/test_hnsw_rabitq_acceptance.cc b/tests/ut/test_hnsw_rabitq_acceptance.cc index d8636199d..687eff90a 100644 --- a/tests/ut/test_hnsw_rabitq_acceptance.cc +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -1027,6 +1027,102 @@ TEST_CASE("RaBitQ range boundaries and range_filter match full-code reference", } } +TEST_CASE("HNSW refined distances by ID match Search and metric references", "[hnsw_rabitq_regression]") { + constexpr int n = 128, d = 32, nq = 3, k = 20; + auto base = GenDataSet(n, d, 1961); + auto query = GenDataSet(nq, d, 1962); + auto* x = const_cast(static_cast(base->GetTensor())); + auto* q = const_cast(static_cast(query->GetTensor())); + // Non-unit vectors expose a raw-IP/cosine mismatch. Use values exactly + // representable in all three refine formats to isolate metric semantics + // from conversion rounding. Also cover the Search norm convention for zeros. + for (int i = 0; i < n; ++i) + for (int j = 0; j < d; ++j) x[i * d + j] = ((i * 17 + j * 13) % 31 - 15) * 0.25f * (1 + i % 7); + std::fill_n(x, 2 * d, 0.0f); + x[0] = 6.0f; + x[1] = 8.0f; + std::fill_n(q, nq * d, 0.0f); + q[0] = 3.0f; + q[1] = 4.0f; + q[d] = -4.0f; + q[d + 1] = 3.0f; + std::vector labels(n); + for (int i = 0; i < n; ++i) labels[i] = n - 1 - i; + + for (const auto* type : {"HNSW_RABITQ", "HNSW_SQ"}) { + for (const auto* metric : {"L2", "IP", "COSINE"}) { + for (const auto* refine : {"FP32", "FP16", "BF16"}) { + CAPTURE(type, metric, refine); + const bool cosine = std::string(metric) == "COSINE"; + const bool l2 = std::string(metric) == "L2"; + knowhere::Json cfg = {{"dim", d}, {"metric_type", metric}, + {"M", 16}, {"efConstruction", 100}, + {"ef", n}, {"k", k}, + {"refine", true}, {"refine_type", refine}, + {"refine_k", 1.3}}; + if (std::string(type) == "HNSW_RABITQ") { + cfg["rbq_bits"] = 4; + cfg["rbq_bits_query"] = 4; + } else { + cfg["sq_type"] = "SQ8"; + } + auto create = [&] { + return knowhere::IndexFactory::Instance() + .Create(type, knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + }; + auto index = create(); + REQUIRE(index.Build(base, cfg) == knowhere::Status::success); + knowhere::BinarySet binary; + REQUIRE(index.Serialize(binary) == knowhere::Status::success); + auto restored = create(); + REQUIRE(restored.Deserialize(binary, cfg) == knowhere::Status::success); + for (const auto* current : {&index, &restored}) { + auto by_storage = + current->Node()->CalcDistByStorageIds(query, {}, labels.data(), labels.size(), cosine); + auto by_public = current->Node()->CalcDistByIDs(query, {}, labels.data(), labels.size(), cosine); + auto search = current->Search(query, cfg, nullptr); + REQUIRE(by_storage.has_value()); + REQUIRE(by_public.has_value()); + REQUIRE(search.has_value()); + for (int qi = 0; qi < nq; ++qi) { + for (int j = 0; j < n; ++j) { + const auto id = labels[j]; + std::vector decoded(d); + for (int c = 0; c < d; ++c) { + const float v = x[id * d + c]; + decoded[c] = std::string(refine) == "FP16" ? float(knowhere::fp16(v)) + : std::string(refine) == "BF16" ? faiss::decode_bf16(faiss::encode_bf16(v)) + : v; + } + float expected = l2 ? faiss::fvec_L2sqr(q + qi * d, decoded.data(), d) + : faiss::fvec_inner_product(q + qi * d, decoded.data(), d); + if (cosine) { + const float qnorm = std::sqrt(faiss::fvec_norm_L2sqr(q + qi * d, d)); + const float xnorm = std::sqrt(faiss::fvec_norm_L2sqr(x + id * d, d)); + expected /= (qnorm > 0 ? qnorm : 1.0f) * (xnorm > 0 ? xnorm : 1.0f); + } + REQUIRE(by_storage.value()->GetDistance()[qi * n + j] == + Catch::Approx(expected).epsilon(1e-5).margin(1e-4)); + REQUIRE(by_public.value()->GetDistance()[qi * n + j] == + by_storage.value()->GetDistance()[qi * n + j]); + } + for (int j = 0; j < k; ++j) { + const auto id = search.value()->GetIds()[qi * k + j]; + REQUIRE(id >= 0); + REQUIRE(id < n); + REQUIRE(search.value()->GetDistance()[qi * k + j] == + Catch::Approx(by_storage.value()->GetDistance()[qi * n + n - 1 - id]) + .epsilon(1e-5) + .margin(1e-4)); + } + } + } + } + } + } +} + TEST_CASE("RaBitQ FP16 BF16 and FP32 refine return refiner distances", "[hnsw_rabitq_acceptance]") { constexpr int n = 128, d = 33, k = 20; auto base = GenDataSet(n, d, 1921); From 821a2439b5896ee394a5cc9608b7d0e7161d38ad Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Fri, 11 Sep 2026 03:15:48 +0000 Subject: [PATCH 7/7] fix: address HNSW RaBitQ clang-tidy findings Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 14 ++++++++++---- src/index/hnsw/impl/IndexHNSWWrapper.cc | 4 +++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 1634a0575..7a1c03927 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -919,7 +919,9 @@ class FaissHnswIterator : public IndexIterator { workspace.qdis.reset(storage_params ? storage_params->storage_distance_computer(index_hnsw) : index_hnsw->get_distance_computer()); if (larger_is_closer) { - workspace.qdis.reset(new faiss::NegativeDistanceComputer(workspace.qdis.release())); + auto negated = std::make_unique(workspace.qdis.get()); + workspace.qdis.release(); + workspace.qdis = std::move(negated); } if (refine_ratio != 0) { @@ -963,7 +965,9 @@ class FaissHnswIterator : public IndexIterator { workspace.qdis.reset(storage_params ? storage_params->storage_distance_computer(index_hnsw) : index_hnsw->get_distance_computer()); if (larger_is_closer) { - workspace.qdis.reset(new faiss::NegativeDistanceComputer(workspace.qdis.release())); + auto negated = std::make_unique(workspace.qdis.get()); + workspace.qdis.release(); + workspace.qdis = std::move(negated); } } @@ -3069,8 +3073,9 @@ class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode Status Deserialize(const BinarySet& binset, std::shared_ptr) override { auto binary = binset.GetByName(Type()); - if (!binary) + if (!binary) { return Status::invalid_binary_set; + } MemoryIOReader reader(binary->data.get(), binary->size); return LoadRaBitQ(reader); } @@ -3137,8 +3142,9 @@ class BaseFaissRegularIndexHNSWRaBitQNode : public BaseFaissRegularIndexHNSWNode const auto* refine = dynamic_cast(loaded.get()); const auto* rbq = dynamic_cast( refine ? refine->base_index : loaded.get()); - if (!rbq) + if (!rbq) { return Status::invalid_serialized_index_type; + } // read_index already checked the storage and HNSW composition. // Only the outer Knowhere type and optional refine relation remain. if (refine) { diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.cc b/src/index/hnsw/impl/IndexHNSWWrapper.cc index a6cef659c..7a2cca881 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWWrapper.cc @@ -206,7 +206,9 @@ IndexHNSWWrapper::range_search(idx_t n, const float* __restrict x, float radius_ std::unique_ptr dis(params ? params->storage_distance_computer(index_hnsw) : index_hnsw->get_distance_computer()); if (faiss::cppcontrib::knowhere::is_similarity_metric(index_hnsw->metric_type)) { - dis.reset(new faiss::NegativeDistanceComputer(dis.release())); + auto negated = std::make_unique(dis.get()); + dis.release(); + dis = std::move(negated); } // radius