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..7a1c03927 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,13 @@ 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) { + auto negated = std::make_unique(workspace.qdis.get()); + workspace.qdis.release(); + workspace.qdis = std::move(negated); + } if (refine_ratio != 0) { // the refine is needed @@ -960,7 +962,13 @@ 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) { + auto negated = std::make_unique(workspace.qdis.get()); + workspace.qdis.release(); + workspace.qdis = std::move(negated); + } } // set query @@ -1155,6 +1163,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 +1372,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_); @@ -1389,9 +1402,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"); @@ -1412,7 +1427,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - knowhere::SearchParametersHNSWWrapper hnsw_search_params; + auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { hnsw_search_params.efSearch = hnsw_cfg.ef.value(); } @@ -1590,7 +1605,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()); } @@ -1660,7 +1688,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) { @@ -1694,9 +1722,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"); @@ -1705,7 +1735,7 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { faiss::Index* index_wrapper_ptr = index_wrapper.get(); // set up faiss search parameters - knowhere::SearchParametersHNSWWrapper hnsw_search_params; + auto& hnsw_search_params = *search_parameters; if (hnsw_cfg.ef.has_value()) { hnsw_search_params.efSearch = hnsw_cfg.ef.value(); @@ -1993,6 +2023,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 +2064,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 +3063,299 @@ 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: + 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(); + 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 + 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; + } + // 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 || + 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(); + 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()) + ->check_cosine_storage_compatibility(); + } else { + index_hnsw_rabitq->check_storage_compatibility(); + } + 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 +3678,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/HnswSearchDispatch.h b/src/index/hnsw/impl/HnswSearchDispatch.h new file mode 100644 index 000000000..4b47b5985 --- /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< + 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); + }; + 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/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/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..c061aa94c --- /dev/null +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.cc @@ -0,0 +1,30 @@ +// Copyright (C) 2026 Zilliz. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "index/hnsw/impl/IndexHNSWRaBitQWrapper.h" + +#include + +#include "index/hnsw/impl/HnswSearchDispatch.h" +#include "index/hnsw/impl/RaBitQSearchParameters.h" + +namespace knowhere { +namespace rabitq_search = faiss::cppcontrib::knowhere::rabitq_search; + +std::unique_ptr +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); + 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..ff4ed41ee --- /dev/null +++ b/src/index/hnsw/impl/IndexHNSWRaBitQWrapper.h @@ -0,0 +1,20 @@ +// 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; + + protected: + std::unique_ptr + 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, + const SearchParametersHNSWWrapper* params) const override; +}; +} // namespace knowhere diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.cc b/src/index/hnsw/impl/IndexHNSWWrapper.cc index 3638fc595..7a2cca881 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.cc +++ b/src/index/hnsw/impl/IndexHNSWWrapper.cc @@ -30,6 +30,7 @@ #include "index/hnsw/impl/DummyVisitor.h" #include "index/hnsw/impl/FederVisitor.h" +#include "index/hnsw/impl/HnswSearchDispatch.h" #include "knowhere/bitsetview.h" #include "knowhere/bitsetview_idselector.h" @@ -39,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 **************************************************************/ @@ -67,6 +50,29 @@ 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::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 +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 { @@ -95,12 +101,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; } // set up hnsw_stats @@ -117,8 +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); - // create a distance computer - std::unique_ptr dis(storage_distance_computer(index_hnsw->storage)); + auto dis = storage_distance_computer(index_hnsw, params); // no parallelism by design for (idx_t i = 0; i < n; i++) { @@ -128,78 +130,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) @@ -271,7 +203,13 @@ 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)) { + auto negated = std::make_unique(dis.get()); + dis.release(); + dis = std::move(negated); + } // radius float radius = radius_in; diff --git a/src/index/hnsw/impl/IndexHNSWWrapper.h b/src/index/hnsw/impl/IndexHNSWWrapper.h index e92e8c9a6..90f899c2a 100644 --- a/src/index/hnsw/impl/IndexHNSWWrapper.h +++ b/src/index/hnsw/impl/IndexHNSWWrapper.h @@ -13,9 +13,13 @@ #include #include +#include +#include +#include #include #include +#include #include "knowhere/feder/HNSW.h" @@ -30,6 +34,15 @@ 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(); + } + + virtual std::unique_ptr + create_hnsw_wrapper(faiss::cppcontrib::knowhere::IndexHNSW* index) const; + inline ~SearchParametersHNSWWrapper() { } }; @@ -51,6 +64,18 @@ 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: + // 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 + 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, + const SearchParametersHNSWWrapper* params) const; }; } // namespace knowhere diff --git a/src/index/hnsw/impl/RaBitQSearchParameters.h b/src/index/hnsw/impl/RaBitQSearchParameters.h new file mode 100644 index 000000000..3b62f09b2 --- /dev/null +++ b/src/index/hnsw/impl/RaBitQSearchParameters.h @@ -0,0 +1,39 @@ +// 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/IndexHNSWRaBitQWrapper.h" + +namespace knowhere { + +// Owns the per-request storage parameters; no shared index state is changed. +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); + 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_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 new file mode 100644 index 000000000..e3d58bc0d --- /dev/null +++ b/tests/ut/test_hnsw_rabitq.cc @@ -0,0 +1,639 @@ +// 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 "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" +#include "utils.h" + +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 threshold heap matches priority queue after every update", "[hnsw_rabitq_core]") { + using faiss::cppcontrib::knowhere::Neighbor; + 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) { + 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") || + !__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); + knowhere::IndexHNSWRaBitQWrapper api(&graph); + knowhere::SearchParametersHNSWWrapper params; + params.efSearch = n; + 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; + 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((similarity ? -1.f : 1.f) * 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()); + // 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()); + 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..687eff90a --- /dev/null +++ b/tests/ut/test_hnsw_rabitq_acceptance.cc @@ -0,0 +1,1172 @@ +// 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 "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" + +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::IndexHNSWRaBitQWrapper 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]); + // 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_NOTHROW(graph->check_storage_compatibility()); + REQUIRE_THROWS(write_decoded()); + --rq->code_size; + auto last = rr->A.back(); + rr->A.pop_back(); + REQUIRE_THROWS(graph->check_storage_compatibility()); + rr->A.push_back(last); + rr->have_bias = true; + 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(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_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(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(); + 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("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); + 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("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); + 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..75c3d2c15 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -0,0 +1,164 @@ +/* + * 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 +#include + +namespace faiss::cppcontrib::knowhere { + + +faiss::DistanceComputer* IndexHNSWRaBitQ::get_staged_distance_computer( + const faiss::RaBitQSearchParameters* params) const { + return new rabitq_search::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::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"); + 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->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"); +} + +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::check_cosine_storage_compatibility() const { + check_storage_compatibility(); + 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"); +} + +} // 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..006666d2f --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h @@ -0,0 +1,87 @@ +/* + * 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; + + /** 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. */ +struct IndexHNSWRaBitQCosine : IndexHNSWRaBitQ, HasInverseL2Norms { + IndexHNSWRaBitQCosine(); + + const float* get_inverse_l2_norms() const override; + 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 70c06d6c3..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 { @@ -51,13 +52,16 @@ constexpr bool track_hnsw_stats = true; /// * 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 < typename DistanceComputerT, typename GraphVisitorT, typename VisitedT, - typename FilterT> + typename FilterT, + typename DistanceEvaluationT = DefaultHnswDistanceEvaluation> struct v2_hnsw_searcher { using storage_idx_t = faiss::cppcontrib::knowhere::HNSW::storage_idx_t; using idx_t = faiss::idx_t; @@ -89,6 +93,8 @@ struct v2_hnsw_searcher { // the pointer is not owned. const faiss::cppcontrib::knowhere::SearchParametersHNSW* params; + DistanceEvaluationT evaluation; + // v2_hnsw_searcher( const faiss::cppcontrib::knowhere::HNSW& hnsw_, @@ -231,62 +237,31 @@ struct v2_hnsw_searcher { ndis += 1; if (counter == 4) { - // 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 = qdis(saved_indices[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; } @@ -326,6 +301,9 @@ struct v2_hnsw_searcher { } } + if (disqualified) { + retset.save_pending(*disqualified); + } // done return stats; } @@ -369,6 +347,8 @@ struct v2_hnsw_searcher { // grab some needed parameters const int efSearch = params ? params->efSearch : hnsw.efSearch; + evaluation.begin(static_cast(k)); + // yes. // greedy search on upper levels. @@ -404,6 +384,8 @@ struct v2_hnsw_searcher { } visited_nodes[nearest] = true; + evaluation.record(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/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/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/RaBitQHnswDistanceEvaluation.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQHnswDistanceEvaluation.h new file mode 100644 index 000000000..b1494d898 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/RaBitQHnswDistanceEvaluation.h @@ -0,0 +1,79 @@ +/* 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 +#include + +namespace faiss::cppcontrib::knowhere::rabitq_search { + +// 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; + + void begin(size_t count) { k = count; results.clear(); results.reserve(k); } + float threshold() const { + 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_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); + 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]); + emit(i, distance); + } + return staged.refine_count - before; + } + 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/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/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..f820b7382 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) @@ -685,6 +684,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->is_trained && 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); @@ -1082,6 +1116,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 +1153,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 +1184,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); @@ -1161,11 +1223,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 +1271,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)->check_storage_compatibility(); + } + if (h == fourcc(kHnswRaBitQCosineFourcc)) { + dynamic_cast(idxhnsw) + ->check_cosine_storage_compatibility(); + } if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { dynamic_cast(idxhnsw->storage)->pq.compute_sdc_table(); } @@ -1222,7 +1302,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); @@ -1484,4 +1564,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..1dc2c8fe0 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) { @@ -502,6 +501,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->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]"); + 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 @@ -542,7 +573,7 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { 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 = @@ -710,6 +741,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 +785,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 +836,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->check_cosine_storage_compatibility(); + } else { + hnsw_rabitq->check_storage_compatibility(); + } + } + 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 +959,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 +1037,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/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 5dac1d60a..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. * @@ -131,6 +145,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 +162,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 +189,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 +216,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 +237,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 +379,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..c80c2b4b3 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,81 @@ 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))}; +} +__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>( @@ -401,6 +476,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 +537,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 +607,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 +652,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 +744,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); } - float result = hsum_avx2(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 = _mm512_reduce_add_ps(acc); result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } @@ -733,13 +835,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) {