From dd3dce0f46304628e6541e4a072d977744395d3a Mon Sep 17 00:00:00 2001 From: foxspy Date: Sat, 29 Aug 2026 10:11:30 +0800 Subject: [PATCH 1/2] update: bump cardinal to v3.0.7 (v2 ref) / v2.5.112 (v1 ref) (#1799) Bump pinned cardinal tags to the freshly cut releases: - v2 line (master): v3.0.6 -> v3.0.7 (includes #922 #928 #931 #929) - v1 line: v2.5.111 -> v2.5.112 (includes #932) Signed-off-by: xianliang.li Co-authored-by: Claude Opus 4.8 --- cmake/libs/cardinal/v1/CMakeLists.txt | 2 +- cmake/libs/cardinal/v2/CMakeLists.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/libs/cardinal/v1/CMakeLists.txt b/cmake/libs/cardinal/v1/CMakeLists.txt index 782ab3c86..0f3fe7b19 100644 --- a/cmake/libs/cardinal/v1/CMakeLists.txt +++ b/cmake/libs/cardinal/v1/CMakeLists.txt @@ -3,7 +3,7 @@ project(knowhere CXX C) # Use short SHA1 as version # currently check out a commit for cardinal v1. TODO: need a tag here -set(CARDINAL_VERSION v2.5.111) +set(CARDINAL_VERSION v2.5.112) set(CARDINAL_REPO_URL "https://github.com/zilliztech/cardinal.git") set(CARDINAL_ROOT "${KNOWHERE_THRID_ROOT}/cardinalv1") diff --git a/cmake/libs/cardinal/v2/CMakeLists.txt b/cmake/libs/cardinal/v2/CMakeLists.txt index cc54e4c9d..aefc347dd 100644 --- a/cmake/libs/cardinal/v2/CMakeLists.txt +++ b/cmake/libs/cardinal/v2/CMakeLists.txt @@ -3,7 +3,7 @@ project(knowhere CXX C) # Use short SHA1 as version # currenly checkout a commit for cardinal v2. TODO: need a tag here -set(CARDINAL_VERSION v3.0.6) +set(CARDINAL_VERSION v3.0.7) set(CARDINAL_REPO_URL "https://github.com/zilliztech/cardinal.git") set(CARDINAL_ROOT "${KNOWHERE_THRID_ROOT}/cardinalv2") From c74c993dd5b124dda5a497e8c6d3acc2883c0934 Mon Sep 17 00:00:00 2001 From: ChenLiqing <23721160+CLiqing@users.noreply.github.com> Date: Wed, 2 Sep 2026 09:24:08 +0000 Subject: [PATCH 2/2] feat: add high-performance HNSW RaBitQ index Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com> --- include/knowhere/comp/index_param.h | 1 + .../index/index_node_data_mock_wrapper.h | 5 + .../index/index_node_thread_pool_wrapper.h | 5 + include/knowhere/index/index_table.h | 4 + src/common/prometheus_client.cc | 1 + src/index/hnsw/faiss_hnsw.cc | 258 ++++- src/index/hnsw/faiss_hnsw_config.h | 54 + tests/faiss/CMakeLists.txt | 1 + tests/ut/test_hnsw_rabitq.cc | 602 +++++++++++ thirdparty/faiss/benchs/bench_rabitq_simd.cpp | 87 +- thirdparty/faiss/faiss/IndexRaBitQ.cpp | 21 +- thirdparty/faiss/faiss/IndexRaBitQ.h | 3 +- .../faiss/faiss/IndexRaBitQFastScan.cpp | 3 + .../cppcontrib/knowhere/IndexHNSWRaBitQ.cpp | 184 ++++ .../cppcontrib/knowhere/IndexHNSWRaBitQ.h | 83 ++ .../cppcontrib/knowhere/impl/index_read.cpp | 187 +++- .../cppcontrib/knowhere/impl/index_write.cpp | 156 ++- .../faiss/faiss/impl/RaBitQuantizer.cpp | 282 +++++- thirdparty/faiss/faiss/impl/RaBitQuantizer.h | 8 +- thirdparty/faiss/faiss/impl/index_read.cpp | 80 +- thirdparty/faiss/faiss/impl/index_write.cpp | 54 +- thirdparty/faiss/faiss/utils/rabitq_simd.h | 210 ++++ .../faiss/utils/simd_impl/rabitq_avx2.cpp | 358 ++++++- .../faiss/utils/simd_impl/rabitq_avx512.cpp | 958 +++++++++++++++++- .../utils/simd_impl/rabitq_avx512_spr.cpp | 112 ++ .../faiss/utils/simd_impl/rabitq_neon.cpp | 19 + .../faiss/utils/simd_impl/rabitq_rvv.cpp | 19 + thirdparty/faiss/faiss/utils/simd_levels.cpp | 22 + thirdparty/faiss/faiss/utils/simd_levels.h | 5 + thirdparty/faiss/perf_tests/CMakeLists.txt | 8 + thirdparty/faiss/tests/test_rabitq_simd.cpp | 174 ++++ 31 files changed, 3824 insertions(+), 140 deletions(-) create mode 100644 tests/ut/test_hnsw_rabitq.cc create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp create mode 100644 thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9d45d6cfb..bfa51331d 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..e39daa53d 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -9,10 +9,13 @@ // 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 @@ -672,31 +675,41 @@ convert_ds_to_float(const DataSetPtr& src, DataFormatEnum data_format) { } Status -add_to_index(faiss::Index* const __restrict index, const DataSetPtr& dataset, const DataFormatEnum data_format) { +add_to_index(faiss::Index* const __restrict index, const DataSetPtr& dataset, const DataFormatEnum data_format, + bool chunk_fp32 = false) { const auto* data = dataset->GetTensor(); const auto rows = dataset->GetRows(); const auto dim = dataset->GetDim(); - if (data_format == DataFormatEnum::fp32) { + if (data_format == DataFormatEnum::fp32 && !chunk_fp32) { // add as is index->add(rows, reinterpret_cast(data)); } else { // convert data into float in pieces and add to the index constexpr int64_t n_tmp_rows = 4096; - std::unique_ptr tmp = std::make_unique(n_tmp_rows * dim); + std::unique_ptr tmp; + if (data_format != DataFormatEnum::fp32) { + tmp = std::make_unique(n_tmp_rows * dim); + } for (int64_t irow = 0; irow < rows; irow += n_tmp_rows) { const int64_t start_row = irow; const int64_t end_row = std::min(rows, start_row + n_tmp_rows); const int64_t count_rows = end_row - start_row; - if (!convert_rows_to_fp32(data, tmp.get(), data_format, start_row, count_rows, dim)) { - LOG_KNOWHERE_ERROR_ << "Unsupported data format"; - return Status::invalid_args; + const float* chunk = nullptr; + if (data_format == DataFormatEnum::fp32) { + chunk = reinterpret_cast(data) + start_row * dim; + } else { + if (!convert_rows_to_fp32(data, tmp.get(), data_format, start_row, count_rows, dim)) { + LOG_KNOWHERE_ERROR_ << "Unsupported data format"; + return Status::invalid_args; + } + chunk = tmp.get(); } // add - index->add(count_rows, tmp.get()); + index->add(count_rows, chunk); } } @@ -3032,6 +3045,235 @@ 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: + BaseFaissRegularIndexHNSWRaBitQNode(const int32_t& version, const Object& object, DataFormatEnum data_format) + : BaseFaissRegularIndexHNSWNode(version, object, data_format) { + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + std::string + Type() const override { + return knowhere::IndexEnum::INDEX_HNSW_RABITQ; + } + + bool + IsAdditionalScalarSupported(bool) const override { + return false; + } + + bool + IsIndexRefineEnabled() const override { + return !indexes.empty() && std::all_of(indexes.begin(), indexes.end(), [](const auto& index) { + return index != nullptr && + dynamic_cast(index.get()) != nullptr; + }); + } + + protected: + std::vector> tmp_index_rabitq; + + Status + TrainInternal(const DataSetPtr dataset, const Config& cfg) override { + const auto rows = dataset->GetRows(); + const auto dim = dataset->GetDim(); + const auto& hnsw_cfg = static_cast(cfg); + + auto metric = Str2FaissMetricType(hnsw_cfg.metric_type.value()); + if (!metric.has_value() || + (metric.value() != faiss::METRIC_L2 && metric.value() != faiss::METRIC_INNER_PRODUCT)) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ only supports L2, IP and COSINE metrics"; + return Status::invalid_metric_type; + } + const bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); + const auto& scalar_info_map = + dataset->Get>>>(meta::SCALAR_INFO); + if (!scalar_info_map.empty()) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ does not support building with scalar info"; + return Status::invalid_args; + } + + auto float_ds_ptr = convert_ds_to_float(dataset, data_format); + if (float_ds_ptr == nullptr) { + LOG_KNOWHERE_ERROR_ << "Unsupported data format"; + return Status::invalid_args; + } + const auto* data = static_cast(float_ds_ptr->GetTensor()); + + try { + std::unique_ptr hnsw_index; + if (is_cosine) { + hnsw_index = + std::make_unique(dim, hnsw_cfg.M.value()); + } else { + hnsw_index = std::make_unique(dim, hnsw_cfg.M.value(), + metric.value()); + } + hnsw_index->hnsw.efConstruction = hnsw_cfg.efConstruction.value(); + + const auto rbq_bits = static_cast(hnsw_cfg.rbq_bits.value()); + auto rabitq_index = std::make_unique(dim, metric.value(), rbq_bits, rbq_bits > 1); + // HNSW_RABITQ V1 always uses an FP32 query. + rabitq_index->qb = 0; + 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"; + // IndexPreTransform materializes its transformed input. Bound that + // temporary allocation instead of rotating the full FP32 dataset + // in one call. + status = add_to_index(tmp_index_rabitq[0].get(), dataset, data_format, true); + if (status != Status::success) { + return status; + } + + faiss::cppcontrib::knowhere::IndexRefine* index_refine = + dynamic_cast(indexes[0].get()); + auto* index_hnsw = index_refine != nullptr + ? dynamic_cast(index_refine->base_index) + : dynamic_cast(indexes[0].get()); + if (index_hnsw == nullptr) { + LOG_KNOWHERE_ERROR_ << "HNSW_RABITQ build produced an unexpected base index"; + return Status::invalid_index_error; + } + + const bool is_cosine = faiss::cppcontrib::knowhere::is_cosine_index(index_hnsw->storage); + std::unique_ptr index_hnsw_rabitq; + if (is_cosine) { + index_hnsw_rabitq = std::make_unique(); + } else { + index_hnsw_rabitq = std::make_unique(); + } + // C++ slicing is intentional: preserve the exact graph while + // changing only the runtime HNSW type and its vector storage. + static_cast(*index_hnsw_rabitq) = + static_cast(*index_hnsw); + + // Validate the replacement before relinquishing either owner so a + // malformed storage cannot leave the exact graph half-finalized. + auto* flat_storage = index_hnsw->storage; + index_hnsw_rabitq->storage = tmp_index_rabitq[0].get(); + index_hnsw_rabitq->own_fields = false; + if (is_cosine) { + dynamic_cast(index_hnsw_rabitq.get()) + ->validate_cosine_storage(); + } else { + index_hnsw_rabitq->validate_storage(); + } + index_hnsw_rabitq->own_fields = true; + tmp_index_rabitq[0].release(); + index_hnsw->storage = nullptr; + delete flat_storage; + + if (index_refine != nullptr) { + delete index_refine->base_index; + index_refine->base_index = index_hnsw_rabitq.release(); + } else { + indexes[0] = std::move(index_hnsw_rabitq); + } + tmp_index_rabitq.clear(); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what(); + return Status::faiss_inner_error; + } + + return Status::success; + } +}; + +template +class BaseFaissRegularIndexHNSWRaBitQNodeTemplate : public BaseFaissRegularIndexHNSWRaBitQNode { + public: + BaseFaissRegularIndexHNSWRaBitQNodeTemplate(const int32_t& version, const Object& object) + : BaseFaissRegularIndexHNSWRaBitQNode(version, object, datatype_v) { + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + const auto& hnsw_cfg = static_cast(config); + return has_lossless_refine_index(hnsw_cfg.refine, hnsw_cfg.refine_type, datatype_v); + } +}; + // this index trains PRQ and HNSW+FLAT separately, then constructs HNSW+PRQ class BaseFaissRegularIndexHNSWPRQNode : public BaseFaissRegularIndexHNSWNode { public: @@ -3354,5 +3596,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..64d35c918 100644 --- a/src/index/hnsw/faiss_hnsw_config.h +++ b/src/index/hnsw/faiss_hnsw_config.h @@ -195,6 +195,60 @@ class FaissHnswPqConfig : public FaissHnswConfig { } }; +class FaissHnswRaBitQConfig : public FaissHnswConfig { + public: + // Number of bits per database vector dimension. + CFG_INT rbq_bits; + // Declared to reject an IVF_RABITQ query-bit parameter explicitly instead + // of silently ignoring the unknown JSON key. HNSW_RABITQ V1 uses FP32 + // queries only. + 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, 8) + .for_train() + .for_static(); + KNOWHERE_CONFIG_DECLARE_FIELD(rbq_bits_query) + .description("not supported on HNSW_RABITQ; must be 0 or omitted") + .set_default(0) + .set_range(0, 8) + .for_search() + .for_range_search(); + } + + 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::SEARCH || param_type == PARAM_TYPE::RANGE_SEARCH) && + rbq_bits_query.value_or(0) != 0) { + return HandleError(err_msg, "HNSW_RABITQ only supports FP32 queries (rbq_bits_query must be 0 or omitted)", + 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/tests/faiss/CMakeLists.txt b/tests/faiss/CMakeLists.txt index 270302429..6b1641665 100644 --- a/tests/faiss/CMakeLists.txt +++ b/tests/faiss/CMakeLists.txt @@ -7,6 +7,7 @@ set(FAISS_TEST_SRCS ../../thirdparty/faiss/tests/test_pq_code_distance.cpp ../../thirdparty/faiss/tests/test_cppcontrib_uintreader.cpp ../../thirdparty/faiss/tests/test_distances_simd.cpp + ../../thirdparty/faiss/tests/test_rabitq_simd.cpp ../../thirdparty/faiss/tests/test_hnsw.cpp ../../thirdparty/faiss/tests/test_binary_flat.cpp ../../thirdparty/faiss/tests/test_cppcontrib_sa_decode.cpp diff --git a/tests/ut/test_hnsw_rabitq.cc b/tests/ut/test_hnsw_rabitq.cc new file mode 100644 index 000000000..35bf32243 --- /dev/null +++ b/tests/ut/test_hnsw_rabitq.cc @@ -0,0 +1,602 @@ +// Copyright (C) 2019-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. + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "catch2/catch_approx.hpp" +#include "catch2/catch_test_macros.hpp" +#include "faiss/IndexRaBitQ.h" +#include "faiss/IndexRaBitQFastScan.h" +#include "faiss/VectorTransform.h" +#include "faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h" +#include "faiss/cppcontrib/knowhere/index_io.h" +#include "faiss/impl/FaissException.h" +#include "faiss/impl/RaBitQUtils.h" +#include "faiss/impl/io.h" +#include "faiss/index_io.h" +#include "knowhere/bitsetview.h" +#include "knowhere/comp/index_param.h" +#include "knowhere/comp/knowhere_check.h" +#include "knowhere/comp/knowhere_config.h" +#include "knowhere/index/index_factory.h" +#include "knowhere/utils.h" +#include "utils.h" + +namespace { + +constexpr int64_t kNb = 256; +constexpr int64_t kNq = 4; +constexpr int64_t kTopk = 8; + +knowhere::Json +MakeHnswRaBitQConfig(int64_t dim, const std::string& metric, int rbq_bits = 1) { + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::ROWS] = kNb; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = kTopk; + json[knowhere::indexparam::HNSW_M] = 12; + json[knowhere::indexparam::EFCONSTRUCTION] = 64; + json[knowhere::indexparam::EF] = 64; + json[knowhere::indexparam::RABITQ_BITS] = rbq_bits; + return json; +} + +bool +SerializedIndexContainsFourcc(const knowhere::BinarySet& binary_set, const std::string& name, const char* fourcc) { + const auto binary = binary_set.GetByName(name); + if (binary == nullptr) { + return false; + } + const auto* begin = binary->data.get(); + const auto* end = begin + binary->size; + return std::search(begin, end, fourcc, fourcc + 4) != end; +} + +void +CheckValidKnnResult(const knowhere::DataSet& result, int64_t nb, int64_t nq, int64_t topk) { + REQUIRE(result.GetRows() == nq); + REQUIRE(result.GetDim() == topk); + const auto* ids = result.GetIds(); + const auto* distances = result.GetDistance(); + for (int64_t i = 0; i < nq * topk; ++i) { + REQUIRE(ids[i] >= 0); + REQUIRE(ids[i] < nb); + REQUIRE(std::isfinite(distances[i])); + } +} + +void +CheckKnnOrder(const knowhere::DataSet& result, const std::string& metric) { + const auto rows = result.GetRows(); + const auto topk = result.GetDim(); + const auto* distances = result.GetDistance(); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 1; j < topk; ++j) { + const auto previous = distances[i * topk + j - 1]; + const auto current = distances[i * topk + j]; + if (knowhere::IsMetricType(metric, knowhere::metric::L2)) { + REQUIRE(previous <= current); + } else { + REQUIRE(previous >= current); + } + } + } +} + +template +void +CheckTypedBuildAndSearch(int64_t dim) { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + auto index = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version); + REQUIRE(index.has_value()); + + const auto train = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kNb, dim, 101)); + const auto query = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kNq, dim, 102)); + auto json = MakeHnswRaBitQConfig(dim, knowhere::metric::L2); + + REQUIRE(index.value().Build(train, json) == knowhere::Status::success); + const auto result = index.value().Search(query, json, nullptr); + REQUIRE(result.has_value()); + CheckValidKnnResult(*result.value(), kNb, kNq, kTopk); +} + +struct SearchSnapshot { + std::vector ids; + std::vector distances; +}; + +SearchSnapshot +TakeSnapshot(const knowhere::DataSet& result) { + const auto count = result.GetRows() * result.GetDim(); + return {std::vector(result.GetIds(), result.GetIds() + count), + std::vector(result.GetDistance(), result.GetDistance() + count)}; +} + +std::vector +ExactTopK(const float* base, const float* queries, int64_t nb, int64_t nq, int64_t dim, int64_t topk, + bool inner_product) { + std::vector result(nq * topk); + std::vector> scored(nb); + for (int64_t q = 0; q < nq; ++q) { + for (int64_t i = 0; i < nb; ++i) { + float value = 0.0f; + for (int64_t j = 0; j < dim; ++j) { + const float bv = base[i * dim + j]; + const float qv = queries[q * dim + j]; + value += inner_product ? bv * qv : (bv - qv) * (bv - qv); + } + scored[i] = {value, i}; + } + std::partial_sort(scored.begin(), scored.begin() + topk, scored.end(), + [inner_product](const auto& lhs, const auto& rhs) { + return inner_product ? lhs.first > rhs.first : lhs.first < rhs.first; + }); + for (int64_t k = 0; k < topk; ++k) { + result[q * topk + k] = scored[k].second; + } + } + return result; +} + +double +RecallAtK(const int64_t* actual, const std::vector& expected, int64_t nq, int64_t topk) { + int64_t matches = 0; + for (int64_t q = 0; q < nq; ++q) { + for (int64_t k = 0; k < topk; ++k) { + const auto begin = expected.begin() + q * topk; + const auto end = begin + topk; + matches += std::find(begin, end, actual[q * topk + k]) != end; + } + } + return static_cast(matches) / static_cast(nq * topk); +} + +} // namespace + +TEST_CASE("HNSW_RABITQ supports the public type and data-type boundary", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + const auto type = knowhere::IndexEnum::INDEX_HNSW_RABITQ; + auto& factory = knowhere::IndexFactory::Instance(); + + REQUIRE(factory.Create(type, version).has_value()); + REQUIRE(factory.Create(type, version).has_value()); + REQUIRE(factory.Create(type, version).has_value()); + REQUIRE_FALSE(factory.Create(type, version).has_value()); + REQUIRE_FALSE(factory.Create(type, version).has_value()); + + REQUIRE(factory.FeatureCheck(type, knowhere::feature::FLOAT32)); + REQUIRE(factory.FeatureCheck(type, knowhere::feature::FP16)); + REQUIRE(factory.FeatureCheck(type, knowhere::feature::BF16)); + REQUIRE_FALSE(factory.FeatureCheck(type, knowhere::feature::INT8)); + REQUIRE_FALSE(factory.FeatureCheck(type, knowhere::feature::BINARY)); + REQUIRE_FALSE(factory.FeatureCheck(type, knowhere::feature::MMAP)); + REQUIRE_FALSE(factory.FeatureCheck(type, knowhere::feature::MV)); + REQUIRE_FALSE(factory.FeatureCheck(type, knowhere::feature::EMB_LIST)); + + REQUIRE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_FLOAT)); + REQUIRE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_FLOAT16)); + REQUIRE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_BFLOAT16)); + REQUIRE_FALSE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_INT8)); + REQUIRE_FALSE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_BINARY)); + REQUIRE_FALSE(knowhere::KnowhereCheck::SupportMmapIndexTypeCheck(type)); + REQUIRE_FALSE(knowhere::KnowhereCheck::IndexTypeAndDataTypeCheck(type, knowhere::VecType::VECTOR_FLOAT, true)); + + auto fp32_index = factory.Create(type, version).value(); + REQUIRE_FALSE(fp32_index.IsAdditionalScalarSupported(false)); + REQUIRE_FALSE(fp32_index.IsAdditionalScalarSupported(true)); + + CheckTypedBuildAndSearch(13); + CheckTypedBuildAndSearch(13); + CheckTypedBuildAndSearch(13); +} + +TEST_CASE("HNSW_RABITQ validates metrics and bit widths", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + const auto train = GenDataSet(kNb, 32, 201); + + for (const int rbq_bits : {0, 9}) { + CAPTURE(rbq_bits); + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + const auto json = MakeHnswRaBitQConfig(32, knowhere::metric::L2, rbq_bits); + REQUIRE(index.Build(train, json) == knowhere::Status::out_of_range_in_json); + } + + { + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + const auto json = MakeHnswRaBitQConfig(32, knowhere::metric::JACCARD); + REQUIRE(index.Build(train, json) == knowhere::Status::invalid_metric_type); + } + + { + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + auto json = MakeHnswRaBitQConfig(32, knowhere::metric::L2, 4); + REQUIRE(index.Build(train, json) == knowhere::Status::success); + const auto query = GenDataSet(kNq, 32, 202); + json[knowhere::indexparam::RABITQ_QUERY_BITS] = 4; + REQUIRE_FALSE(index.Search(query, json, nullptr).has_value()); + } +} + +TEST_CASE("HNSW_RABITQ searches, ranges, iterates, and round-trips", "[hnsw_rabitq]") { + struct Scenario { + const char* metric; + int64_t dim; + int rbq_bits; + }; + const std::vector scenarios = { + {knowhere::metric::L2, 127, 1}, {knowhere::metric::IP, 129, 1}, {knowhere::metric::L2, 129, 2}, + {knowhere::metric::IP, 65, 3}, {knowhere::metric::L2, 64, 4}, {knowhere::metric::IP, 65, 5}, + {knowhere::metric::L2, 64, 6}, {knowhere::metric::COSINE, 65, 4}, {knowhere::metric::COSINE, 128, 8}, + {knowhere::metric::L2, 128, 8}, + }; + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + + for (const auto& scenario : scenarios) { + CAPTURE(scenario.metric, scenario.dim, scenario.rbq_bits); + const auto train = GenDataSet(kNb, scenario.dim, 301 + scenario.dim + scenario.rbq_bits); + const auto query = GenDataSet(kNq, scenario.dim, 401 + scenario.dim + scenario.rbq_bits); + auto json = MakeHnswRaBitQConfig(scenario.dim, scenario.metric, scenario.rbq_bits); + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(train, json) == knowhere::Status::success); + REQUIRE(index.Count() == kNb); + + const auto before = index.Search(query, json, nullptr); + REQUIRE(before.has_value()); + CheckValidKnnResult(*before.value(), kNb, kNq, kTopk); + CheckKnnOrder(*before.value(), scenario.metric); + + auto range_json = json; + if (knowhere::IsMetricType(scenario.metric, knowhere::metric::L2)) { + range_json[knowhere::meta::RADIUS] = std::numeric_limits::max(); + range_json[knowhere::meta::RANGE_FILTER] = 0.0f; + } else { + range_json[knowhere::meta::RADIUS] = -std::numeric_limits::max(); + range_json[knowhere::meta::RANGE_FILTER] = std::numeric_limits::max(); + } + range_json[knowhere::meta::RANGE_SEARCH_K] = 32; + const auto range_result = index.RangeSearch(query, range_json, nullptr); + REQUIRE(range_result.has_value()); + REQUIRE(range_result.value()->GetLims()[kNq] > 0); + + const auto iterators = index.AnnIterator(query, json, nullptr, false); + REQUIRE(iterators.has_value()); + REQUIRE(iterators.value().size() == kNq); + for (auto& iterator : iterators.value()) { + REQUIRE(iterator->HasNext().value()); + const auto next = iterator->Next(); + REQUIRE(next.has_value()); + REQUIRE(next.value().first >= 0); + REQUIRE(next.value().first < kNb); + REQUIRE(std::isfinite(next.value().second)); + } + + const auto* labels = before.value()->GetIds(); + const auto one_query = knowhere::GenDataSet(1, scenario.dim, query->GetTensor()); + const bool is_cosine = knowhere::IsMetricType(scenario.metric, knowhere::metric::COSINE); + const auto distances = index.CalcDistByIDs(one_query, nullptr, labels, kTopk, is_cosine); + REQUIRE(distances.has_value()); + for (int64_t i = 0; i < kTopk; ++i) { + REQUIRE(distances.value()->GetDistance()[i] == + Catch::Approx(before.value()->GetDistance()[i]).epsilon(1e-5)); + } + + knowhere::BinarySet binary_set; + REQUIRE(index.Serialize(binary_set) == knowhere::Status::success); + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), is_cosine ? "IHRC" : "IHRK")); + REQUIRE_FALSE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IHNr")); + if (is_cosine) { + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IRKC")); + } + const char* storage_fourcc = scenario.rbq_bits == 1 ? "Ixrq" : "Ixrd"; + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), storage_fourcc)); + + const auto serialized = binary_set.GetByName(index.Type()); + REQUIRE(serialized != nullptr); + faiss::VectorIOReader standard_reader; + standard_reader.data.assign(serialized->data.get(), serialized->data.get() + serialized->size); + std::unique_ptr standard_loaded(faiss::read_index(&standard_reader)); + REQUIRE(dynamic_cast(standard_loaded.get()) != nullptr); + if (is_cosine) { + REQUIRE(dynamic_cast(standard_loaded.get()) != + nullptr); + } + const auto* standard_storage = dynamic_cast( + dynamic_cast(standard_loaded.get())->storage); + REQUIRE(standard_storage != nullptr); + const auto* standard_rabitq = dynamic_cast(standard_storage->index); + REQUIRE(standard_rabitq != nullptr); + REQUIRE(standard_rabitq->rabitq.dense_layout == (scenario.rbq_bits > 1)); + + faiss::VectorIOWriter standard_writer; + faiss::write_index(standard_loaded.get(), &standard_writer); + faiss::VectorIOReader contrib_reader; + contrib_reader.data = standard_writer.data; + std::unique_ptr contrib_loaded(faiss::cppcontrib::knowhere::read_index(&contrib_reader)); + REQUIRE(dynamic_cast(contrib_loaded.get()) != nullptr); + + auto loaded = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(loaded.Deserialize(binary_set, json) == knowhere::Status::success); + const auto after = loaded.Search(query, json, nullptr); + REQUIRE(after.has_value()); + const auto before_snapshot = TakeSnapshot(*before.value()); + const auto after_snapshot = TakeSnapshot(*after.value()); + REQUIRE(after_snapshot.ids == before_snapshot.ids); + REQUIRE(after_snapshot.distances == before_snapshot.distances); + + REQUIRE(index.Add(GenDataSet(8, scenario.dim, 501), json) == knowhere::Status::not_implemented); + } +} + +TEST_CASE("HNSW_RABITQ agrees with an independent FP32 oracle", "[hnsw_rabitq]") { + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + for (const auto& [metric, dim, bits] : + std::vector>{{knowhere::metric::L2, 64, 1}, + {knowhere::metric::IP, 65, 1}, + {knowhere::metric::L2, 65, 4}, + {knowhere::metric::IP, 64, 8}}) { + CAPTURE(metric, dim, bits); + const auto train = GenDataSet(kNb, dim, 2101 + dim + bits); + const auto query = GenDataSet(kNq, dim, 2201 + dim + bits); + auto json = MakeHnswRaBitQConfig(dim, metric, bits); + json[knowhere::indexparam::EF] = kNb; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(train, json) == knowhere::Status::success); + const auto actual = index.Search(query, json, nullptr); + REQUIRE(actual.has_value()); + + const bool is_ip = knowhere::IsMetricType(metric, knowhere::metric::IP); + const auto expected = ExactTopK(static_cast(train->GetTensor()), + static_cast(query->GetTensor()), kNb, kNq, dim, kTopk, is_ip); + const double recall = RecallAtK(actual.value()->GetIds(), expected, kNq, kTopk); + REQUIRE(recall >= (bits == 1 ? 0.20 : 0.75)); + } +} + +TEST_CASE("RaBitQ dense layout preserves packed distances", "[hnsw_rabitq]") { + constexpr int64_t nb = 64; + constexpr int64_t nq = 2; + for (const int64_t dim : {13, 64, 65, 128, 129}) { + const auto base = GenDataSet(nb, dim, 1701 + dim); + const auto query = GenDataSet(nq, dim, 1801 + dim); + const auto* base_data = static_cast(base->GetTensor()); + const auto* query_data = static_cast(query->GetTensor()); + + for (const uint8_t bits : {2, 3, 4, 5, 6, 7, 8}) { + const size_t ex_bits = bits - 1; + for (const auto metric : {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}) { + faiss::IndexRaBitQ packed(dim, metric, bits, false); + faiss::IndexRaBitQ dense(dim, metric, bits, true); + packed.qb = 0; + REQUIRE(dense.qb == 0); + packed.train(nb, base_data); + dense.train(nb, base_data); + packed.add(nb, base_data); + dense.add(nb, base_data); + REQUIRE_THROWS(faiss::IndexRaBitQFastScan(dense)); + + for (int64_t i = 0; i < nb; ++i) { + const uint8_t* packed_code = packed.codes.data() + i * packed.code_size; + const uint8_t* dense_code = dense.codes.data() + i * dense.code_size; + const uint8_t* extra_code = + packed_code + (dim + 7) / 8 + sizeof(faiss::rabitq_utils::SignBitFactorsWithError); + for (int64_t j = 0; j < dim; ++j) { + const uint16_t sign = + (packed_code[j / 8] & (1u << (j % 8))) != 0 ? static_cast(1u << ex_bits) : 0; + const uint16_t extra = + static_cast(faiss::rabitq_utils::extract_code_inline(extra_code, j, ex_bits)); + REQUIRE(faiss::rabitq_utils::extract_code_inline(dense_code, j, bits) == + static_cast(sign | extra)); + } + } + + std::unique_ptr packed_dc(packed.get_distance_computer()); + std::unique_ptr dense_dc(dense.get_distance_computer()); + for (int64_t q = 0; q < nq; ++q) { + packed_dc->set_query(query_data + q * dim); + dense_dc->set_query(query_data + q * dim); + for (int64_t i = 0; i < nb; ++i) { + REQUIRE((*dense_dc)(i) == Catch::Approx((*packed_dc)(i)).epsilon(2e-5)); + } + for (int64_t i = 0; i + 3 < nb; i += 4) { + float batch[4]; + dense_dc->distances_batch_4(i, i + 1, i + 2, i + 3, batch[0], batch[1], batch[2], batch[3]); + for (int64_t j = 0; j < 4; ++j) { + REQUIRE(batch[j] == Catch::Approx((*packed_dc)(i + j)).epsilon(2e-5)); + } + } + } + + constexpr int64_t topk = 8; + std::vector packed_distances(nq * topk); + std::vector dense_distances(nq * topk); + std::vector packed_labels(nq * topk); + std::vector dense_labels(nq * topk); + packed.search(nq, query_data, topk, packed_distances.data(), packed_labels.data()); + dense.search(nq, query_data, topk, dense_distances.data(), dense_labels.data()); + REQUIRE(dense_labels == packed_labels); + for (size_t i = 0; i < dense_distances.size(); ++i) { + REQUIRE(dense_distances[i] == Catch::Approx(packed_distances[i]).epsilon(2e-5)); + } + } + } + } +} + +TEST_CASE("RaBitQ cosine storage applies post-norm correction", "[hnsw_rabitq]") { + constexpr int64_t dim = 13; + constexpr int64_t nb = 8; + const auto generated_base = GenDataSet(nb, dim, 1901); + auto query = GenDataSet(2, dim, 1902); + const auto* generated_base_data = static_cast(generated_base->GetTensor()); + std::vector base_values(generated_base_data, generated_base_data + nb * dim); + auto* base_data = base_values.data(); + const auto* query_data = static_cast(query->GetTensor()); + + for (int64_t i = 0; i < nb; ++i) { + const float scale = static_cast(i + 1); + for (int64_t j = 0; j < dim; ++j) { + base_data[i * dim + j] *= scale; + } + } + std::fill(base_data, base_data + dim, 0.0f); + + auto* rotation = new faiss::RandomRotationMatrix(dim, dim); + auto* rabitq = new faiss::IndexRaBitQ(dim, faiss::METRIC_INNER_PRODUCT, 4, true); + rabitq->qb = 0; + rabitq->centered = false; + faiss::cppcontrib::knowhere::IndexPreTransformRaBitQCosine storage(rotation, rabitq); + storage.own_fields = true; + storage.train(nb, base_data); + storage.add(nb, base_data); + storage.validate_norms(); + + std::unique_ptr raw_dc(storage.faiss::IndexPreTransform::get_distance_computer()); + std::unique_ptr cosine_dc(storage.get_distance_computer()); + const auto* inverse_norms = storage.get_inverse_l2_norms(); + + for (int64_t q = 0; q < 2; ++q) { + const float* current_query = query_data + q * dim; + float query_norm_sqr = 0.0f; + for (int64_t j = 0; j < dim; ++j) { + query_norm_sqr += current_query[j] * current_query[j]; + } + const float inverse_query_norm = query_norm_sqr > 0.0f ? 1.0f / std::sqrt(query_norm_sqr) : 1.0f; + raw_dc->set_query(current_query); + cosine_dc->set_query(current_query); + for (int64_t i = 0; i < nb; ++i) { + const float expected = (*raw_dc)(i)*inverse_norms[i] * inverse_query_norm; + REQUIRE((*cosine_dc)(i) == Catch::Approx(expected).epsilon(1e-6)); + } + + float raw_batch[4]; + float cosine_batch[4]; + raw_dc->distances_batch_4(0, 1, 2, 3, raw_batch[0], raw_batch[1], raw_batch[2], raw_batch[3]); + cosine_dc->distances_batch_4(0, 1, 2, 3, cosine_batch[0], cosine_batch[1], cosine_batch[2], cosine_batch[3]); + for (int64_t i = 0; i < 4; ++i) { + REQUIRE(cosine_batch[i] == + Catch::Approx(raw_batch[i] * inverse_norms[i] * inverse_query_norm).epsilon(1e-6)); + } + } +} + +TEST_CASE("HNSW_RABITQ supports optional refinement", "[hnsw_rabitq]") { + constexpr int64_t dim = 32; + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + const auto train = GenDataSet(kNb, dim, 701); + const auto query = GenDataSet(kNq, dim, 702); + auto json = MakeHnswRaBitQConfig(dim, knowhere::metric::L2); + json[knowhere::indexparam::REFINE] = true; + json[knowhere::indexparam::REFINE_TYPE] = "FLAT"; + json[knowhere::indexparam::REFINE_K] = 2.0f; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(train, json) == knowhere::Status::success); + REQUIRE(index.IsIndexRefineEnabled()); + REQUIRE(index.HasRawData(knowhere::metric::L2)); + + const auto result = index.Search(query, json, nullptr); + REQUIRE(result.has_value()); + CheckValidKnnResult(*result.value(), kNb, kNq, kTopk); + + knowhere::BinarySet binary_set; + REQUIRE(index.Serialize(binary_set) == knowhere::Status::success); + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IHRK")); + REQUIRE_FALSE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IHNr")); + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), "Ixrq")); + + auto loaded = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(loaded.Deserialize(binary_set, json) == knowhere::Status::success); + const auto loaded_result = loaded.Search(query, json, nullptr); + REQUIRE(loaded_result.has_value()); + REQUIRE(TakeSnapshot(*loaded_result.value()).ids == TakeSnapshot(*result.value()).ids); +} + +TEST_CASE("HNSW_RABITQ cosine refinement returns exact cosine", "[hnsw_rabitq]") { + constexpr int64_t dim = 32; + const auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); + const auto train = GenDataSet(kNb, dim, 2001); + const auto query = GenDataSet(kNq, dim, 2002); + auto json = MakeHnswRaBitQConfig(dim, knowhere::metric::COSINE, 4); + json[knowhere::indexparam::REFINE] = true; + json[knowhere::indexparam::REFINE_TYPE] = "FLAT"; + json[knowhere::indexparam::REFINE_K] = 2.0f; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(index.Build(train, json) == knowhere::Status::success); + const auto result = index.Search(query, json, nullptr); + REQUIRE(result.has_value()); + CheckValidKnnResult(*result.value(), kNb, kNq, kTopk); + CheckKnnOrder(*result.value(), knowhere::metric::COSINE); + + const auto* base_data = static_cast(train->GetTensor()); + const auto* query_data = static_cast(query->GetTensor()); + for (int64_t q = 0; q < kNq; ++q) { + for (int64_t k = 0; k < kTopk; ++k) { + const int64_t id = result.value()->GetIds()[q * kTopk + k]; + float dot = 0.0f; + float base_norm_sqr = 0.0f; + float query_norm_sqr = 0.0f; + for (int64_t j = 0; j < dim; ++j) { + const float bv = base_data[id * dim + j]; + const float qv = query_data[q * dim + j]; + dot += bv * qv; + base_norm_sqr += bv * bv; + query_norm_sqr += qv * qv; + } + const float expected = (base_norm_sqr == 0.0f || query_norm_sqr == 0.0f) + ? 0.0f + : dot / std::sqrt(base_norm_sqr * query_norm_sqr); + REQUIRE(result.value()->GetDistance()[q * kTopk + k] == Catch::Approx(expected).epsilon(2e-5)); + } + } + + knowhere::BinarySet binary_set; + REQUIRE(index.Serialize(binary_set) == knowhere::Status::success); + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IHRC")); + REQUIRE(SerializedIndexContainsFourcc(binary_set, index.Type(), "IRKC")); + auto loaded = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_RABITQ, version) + .value(); + REQUIRE(loaded.Deserialize(binary_set, json) == knowhere::Status::success); + const auto loaded_result = loaded.Search(query, json, nullptr); + REQUIRE(loaded_result.has_value()); + REQUIRE(TakeSnapshot(*loaded_result.value()).ids == TakeSnapshot(*result.value()).ids); + REQUIRE(TakeSnapshot(*loaded_result.value()).distances == TakeSnapshot(*result.value()).distances); +} diff --git a/thirdparty/faiss/benchs/bench_rabitq_simd.cpp b/thirdparty/faiss/benchs/bench_rabitq_simd.cpp index 01b136c70..4506249a1 100644 --- a/thirdparty/faiss/benchs/bench_rabitq_simd.cpp +++ b/thirdparty/faiss/benchs/bench_rabitq_simd.cpp @@ -9,9 +9,22 @@ #include #include #include +#include namespace faiss { +template +bool check_simd_available(benchmark::State& state) { + if constexpr (SL != SIMDLevel::NONE) { + if (!SIMDConfig::is_simd_level_available(SL)) { + state.SkipWithError( + "requested SIMD level is unavailable on this CPU"); + return false; + } + } + return true; +} + const auto& randomData() { static auto data = [] { AlignedTable x(10 << 20); // 10 MiB @@ -96,6 +109,9 @@ void bench_rabitq_and_dot_product_with_sum_fused(benchmark::State& state) { template void bench_rabitq_rearrange_impl(benchmark::State& state) { + if (!check_simd_available(state)) { + return; + } size_t qb = state.range(0); size_t d = state.range(1); size_t offset = (d + 7) / 8; @@ -126,8 +142,57 @@ void bench_rabitq_rearrange_avx2(benchmark::State& state) { } #endif +template +void bench_rabitq_selected_float_sum(benchmark::State& state) { + if (!check_simd_available(state)) { + return; + } + const size_t d = state.range(0); + AlignedTable sign_bits((d + 7) / 8); + AlignedTable values(d); + byte_rand(sign_bits.data(), sign_bits.size(), 721); + float_rand(values.data(), values.size(), 812); + + for (auto _ : state) { + float result = rabitq::selected_float_sum( + sign_bits.data(), values.data(), d); + benchmark::DoNotOptimize(result); + } + state.SetItemsProcessed(state.iterations()); +} + +template +void bench_rabitq_multibit_inner_product(benchmark::State& state) { + if (!check_simd_available(state)) { + return; + } + const size_t ex_bits = state.range(0); + const size_t d = state.range(1); + AlignedTable sign_bits((d + 7) / 8); + // Padding keeps the benchmark input valid for kernels that load a complete + // machine word at the end of a bit-packed code. + AlignedTable extra_code((d * ex_bits + 7) / 8 + 8); + AlignedTable query(d); + byte_rand(sign_bits.data(), sign_bits.size(), 913); + byte_rand(extra_code.data(), extra_code.size(), 114); + float_rand(query.data(), query.size(), 315); + + for (auto _ : state) { + float result = rabitq::multibit::compute_inner_product( + sign_bits.data(), + extra_code.data(), + query.data(), + d, + ex_bits, + -3.25f); + benchmark::DoNotOptimize(result); + } + state.SetItemsProcessed(state.iterations()); +} + const std::vector qbs{1, 2, 4, 8}; const std::vector dims{64, 100, 256, 512, 1000, 1024, 3072}; +const std::vector scorer_dims{128, 768, 1537}; BENCHMARK(bench_rabitq_sum)->ArgsProduct({{0}, dims})->ArgNames({"qb", "d"}); BENCHMARK(bench_rabitq_and_dot_product) @@ -149,7 +214,27 @@ BENCHMARK(bench_rabitq_rearrange_scalar) BENCHMARK(bench_rabitq_rearrange_avx2) ->ArgsProduct({qbs, dims}) ->ArgNames({"qb", "d"}); + +BENCHMARK_TEMPLATE(bench_rabitq_selected_float_sum, SIMDLevel::NONE) + ->ArgsProduct({scorer_dims}) + ->ArgNames({"d"}); +BENCHMARK_TEMPLATE(bench_rabitq_selected_float_sum, SIMDLevel::AVX2) + ->ArgsProduct({scorer_dims}) + ->ArgNames({"d"}); +BENCHMARK_TEMPLATE(bench_rabitq_selected_float_sum, SIMDLevel::AVX512) + ->ArgsProduct({scorer_dims}) + ->ArgNames({"d"}); +BENCHMARK_TEMPLATE(bench_rabitq_multibit_inner_product, SIMDLevel::NONE) + ->ArgsProduct({{1, 3, 7}, scorer_dims}) + ->ArgNames({"extra_bits", "d"}); +BENCHMARK_TEMPLATE(bench_rabitq_multibit_inner_product, SIMDLevel::AVX2) + ->ArgsProduct({{1, 3, 7}, scorer_dims}) + ->ArgNames({"extra_bits", "d"}); +BENCHMARK_TEMPLATE(bench_rabitq_multibit_inner_product, SIMDLevel::AVX512) + ->ArgsProduct({{1, 3, 7}, scorer_dims}) + ->ArgNames({"extra_bits", "d"}); #endif -BENCHMARK_MAIN(); } // namespace faiss + +BENCHMARK_MAIN(); diff --git a/thirdparty/faiss/faiss/IndexRaBitQ.cpp b/thirdparty/faiss/faiss/IndexRaBitQ.cpp index a41133fa4..e40ec5d0f 100644 --- a/thirdparty/faiss/faiss/IndexRaBitQ.cpp +++ b/thirdparty/faiss/faiss/IndexRaBitQ.cpp @@ -21,11 +21,23 @@ using rabitq_utils::SignBitFactorsWithError; IndexRaBitQ::IndexRaBitQ() = default; -IndexRaBitQ::IndexRaBitQ(idx_t d_in, MetricType metric, uint8_t nb_bits_in) - : IndexFlatCodes(0, d_in, metric), rabitq(d_in, metric, nb_bits_in) { +IndexRaBitQ::IndexRaBitQ( + idx_t d_in, + MetricType metric, + uint8_t nb_bits_in, + bool dense_layout) + : IndexFlatCodes(0, d_in, metric), + rabitq(d_in, metric, nb_bits_in, dense_layout) { // Update code size based on nb_bits code_size = rabitq.code_size; + // Dense codes currently use the FP32-query distance computer. Keep the + // constructor valid on its own instead of requiring every caller to + // override the legacy qb=4 default before the first search. + if (dense_layout) { + qb = 0; + } + is_trained = false; } @@ -137,7 +149,10 @@ struct Run_search_with_dc_res { float est_distance = dc->distance_to_code_1bit(code); - size_t code_size_base = (index->d + 7) / 8; + const size_t code_size_base = + index->rabitq.dense_layout + ? (index->d * index->rabitq.nb_bits + 7) / 8 + : (index->d + 7) / 8; const rabitq_utils::SignBitFactorsWithError* base_fac = reinterpret_cast< const rabitq_utils:: diff --git a/thirdparty/faiss/faiss/IndexRaBitQ.h b/thirdparty/faiss/faiss/IndexRaBitQ.h index e6a48a6e0..4d75ec49f 100644 --- a/thirdparty/faiss/faiss/IndexRaBitQ.h +++ b/thirdparty/faiss/faiss/IndexRaBitQ.h @@ -37,7 +37,8 @@ struct IndexRaBitQ : IndexFlatCodes { explicit IndexRaBitQ( idx_t d, MetricType metric = METRIC_L2, - uint8_t nb_bits = 1); + uint8_t nb_bits = 1, + bool dense_layout = false); void train(idx_t n, const float* x) override; diff --git a/thirdparty/faiss/faiss/IndexRaBitQFastScan.cpp b/thirdparty/faiss/faiss/IndexRaBitQFastScan.cpp index 8dacbd6a5..3831ccf6f 100644 --- a/thirdparty/faiss/faiss/IndexRaBitQFastScan.cpp +++ b/thirdparty/faiss/faiss/IndexRaBitQFastScan.cpp @@ -111,6 +111,9 @@ IndexRaBitQFastScan::IndexRaBitQFastScan(const IndexRaBitQ& orig, int bbs_in) orig.metric_type == METRIC_L2 || orig.metric_type == METRIC_INNER_PRODUCT, "RaBitQ FastScan only supports L2 and Inner Product metrics"); + FAISS_THROW_IF_NOT_MSG( + !orig.rabitq.dense_layout, + "RaBitQ FastScan conversion does not support dense codes"); // RaBitQ uses 1 bit per dimension packed into 4-bit FastScan sub-quantizers // Each FastScan sub-quantizer handles 4 RaBitQ dimensions diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp new file mode 100644 index 000000000..d9c24311a --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.cpp @@ -0,0 +1,184 @@ +/* + * 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 + +namespace faiss::cppcontrib::knowhere { + +IndexPreTransformRaBitQCosine::IndexPreTransformRaBitQCosine() = default; + +IndexPreTransformRaBitQCosine::IndexPreTransformRaBitQCosine( + faiss::VectorTransform* transform, + faiss::Index* index_in) + : faiss::IndexPreTransform(transform, index_in) {} + +void IndexPreTransformRaBitQCosine::add(idx_t n, const float* x) { + faiss::IndexPreTransform::add(n, x); + inverse_norms_storage.add(x, n, d); +} + +void IndexPreTransformRaBitQCosine::reset() { + faiss::IndexPreTransform::reset(); + inverse_norms_storage.reset(); +} + +faiss::DistanceComputer* IndexPreTransformRaBitQCosine::get_distance_computer() + const { + FAISS_THROW_IF_NOT_MSG( + inverse_norms_storage.inverse_l2_norms.size() == + static_cast(ntotal), + "cosine RaBitQ inverse norm count must match ntotal"); + return new WithCosineNormDistanceComputer( + get_inverse_l2_norms(), + d, + std::unique_ptr( + faiss::IndexPreTransform::get_distance_computer())); +} + +const float* IndexPreTransformRaBitQCosine::get_inverse_l2_norms() const { + return inverse_norms_storage.inverse_l2_norms.data(); +} + +void IndexPreTransformRaBitQCosine::validate_norms() const { + FAISS_THROW_IF_NOT_MSG( + inverse_norms_storage.inverse_l2_norms.size() == + static_cast(ntotal), + "cosine RaBitQ inverse norm count must match ntotal"); + for (const float inverse_norm : inverse_norms_storage.inverse_l2_norms) { + FAISS_THROW_IF_NOT_MSG( + std::isfinite(inverse_norm) && inverse_norm > 0.0f, + "cosine RaBitQ inverse norms must be finite and positive"); + } +} + +IndexHNSWRaBitQ::IndexHNSWRaBitQ() = default; + +void IndexHNSWRaBitQ::add(idx_t, const float*) { + FAISS_THROW_MSG( + "IndexHNSWRaBitQ does not support incremental add: build the " + "HNSW graph with exact storage before attaching RaBitQ storage"); +} + +const faiss::IndexPreTransform* IndexHNSWRaBitQ::pretransform_index() const { + return dynamic_cast(storage); +} + +const faiss::IndexRaBitQ* IndexHNSWRaBitQ::rabitq_index() const { + const auto* pretransform = pretransform_index(); + return pretransform + ? dynamic_cast(pretransform->index) + : nullptr; +} + +void IndexHNSWRaBitQ::validate_storage() const { + FAISS_THROW_IF_NOT_MSG( + metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT, + "IndexHNSWRaBitQ only supports L2 and inner product metrics"); + FAISS_THROW_IF_NOT_MSG( + storage != nullptr, "IndexHNSWRaBitQ requires non-null storage"); + + const auto* pretransform = pretransform_index(); + FAISS_THROW_IF_NOT_MSG( + pretransform != nullptr, + "IndexHNSWRaBitQ storage must be IndexPreTransform"); + FAISS_THROW_IF_NOT_MSG( + pretransform->chain.size() == 1, + "IndexHNSWRaBitQ storage must contain exactly one transform"); + + const auto* rotation = dynamic_cast( + pretransform->chain[0]); + FAISS_THROW_IF_NOT_MSG( + rotation != nullptr, + "IndexHNSWRaBitQ transform must be RandomRotationMatrix"); + + const auto* rabitq = rabitq_index(); + FAISS_THROW_IF_NOT_MSG( + rabitq != nullptr, + "IndexHNSWRaBitQ pretransform leaf must be IndexRaBitQ"); + + FAISS_THROW_IF_NOT_MSG( + d == pretransform->d && metric_type == pretransform->metric_type && + ntotal == pretransform->ntotal, + "IndexHNSWRaBitQ outer index and pretransform metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + is_trained && pretransform->is_trained && rotation->is_trained && + rabitq->is_trained, + "IndexHNSWRaBitQ requires fully trained storage"); + FAISS_THROW_IF_NOT_MSG( + pretransform->index != nullptr && + pretransform->ntotal == rabitq->ntotal && + pretransform->metric_type == rabitq->metric_type, + "IndexHNSWRaBitQ pretransform and RaBitQ metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + rotation->d_in == d && rotation->d_out == rabitq->d && + rotation->d_in == rotation->d_out, + "IndexHNSWRaBitQ requires a square rotation matching index dimensions"); + FAISS_THROW_IF_NOT_MSG( + rotation->is_orthonormal && !rotation->have_bias && + rotation->b.empty() && + rotation->A.size() == + static_cast(rotation->d_in) * + rotation->d_out, + "IndexHNSWRaBitQ rotation matrix has invalid storage"); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.d == static_cast(rabitq->d) && + rabitq->rabitq.metric_type == rabitq->metric_type, + "IndexHNSWRaBitQ RaBitQ quantizer metadata mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.nb_bits >= 1 && rabitq->rabitq.nb_bits <= 9, + "IndexHNSWRaBitQ RaBitQ nb_bits must be in [1, 9]"); + + const size_t expected_code_size = + rabitq->rabitq.compute_code_size(rabitq->d, rabitq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.code_size == expected_code_size && + rabitq->code_size == expected_code_size, + "IndexHNSWRaBitQ RaBitQ code size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->codes.size() == + static_cast(rabitq->ntotal) * expected_code_size, + "IndexHNSWRaBitQ RaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->center.size() == static_cast(rabitq->d), + "IndexHNSWRaBitQ RaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_MSG( + rabitq->qb <= 8, "IndexHNSWRaBitQ RaBitQ qb must be in [0, 8]"); + FAISS_THROW_IF_NOT_MSG( + rabitq->rabitq.nb_bits == 1 || rabitq->qb == 0, + "IndexHNSWRaBitQ requires qb=0 when nb_bits > 1"); + FAISS_THROW_IF_NOT_MSG( + !rabitq->centered, "IndexHNSWRaBitQ V1 requires centered=false"); +} + +IndexHNSWRaBitQCosine::IndexHNSWRaBitQCosine() = default; + +const float* IndexHNSWRaBitQCosine::get_inverse_l2_norms() const { + const auto* cosine_storage = + dynamic_cast(storage); + return cosine_storage ? cosine_storage->get_inverse_l2_norms() : nullptr; +} + +void IndexHNSWRaBitQCosine::validate_cosine_storage() const { + validate_storage(); + const auto* cosine_storage = + dynamic_cast(storage); + FAISS_THROW_IF_NOT_MSG( + cosine_storage != nullptr, + "IndexHNSWRaBitQCosine requires cosine-aware pretransform storage"); + FAISS_THROW_IF_NOT_MSG( + metric_type == METRIC_INNER_PRODUCT, + "IndexHNSWRaBitQCosine requires inner product storage"); + cosine_storage->validate_norms(); +} + +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h new file mode 100644 index 000000000..39f864396 --- /dev/null +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/IndexHNSWRaBitQ.h @@ -0,0 +1,83 @@ +/* + * 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[] = "IHRK"; +inline constexpr char kHnswRaBitQCosineFourcc[] = "IHRC"; +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; + + /** Validate the complete runtime/storage shape and serialized invariants. + * Throws FaissException on malformed state. */ + void validate_storage() const; +}; + +/** Cosine runtime marker for HNSW backed by cosine-aware RaBitQ storage. */ +struct IndexHNSWRaBitQCosine : IndexHNSWRaBitQ, HasInverseL2Norms { + IndexHNSWRaBitQCosine(); + + const float* get_inverse_l2_norms() const override; + void validate_cosine_storage() const; +}; + +} // namespace faiss::cppcontrib::knowhere diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp index ceaf05193..20c3013f6 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_read.cpp @@ -15,36 +15,38 @@ #include #include +#include #include #include #include -#include #include #include #include -#include #include -#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include +#include #include #include #include #include #include -#include -#include #include #include #include #include -#include -#include +#include #include #include @@ -58,11 +60,8 @@ #include #include - - namespace faiss::cppcontrib::knowhere { - uint32_t read_value(IOReader* f) { uint32_t h; READ1(h) @@ -216,7 +215,10 @@ void read_xb_vector(VectorT& target, IOReader* f) { * Read **************************************************************/ -static void read_index_header(Index* idx, IOReader* f, bool* is_cosine_out = nullptr) { +static void read_index_header( + Index* idx, + IOReader* f, + bool* is_cosine_out = nullptr) { READ1(idx->d); READ1(idx->ntotal); @@ -362,24 +364,27 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { "read_InvertedLists:" " WARN! inverted lists not stored with IVF object\n"); return nullptr; - } else if (h == fourcc ("iloa") && !(io_flags & IO_FLAG_MMAP)) { + } else if (h == fourcc("iloa") && !(io_flags & IO_FLAG_MMAP)) { size_t nlist; size_t code_size; - std::vector list_length; + std::vector list_length; READ1(nlist); READ1(code_size); READVECTOR(list_length); - auto ails = new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); + auto ails = + new ReadOnlyArrayInvertedLists(nlist, code_size, list_length); size_t n; READ1(n); #ifdef USE_GPU - ails->pin_readonly_ids = std::make_shared(n * sizeof(idx_t)); - ails->pin_readonly_codes = std::make_shared(n * code_size * sizeof(uint8_t)); - READANDCHECK((idx_t *) ails->pin_readonly_ids->data, n); - READANDCHECK((uint8_t *) ails->pin_readonly_codes->data, n * code_size); + ails->pin_readonly_ids = + std::make_shared(n * sizeof(idx_t)); + ails->pin_readonly_codes = std::make_shared( + n * code_size * sizeof(uint8_t)); + READANDCHECK((idx_t*)ails->pin_readonly_ids->data, n); + READANDCHECK((uint8_t*)ails->pin_readonly_codes->data, n * code_size); #else ails->readonly_ids.resize(n); - ails->readonly_codes.resize(n*code_size); + ails->readonly_codes.resize(n * code_size); READANDCHECK(ails->readonly_ids.data(), n); READANDCHECK(ails->readonly_codes.data(), n * code_size); #endif @@ -391,7 +396,8 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { READ1(segment_size); bool save_norm = io_flags & IO_FLAG_WITH_NORM; - auto lca = new ConcurrentArrayInvertedLists(nlist, code_size, segment_size, save_norm); + auto lca = new ConcurrentArrayInvertedLists( + nlist, code_size, segment_size, save_norm); std::vector sizes(nlist); read_ArrayInvertedLists_sizes(f, sizes); for (size_t i = 0; i < lca->nlist; i++) { @@ -402,12 +408,15 @@ ::faiss::InvertedLists* read_InvertedLists(IOReader* f, int io_flags) { if (n > 0) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { - size_t seg_size = lca->get_segment_size(i , j); + size_t seg_size = lca->get_segment_size(i, j); size_t seg_off = lca->get_segment_offset(i, j); - READANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); + READANDCHECK( + lca->codes[i][j].data_.data(), + seg_size * lca->code_size); READANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (save_norm) { - READANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); + READANDCHECK( + lca->code_norms[i][j].data_.data(), seg_size); } } } @@ -585,9 +594,7 @@ static void read_ProductLocalSearchQuantizer( } } -static void read_ScalarQuantizer( - ::faiss::ScalarQuantizer* ivsc, - IOReader* f) { +static void read_ScalarQuantizer(::faiss::ScalarQuantizer* ivsc, IOReader* f) { READ1(ivsc->qtype); READ1(ivsc->rangestat); READ1(ivsc->rangestat_arg); @@ -685,6 +692,44 @@ 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]"); + FAISS_THROW_IF_NOT_MSG( + !idxq->rabitq.dense_layout || idxq->rabitq.nb_bits > 1, + "IndexRaBitQ dense layout requires nb_bits>1"); + + const size_t expected_code_size = + idxq->rabitq.compute_code_size(idxq->d, idxq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.code_size == expected_code_size, + "IndexRaBitQ quantizer code size mismatch"); + idxq->code_size = expected_code_size; + FAISS_THROW_IF_NOT_MSG( + idxq->codes.size() == + static_cast(idxq->ntotal) * expected_code_size, + "IndexRaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->center.empty() || + idxq->center.size() == static_cast(idxq->d), + "IndexRaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_FMT( + idxq->qb <= 8, + "invalid RaBitQ qb=%d (must be in [0, 8])", + idxq->qb); + // The V1 cppcontrib wire format intentionally has no centered field. + idxq->centered = false; +} + static void read_direct_map(DirectMap* dm, IOReader* f) { char maintain_direct_map; READ1(maintain_direct_map); @@ -699,11 +744,11 @@ static void read_direct_map(DirectMap* dm, IOReader* f) { map[it.first] = it.second; } } - // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` - // read branch is gone — see the symmetric comment in index_write.cpp. - // Old files (if any) with `type == 3` would fail to round-trip here - // since the enum value no longer exists; in practice CC indexes - // were never written through this path. + // Path-D step 10.9: the former `if (dm->type == + // DirectMap::ConcurrentArray)` read branch is gone — see the symmetric + // comment in index_write.cpp. Old files (if any) with `type == 3` would + // fail to round-trip here since the enum value no longer exists; in + // practice CC indexes were never written through this path. } static void read_ivf_header( @@ -791,7 +836,8 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(wire_l2_norms); // reconstruct inverse norms from wire L2 norms - idxf->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_l2_norms); + idxf->inverse_norms_storage = + L2NormsStorage::from_l2_norms(wire_l2_norms); FAISS_THROW_IF_NOT( idxf->codes.size() == idxf->ntotal * idxf->code_size); @@ -828,7 +874,8 @@ Index* read_index(IOReader* f, int io_flags) { idxfc->code_size = idxf->code_size; idxfc->codes = std::move(idxf->codes); // reconstruct inverse norms from wire L2 norms - idxfc->inverse_norms_storage = L2NormsStorage::from_l2_norms(wire_code_norms); + idxfc->inverse_norms_storage = + L2NormsStorage::from_l2_norms(wire_code_norms); delete idxf; idxf = idxfc; } @@ -849,7 +896,7 @@ Index* read_index(IOReader* f, int io_flags) { READVECTOR(idxp->inverse_norms_storage.inverse_l2_norms); if (!(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { - idxp->pq.compute_sdc_table (); + idxp->pq.compute_sdc_table(); } idx = idxp; @@ -875,7 +922,7 @@ Index* read_index(IOReader* f, int io_flags) { // the following "if" block is Knowhere-specific if (h == fourcc("IxPq")) { - idxp->pq.compute_sdc_table (); + idxp->pq.compute_sdc_table(); } idx = idxp; @@ -1070,8 +1117,7 @@ Index* read_index(IOReader* f, int io_flags) { // either enum name, and route legacy data to // IndexBinaryScalarQuantizer. const int legacy_qt_1bit_direct_marker = 9; - if (static_cast(idxs->sq.qtype) == - legacy_qt_1bit_direct_marker) { + if (static_cast(idxs->sq.qtype) == legacy_qt_1bit_direct_marker) { IndexBinaryScalarQuantizer* bsq = new IndexBinaryScalarQuantizer( static_cast(idxs->d), idxs->metric_type); bsq->ntotal = idxs->ntotal; @@ -1082,6 +1128,22 @@ Index* read_index(IOReader* f, int io_flags) { } else { idx = idxs; } + } else if ( + h == fourcc("Ixrq") || h == fourcc("Ixrr") || h == fourcc("Ixrb") || + h == fourcc("Ixrd")) { + auto idxq = std::make_unique<::faiss::IndexRaBitQ>(); + read_index_header(idxq.get(), f); + read_RaBitQuantizer( + &idxq->rabitq, f, /*multi_bit=*/h != fourcc("Ixrq")); + idxq->rabitq.dense_layout = h == fourcc("Ixrb") || h == fourcc("Ixrd"); + FAISS_THROW_IF_NOT_MSG( + h != fourcc("Ixrb") || idxq->rabitq.nb_bits == 8, + "legacy IndexRaBitQ byte layout requires nb_bits=8"); + 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,6 +1170,21 @@ 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* ixpt = new IndexPreTransformRaBitQCosine(); + 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 = ixpt; } else if (h == fourcc("IxPT")) { IndexPreTransform* ixpt = new IndexPreTransform(); ixpt->own_fields = true; @@ -1151,7 +1228,8 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxrf->k_factor); if (dynamic_cast<::faiss::IndexFlat*>(idxrf->refine_index)) { // then make a RefineFlat with it. Refine index may be a baseline - // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware subclass. + // ::faiss::IndexFlat{,IP,L2} or the knowhere Jaccard-aware + // subclass. IndexRefine* idxrf_old = idxrf; idxrf = new IndexRefineFlat(); *idxrf = *idxrf_old; @@ -1161,11 +1239,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 +1287,13 @@ Index* read_index(IOReader* f, int io_flags) { read_HNSW(&idxhnsw->hnsw, f); idxhnsw->storage = read_index(f, io_flags); idxhnsw->own_fields = idxhnsw->storage != nullptr; + if (h == fourcc(kHnswRaBitQFourcc)) { + dynamic_cast(idxhnsw)->validate_storage(); + } + if (h == fourcc(kHnswRaBitQCosineFourcc)) { + dynamic_cast(idxhnsw) + ->validate_cosine_storage(); + } if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) { dynamic_cast(idxhnsw->storage)->pq.compute_sdc_table(); } @@ -1222,7 +1318,7 @@ Index* read_index(IOReader* f, int io_flags) { delete idxhnsw; idxhnsw = newh; } - idx = idxhnsw; + idx = idxhnsw_rabitq_owner ? idxhnsw_rabitq_owner.release() : idxhnsw; } else if (h == fourcc("IwPf")) { ::faiss::IndexIVFPQFastScan* ivpq = new ::faiss::IndexIVFPQFastScan(); read_ivf_header(ivpq, f); @@ -1333,7 +1429,8 @@ Index* read_index(IOReader* f, int io_flags) { // field); Iwrr is baseline multi-bit and does serialize nb_bits. auto ivrq = new IndexIVFRaBitQ(); read_ivf_header(ivrq, f); - read_RaBitQuantizer(&ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); + read_RaBitQuantizer( + &ivrq->rabitq, f, /*multi_bit=*/h == fourcc("Iwrr")); READ1(ivrq->code_size); READ1(ivrq->by_residual); READ1(ivrq->qb); @@ -1484,4 +1581,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..ef66e1b85 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/impl/index_write.cpp @@ -18,30 +18,32 @@ #include #include +#include #include #include -#include #include -#include #include -#include +#include +#include +#include +#include +#include +#include #include #include #include #include #include +#include #include #include #include -#include -#include #include #include #include #include -#include -#include +#include #include #include @@ -69,11 +71,8 @@ * leak memory. **************************************************************/ - - namespace faiss::cppcontrib::knowhere { - /************************************************************* * Write **************************************************************/ @@ -94,7 +93,7 @@ static void write_index_header(const Index* idx, IOWriter* f) { WRITE1(dummy32); idx_t dummy = 0; WRITE1(dummy); - + WRITE1(idx->is_trained); WRITE1(idx->metric_type); if (idx->metric_type > 1) { @@ -305,8 +304,9 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { } } } - } else if (const auto & lca = - dynamic_cast(ils)) { + } else if ( + const auto& lca = + dynamic_cast(ils)) { uint32_t h = fourcc("ilca"); WRITE1(h); WRITE1(lca->nlist); @@ -350,16 +350,20 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t seg_num = lca->get_segment_num(i); for (size_t j = 0; j < seg_num; j++) { size_t seg_size = lca->get_segment_size(i, j); - WRITEANDCHECK(lca->codes[i][j].data_.data(), seg_size * lca->code_size); + WRITEANDCHECK( + lca->codes[i][j].data_.data(), + seg_size * lca->code_size); WRITEANDCHECK(lca->ids[i][j].data_.data(), seg_size); if (lca->save_norm) { - WRITEANDCHECK(lca->code_norms[i][j].data_.data(), seg_size); + WRITEANDCHECK( + lca->code_norms[i][j].data_.data(), seg_size); } } } } - } else if (const auto & oa = - dynamic_cast(ils)) { + } else if ( + const auto& oa = + dynamic_cast(ils)) { uint32_t h = fourcc("iloa"); WRITE1(h); WRITE1(oa->nlist); @@ -369,16 +373,16 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { size_t n = oa->pin_readonly_ids->size() / sizeof(InvertedLists::idx_t); WRITE1(n); WRITEANDCHECK((InvertedLists::idx_t*)oa->pin_readonly_ids->data, n); - WRITEANDCHECK((uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); + WRITEANDCHECK( + (uint8_t*)oa->pin_readonly_codes->data, n * oa->code_size); #else size_t n = oa->readonly_ids.size(); WRITE1(n); WRITEANDCHECK(oa->readonly_ids.data(), n); WRITEANDCHECK(oa->readonly_codes.data(), n * oa->code_size); #endif - } else if (const auto & od = - dynamic_cast(ils)) { - uint32_t h = fourcc ("ilod"); + } else if (const auto& od = dynamic_cast(ils)) { + uint32_t h = fourcc("ilod"); WRITE1(h); WRITE1(ils->nlist); WRITE1(ils->code_size); @@ -387,7 +391,7 @@ void write_InvertedLists(const ::faiss::InvertedLists* ils, IOWriter* f) { { std::vector v( - od->slots.begin(), od->slots.end()); + od->slots.begin(), od->slots.end()); WRITEVECTOR(v); } { @@ -502,6 +506,41 @@ 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]"); + FAISS_THROW_IF_NOT_MSG( + !idxq->rabitq.dense_layout || idxq->rabitq.nb_bits > 1, + "IndexRaBitQ dense layout requires nb_bits>1"); + const size_t expected_code_size = + idxq->rabitq.compute_code_size(idxq->d, idxq->rabitq.nb_bits); + FAISS_THROW_IF_NOT_MSG( + idxq->rabitq.code_size == expected_code_size && + idxq->code_size == expected_code_size, + "IndexRaBitQ code size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->codes.size() == + static_cast(idxq->ntotal) * expected_code_size, + "IndexRaBitQ codes size mismatch"); + FAISS_THROW_IF_NOT_MSG( + idxq->center.empty() || + idxq->center.size() == static_cast(idxq->d), + "IndexRaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_MSG(idxq->qb <= 8, "IndexRaBitQ qb must be in [0, 8]"); + FAISS_THROW_IF_NOT_MSG( + !idxq->centered, + "cppcontrib IndexRaBitQ V1 serialization requires centered=false"); +} + static void write_direct_map(const DirectMap* dm, IOWriter* f) { char maintain_direct_map = (char)dm->type; // for backwards compatibility with bool @@ -514,11 +553,12 @@ static void write_direct_map(const DirectMap* dm, IOWriter* f) { std::copy(map.begin(), map.end(), v.begin()); WRITEVECTOR(v); } - // Path-D step 10.9: the former `if (dm->type == DirectMap::ConcurrentArray)` - // write branch is gone — fork DirectMap no longer supports that - // variant. CC indexes now carry their own `cc_direct_map` member - // (ConcurrentDirectMap) which is not serialized through this path - // (CC indexes have no serialize stage; see ivf.cc:619 comment). + // Path-D step 10.9: the former `if (dm->type == + // DirectMap::ConcurrentArray)` write branch is gone — fork DirectMap no + // longer supports that variant. CC indexes now carry their own + // `cc_direct_map` member (ConcurrentDirectMap) which is not serialized + // through this path (CC indexes have no serialize stage; see ivf.cc:619 + // comment). } static void write_ivf_header(const IndexIVF* ivf, IOWriter* f) { @@ -536,13 +576,15 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { // eg. for a storage component of HNSW that is set to nullptr uint32_t h = fourcc("null"); WRITE1(h); - } else if (const IndexFlatCosine* idxf = dynamic_cast(idx)) { + } else if ( + const IndexFlatCosine* idxf = + dynamic_cast(idx)) { uint32_t h = fourcc("IxF9"); WRITE1(h); write_index_header(idx, f); WRITEXBVECTOR(idxf->codes); // we're storing real l2 norms, because of - // backward compatibility issues. + // backward compatibility issues. WRITEVECTOR(idxf->inverse_norms_storage.as_l2_norms()); } else if ( const ::faiss::IndexFlat* idxf = @@ -556,7 +598,9 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITE1(h); write_index_header(idx, f); WRITEXBVECTOR(idxf->codes); - } else if (const IndexPQCosine* idxp = dynamic_cast(idx)) { + } else if ( + const IndexPQCosine* idxp = + dynamic_cast(idx)) { uint32_t h = fourcc("IxP7"); WRITE1(h); write_index_header(idx, f); @@ -598,7 +642,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { WRITEVECTOR(idxr_2->codes); } else if ( const IndexProductResidualQuantizerCosine* idxpr = - dynamic_cast(idx)) { + dynamic_cast( + idx)) { uint32_t h = fourcc("IxP5"); WRITE1(h); write_index_header(idx, f); @@ -710,6 +755,20 @@ 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 = idxq->rabitq.dense_layout + ? fourcc("Ixrd") + : (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 +801,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 +852,23 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index(idxrf->refine_index, f); WRITE1(idxrf->k_factor); } else if (const IndexHNSW* idxhnsw = dynamic_cast(idx)) { - uint32_t h = dynamic_cast(idx) ? fourcc("IHNf") - : dynamic_cast(idx) ? fourcc("IHNp") + const auto* hnsw_rabitq_cosine = + dynamic_cast(idx); + const auto* hnsw_rabitq = dynamic_cast(idx); + if (hnsw_rabitq) { + FAISS_THROW_IF_NOT_MSG( + !(io_flags & IO_FLAG_SKIP_STORAGE), + "IndexHNSWRaBitQ cannot be serialized without its RaBitQ storage"); + if (hnsw_rabitq_cosine) { + hnsw_rabitq_cosine->validate_cosine_storage(); + } else { + hnsw_rabitq->validate_storage(); + } + } + uint32_t h = hnsw_rabitq_cosine ? fourcc(kHnswRaBitQCosineFourcc) + : hnsw_rabitq ? fourcc(kHnswRaBitQFourcc) + : dynamic_cast(idx) ? fourcc("IHNf") + : dynamic_cast(idx) ? fourcc("IHNp") // IndexHNSWBinary reuses the legacy IHNs fourcc so // on-disk bytes match what IndexHNSWSQ(QT_1bit_direct, // metric) used to produce. Readers dispatch to @@ -888,7 +975,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 +1053,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..4e46f7f29 100644 --- a/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp +++ b/thirdparty/faiss/faiss/impl/RaBitQuantizer.cpp @@ -33,12 +33,15 @@ using rabitq_utils::SignBitFactorsWithError; RaBitQuantizer::RaBitQuantizer( size_t d_in, MetricType metric, - size_t nb_bits_in) + size_t nb_bits_in, + bool dense_layout_in) : Quantizer(d_in, 0), // code_size will be set below metric_type{metric}, - nb_bits{nb_bits_in} { + nb_bits{nb_bits_in}, + dense_layout{dense_layout_in} { // Validate nb_bits range FAISS_THROW_IF_NOT(nb_bits >= 1 && nb_bits <= 9); + FAISS_THROW_IF_NOT(!dense_layout || nb_bits > 1); // Set code_size using compute_code_size code_size = compute_code_size(d, nb_bits); @@ -50,6 +53,13 @@ size_t RaBitQuantizer::compute_code_size(size_t d_in, size_t num_bits) const { size_t ex_bits = num_bits - 1; + if (dense_layout) { + FAISS_THROW_IF_NOT(num_bits > 1); + return (d_in * num_bits + 7) / 8 + + sizeof(SignBitFactorsWithError) + + sizeof(ExtraBitsFactors); + } + // Base: 1-bit codes + base factors // Layout for 1-bit: [binary_code: (d+7)/8 bytes][SignBitFactors: 8 bytes] // base_factors = or_minus_c_l2sqr (4) + dp_multiplier (4) @@ -113,6 +123,9 @@ void RaBitQuantizer::compute_codes_core( // 12 bytes] // [ex_code: (d*ex_bits+7)/8 bytes][ex_factors: 8 bytes] uint8_t* binary_code = code; + const size_t base_code_size = dense_layout + ? (d * nb_bits + 7) / 8 + : (d + 7) / 8; // Step 1: Compute 1-bit quantization and base factors // Store residual for potential ex-bits quantization @@ -134,7 +147,7 @@ void RaBitQuantizer::compute_codes_core( // For multi-bit: write full SignBitFactorsWithError (12 bytes) SignBitFactorsWithError* full_factors = reinterpret_cast( - code + (d + 7) / 8); + code + base_code_size); *full_factors = factors_data; } @@ -149,7 +162,7 @@ void RaBitQuantizer::compute_codes_core( const bool xb = (or_minus_c > 0.0f); // Store the 1-bit sign code - if (xb) { + if (xb && !dense_layout) { rabitq_utils::set_bit_standard(binary_code, j); } } @@ -158,10 +171,18 @@ void RaBitQuantizer::compute_codes_core( if (ex_bits > 0) { // Pointer to ex-bit code section uint8_t* ex_code = - code + (d + 7) / 8 + sizeof(SignBitFactorsWithError); + code + base_code_size + sizeof(SignBitFactorsWithError); + std::vector packed_ex_code; + if (dense_layout) { + packed_ex_code.resize((d * ex_bits + 7) / 8); + ex_code = packed_ex_code.data(); + } // Pointer to ex-factors section - ExtraBitsFactors* ex_factors = reinterpret_cast( - ex_code + (d * ex_bits + 7) / 8); + ExtraBitsFactors byte_ex_factors; + ExtraBitsFactors* ex_factors = dense_layout + ? &byte_ex_factors + : reinterpret_cast( + ex_code + (d * ex_bits + 7) / 8); // Quantize residual to ex-bits (pass centroid for IP metric) rabitq_multibit::quantize_ex_bits( @@ -172,6 +193,31 @@ void RaBitQuantizer::compute_codes_core( *ex_factors, metric_type, centroid_in); + + if (dense_layout) { + for (size_t j = 0; j < d; j++) { + const uint16_t sign = residual[j] > 0.0f + ? static_cast(1u << ex_bits) + : 0; + const uint16_t value = sign | static_cast( + rabitq_utils::extract_code_inline( + ex_code, j, ex_bits)); + const size_t bit_pos = j * nb_bits; + const size_t byte_pos = bit_pos / 8; + const size_t shift = bit_pos % 8; + const uint32_t shifted = + static_cast(value) << shift; + const size_t nbytes = (shift + nb_bits + 7) / 8; + for (size_t b = 0; b < nbytes; b++) { + binary_code[byte_pos + b] |= + static_cast(shifted >> (8 * b)); + } + } + memcpy(code + base_code_size + + sizeof(SignBitFactorsWithError), + ex_factors, + sizeof(ExtraBitsFactors)); + } } } } @@ -204,10 +250,13 @@ void RaBitQuantizer::decode_core( // For 1-bit: use SignBitFactors (8 bytes) // For multi-bit: use SignBitFactorsWithError (12 bytes, but only first // 8 bytes used for decode) + const size_t base_code_size = dense_layout + ? (d * nb_bits + 7) / 8 + : (d + 7) / 8; const SignBitFactors* fac = (ex_bits == 0) ? reinterpret_cast(code + (d + 7) / 8) : reinterpret_cast( - code + (d + 7) / 8); + code + base_code_size); // this is the baseline code // @@ -215,7 +264,14 @@ void RaBitQuantizer::decode_core( for (size_t j = 0; j < d; j++) { // extract i-th bit const uint8_t masker = (1 << (j % 8)); - const float bit = ((binary_data[j / 8] & masker) == masker) ? 1 : 0; + const float bit = dense_layout + ? ((rabitq_utils::extract_code_inline( + binary_data, j, nb_bits) & + (1u << ex_bits)) != 0 + ? 1.0f + : 0.0f) + : (((binary_data[j / 8] & masker) == masker) ? 1.0f + : 0.0f); // compute the output code x[i * d + j] = (bit - 0.5f) * fac->dp_multiplier * 2 * inv_d_sqrt + @@ -226,6 +282,65 @@ void RaBitQuantizer::decode_core( namespace { +template +void distance_to_code_full_batch_4_impl( + const uint8_t* const codes[4], + size_t d, + size_t nb_bits, + const float* rotated_q, + float qr_base, + MetricType metric_type, + bool dense_layout, + float out[4]) { + const size_t ex_bits = nb_bits - 1; + if (ex_bits == 0) { + FAISS_THROW_MSG("multi-bit batch helper requires extra bits"); + } + + const size_t code_size_base = dense_layout + ? (d * nb_bits + 7) / 8 + : (d + 7) / 8; + const size_t ex_offset = + code_size_base + sizeof(SignBitFactorsWithError); + const size_t ex_code_size = (d * ex_bits + 7) / 8; + const uint8_t* sign_bits[4]; + const uint8_t* ex_codes[4]; + const ExtraBitsFactors* ex_factors[4]; + for (size_t i = 0; i < 4; i++) { + sign_bits[i] = codes[i]; + ex_codes[i] = dense_layout ? codes[i] : codes[i] + ex_offset; + ex_factors[i] = reinterpret_cast( + dense_layout + ? codes[i] + code_size_base + + sizeof(SignBitFactorsWithError) + : ex_codes[i] + ex_code_size); + } + + const float cb = -(static_cast(1 << ex_bits) - 0.5f); + float inner_products[4]; + if (dense_layout) { + rabitq::multibit::compute_inner_product_dense_batch_4( + ex_codes, rotated_q, d, nb_bits, cb, inner_products); + } else { + rabitq::multibit::compute_inner_product_batch_4( + sign_bits, + ex_codes, + rotated_q, + d, + ex_bits, + cb, + inner_products); + } + + for (size_t i = 0; i < 4; i++) { + float distance = qr_base + ex_factors[i]->f_add_ex + + ex_factors[i]->f_rescale_ex * inner_products[i]; + out[i] = metric_type == MetricType::METRIC_L2 + ? std::max(0.0f, distance) + : distance; + } +} + // Distance computers templatized on SIMDLevel to avoid per-call dynamic // dispatch. The SIMDLevel is baked in at construction time via // get_distance_computer, so virtual calls through the base class go @@ -237,6 +352,7 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { std::vector rotated_q; // some additional numbers for the query QueryFactorsData query_fac; + bool dense_layout = false; RaBitQDistanceComputerNotQ() = default; @@ -247,24 +363,11 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { // this is the baseline code // // compute using floats - float dot_qo = 0; - // It was a willful decision (after the discussion) to not to pre-cache - // the sum of all bits, just in order to reduce the overhead per - // vector. - uint64_t sum_q = 0; - - for (size_t i = 0; i < d; i++) { - // Extract i-th bit - bool bit = rabitq_utils::extract_bit_standard(binary_data, i); - // accumulate dp - dot_qo += bit ? rotated_q[i] : 0; - // accumulate sum-of-bits - sum_q += bit ? 1 : 0; - } + const float dot_qo = rabitq::selected_float_sum( + binary_data, rotated_q.data(), d); // Apply query factors - float final_dot = - query_fac.c1 * dot_qo + query_fac.c2 * sum_q - query_fac.c34; + float final_dot = query_fac.c1 * dot_qo - query_fac.c34; // pre_dist = ||or - c||^2 + ||qr - c||^2 - // 2 * ||or - c|| * ||qr - c|| * - (IP ? ||or||^2 : 0) @@ -288,13 +391,33 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { metric_type == MetricType::METRIC_INNER_PRODUCT)); FAISS_ASSERT(rotated_q.size() == d); - const size_t code_size_base = (d + 7) / 8; + const size_t code_size_base = dense_layout + ? (d * nb_bits + 7) / 8 + : (d + 7) / 8; const size_t ex_bits = nb_bits - 1; const SignBitFactors* base_fac = (ex_bits == 0) ? reinterpret_cast(code + code_size_base) : reinterpret_cast( code + code_size_base); - return distance_to_code_1bit_impl(code, base_fac); + if (!dense_layout) { + return distance_to_code_1bit_impl(code, base_fac); + } + + float dot_qo = 0.0f; + for (size_t i = 0; i < d; i++) { + if ((rabitq_utils::extract_code_inline( + code, i, nb_bits) & + (1u << (nb_bits - 1))) != 0) { + dot_qo += rotated_q[i]; + } + } + const float final_dot = query_fac.c1 * dot_qo - query_fac.c34; + const float pre_dist = base_fac->or_minus_c_l2sqr + + query_fac.qr_to_c_L2sqr - + 2 * base_fac->dp_multiplier * final_dot; + return metric_type == MetricType::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) @@ -314,15 +437,32 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { // Extract pointers to code sections const uint8_t* binary_data = code; - size_t offset = (d + 7) / 8 + sizeof(SignBitFactorsWithError); - const uint8_t* ex_code = code + offset; + const size_t dense_code_size = (d * nb_bits + 7) / 8; + size_t offset = (dense_layout ? dense_code_size : (d + 7) / 8) + + sizeof(SignBitFactorsWithError); + const uint8_t* ex_code = dense_layout ? code : code + offset; const ExtraBitsFactors* ex_fac = reinterpret_cast( - ex_code + (d * ex_bits + 7) / 8); + dense_layout + ? code + dense_code_size + + sizeof(SignBitFactorsWithError) + : ex_code + (d * ex_bits + 7) / 8); float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT) ? query_fac.q_dot_c : query_fac.qr_to_c_L2sqr; + if (dense_layout) { + const float cb = + -(static_cast(1 << ex_bits) - 0.5f); + const float ex_ip = + rabitq::multibit::compute_inner_product_dense( + ex_code, rotated_q.data(), d, nb_bits, cb); + const float distance = qr_base + ex_fac->f_add_ex + + ex_fac->f_rescale_ex * ex_ip; + return metric_type == MetricType::METRIC_L2 + ? std::max(0.0f, distance) + : distance; + } return rabitq_utils::compute_full_multibit_distance( binary_data, ex_code, @@ -334,6 +474,42 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { metric_type); } + void distance_to_code_batch_4( + const uint8_t* code0, + const uint8_t* code1, + const uint8_t* code2, + const uint8_t* code3, + float& dis0, + float& dis1, + float& dis2, + float& dis3) final { + if (nb_bits == 1) { + dis0 = distance_to_code_full(code0); + dis1 = distance_to_code_full(code1); + dis2 = distance_to_code_full(code2); + dis3 = distance_to_code_full(code3); + return; + } + const uint8_t* codes[4] = {code0, code1, code2, code3}; + float distances[4]; + const float qr_base = metric_type == MetricType::METRIC_INNER_PRODUCT + ? query_fac.q_dot_c + : query_fac.qr_to_c_L2sqr; + distance_to_code_full_batch_4_impl( + codes, + d, + nb_bits, + rotated_q.data(), + qr_base, + metric_type, + dense_layout, + distances); + dis0 = distances[0]; + dis1 = distances[1]; + dis2 = distances[2]; + dis3 = distances[3]; + } + void set_query(const float* x) final { q = x; FAISS_ASSERT(x != nullptr); @@ -386,7 +562,9 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { const IDSelector* sel, bool keep_max, ResultHandler& handler) final { - const size_t code_size_base = (d + 7) / 8; + const size_t code_size_base = dense_layout + ? (d * nb_bits + 7) / 8 + : (d + 7) / 8; const size_t ex_bits = nb_bits - 1; FAISS_ASSERT(ex_bits > 0); @@ -404,7 +582,7 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { reinterpret_cast( codes + code_size_base); const float est_distance = - distance_to_code_1bit_impl(codes, base_fac); + distance_to_code_1bit(codes); const bool should_refine = rabitq_utils::should_refine_candidate( est_distance, @@ -547,6 +725,42 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer { metric_type); } + void distance_to_code_batch_4( + const uint8_t* code0, + const uint8_t* code1, + const uint8_t* code2, + const uint8_t* code3, + float& dis0, + float& dis1, + float& dis2, + float& dis3) final { + if (nb_bits == 1) { + dis0 = distance_to_code_full(code0); + dis1 = distance_to_code_full(code1); + dis2 = distance_to_code_full(code2); + dis3 = distance_to_code_full(code3); + return; + } + const uint8_t* codes[4] = {code0, code1, code2, code3}; + float distances[4]; + const float qr_base = metric_type == MetricType::METRIC_INNER_PRODUCT + ? query_fac.q_dot_c + : query_fac.qr_to_c_L2sqr; + distance_to_code_full_batch_4_impl( + codes, + d, + nb_bits, + rotated_q.data(), + qr_base, + metric_type, + false, + distances); + dis0 = distances[0]; + dis1 = distances[1]; + dis2 = distances[2]; + dis3 = distances[3]; + } + void set_query(const float* x) final { q = x; FAISS_ASSERT(x != nullptr); @@ -667,9 +881,13 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer( dc->d = d; dc->centroid = centroid_in; dc->nb_bits = nb_bits; + dc->dense_layout = dense_layout; return dc.release(); } else { + FAISS_THROW_IF_NOT_MSG( + !dense_layout, + "dense-layout RaBitQ currently requires qb=0"); auto dc = std::make_unique>(); dc->metric_type = metric_type; dc->d = d; diff --git a/thirdparty/faiss/faiss/impl/RaBitQuantizer.h b/thirdparty/faiss/faiss/impl/RaBitQuantizer.h index 7766df35c..bb25c33f8 100644 --- a/thirdparty/faiss/faiss/impl/RaBitQuantizer.h +++ b/thirdparty/faiss/faiss/impl/RaBitQuantizer.h @@ -47,10 +47,16 @@ struct RaBitQuantizer : Quantizer { // - nb_bits = 2-9: multi-bit RaBitQ (1 sign bit + ex_bits extra bits) size_t nb_bits = 1; + // Store one complete scalar code per dimension in a dense n-bit stream. + // This combines the sign bit and extra bits at build time while retaining + // the exact nbits/dimension budget. For nbits=8 this is one byte/dimension. + bool dense_layout = false; + RaBitQuantizer( size_t d = 0, MetricType metric = MetricType::METRIC_L2, - size_t nb_bits = 1); + size_t nb_bits = 1, + bool dense_layout = false); // Compute code size based on dimensionality and number of bits // Returns: size in bytes for one encoded vector diff --git a/thirdparty/faiss/faiss/impl/index_read.cpp b/thirdparty/faiss/faiss/impl/index_read.cpp index 29151ef50..92d79ba5a 100644 --- a/thirdparty/faiss/faiss/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/impl/index_read.cpp @@ -57,6 +57,7 @@ #include #include #include +#include #ifdef FAISS_ENABLE_SVS #include #include @@ -1305,6 +1306,19 @@ static void read_HNSW(HNSW& hnsw, IOReader* f) { validate_HNSW(hnsw); } +static void read_knowhere_HNSW(cppcontrib::knowhere::HNSW& hnsw, IOReader* f) { + READVECTOR(hnsw.assign_probas); + READVECTOR(hnsw.cum_nneighbor_per_level); + READVECTOR(hnsw.levels); + READVECTOR(hnsw.offsets); + read_vector(hnsw.neighbors, f); + READ1(hnsw.entry_point); + READ1(hnsw.max_level); + READ1(hnsw.efConstruction); + READ1(hnsw.efSearch); + READ1(hnsw.upper_beam); +} + static void read_NSG(NSG& nsg, IOReader* f) { READ1(nsg.ntotal); READ1(nsg.R); @@ -1420,7 +1434,8 @@ static void read_RaBitQuantizer( RaBitQuantizer& rabitq, IOReader* f, int expected_d, - bool multi_bit = true) { + bool multi_bit = true, + bool dense_layout = false) { READ1(rabitq.d); READ1(rabitq.code_size); int metric_type_int; @@ -1432,12 +1447,21 @@ static void read_RaBitQuantizer( } else { rabitq.nb_bits = 1; } + rabitq.dense_layout = dense_layout; FAISS_THROW_IF_NOT_FMT( rabitq.d == static_cast(expected_d), "RaBitQuantizer dimension mismatch: rabitq.d=%zu vs index d=%d", rabitq.d, expected_d); + + const size_t expected_code_size = + rabitq.compute_code_size(rabitq.d, rabitq.nb_bits); + FAISS_THROW_IF_NOT_FMT( + rabitq.code_size == expected_code_size, + "RaBitQuantizer code size mismatch: stored=%zu expected=%zu", + rabitq.code_size, + expected_code_size); } static void read_EDENScalarQuantizer( @@ -2351,6 +2375,24 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { READ1(index_ivfpq->use_precomputed_table); } idx = std::move(indep); + } else if ( + h == + fourcc(cppcontrib::knowhere::kRaBitQPreTransformCosineFourcc)) { + auto ixpt = std::make_unique< + cppcontrib::knowhere::IndexPreTransformRaBitQCosine>(); + 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 = std::move(ixpt); } else if (h == fourcc("IxPT")) { auto ixpt = std::make_unique(); ixpt->own_fields = true; @@ -2494,6 +2536,30 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxp->code_size, "Index2Layer codes")); idx = std::move(idxp); + } else if ( + h == fourcc(cppcontrib::knowhere::kHnswRaBitQFourcc) || + h == fourcc(cppcontrib::knowhere::kHnswRaBitQCosineFourcc)) { + const bool is_cosine = + h == fourcc(cppcontrib::knowhere::kHnswRaBitQCosineFourcc); + std::unique_ptr idxhnsw; + if (is_cosine) { + idxhnsw = std::make_unique< + cppcontrib::knowhere::IndexHNSWRaBitQCosine>(); + } else { + idxhnsw = std::make_unique(); + } + read_index_header(*idxhnsw, f); + read_knowhere_HNSW(idxhnsw->hnsw, f); + idxhnsw->storage = read_index(f, io_flags); + idxhnsw->own_fields = idxhnsw->storage != nullptr; + if (is_cosine) { + dynamic_cast( + idxhnsw.get()) + ->validate_cosine_storage(); + } else { + idxhnsw->validate_storage(); + } + idx = std::move(idxhnsw); } else if ( h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") || h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHc2") || @@ -2851,12 +2917,12 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { // rabitq.nb_bits is already set to 1 by read_RaBitQuantizer idxq->code_size = idxq->rabitq.code_size; idx = std::move(idxq); - } else if (h == fourcc("Ixrr")) { - // Ixrr = multi-bit format (new) + } else if (h == fourcc("Ixrr") || h == fourcc("Ixrd")) { + // Ixrr = split multi-bit format; Ixrd = dense multi-bit format. auto idxq = std::make_unique(); read_index_header(*idxq, f); read_RaBitQuantizer( - idxq->rabitq, f, idxq->d, true); // Reads nb_bits from file + idxq->rabitq, f, idxq->d, true, h == fourcc("Ixrd")); READVECTOR(idxq->codes); READVECTOR(idxq->center); READ1(idxq->qb); @@ -2868,6 +2934,12 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxq->qb); idxq->code_size = idxq->rabitq.code_size; + FAISS_THROW_IF_NOT_FMT( + idxq->codes.size() == + static_cast(idxq->ntotal) * idxq->code_size, + "IndexRaBitQ codes size mismatch: stored=%zu expected=%zu", + idxq->codes.size(), + static_cast(idxq->ntotal) * idxq->code_size); idx = std::move(idxq); } else if (h == fourcc("Iwrq")) { auto ivrq = std::make_unique(); diff --git a/thirdparty/faiss/faiss/impl/index_write.cpp b/thirdparty/faiss/faiss/impl/index_write.cpp index 873fc0f9c..f0def66ee 100644 --- a/thirdparty/faiss/faiss/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/impl/index_write.cpp @@ -52,6 +52,7 @@ #include #include #include +#include #ifdef FAISS_ENABLE_SVS #include #include @@ -372,6 +373,21 @@ static void write_HNSW(const HNSW* hnsw, IOWriter* f) { WRITE1(tmp_upper_beam); } +static void write_knowhere_HNSW( + const cppcontrib::knowhere::HNSW* hnsw, + IOWriter* f) { + WRITEVECTOR(hnsw->assign_probas); + WRITEVECTOR(hnsw->cum_nneighbor_per_level); + WRITEVECTOR(hnsw->levels); + WRITEVECTOR(hnsw->offsets); + WRITEVECTOR(hnsw->neighbors); + WRITE1(hnsw->entry_point); + WRITE1(hnsw->max_level); + WRITE1(hnsw->efConstruction); + WRITE1(hnsw->efSearch); + WRITE1(hnsw->upper_beam); +} + static void write_NSG(const NSG* nsg, IOWriter* f) { WRITE1(nsg->ntotal); WRITE1(nsg->R); @@ -827,6 +843,22 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { if (auto index_ivfpq = dynamic_cast(indep->index_ivf)) { WRITE1(index_ivfpq->use_precomputed_table); } + } else if ( + const auto* cosine_rabitq = dynamic_cast< + const cppcontrib::knowhere::IndexPreTransformRaBitQCosine*>( + idx)) { + cosine_rabitq->validate_norms(); + uint32_t h = + fourcc(cppcontrib::knowhere::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, io_flags); + WRITEVECTOR(cosine_rabitq->inverse_norms_storage.inverse_l2_norms); } else if ( const IndexPreTransform* ixpt = dynamic_cast(idx)) { @@ -865,6 +897,25 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index_header(idxmap, f); write_index(idxmap->index, f); WRITEVECTOR(idxmap->id_map); + } else if ( + const auto* idxhnsw = + dynamic_cast( + idx)) { + FAISS_THROW_IF_NOT_MSG( + !(io_flags & IO_FLAG_SKIP_STORAGE), + "IndexHNSWRaBitQ cannot be serialized without its RaBitQ storage"); + idxhnsw->validate_storage(); + const bool is_cosine = + dynamic_cast< + const cppcontrib::knowhere::IndexHNSWRaBitQCosine*>( + idxhnsw) != nullptr; + uint32_t h = + fourcc(is_cosine ? cppcontrib::knowhere::kHnswRaBitQCosineFourcc + : cppcontrib::knowhere::kHnswRaBitQFourcc); + WRITE1(h); + write_index_header(idxhnsw, f); + write_knowhere_HNSW(&idxhnsw->hnsw, f); + write_index(idxhnsw->storage, f, io_flags); } else if (const IndexHNSW* idxhnsw = dynamic_cast(idx)) { uint32_t h = dynamic_cast(idx) ? fourcc("IHfP") @@ -1024,7 +1075,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { write_index_header(idx, f); write_RaBitQuantizer(&idxq->rabitq, f, false); } else { - uint32_t h = fourcc("Ixrr"); // multi-bit (new format) + uint32_t h = + idxq->rabitq.dense_layout ? fourcc("Ixrd") : fourcc("Ixrr"); WRITE1(h); write_index_header(idx, f); write_RaBitQuantizer(&idxq->rabitq, f, true); diff --git a/thirdparty/faiss/faiss/utils/rabitq_simd.h b/thirdparty/faiss/faiss/utils/rabitq_simd.h index 5dac1d60a..1f90c60ab 100644 --- a/thirdparty/faiss/faiss/utils/rabitq_simd.h +++ b/thirdparty/faiss/faiss/utils/rabitq_simd.h @@ -79,6 +79,13 @@ uint64_t bitwise_xor_dot_product( template uint64_t popcount(const uint8_t* data, size_t size); +/** Sum float values selected by one packed sign bit per dimension. */ +template +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d); + /** * Rearrange per-dimension quantized query codes into bit-plane layout. * @@ -225,6 +232,20 @@ inline uint64_t popcount(const uint8_t* data, size_t size) { return sum; } +template <> +inline float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + float sum = 0.0f; + for (size_t i = 0; i < d; i++) { + if ((sign_bits[i / 8] >> (i % 8)) & 1) { + sum += values[i]; + } + } + return sum; +} + template <> inline void rearrange_bit_planes( const uint8_t* rotated_qq, @@ -345,6 +366,144 @@ inline void quantize_query_values( *********************************************************/ namespace faiss::rabitq::multibit { +inline float ip_byte_scalar( + const uint8_t* __restrict code, + const float* __restrict query, + size_t start, + size_t d, + float cb) { + float result = 0.0f; + for (size_t i = start; i < d; i++) { + result += query[i] * (static_cast(code[i]) + cb); + } + return result; +} + +template +float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb); + +/** Dot product between an FP32 query and dense packed n-bit scalar codes. */ +template +float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb); + +template +inline void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product_byte(codes[i], query, d, cb); + } +} + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]); + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]); + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]); + +template +inline void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + size_t nbits, + float cb, + float out[4]) { + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product_dense( + codes[i], query, d, nbits, cb); + } +} + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]); + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]); + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]); + +template <> +inline float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + return ip_byte_scalar(code, query, 0, d, cb); +} + +template <> +inline float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb) { + float result = 0.0f; + const uint64_t mask = (uint64_t{1} << nbits) - 1; + for (size_t i = 0; i < d; i++) { + const size_t bit_pos = i * nbits; + const size_t byte_pos = bit_pos / 8; + const size_t shift = bit_pos % 8; + uint32_t window = code[byte_pos]; + if (shift + nbits > 8) { + window |= uint32_t(code[byte_pos + 1]) << 8; + } + result += query[i] * + (static_cast((window >> shift) & mask) + cb); + } + return result; +} + /// Scalar inner product for multi-bit RaBitQ. /// Extracts each code value in O(1) via 64-bit window read + shift + mask. /// Also serves as the tail handler for SIMD kernels via the @p start parameter. @@ -393,6 +552,57 @@ float compute_inner_product( size_t ex_bits, float cb); +/** Compute four independent multi-bit inner products against one query. + * + * HNSW evaluates graph neighbors in groups of four. Keeping the four + * accumulators in one dimension loop amortizes query loads and exposes ILP + * without requiring the database codes to be contiguous. + */ +template +inline void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product( + sign_bits[i], ex_codes[i], rotated_q, d, ex_bits, cb); + } +} + +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]); + +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]); + +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]); + // NONE specialization — pure scalar template <> inline float compute_inner_product( diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp index 0d63504dc..68cd58b77 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp @@ -437,6 +437,33 @@ uint64_t popcount(const uint8_t* data, size_t size) { return sum; } +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + const __m256i bit_positions = + _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + __m256 sum = _mm256_setzero_ps(); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256i packed = _mm256_set1_epi32(sign_bits[i / 8]); + const __m256i selected = _mm256_cmpeq_epi32( + _mm256_and_si256(packed, bit_positions), bit_positions); + const __m256 values_i = _mm256_loadu_ps(values + i); + sum = _mm256_add_ps( + sum, + _mm256_and_ps(values_i, _mm256_castsi256_ps(selected))); + } + alignas(32) float lanes[8]; + _mm256_store_ps(lanes, sum); + float result = lanes[0] + lanes[1] + lanes[2] + lanes[3] + lanes[4] + + lanes[5] + lanes[6] + lanes[7]; + result += selected_float_sum( + sign_bits + i / 8, values + i, d - i); + return result; +} + template <> void rearrange_bit_planes( const uint8_t* rotated_qq, @@ -471,6 +498,221 @@ namespace faiss::rabitq::multibit { namespace { +template +inline __m256 dense_decode_8_avx2(const uint8_t* code) { + static_assert(NBITS >= 2 && NBITS <= 8); + uint64_t packed = 0; + memcpy(&packed, code, NBITS); + constexpr uint64_t mask = (uint64_t{1} << NBITS) - 1; + const __m256i values = _mm256_setr_epi32( + (packed >> (0 * NBITS)) & mask, + (packed >> (1 * NBITS)) & mask, + (packed >> (2 * NBITS)) & mask, + (packed >> (3 * NBITS)) & mask, + (packed >> (4 * NBITS)) & mask, + (packed >> (5 * NBITS)) & mask, + (packed >> (6 * NBITS)) & mask, + (packed >> (7 * NBITS)) & mask); + return _mm256_cvtepi32_ps(values); +} + +template +float ip_dense_avx2( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + __m256 acc = _mm256_setzero_ps(); + const __m256 bias = _mm256_set1_ps(cb); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 values = dense_decode_8_avx2( + code + (i * NBITS) / 8); + acc = _mm256_fmadd_ps( + _mm256_loadu_ps(query + i), + _mm256_add_ps(values, bias), + acc); + } + float lanes[8]; + _mm256_storeu_ps(lanes, acc); + float result = lanes[0] + lanes[1] + lanes[2] + lanes[3] + lanes[4] + + lanes[5] + lanes[6] + lanes[7]; + return result + compute_inner_product_dense( + code + (i * NBITS) / 8, + query + i, + d - i, + NBITS, + cb); +} + +template +void ip_dense_batch_4_avx2( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + __m256 acc[4] = { + _mm256_setzero_ps(), + _mm256_setzero_ps(), + _mm256_setzero_ps(), + _mm256_setzero_ps()}; + const __m256 bias = _mm256_set1_ps(cb); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 q = _mm256_loadu_ps(query + i); + for (size_t j = 0; j < 4; j++) { + const __m256 values = dense_decode_8_avx2( + codes[j] + (i * NBITS) / 8); + acc[j] = _mm256_fmadd_ps( + q, _mm256_add_ps(values, bias), acc[j]); + } + } + for (size_t j = 0; j < 4; j++) { + float lanes[8]; + _mm256_storeu_ps(lanes, acc[j]); + out[j] = lanes[0] + lanes[1] + lanes[2] + lanes[3] + lanes[4] + + lanes[5] + lanes[6] + lanes[7] + + compute_inner_product_dense( + codes[j] + (i * NBITS) / 8, + query + i, + d - i, + NBITS, + cb); + } +} + +} // namespace + +template <> +float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb) { +#define FAISS_RABITQ_DENSE_CASE(N) \ + case N: \ + return ip_dense_avx2(code, query, d, cb) + switch (nbits) { + FAISS_RABITQ_DENSE_CASE(2); + FAISS_RABITQ_DENSE_CASE(3); + FAISS_RABITQ_DENSE_CASE(4); + FAISS_RABITQ_DENSE_CASE(5); + FAISS_RABITQ_DENSE_CASE(6); + FAISS_RABITQ_DENSE_CASE(7); + FAISS_RABITQ_DENSE_CASE(8); + default: + return compute_inner_product_dense( + code, query, d, nbits, cb); + } +#undef FAISS_RABITQ_DENSE_CASE +} + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]) { +#define FAISS_RABITQ_DENSE_CASE(N) \ + case N: \ + return ip_dense_batch_4_avx2(codes, query, d, cb, out) + switch (nbits) { + FAISS_RABITQ_DENSE_CASE(2); + FAISS_RABITQ_DENSE_CASE(3); + FAISS_RABITQ_DENSE_CASE(4); + FAISS_RABITQ_DENSE_CASE(5); + FAISS_RABITQ_DENSE_CASE(6); + FAISS_RABITQ_DENSE_CASE(7); + FAISS_RABITQ_DENSE_CASE(8); + default: + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product_dense( + codes[i], query, d, nbits, cb); + } + } +#undef FAISS_RABITQ_DENSE_CASE +} + +template <> +float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + __m256 acc = _mm256_setzero_ps(); + const __m256 bias = _mm256_set1_ps(cb); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m128i bytes = _mm_loadl_epi64( + reinterpret_cast(code + i)); + const __m256 values = _mm256_cvtepi32_ps( + _mm256_cvtepu8_epi32(bytes)); + acc = _mm256_fmadd_ps( + _mm256_loadu_ps(query + i), + _mm256_add_ps(values, bias), + acc); + } + float lanes[8]; + _mm256_storeu_ps(lanes, acc); + float result = lanes[0] + lanes[1] + lanes[2] + lanes[3] + lanes[4] + + lanes[5] + lanes[6] + lanes[7]; + return result + ip_byte_scalar(code, query, i, d, cb); +} + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + __m256 acc[4] = { + _mm256_setzero_ps(), + _mm256_setzero_ps(), + _mm256_setzero_ps(), + _mm256_setzero_ps()}; + const __m256 bias = _mm256_set1_ps(cb); + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 q = _mm256_loadu_ps(query + i); + for (size_t j = 0; j < 4; j++) { + const __m128i bytes = _mm_loadl_epi64( + reinterpret_cast(codes[j] + i)); + const __m256 values = _mm256_cvtepi32_ps( + _mm256_cvtepu8_epi32(bytes)); + acc[j] = _mm256_fmadd_ps( + q, _mm256_add_ps(values, bias), acc[j]); + } + } + for (size_t j = 0; j < 4; j++) { + float lanes[8]; + _mm256_storeu_ps(lanes, acc[j]); + out[j] = lanes[0] + lanes[1] + lanes[2] + lanes[3] + lanes[4] + + lanes[5] + lanes[6] + lanes[7] + + ip_byte_scalar(codes[j], query, i, d, cb); + } +} + +namespace { + +#if (defined(__GNUC__) || defined(__clang__)) && \ + (defined(__x86_64__) || defined(__i386__)) +#define FAISS_RABITQ_HAS_BMI2_TARGET 1 +#define FAISS_RABITQ_TARGET_BMI2 __attribute__((target("bmi2"))) + +inline bool cpu_supports_fast_bmi2() { + static const bool supported = __builtin_cpu_supports("bmi2"); + return supported && SIMDConfig::bmi2_fast; +} +#else +#define FAISS_RABITQ_HAS_BMI2_TARGET 0 +#define FAISS_RABITQ_TARGET_BMI2 +#endif + inline float hsum_avx2(__m256 v) { __m128 hi = _mm256_extractf128_ps(v, 1); __m128 lo = _mm256_castps256_ps128(v); @@ -517,8 +759,8 @@ inline float ip_1exbit_avx2( return result; } -#ifdef __BMI2__ -inline float ip_bitplane_avx2( +#if FAISS_RABITQ_HAS_BMI2_TARGET +FAISS_RABITQ_TARGET_BMI2 inline float ip_bitplane_avx2( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, @@ -571,7 +813,88 @@ inline float ip_bitplane_avx2( result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } -#endif // __BMI2__ + +FAISS_RABITQ_TARGET_BMI2 inline void ip_bitplane_batch_4_avx2( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _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); + + uint64_t pext_masks[7]; + __m256 v_weights[8]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t mask = 0; + for (int j = 0; j < 8; j++) { + mask |= (1ULL << (b + j * ex_bits)); + } + pext_masks[b] = mask; + v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); + } + v_weights[ex_bits] = + _mm256_set1_ps(static_cast(1u << ex_bits)); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 query = _mm256_loadu_ps(rotated_q + i); + + auto reconstruct = [&](size_t code_index) FAISS_RABITQ_TARGET_BMI2 { + const __m256i sign_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(sign_bits[code_index][i / 8]), + bit_pos), + zero); + __m256 reconstruction = _mm256_mul_ps( + _mm256_and_ps(_mm256_castsi256_ps(sign_cmp), v_one), + v_weights[ex_bits]); + + uint64_t extra_bits = 0; + memcpy( + &extra_bits, + ex_codes[code_index] + (i / 8) * ex_bits, + sizeof(extra_bits)); + for (size_t b = 0; b < ex_bits; b++) { + const auto plane = static_cast( + _pext_u64(extra_bits, pext_masks[b])); + const __m256i plane_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(plane), bit_pos), + zero); + const __m256 plane_values = _mm256_and_ps( + _mm256_castsi256_ps(plane_cmp), v_one); + reconstruction = _mm256_fmadd_ps( + plane_values, v_weights[b], reconstruction); + } + return _mm256_add_ps(reconstruction, v_cb); + }; + + acc0 = _mm256_fmadd_ps(query, reconstruct(0), acc0); + acc1 = _mm256_fmadd_ps(query, reconstruct(1), acc1); + acc2 = _mm256_fmadd_ps(query, reconstruct(2), acc2); + acc3 = _mm256_fmadd_ps(query, reconstruct(3), acc3); + } + + out[0] = hsum_avx2(acc0) + + ip_scalar(sign_bits[0], ex_codes[0], rotated_q, i, d, ex_bits, cb); + out[1] = hsum_avx2(acc1) + + ip_scalar(sign_bits[1], ex_codes[1], rotated_q, i, d, ex_bits, cb); + out[2] = hsum_avx2(acc2) + + ip_scalar(sign_bits[2], ex_codes[2], rotated_q, i, d, ex_bits, cb); + out[3] = hsum_avx2(acc3) + + ip_scalar(sign_bits[3], ex_codes[3], rotated_q, i, d, ex_bits, cb); +} +#endif } // namespace @@ -587,14 +910,39 @@ float compute_inner_product( return ip_1exbit_avx2(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ - if (ex_bits <= 7) { +#if FAISS_RABITQ_HAS_BMI2_TARGET + if (ex_bits <= 7 && cpu_supports_fast_bmi2()) { return ip_bitplane_avx2(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); } +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { +#if FAISS_RABITQ_HAS_BMI2_TARGET + if (ex_bits <= 7 && cpu_supports_fast_bmi2()) { + ip_bitplane_batch_4_avx2( + sign_bits, ex_codes, rotated_q, d, ex_bits, cb, out); + return; + } +#endif + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product( + sign_bits[i], ex_codes[i], rotated_q, d, ex_bits, cb); + } +} + +#undef FAISS_RABITQ_TARGET_BMI2 +#undef FAISS_RABITQ_HAS_BMI2_TARGET + } // namespace faiss::rabitq::multibit #endif // COMPILE_SIMD_AVX2 diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp index 991255f77..eeb2f45bc 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp @@ -582,6 +582,27 @@ uint64_t popcount(const uint8_t* data, size_t size) { return sum; } +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + __m512 sum = _mm512_setzero_ps(); + size_t i = 0; + for (; i + 16 <= d; i += 16) { + uint16_t packed = 0; + memcpy(&packed, sign_bits + i / 8, sizeof(packed)); + sum = _mm512_add_ps( + sum, + _mm512_maskz_loadu_ps( + static_cast<__mmask16>(packed), values + i)); + } + float result = _mm512_reduce_add_ps(sum); + result += selected_float_sum( + sign_bits + i / 8, values + i, d - i); + return result; +} + template <> void rearrange_bit_planes( const uint8_t* rotated_qq, @@ -623,6 +644,388 @@ namespace faiss::rabitq::multibit { namespace { +template +inline __m256i dense_decode_8_i32_avx512(const uint8_t* code) { + static_assert(NBITS >= 2 && NBITS <= 8); + uint64_t packed = 0; + memcpy(&packed, code, NBITS); + constexpr uint64_t mask = (uint64_t{1} << NBITS) - 1; + return _mm256_setr_epi32( + (packed >> (0 * NBITS)) & mask, + (packed >> (1 * NBITS)) & mask, + (packed >> (2 * NBITS)) & mask, + (packed >> (3 * NBITS)) & mask, + (packed >> (4 * NBITS)) & mask, + (packed >> (5 * NBITS)) & mask, + (packed >> (6 * NBITS)) & mask, + (packed >> (7 * NBITS)) & mask); +} + +template +inline __m512 dense_decode_16_avx512(const uint8_t* code) { + const __m256i lo = dense_decode_8_i32_avx512(code); + const __m256i hi = dense_decode_8_i32_avx512(code + NBITS); + const __m512i values = _mm512_inserti32x8( + _mm512_castsi256_si512(lo), hi, 1); + return _mm512_cvtepi32_ps(values); +} + +inline __m512i dense_decode_3_64_u8_avx512(const uint8_t* code) { + const __m256i shuf_0 = _mm256_setr_epi8( + 0, -1, 0, 1, 1, -1, 2, -1, 3, -1, 3, 4, 4, -1, 5, -1, + 6, -1, 6, 7, 7, -1, 8, -1, 9, -1, 9, 10, 10, -1, 11, -1); + const __m256i shuf_1 = _mm256_setr_epi8( + 0, -1, 1, -1, 1, 2, 2, -1, 3, -1, 4, -1, 4, 5, 5, -1, + 6, -1, 7, -1, 7, 8, 8, -1, 9, -1, 10, -1, 10, 11, 11, -1); + const __m256i shuf_2 = _mm256_setr_epi8( + 12, -1, 12, 13, 13, -1, 14, -1, 15, -1, 15, 0, 0, -1, 1, -1, + 2, -1, 2, 3, 3, -1, 4, -1, 5, -1, 5, 6, 6, -1, 7, -1); + const __m256i shuf_3 = _mm256_setr_epi8( + 12, -1, 13, -1, 13, 14, 14, -1, 15, -1, 0, -1, 0, 1, 1, -1, + 2, -1, 3, -1, 3, 4, 4, -1, 5, -1, 6, -1, 6, 7, 7, -1); + const __m512i shuf_02 = _mm512_inserti32x8( + _mm512_castsi256_si512(shuf_0), shuf_2, 1); + const __m512i shuf_13 = _mm512_inserti32x8( + _mm512_castsi256_si512(shuf_1), shuf_3, 1); + const __m256i shifts_left = _mm256_setr_epi16( + 5, 7, 1, 3, 5, 7, 1, 3, 5, 7, 1, 3, 5, 7, 1, 3); + const __m256i shifts_right = _mm256_setr_epi16( + 0, 6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2, 0, 6, 4, 2); + const __m512i v_shl = _mm512_inserti32x8( + _mm512_castsi256_si512(shifts_left), shifts_left, 1); + const __m512i v_shr = _mm512_inserti32x8( + _mm512_castsi256_si512(shifts_right), shifts_right, 1); + + const __m128i raw_0 = _mm_loadu_si128( + reinterpret_cast(code)); + const __m128i raw_2_low = _mm_loadl_epi64( + reinterpret_cast(code + 16)); + const __m128i raw_2 = _mm_blend_epi16(raw_0, raw_2_low, 0x0f); + const __m256i raw_01 = _mm256_inserti32x4( + _mm256_castsi128_si256(raw_0), raw_0, 1); + const __m256i raw_23 = _mm256_inserti32x4( + _mm256_castsi128_si256(raw_2), raw_2, 1); + const __m512i raw = _mm512_inserti32x8( + _mm512_castsi256_si512(raw_01), raw_23, 1); + const __m512i right = _mm512_srlv_epi16( + _mm512_shuffle_epi8(raw, shuf_02), v_shr); + const __m512i left = _mm512_sllv_epi16( + _mm512_shuffle_epi8(raw, shuf_13), v_shl); + return _mm512_and_si512( + _mm512_mask_blend_epi8(0xaaaaaaaaaaaaaaaaULL, right, left), + _mm512_set1_epi8(7)); +} + +inline __m512i dense_decode_4_64_u8_avx512(const uint8_t* code) { + const __m256i packed = _mm256_loadu_si256( + reinterpret_cast(code)); + const __m512i widened = _mm512_cvtepu8_epi16(packed); + return _mm512_and_si512( + _mm512_or_si512( + widened, + _mm512_slli_epi16(_mm512_srli_epi16(widened, 4), 8)), + _mm512_set1_epi16(0x0f0f)); +} + +inline __m512i dense_decode_5_64_u8_avx512(const uint8_t* code) { + const __m512i low = _mm512_zextsi256_si512(_mm256_loadu_si256( + reinterpret_cast(code))); + const __m128i high = _mm_loadl_epi64( + reinterpret_cast(code + 32)); + const __m512i raw = _mm512_inserti32x4(low, high, 2); + const __m512i spread = _mm512_permutexvar_epi64( + _mm512_setr_epi64(0, 1, 1, 2, 2, 3, 3, 4), raw); + const __m512i shuf_a = _mm512_setr_epi64( + 0x04030302ff01ff00ULL, 0x09080807ff06ff05ULL, + 0x06050504ff03ff02ULL, 0x0b0a0a09ff08ff07ULL, + 0x08070706ff05ff04ULL, 0x0d0c0c0bff0aff09ULL, + 0x0a090908ff07ff06ULL, 0x0f0e0e0dff0cff0bULL); + const __m512i shuf_b = _mm512_setr_epi64( + 0xff04ff0302010100ULL, 0xff09ff0807060605ULL, + 0xff06ff0504030302ULL, 0xff0bff0a09080807ULL, + 0xff08ff0706050504ULL, 0xff0dff0c0b0a0a09ULL, + 0xff0aff0908070706ULL, 0xff0fff0e0d0c0c0bULL); + const __m512i right = _mm512_srlv_epi16( + _mm512_shuffle_epi8(spread, shuf_a), + _mm512_set1_epi64(0x0006000400020000ULL)); + const __m512i left = _mm512_sllv_epi16( + _mm512_shuffle_epi8(spread, shuf_b), + _mm512_set1_epi64(0x0005000700010003ULL)); + return _mm512_and_si512( + _mm512_mask_blend_epi8(0xaaaaaaaaaaaaaaaaULL, right, left), + _mm512_set1_epi8(0x1f)); +} + +inline __m512i dense_decode_6_64_u8_avx512(const uint8_t* code) { + const __m512i packed = _mm512_maskz_loadu_epi8( + (uint64_t{1} << 48) - 1, code); + const __m512i expanded = _mm512_permutexvar_epi32( + _mm512_setr_epi32( + 0, 1, 2, -1, 3, 4, 5, -1, + 6, 7, 8, -1, 9, 10, 11, -1), + packed); + const __m512i shuf_0 = _mm512_broadcast_i32x4(_mm_setr_epi8( + 0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 8, 9, 10, 10, 11)); + const __m512i shuf_1 = _mm512_broadcast_i32x4(_mm_setr_epi8( + 0, 1, 2, -1, 3, 4, 5, -1, 6, 7, 8, -1, 9, 10, 11, -1)); + const __m512i left = _mm512_sllv_epi16( + _mm512_shuffle_epi8(expanded, shuf_1), + _mm512_set1_epi32(0x00060002)); + const __m512i right = _mm512_srlv_epi16( + _mm512_shuffle_epi8(expanded, shuf_0), + _mm512_set1_epi32(0x00040000)); + return _mm512_and_si512( + _mm512_mask_blend_epi8(0x5555555555555555ULL, left, right), + _mm512_set1_epi8(0x3f)); +} + +template +inline __m512i dense_decode_64_u8_avx512(const uint8_t* code) { + if constexpr (NBITS == 3) { + return dense_decode_3_64_u8_avx512(code); + } else if constexpr (NBITS == 4) { + return dense_decode_4_64_u8_avx512(code); + } else if constexpr (NBITS == 5) { + return dense_decode_5_64_u8_avx512(code); + } else if constexpr (NBITS == 6) { + return dense_decode_6_64_u8_avx512(code); + } else { + static_assert(NBITS == 8); + return _mm512_loadu_si512(code); + } +} + +inline void dense_fma_64_avx512( + __m512& acc, + __m512i decoded, + const float* query, + __m512 bias) { +#define FAISS_RABITQ_DENSE_FMA_QUARTER(Q, BYTES) \ + do { \ + const __m512 values = _mm512_cvtepi32_ps( \ + _mm512_cvtepu8_epi32(BYTES)); \ + acc = _mm512_fmadd_ps( \ + _mm512_loadu_ps(query + 16 * (Q)), \ + _mm512_add_ps(values, bias), \ + acc); \ + } while (false) + FAISS_RABITQ_DENSE_FMA_QUARTER(0, _mm512_castsi512_si128(decoded)); + FAISS_RABITQ_DENSE_FMA_QUARTER( + 1, _mm512_extracti32x4_epi32(decoded, 1)); + FAISS_RABITQ_DENSE_FMA_QUARTER( + 2, _mm512_extracti32x4_epi32(decoded, 2)); + FAISS_RABITQ_DENSE_FMA_QUARTER( + 3, _mm512_extracti32x4_epi32(decoded, 3)); +#undef FAISS_RABITQ_DENSE_FMA_QUARTER +} + +template +float ip_dense_avx512( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + __m512 acc = _mm512_setzero_ps(); + const __m512 bias = _mm512_set1_ps(cb); + size_t i = 0; + if constexpr ((NBITS >= 3 && NBITS <= 6) || NBITS == 8) { + for (; i + 64 <= d; i += 64) { + dense_fma_64_avx512( + acc, + dense_decode_64_u8_avx512( + code + (i * NBITS) / 8), + query + i, + bias); + } + } + for (; i + 16 <= d; i += 16) { + const __m512 values = dense_decode_16_avx512( + code + (i * NBITS) / 8); + acc = _mm512_fmadd_ps( + _mm512_loadu_ps(query + i), + _mm512_add_ps(values, bias), + acc); + } + return _mm512_reduce_add_ps(acc) + + compute_inner_product_dense( + code + (i * NBITS) / 8, + query + i, + d - i, + NBITS, + cb); +} + +template +void ip_dense_batch_4_avx512( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + __m512 acc[4] = { + _mm512_setzero_ps(), + _mm512_setzero_ps(), + _mm512_setzero_ps(), + _mm512_setzero_ps()}; + const __m512 bias = _mm512_set1_ps(cb); + size_t i = 0; + if constexpr ((NBITS >= 3 && NBITS <= 6) || NBITS == 8) { + for (; i + 64 <= d; i += 64) { + for (size_t j = 0; j < 4; j++) { + dense_fma_64_avx512( + acc[j], + dense_decode_64_u8_avx512( + codes[j] + (i * NBITS) / 8), + query + i, + bias); + } + } + } + for (; i + 16 <= d; i += 16) { + const __m512 q = _mm512_loadu_ps(query + i); + for (size_t j = 0; j < 4; j++) { + const __m512 values = dense_decode_16_avx512( + codes[j] + (i * NBITS) / 8); + acc[j] = _mm512_fmadd_ps( + q, _mm512_add_ps(values, bias), acc[j]); + } + } + for (size_t j = 0; j < 4; j++) { + out[j] = _mm512_reduce_add_ps(acc[j]) + + compute_inner_product_dense( + codes[j] + (i * NBITS) / 8, + query + i, + d - i, + NBITS, + cb); + } +} + +} // namespace + +template <> +float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb) { +#define FAISS_RABITQ_DENSE_CASE(N) \ + case N: \ + return ip_dense_avx512(code, query, d, cb) + switch (nbits) { + FAISS_RABITQ_DENSE_CASE(2); + FAISS_RABITQ_DENSE_CASE(3); + FAISS_RABITQ_DENSE_CASE(4); + FAISS_RABITQ_DENSE_CASE(5); + FAISS_RABITQ_DENSE_CASE(6); + FAISS_RABITQ_DENSE_CASE(7); + FAISS_RABITQ_DENSE_CASE(8); + default: + return compute_inner_product_dense( + code, query, d, nbits, cb); + } +#undef FAISS_RABITQ_DENSE_CASE +} + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]) { +#define FAISS_RABITQ_DENSE_CASE(N) \ + case N: \ + return ip_dense_batch_4_avx512(codes, query, d, cb, out) + switch (nbits) { + FAISS_RABITQ_DENSE_CASE(2); + FAISS_RABITQ_DENSE_CASE(3); + FAISS_RABITQ_DENSE_CASE(4); + FAISS_RABITQ_DENSE_CASE(5); + FAISS_RABITQ_DENSE_CASE(6); + FAISS_RABITQ_DENSE_CASE(7); + FAISS_RABITQ_DENSE_CASE(8); + default: + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product_dense( + codes[i], query, d, nbits, cb); + } + } +#undef FAISS_RABITQ_DENSE_CASE +} + +template <> +float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + __m512 acc = _mm512_setzero_ps(); + const __m512 bias = _mm512_set1_ps(cb); + size_t i = 0; + for (; i + 16 <= d; i += 16) { + const __m128i bytes = _mm_loadu_si128( + reinterpret_cast(code + i)); + const __m512 values = _mm512_cvtepi32_ps( + _mm512_cvtepu8_epi32(bytes)); + acc = _mm512_fmadd_ps( + _mm512_loadu_ps(query + i), + _mm512_add_ps(values, bias), + acc); + } + return _mm512_reduce_add_ps(acc) + + ip_byte_scalar(code, query, i, d, cb); +} + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + __m512 acc[4] = { + _mm512_setzero_ps(), + _mm512_setzero_ps(), + _mm512_setzero_ps(), + _mm512_setzero_ps()}; + const __m512 bias = _mm512_set1_ps(cb); + size_t i = 0; + for (; i + 16 <= d; i += 16) { + const __m512 q = _mm512_loadu_ps(query + i); + for (size_t j = 0; j < 4; j++) { + const __m128i bytes = _mm_loadu_si128( + reinterpret_cast(codes[j] + i)); + const __m512 values = _mm512_cvtepi32_ps( + _mm512_cvtepu8_epi32(bytes)); + acc[j] = _mm512_fmadd_ps( + q, _mm512_add_ps(values, bias), acc[j]); + } + } + for (size_t j = 0; j < 4; j++) { + out[j] = _mm512_reduce_add_ps(acc[j]) + + ip_byte_scalar(codes[j], query, i, d, cb); + } +} + +namespace { + +#if (defined(__GNUC__) || defined(__clang__)) && \ + (defined(__x86_64__) || defined(__i386__)) +#define FAISS_RABITQ_HAS_BMI2_TARGET 1 +#define FAISS_RABITQ_TARGET_BMI2 __attribute__((target("bmi2"))) + +inline bool cpu_supports_fast_bmi2() { + static const bool supported = __builtin_cpu_supports("bmi2"); + return supported && SIMDConfig::bmi2_fast; +} +#else +#define FAISS_RABITQ_HAS_BMI2_TARGET 0 +#define FAISS_RABITQ_TARGET_BMI2 +#endif + inline float hsum_avx2(__m256 v) { __m128 hi = _mm256_extractf128_ps(v, 1); __m128 lo = _mm256_castps256_ps128(v); @@ -667,8 +1070,8 @@ inline float ip_1exbit_avx512( // 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( +#if FAISS_RABITQ_HAS_BMI2_TARGET +FAISS_RABITQ_TARGET_BMI2 inline float ip_bitplane_avx2( const uint8_t* __restrict sign_bits, const uint8_t* __restrict ex_code, const float* __restrict rotated_q, @@ -721,7 +1124,488 @@ inline float ip_bitplane_avx2( result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); return result; } -#endif // __BMI2__ + +// The 16-lane AVX-512 bit extraction below needs more than one 64-bit PEXT +// window once ex_bits exceeds four. Keep a genuine four-code path for those +// wider codes by sharing each query load across four 8-lane reconstructions. +FAISS_RABITQ_TARGET_BMI2 inline void ip_bitplane_batch_4_avx2( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _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); + + uint64_t pext_masks[7]; + __m256 v_weights[8]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t mask = 0; + for (int j = 0; j < 8; j++) { + mask |= (1ULL << (b + j * ex_bits)); + } + pext_masks[b] = mask; + v_weights[b] = _mm256_set1_ps(static_cast(1u << b)); + } + v_weights[ex_bits] = + _mm256_set1_ps(static_cast(1u << ex_bits)); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 query = _mm256_loadu_ps(rotated_q + i); + + auto reconstruct = [&](size_t code_index) FAISS_RABITQ_TARGET_BMI2 { + const __m256i sign_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(sign_bits[code_index][i / 8]), + bit_pos), + zero); + __m256 reconstruction = _mm256_mul_ps( + _mm256_and_ps(_mm256_castsi256_ps(sign_cmp), v_one), + v_weights[ex_bits]); + + uint64_t extra_bits = 0; + memcpy( + &extra_bits, + ex_codes[code_index] + (i / 8) * ex_bits, + sizeof(extra_bits)); + for (size_t b = 0; b < ex_bits; b++) { + const auto plane = static_cast( + _pext_u64(extra_bits, pext_masks[b])); + const __m256i plane_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(plane), bit_pos), + zero); + const __m256 plane_values = _mm256_and_ps( + _mm256_castsi256_ps(plane_cmp), v_one); + reconstruction = _mm256_fmadd_ps( + plane_values, v_weights[b], reconstruction); + } + return _mm256_add_ps(reconstruction, v_cb); + }; + + acc0 = _mm256_fmadd_ps(query, reconstruct(0), acc0); + acc1 = _mm256_fmadd_ps(query, reconstruct(1), acc1); + acc2 = _mm256_fmadd_ps(query, reconstruct(2), acc2); + acc3 = _mm256_fmadd_ps(query, reconstruct(3), acc3); + } + + out[0] = hsum_avx2(acc0) + + ip_scalar(sign_bits[0], ex_codes[0], rotated_q, i, d, ex_bits, cb); + out[1] = hsum_avx2(acc1) + + ip_scalar(sign_bits[1], ex_codes[1], rotated_q, i, d, ex_bits, cb); + out[2] = hsum_avx2(acc2) + + ip_scalar(sign_bits[2], ex_codes[2], rotated_q, i, d, ex_bits, cb); + out[3] = hsum_avx2(acc3) + + ip_scalar(sign_bits[3], ex_codes[3], rotated_q, i, d, ex_bits, cb); +} +#endif + +// For five to seven extra bits, eight packed coordinates still fit in one +// uint64_t. Extract the per-coordinate integers directly with AVX-512 +// variable shifts instead of rebuilding them one bitplane at a time. +template +inline float ip_packed_8_avx512( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + float cb) { + static_assert(ExBits >= 5 && ExBits <= 7); + __m256 acc = _mm256_setzero_ps(); + 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); + const __m256 v_sign_weight = + _mm256_set1_ps(static_cast(1u << ExBits)); + const __m512i shifts = _mm512_setr_epi64( + 0, + ExBits, + 2 * ExBits, + 3 * ExBits, + 4 * ExBits, + 5 * ExBits, + 6 * ExBits, + 7 * ExBits); + const __m512i value_mask = _mm512_set1_epi64((1u << ExBits) - 1); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + uint64_t packed = 0; + memcpy(&packed, ex_code + (i / 8) * ExBits, ExBits); + const __m512i values = _mm512_and_si512( + _mm512_srlv_epi64(_mm512_set1_epi64(packed), shifts), + value_mask); + const __m256 extra_values = + _mm256_cvtepi32_ps(_mm512_cvtepi64_epi32(values)); + + const __m256i sign_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(sign_bits[i / 8]), bit_pos), + zero); + const __m256 sign_values = _mm256_and_ps( + _mm256_castsi256_ps(sign_cmp), v_sign_weight); + const __m256 reconstruction = _mm256_add_ps( + _mm256_add_ps(extra_values, sign_values), v_cb); + acc = _mm256_fmadd_ps( + _mm256_loadu_ps(rotated_q + i), reconstruction, acc); + } + + return hsum_avx2(acc) + + ip_scalar( + sign_bits, ex_code, rotated_q, i, d, ExBits, cb); +} + +template +inline void ip_packed_8_batch_4_avx512( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + float cb, + float out[4]) { + static_assert(ExBits >= 5 && ExBits <= 7); + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _mm256_setzero_ps(); + 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); + const __m256 v_sign_weight = + _mm256_set1_ps(static_cast(1u << ExBits)); + const __m512i shifts = _mm512_setr_epi64( + 0, + ExBits, + 2 * ExBits, + 3 * ExBits, + 4 * ExBits, + 5 * ExBits, + 6 * ExBits, + 7 * ExBits); + const __m512i value_mask = _mm512_set1_epi64((1u << ExBits) - 1); + + size_t i = 0; + for (; i + 8 <= d; i += 8) { + const __m256 query = _mm256_loadu_ps(rotated_q + i); + + auto reconstruct = [&](size_t code_index) { + uint64_t packed = 0; + memcpy( + &packed, + ex_codes[code_index] + (i / 8) * ExBits, + ExBits); + const __m512i values = _mm512_and_si512( + _mm512_srlv_epi64(_mm512_set1_epi64(packed), shifts), + value_mask); + const __m256 extra_values = + _mm256_cvtepi32_ps(_mm512_cvtepi64_epi32(values)); + + const __m256i sign_cmp = _mm256_cmpgt_epi32( + _mm256_and_si256( + _mm256_set1_epi32(sign_bits[code_index][i / 8]), + bit_pos), + zero); + const __m256 sign_values = _mm256_and_ps( + _mm256_castsi256_ps(sign_cmp), v_sign_weight); + return _mm256_add_ps( + _mm256_add_ps(extra_values, sign_values), v_cb); + }; + + acc0 = _mm256_fmadd_ps(query, reconstruct(0), acc0); + acc1 = _mm256_fmadd_ps(query, reconstruct(1), acc1); + acc2 = _mm256_fmadd_ps(query, reconstruct(2), acc2); + acc3 = _mm256_fmadd_ps(query, reconstruct(3), acc3); + } + + out[0] = hsum_avx2(acc0) + + ip_scalar(sign_bits[0], ex_codes[0], rotated_q, i, d, ExBits, cb); + out[1] = hsum_avx2(acc1) + + ip_scalar(sign_bits[1], ex_codes[1], rotated_q, i, d, ExBits, cb); + out[2] = hsum_avx2(acc2) + + ip_scalar(sign_bits[2], ex_codes[2], rotated_q, i, d, ExBits, cb); + out[3] = hsum_avx2(acc3) + + ip_scalar(sign_bits[3], ex_codes[3], rotated_q, i, d, ExBits, cb); +} + +template +inline __m256i unpack_packed_8_avx512( + const uint8_t* __restrict ex_code, + const __m512i shifts, + const __m512i value_mask) { + uint64_t packed = 0; + memcpy(&packed, ex_code, ExBits); + const __m512i values = _mm512_and_si512( + _mm512_srlv_epi64(_mm512_set1_epi64(packed), shifts), + value_mask); + return _mm512_cvtepi64_epi32(values); +} + +template +inline float ip_packed_16_avx512( + const uint8_t* __restrict sign_bits, + const uint8_t* __restrict ex_code, + const float* __restrict rotated_q, + size_t d, + float cb) { + static_assert(ExBits >= 5 && ExBits <= 7); + __m512 acc = _mm512_setzero_ps(); + const __m512 v_one = _mm512_set1_ps(1.0f); + const __m512 v_cb = _mm512_set1_ps(cb); + const __m512 v_sign_weight = + _mm512_set1_ps(static_cast(1u << ExBits)); + const __m512i shifts = _mm512_setr_epi64( + 0, + ExBits, + 2 * ExBits, + 3 * ExBits, + 4 * ExBits, + 5 * ExBits, + 6 * ExBits, + 7 * ExBits); + const __m512i value_mask = _mm512_set1_epi64((1u << ExBits) - 1); + + size_t i = 0; + for (; i + 16 <= d; i += 16) { + const uint8_t* block = ex_code + (i / 8) * ExBits; + const __m256i lo = + unpack_packed_8_avx512(block, shifts, value_mask); + const __m256i hi = unpack_packed_8_avx512( + block + ExBits, shifts, value_mask); + const __m512i values = _mm512_inserti64x4( + _mm512_castsi256_si512(lo), hi, 1); + const __m512 extra_values = _mm512_cvtepi32_ps(values); + + uint16_t sign_plane = 0; + memcpy(&sign_plane, sign_bits + i / 8, sizeof(sign_plane)); + const __m512 sign_values = _mm512_mul_ps( + _mm512_maskz_mov_ps(_cvtu32_mask16(sign_plane), v_one), + v_sign_weight); + const __m512 reconstruction = _mm512_add_ps( + _mm512_add_ps(extra_values, sign_values), v_cb); + acc = _mm512_fmadd_ps( + _mm512_loadu_ps(rotated_q + i), reconstruction, acc); + } + + return _mm512_reduce_add_ps(acc) + + ip_scalar( + sign_bits, ex_code, rotated_q, i, d, ExBits, cb); +} + +template +inline void ip_packed_16_batch_4_avx512( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + float cb, + float out[4]) { + static_assert(ExBits >= 5 && ExBits <= 7); + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + const __m512 v_one = _mm512_set1_ps(1.0f); + const __m512 v_cb = _mm512_set1_ps(cb); + const __m512 v_sign_weight = + _mm512_set1_ps(static_cast(1u << ExBits)); + const __m512i shifts = _mm512_setr_epi64( + 0, + ExBits, + 2 * ExBits, + 3 * ExBits, + 4 * ExBits, + 5 * ExBits, + 6 * ExBits, + 7 * ExBits); + const __m512i value_mask = _mm512_set1_epi64((1u << ExBits) - 1); + + size_t i = 0; + for (; i + 16 <= d; i += 16) { + const __m512 query = _mm512_loadu_ps(rotated_q + i); + + auto reconstruct = [&](size_t code_index) { + const uint8_t* block = + ex_codes[code_index] + (i / 8) * ExBits; + const __m256i lo = + unpack_packed_8_avx512(block, shifts, value_mask); + const __m256i hi = unpack_packed_8_avx512( + block + ExBits, shifts, value_mask); + const __m512i values = _mm512_inserti64x4( + _mm512_castsi256_si512(lo), hi, 1); + const __m512 extra_values = _mm512_cvtepi32_ps(values); + + uint16_t sign_plane = 0; + memcpy( + &sign_plane, + sign_bits[code_index] + i / 8, + sizeof(sign_plane)); + const __m512 sign_values = _mm512_mul_ps( + _mm512_maskz_mov_ps( + _cvtu32_mask16(sign_plane), v_one), + v_sign_weight); + return _mm512_add_ps( + _mm512_add_ps(extra_values, sign_values), v_cb); + }; + + acc0 = _mm512_fmadd_ps(query, reconstruct(0), acc0); + acc1 = _mm512_fmadd_ps(query, reconstruct(1), acc1); + acc2 = _mm512_fmadd_ps(query, reconstruct(2), acc2); + acc3 = _mm512_fmadd_ps(query, reconstruct(3), acc3); + } + + out[0] = _mm512_reduce_add_ps(acc0) + + ip_scalar(sign_bits[0], ex_codes[0], rotated_q, i, d, ExBits, cb); + out[1] = _mm512_reduce_add_ps(acc1) + + ip_scalar(sign_bits[1], ex_codes[1], rotated_q, i, d, ExBits, cb); + out[2] = _mm512_reduce_add_ps(acc2) + + ip_scalar(sign_bits[2], ex_codes[2], rotated_q, i, d, ExBits, cb); + out[3] = _mm512_reduce_add_ps(acc3) + + ip_scalar(sign_bits[3], ex_codes[3], rotated_q, i, d, ExBits, cb); +} + +#if FAISS_RABITQ_HAS_BMI2_TARGET +FAISS_RABITQ_TARGET_BMI2 inline 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) { + __m512 acc = _mm512_setzero_ps(); + const __m512 v_one = _mm512_set1_ps(1.0f); + const __m512 v_cb = _mm512_set1_ps(cb); + const __m512 v_sign_weight = + _mm512_set1_ps(static_cast(1u << ex_bits)); + + uint64_t pext_masks[4]; + __m512 v_weights[4]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t mask = 0; + for (int j = 0; j < 16; j++) { + mask |= (1ULL << (j * ex_bits + b)); + } + pext_masks[b] = mask; + v_weights[b] = _mm512_set1_ps(static_cast(1u << b)); + } + + size_t i = 0; + for (; i + 16 <= d; i += 16) { + uint16_t sign_plane = 0; + memcpy(&sign_plane, sign_bits + i / 8, sizeof(sign_plane)); + __m512 reconstruction = _mm512_mul_ps( + _mm512_maskz_mov_ps(_cvtu32_mask16(sign_plane), v_one), + v_sign_weight); + + uint64_t extra_bits = 0; + memcpy( + &extra_bits, + ex_code + (i / 8) * ex_bits, + sizeof(extra_bits)); + for (size_t b = 0; b < ex_bits; b++) { + const uint16_t plane = static_cast( + _pext_u64(extra_bits, pext_masks[b])); + const __m512 plane_values = + _mm512_maskz_mov_ps(_cvtu32_mask16(plane), v_one); + reconstruction = _mm512_fmadd_ps( + plane_values, v_weights[b], reconstruction); + } + + const __m512 query = _mm512_loadu_ps(rotated_q + i); + acc = _mm512_fmadd_ps( + query, _mm512_add_ps(reconstruction, v_cb), acc); + } + + float result = _mm512_reduce_add_ps(acc); + result += ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb); + return result; +} + +FAISS_RABITQ_TARGET_BMI2 inline void ip_bitplane_batch_4_avx512( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + const __m512 v_one = _mm512_set1_ps(1.0f); + const __m512 v_cb = _mm512_set1_ps(cb); + const __m512 v_sign_weight = + _mm512_set1_ps(static_cast(1u << ex_bits)); + + uint64_t pext_masks[4]; + __m512 v_weights[4]; + for (size_t b = 0; b < ex_bits; b++) { + uint64_t mask = 0; + for (int j = 0; j < 16; j++) { + mask |= (1ULL << (j * ex_bits + b)); + } + pext_masks[b] = mask; + v_weights[b] = _mm512_set1_ps(static_cast(1u << b)); + } + + size_t i = 0; + for (; i + 16 <= d; i += 16) { + const __m512 query = _mm512_loadu_ps(rotated_q + i); + + auto reconstruct = [&](size_t code_index) FAISS_RABITQ_TARGET_BMI2 { + uint16_t sign_plane = 0; + memcpy( + &sign_plane, + sign_bits[code_index] + i / 8, + sizeof(sign_plane)); + __m512 reconstruction = _mm512_mul_ps( + _mm512_maskz_mov_ps(_cvtu32_mask16(sign_plane), v_one), + v_sign_weight); + + uint64_t extra_bits = 0; + memcpy( + &extra_bits, + ex_codes[code_index] + (i / 8) * ex_bits, + sizeof(extra_bits)); + for (size_t b = 0; b < ex_bits; b++) { + const uint16_t plane = static_cast( + _pext_u64(extra_bits, pext_masks[b])); + const __m512 plane_values = + _mm512_maskz_mov_ps(_cvtu32_mask16(plane), v_one); + reconstruction = _mm512_fmadd_ps( + plane_values, v_weights[b], reconstruction); + } + return _mm512_add_ps(reconstruction, v_cb); + }; + + acc0 = _mm512_fmadd_ps(query, reconstruct(0), acc0); + acc1 = _mm512_fmadd_ps(query, reconstruct(1), acc1); + acc2 = _mm512_fmadd_ps(query, reconstruct(2), acc2); + acc3 = _mm512_fmadd_ps(query, reconstruct(3), acc3); + } + + out[0] = _mm512_reduce_add_ps(acc0) + + ip_scalar(sign_bits[0], ex_codes[0], rotated_q, i, d, ex_bits, cb); + out[1] = _mm512_reduce_add_ps(acc1) + + ip_scalar(sign_bits[1], ex_codes[1], rotated_q, i, d, ex_bits, cb); + out[2] = _mm512_reduce_add_ps(acc2) + + ip_scalar(sign_bits[2], ex_codes[2], rotated_q, i, d, ex_bits, cb); + out[3] = _mm512_reduce_add_ps(acc3) + + ip_scalar(sign_bits[3], ex_codes[3], rotated_q, i, d, ex_bits, cb); +} +#endif } // namespace @@ -737,14 +1621,78 @@ float compute_inner_product( return ip_1exbit_avx512(sign_bits, ex_code, rotated_q, d, cb); } -#ifdef __BMI2__ - if (ex_bits <= 7) { + switch (ex_bits) { + case 5: + return ip_packed_16_avx512<5>( + sign_bits, ex_code, rotated_q, d, cb); + case 6: + return ip_packed_16_avx512<6>( + sign_bits, ex_code, rotated_q, d, cb); + case 7: + return ip_packed_16_avx512<7>( + sign_bits, ex_code, rotated_q, d, cb); + default: + break; + } + +#if FAISS_RABITQ_HAS_BMI2_TARGET + if (ex_bits <= 4 && cpu_supports_fast_bmi2()) { + return ip_bitplane_avx512( + sign_bits, ex_code, rotated_q, d, ex_bits, cb); + } + if (ex_bits <= 7 && cpu_supports_fast_bmi2()) { return ip_bitplane_avx2(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); } +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + switch (ex_bits) { + case 5: + ip_packed_16_batch_4_avx512<5>( + sign_bits, ex_codes, rotated_q, d, cb, out); + return; + case 6: + ip_packed_16_batch_4_avx512<6>( + sign_bits, ex_codes, rotated_q, d, cb, out); + return; + case 7: + ip_packed_16_batch_4_avx512<7>( + sign_bits, ex_codes, rotated_q, d, cb, out); + return; + default: + break; + } +#if FAISS_RABITQ_HAS_BMI2_TARGET + if (ex_bits <= 4 && cpu_supports_fast_bmi2()) { + ip_bitplane_batch_4_avx512( + sign_bits, ex_codes, rotated_q, d, ex_bits, cb, out); + return; + } + if (ex_bits <= 7 && cpu_supports_fast_bmi2()) { + ip_bitplane_batch_4_avx2( + sign_bits, ex_codes, rotated_q, d, ex_bits, cb, out); + return; + } +#endif + for (size_t i = 0; i < 4; i++) { + out[i] = compute_inner_product( + sign_bits[i], ex_codes[i], rotated_q, d, ex_bits, cb); + } +} + +#undef FAISS_RABITQ_TARGET_BMI2 +#undef FAISS_RABITQ_HAS_BMI2_TARGET + } // namespace faiss::rabitq::multibit #endif // COMPILE_SIMD_AVX512 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..0205048cc 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp @@ -65,6 +65,11 @@ uint64_t bitwise_xor_dot_product( size_t qb); template <> uint64_t popcount(const uint8_t* data, size_t size); +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d); namespace { @@ -402,10 +407,94 @@ uint64_t popcount(const uint8_t* data, size_t size) { return sum; } +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + return selected_float_sum(sign_bits, values, d); +} + } // namespace faiss::rabitq namespace faiss::rabitq::multibit { +template <> +float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb); + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]); + +template <> +float compute_inner_product_dense( + const uint8_t* code, + const float* query, + size_t d, + size_t nbits, + float cb); + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]); + +template <> +float compute_inner_product_byte( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + float cb) { + return compute_inner_product_byte( + code, query, d, cb); +} + +template <> +void compute_inner_product_byte_batch_4( + const uint8_t* const codes[4], + const float* __restrict query, + size_t d, + float cb, + float out[4]) { + compute_inner_product_byte_batch_4( + codes, query, d, cb, out); +} + +template <> +float compute_inner_product_dense( + const uint8_t* code, + const float* query, + size_t d, + size_t nbits, + float cb) { + return compute_inner_product_dense( + code, query, d, nbits, cb); +} + +template <> +void compute_inner_product_dense_batch_4( + const uint8_t* const codes[4], + const float* query, + size_t d, + size_t nbits, + float cb, + float out[4]) { + compute_inner_product_dense_batch_4( + codes, query, d, nbits, cb, out); +} + // Forward-declare the AVX512 floating-point inner-product kernel. // VPOPCNTDQ does not help this kernel (it operates on FP32), so we // reuse the AVX512 implementation rather than duplicate it. @@ -418,6 +507,16 @@ float compute_inner_product( size_t ex_bits, float cb); +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]); + template <> float compute_inner_product( const uint8_t* __restrict sign_bits, @@ -430,6 +529,19 @@ float compute_inner_product( sign_bits, ex_code, rotated_q, d, ex_bits, cb); } +template <> +void compute_inner_product_batch_4( + const uint8_t* const sign_bits[4], + const uint8_t* const ex_codes[4], + const float* __restrict rotated_q, + size_t d, + size_t ex_bits, + float cb, + float out[4]) { + compute_inner_product_batch_4( + sign_bits, ex_codes, rotated_q, d, ex_bits, cb, out); +} + } // namespace faiss::rabitq::multibit #endif // COMPILE_SIMD_AVX512_SPR diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_neon.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_neon.cpp index 6294aa5f2..09f68a41f 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_neon.cpp @@ -45,6 +45,14 @@ uint64_t popcount(const uint8_t* data, size_t size) { return popcount(data, size); } +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + return selected_float_sum(sign_bits, values, d); +} + } // namespace faiss::rabitq namespace faiss::rabitq::multibit { @@ -61,6 +69,17 @@ float compute_inner_product( sign_bits, ex_code, rotated_q, d, ex_bits, cb); } +template <> +float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb) { + return compute_inner_product_dense( + code, query, d, nbits, cb); +} + } // namespace faiss::rabitq::multibit #endif // COMPILE_SIMD_ARM_NEON diff --git a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp index 1850a6d86..c673f7e16 100644 --- a/thirdparty/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +++ b/thirdparty/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp @@ -137,6 +137,14 @@ uint64_t popcount(const uint8_t* data, size_t size) { return __riscv_vmv_x_s_u32m1_u32(red); } +template <> +float selected_float_sum( + const uint8_t* sign_bits, + const float* values, + size_t d) { + return selected_float_sum(sign_bits, values, d); +} + } // namespace faiss::rabitq namespace faiss::rabitq::multibit { @@ -187,6 +195,17 @@ compute_inner_product( return ip_scalar(sign_bits, ex_code, rotated_q, 0, d, ex_bits, cb); } +template <> +float compute_inner_product_dense( + const uint8_t* __restrict code, + const float* __restrict query, + size_t d, + size_t nbits, + float cb) { + return compute_inner_product_dense( + code, query, d, nbits, cb); +} + } // namespace faiss::rabitq::multibit #endif // COMPILE_SIMD_RISCV_RVV diff --git a/thirdparty/faiss/faiss/utils/simd_levels.cpp b/thirdparty/faiss/faiss/utils/simd_levels.cpp index e17796bae..32d980d2c 100644 --- a/thirdparty/faiss/faiss/utils/simd_levels.cpp +++ b/thirdparty/faiss/faiss/utils/simd_levels.cpp @@ -23,6 +23,7 @@ uint64_t SIMDConfig::supported_simd_levels = 0; // Microarchitecture flags (x86). Default false; set by // detect_x86_uarch_flags() at load time. bool SIMDConfig::avx512_split = false; +bool SIMDConfig::bmi2_fast = false; // ARM SVE runtime detection #if defined(__aarch64__) || defined(_M_ARM64) @@ -71,9 +72,27 @@ void detect_x86_uarch_flags() { asm volatile("cpuid" : "=a"(eax), "=b"(ebx), "=c"(ecx), "=d"(edx) : "a"(eax), "c"(ecx)); + const unsigned int max_basic_leaf = eax; const bool is_amd = ebx == 0x68747541u && edx == 0x69746e65u && ecx == 0x444d4163u; + // CPUID.7.0: EBX bit 8 is BMI2 and bit 16 is AVX-512F. All AMD Zen 4+ + // CPUs implement AVX-512F, whereas AMD CPUs before Zen 4 do not. Use the + // latter as a robust generation gate rather than family/model tables, + // because family 0x19 contains both Zen 3 and Zen 4 model ranges. + bool has_bmi2 = false; + bool has_avx512f = false; + if (max_basic_leaf >= 7) { + eax = 7; + ecx = 0; + asm volatile("cpuid" + : "=a"(eax), "=b"(ebx), "=c"(ecx), "=d"(edx) + : "a"(eax), "c"(ecx)); + has_bmi2 = (ebx & (1u << 8)) != 0; + has_avx512f = (ebx & (1u << 16)) != 0; + } + SIMDConfig::bmi2_fast = has_bmi2 && (!is_amd || has_avx512f); + // Family/model (CPUID.1 EAX). eax = 1; ecx = 0; @@ -105,6 +124,9 @@ void detect_x86_uarch_flags() {} static SIMDConfig simd_config_initializer; SIMDConfig::SIMDConfig(const char** faiss_simd_level_env) { + // Initialize microarchitecture flags even when the SIMD level is forced + // through FAISS_SIMD_LEVEL. + detect_x86_uarch_flags(); // Support dependency injection for testing const char* env_var = faiss_simd_level_env ? *faiss_simd_level_env : getenv("FAISS_SIMD_LEVEL"); diff --git a/thirdparty/faiss/faiss/utils/simd_levels.h b/thirdparty/faiss/faiss/utils/simd_levels.h index 9c68c4262..f8370d7d2 100644 --- a/thirdparty/faiss/faiss/utils/simd_levels.h +++ b/thirdparty/faiss/faiss/utils/simd_levels.h @@ -167,6 +167,11 @@ struct FAISS_API SIMDConfig { /// 256-bit kernel. static bool avx512_split; + /// BMI2 PEXT is a fast hardware instruction on Intel and AMD Zen 4+. + /// AMD CPUs before Zen 4 advertise BMI2 but execute PEXT very slowly, so + /// RaBitQ bit-plane kernels must use their non-BMI2 fallback there. + static bool bmi2_fast; + static SIMDLevel auto_detect_simd_level(); static constexpr bool has_dynamic_dispatch() { diff --git a/thirdparty/faiss/perf_tests/CMakeLists.txt b/thirdparty/faiss/perf_tests/CMakeLists.txt index db79f59c9..8ce5a6f61 100644 --- a/thirdparty/faiss/perf_tests/CMakeLists.txt +++ b/thirdparty/faiss/perf_tests/CMakeLists.txt @@ -31,6 +31,14 @@ include(../cmake/link_to_faiss_lib.cmake) link_to_faiss_lib(faiss_perf_tests_utils) +if(FAISS_OPT_LEVEL STREQUAL "dd" AND CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|amd64|AMD64)") + # The benchmark directly instantiates DD-only AVX2/AVX-512 symbols. Keep it + # opt-in so generic/avx2 default builds neither link nor run unavailable ISA. + add_executable(bench_rabitq_simd EXCLUDE_FROM_ALL ../benchs/bench_rabitq_simd.cpp) + link_to_faiss_lib(bench_rabitq_simd) + target_link_libraries(bench_rabitq_simd PRIVATE benchmark::benchmark) +endif() + set(FAISS_PERF_TEST_SRC bench_no_multithreading_rcq_search.cpp bench_scalar_quantizer_accuracy.cpp diff --git a/thirdparty/faiss/tests/test_rabitq_simd.cpp b/thirdparty/faiss/tests/test_rabitq_simd.cpp index 35c956994..b9a24e537 100644 --- a/thirdparty/faiss/tests/test_rabitq_simd.cpp +++ b/thirdparty/faiss/tests/test_rabitq_simd.cpp @@ -7,6 +7,8 @@ #include +#include +#include #include #include #include @@ -377,3 +379,175 @@ TEST(RaBitQBitwiseAndDotProductWithPopcount, Avx2MatchesScalar) { } } } + +template +static void check_selected_float_sum_matches_scalar() { + std::mt19937 rng(20260821); + std::uniform_real_distribution value_dist(-10.0f, 10.0f); + for (size_t d : kDims) { + const auto sign_bits = random_bytes((d + 7) / 8, 81723 + d); + std::vector values(d); + for (float& value : values) { + value = value_dist(rng); + } + const float expected = + faiss::rabitq::selected_float_sum( + sign_bits.data(), values.data(), d); + const float actual = faiss::rabitq::selected_float_sum( + sign_bits.data(), values.data(), d); + EXPECT_NEAR(actual, expected, std::max(1e-5f, std::abs(expected) * 1e-5f)) + << "d=" << d; + } +} + +template +static void check_multibit_inner_product_matches_scalar() { + std::mt19937 rng(91827); + std::uniform_real_distribution value_dist(-2.0f, 2.0f); + for (size_t d : kDims) { + const auto sign_bits = random_bytes((d + 7) / 8, 71237 + d); + std::vector query(d); + for (float& value : query) { + value = value_dist(rng); + } + for (size_t extra_bits = 1; extra_bits <= 7; extra_bits++) { + const auto extra_code = random_bytes( + (d * extra_bits + 7) / 8 + 8, + 51971 + d + extra_bits); + constexpr float cb = -3.25f; + const float expected = + faiss::rabitq::multibit::compute_inner_product< + SIMDLevel::NONE>( + sign_bits.data(), + extra_code.data(), + query.data(), + d, + extra_bits, + cb); + const float actual = + faiss::rabitq::multibit::compute_inner_product( + sign_bits.data(), + extra_code.data(), + query.data(), + d, + extra_bits, + cb); + EXPECT_NEAR( + actual, + expected, + std::max(1e-4f, std::abs(expected) * 1e-5f)) + << "d=" << d << " extra_bits=" << extra_bits; + } + } +} + +template +static void check_multibit_inner_product_batch_4_matches_single() { + std::mt19937 rng(20260824); + std::uniform_real_distribution value_dist(-2.0f, 2.0f); + for (size_t d : kDims) { + std::vector query(d); + for (float& value : query) { + value = value_dist(rng); + } + for (size_t extra_bits = 1; extra_bits <= 7; extra_bits++) { + std::vector> sign_storage(4); + std::vector> extra_storage(4); + const uint8_t* sign_bits[4]; + const uint8_t* extra_codes[4]; + for (size_t i = 0; i < 4; i++) { + sign_storage[i] = random_bytes( + (d + 7) / 8, 81237 + d + extra_bits * 17 + i); + extra_storage[i] = random_bytes( + (d * extra_bits + 7) / 8 + 8, + 61971 + d + extra_bits * 17 + i); + sign_bits[i] = sign_storage[i].data(); + extra_codes[i] = extra_storage[i].data(); + } + + constexpr float cb = -3.25f; + float actual[4]; + faiss::rabitq::multibit::compute_inner_product_batch_4( + sign_bits, + extra_codes, + query.data(), + d, + extra_bits, + cb, + actual); + for (size_t i = 0; i < 4; i++) { + const float expected = + faiss::rabitq::multibit::compute_inner_product( + sign_bits[i], + extra_codes[i], + query.data(), + d, + extra_bits, + cb); + EXPECT_NEAR( + actual[i], + expected, + std::max(1e-4f, std::abs(expected) * 1e-5f)) + << "d=" << d << " extra_bits=" << extra_bits + << " code=" << i; + } + } + } +} + +TEST(RaBitQSelectedFloatSum, Avx2MatchesScalar) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX2)) { + GTEST_SKIP() << "AVX2 is not available on this CPU"; + } + check_selected_float_sum_matches_scalar(); +} + +TEST(RaBitQMultiBitInnerProduct, Avx2MatchesScalar) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX2)) { + GTEST_SKIP() << "AVX2 is not available on this CPU"; + } + check_multibit_inner_product_matches_scalar(); +} + +TEST(RaBitQMultiBitInnerProductBatch4, Avx2MatchesSingle) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX2)) { + GTEST_SKIP() << "AVX2 is not available on this CPU"; + } + check_multibit_inner_product_batch_4_matches_single(); +} + +TEST(RaBitQMultiBitInnerProduct, NonBmi2FallbackMatchesScalar) { + const bool original_bmi2_fast = faiss::SIMDConfig::bmi2_fast; + faiss::SIMDConfig::bmi2_fast = false; + + if (faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX2)) { + check_multibit_inner_product_matches_scalar(); + check_multibit_inner_product_batch_4_matches_single(); + } + if (faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX512)) { + check_multibit_inner_product_matches_scalar(); + check_multibit_inner_product_batch_4_matches_single(); + } + + faiss::SIMDConfig::bmi2_fast = original_bmi2_fast; +} + +TEST(RaBitQSelectedFloatSum, Avx512MatchesScalar) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX512)) { + GTEST_SKIP() << "AVX512 is not available on this CPU"; + } + check_selected_float_sum_matches_scalar(); +} + +TEST(RaBitQMultiBitInnerProduct, Avx512MatchesScalar) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX512)) { + GTEST_SKIP() << "AVX512 is not available on this CPU"; + } + check_multibit_inner_product_matches_scalar(); +} +TEST(RaBitQMultiBitInnerProductBatch4, Avx512MatchesSingle) { + if (!faiss::SIMDConfig::is_simd_level_available(SIMDLevel::AVX512)) { + GTEST_SKIP() << "AVX512 is not available on this CPU"; + } + check_multibit_inner_product_batch_4_matches_single(); +}