diff --git a/CMakeLists.txt b/CMakeLists.txt index 99de33262..d42377e9a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -103,6 +103,8 @@ endif() find_package(xxHash REQUIRED) include_directories(${xxHash_INCLUDE_DIRS}) find_package(simde REQUIRED) +find_package(roaring REQUIRED) +include_directories(${roaring_INCLUDE_DIRS}) if(NOT WITH_LIGHT) find_package(opentelemetry-cpp REQUIRED) @@ -114,6 +116,7 @@ if(NOT WITH_LIGHT) endif() if(APPLE) + add_definitions(-D_DARWIN_C_SOURCE) # Prevent SDK sysroot include dir from appearing as -I/-isystem in compile # commands. Transitive find_package calls (CURL, ZLIB, etc.) resolve headers # to the SDK sysroot, and the explicit -I flag breaks C++ wrapper headers @@ -125,7 +128,7 @@ if(APPLE) endif() set(CMAKE_CXX_STANDARD 20) set(CMAKE_OSX_DEPLOYMENT_TARGET - "10.15" + "11.0" CACHE STRING "Minimum OS X deployment version" FORCE) if(OPENMP_FOUND) @@ -139,6 +142,7 @@ include(cmake/utils/compile_flags.cmake) find_package(milvus-common REQUIRED) include(cmake/libs/libhnsw.cmake) include(cmake/libs/libfaiss.cmake) +target_link_libraries(faiss PUBLIC roaring::roaring) if(WITH_COVERAGE) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage") @@ -268,6 +272,7 @@ list(APPEND KNOWHERE_LINKER_LIBS fmt::fmt) list(APPEND KNOWHERE_LINKER_LIBS Folly::folly) list(APPEND KNOWHERE_LINKER_LIBS milvus-common::milvus-common) list(APPEND KNOWHERE_LINKER_LIBS simde::simde) +list(APPEND KNOWHERE_LINKER_LIBS roaring::roaring) add_library(knowhere SHARED ${KNOWHERE_SRCS}) add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS}) diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake index 9c6ce5eb3..150dba9c2 100644 --- a/cmake/libs/libdiskann.cmake +++ b/cmake/libs/libdiskann.cmake @@ -29,7 +29,7 @@ set(DISKANN_SOURCES find_package(folly REQUIRED) set(DISKANN_LINKER_LIBS PUBLIC ${AIO_LIBRARIES} ${DISKANN_BOOST_PROGRAM_OPTIONS_LIB} nlohmann_json::nlohmann_json - Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog) + Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog roaring::roaring) if (WITH_CUVS) list(APPEND DISKANN_LINKER_LIBS PRIVATE cuvs::cuvs) list(APPEND DISKANN_SOURCES thirdparty/DiskANN/src/diskann_gpu.cpp) @@ -49,7 +49,8 @@ target_link_libraries( fmt::fmt prometheus-cpp::core prometheus-cpp::push - glog::glog) + glog::glog + roaring::roaring) if(__X86_64) target_compile_options( diskann PRIVATE -fno-builtin-malloc -fno-builtin-calloc diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 056b530e6..61944a2ff 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -384,7 +384,7 @@ if(__X86_64) -mavx2 -mfma -mf16c -mpopcnt>) target_compile_definitions(faiss_avx2 PRIVATE COMPILE_SIMD_AVX2) target_include_directories(faiss_avx2 PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common roaring::roaring) add_library(faiss_avx512 OBJECT ${FAISS_AVX512_SRCS}) target_compile_options( faiss_avx512 @@ -400,7 +400,7 @@ if(__X86_64) -mpopcnt>) target_compile_definitions(faiss_avx512 PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512) target_include_directories(faiss_avx512 PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common roaring::roaring) if(FAISS_ENABLE_AVX512_SPR) add_library(faiss_avx512_spr OBJECT ${FAISS_AVX512_SPR_SRCS}) @@ -424,7 +424,7 @@ if(__X86_64) target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR) target_include_directories(faiss_avx512_spr PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common roaring::roaring) endif() add_library(faiss STATIC ${FAISS_SRCS}) diff --git a/conanfile.py b/conanfile.py index 41e0edc6c..4fdb3a046 100644 --- a/conanfile.py +++ b/conanfile.py @@ -137,6 +137,7 @@ def requirements(self): self.requires("libcurl/8.10.1#a3113369c86086b0e84231844e7ed0a9", force=True, override=True) self.requires("simde/0.8.2#5e1edfd5cba92f25d79bf6ef4616b972") self.requires("xxhash/0.8.3#caa6d0af1b951c247922e38fbcebdbe6") + self.requires("roaring/3.0.0") if self.settings.os == "Linux": self.requires("openblas/0.3.30") if not self.options.with_light: @@ -283,6 +284,7 @@ def package_info(self): "milvus-common::milvus-common", "prometheus-cpp::core", "prometheus-cpp::push", + "roaring::roaring", ] self.cpp_info.components["libknowhere"].set_property( diff --git a/include/knowhere/bitsetview.h b/include/knowhere/bitsetview.h index ee6db3a47..315f12ec0 100644 --- a/include/knowhere/bitsetview.h +++ b/include/knowhere/bitsetview.h @@ -12,11 +12,15 @@ #ifndef BITSET_H #define BITSET_H +#include + #include #include +#include #include #include #include +#include namespace knowhere { class BitsetView { @@ -25,12 +29,38 @@ class BitsetView { ~BitsetView() = default; BitsetView(const uint8_t* data, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) - : bits_(data), num_bits_(num_bits), num_filtered_out_bits_(num_filtered_out_bits), id_offset_(id_offset) { + : kind_(Kind::Dense), + bits_(data), + num_bits_(num_bits), + num_filtered_out_bits_(num_filtered_out_bits), + id_offset_(id_offset) { + } + + BitsetView(const roaring_bitmap_t* bitmap, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) + : kind_(Kind::Roaring), + roaring_(bitmap), + num_bits_(num_bits), + num_filtered_out_bits_(num_filtered_out_bits), + id_offset_(id_offset) { } BitsetView(const std::nullptr_t) : BitsetView() { } + BitsetView(const std::nullptr_t, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) + : BitsetView(static_cast(nullptr), num_bits, num_filtered_out_bits, id_offset) { + } + + static BitsetView + FromFrozenRoaring(const void* data, size_t byte_size, size_t num_bits, size_t num_filtered_out_bits = 0, + size_t id_offset = 0) { + const auto* bitmap = roaring_bitmap_frozen_view(static_cast(data), byte_size); + BitsetView bitset(bitmap, num_bits, num_filtered_out_bits, id_offset); + bitset.owned_roaring_ = std::shared_ptr( + bitmap, [](const roaring_bitmap_t* p) { roaring_bitmap_free(const_cast(p)); }); + return bitset; + } + bool empty() const { return num_bits_ == 0; @@ -54,6 +84,15 @@ class BitsetView { return num_filtered_out_bits_; } + void + set_count(size_t num_filtered_out_bits) { + if (out_ids_ != nullptr) { + num_filtered_out_ids_ = num_filtered_out_bits; + return; + } + num_filtered_out_bits_ = num_filtered_out_bits; + } + size_t byte_size() const { return (num_bits_ + 8 - 1) >> 3; @@ -64,6 +103,42 @@ class BitsetView { return bits_; } + bool + is_dense() const { + return kind_ == Kind::Dense; + } + + bool + is_roaring() const { + return kind_ == Kind::Roaring; + } + + const roaring_bitmap_t* + roaring() const { + return roaring_; + } + + size_t + id_offset() const { + return id_offset_; + } + + bool + can_iterate_roaring_without_mapping() const { + return kind_ == Kind::Roaring && out_ids_ == nullptr; + } + + std::vector + ToDense() const { + std::vector dense(byte_size(), 0); + for (size_t i = 0; i < num_bits_; ++i) { + if (test(i)) { + dense[i >> 3] |= 0x1 << (i & 0x7); + } + } + return dense; + } + bool has_out_ids() const { return out_ids_ != nullptr; @@ -103,7 +178,13 @@ class BitsetView { out_id = out_ids_[out_id]; } // when index is larger than the max_offset, ignore it - return (out_id >= static_cast(num_bits_)) || (bits_[out_id >> 3] & (0x1 << (out_id & 0x7))); + if (out_id >= static_cast(num_bits_)) { + return true; + } + if (kind_ == Kind::Roaring) { + return roaring_bitmap_contains(roaring_, static_cast(out_id)); + } + return bits_[out_id >> 3] & (0x1 << (out_id & 0x7)); } // return the filtered ratio. if with id mapping, calculated by internal_ids rather than bits. float @@ -126,6 +207,9 @@ class BitsetView { } return count; } + if (kind_ == Kind::Roaring) { + return roaring_bitmap_get_cardinality(roaring_); + } // if without id mapping, use a better algorithm to calculate the number of filtered out bits. size_t ret = 0; auto len_uint8 = byte_size(); @@ -166,6 +250,14 @@ class BitsetView { } return num_internal_ids_; } + if (kind_ == Kind::Roaring) { + for (size_t i = 0; i < num_bits_; i++) { + if (!test(i)) { + return i; + } + } + return num_bits_; + } // if without id mapping, use a better algorithm to find the first valid index. size_t ret = 0; auto len_uint8 = byte_size(); @@ -211,7 +303,12 @@ class BitsetView { } private: + enum class Kind { Dense, Roaring }; + + Kind kind_ = Kind::Dense; const uint8_t* bits_ = nullptr; + const roaring_bitmap_t* roaring_ = nullptr; + std::shared_ptr owned_roaring_; size_t num_bits_ = 0; size_t num_filtered_out_bits_ = 0; diff --git a/src/index/index.cc b/src/index/index.cc index 70f615494..268a95e2e 100644 --- a/src/index/index.cc +++ b/src/index/index.cc @@ -154,7 +154,8 @@ Index::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b if (bitset_.count() == 0) { // traverse bitset to get the filtered out num auto filtered_out_num = bitset_.get_filtered_out_num_(); - bitset = BitsetView(bitset_.data(), bitset_.size(), filtered_out_num); + bitset = bitset_; + bitset.set_count(filtered_out_num); } else { // if bitset has filtered out num, use it bitset = bitset_; @@ -212,7 +213,10 @@ Index::AnnIterator(const DataSetPtr dataset, const Json& json, const BitsetVi return expected>>::Err(Status::invalid_args, msg); } - const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_()); + auto bitset = bitset_; + if (bitset.count() == 0) { + bitset.set_count(bitset.get_filtered_out_num_()); + } #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) // note that this time includes only the initial search phase of iterator. @@ -252,7 +256,10 @@ Index::RangeSearch(const DataSetPtr dataset, const Json& json, const BitsetVi return expected::Err(Status::invalid_args, msg); } - const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_()); + auto bitset = bitset_; + if (bitset.count() == 0) { + bitset.set_count(bitset.get_filtered_out_num_()); + } #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) const BaseConfig& b_cfg = static_cast(*cfg); diff --git a/tests/ut/test_distances.cc b/tests/ut/test_distances.cc index f815b346e..0fe506e06 100644 --- a/tests/ut/test_distances.cc +++ b/tests/ut/test_distances.cc @@ -12,10 +12,77 @@ #include #include #include +#include +#include +#include +#include +#include #include +#include +#include +#include +#include "knowhere/operands.h" #include "simd/distances_ref.h" #include "simd/hook.h" + +namespace { + +template +void +CheckByIdxDistancePositions(ByIdxDistance by_idx_distance, DirectDistance direct_distance) { + constexpr size_t dim = 3; + constexpr size_t database_size = 32; + const std::array selected_ids = {5, 9, 17, 23, 31}; + const std::array query = {DataType(2.0f), DataType(-3.0f), DataType(5.0f)}; + std::vector database(database_size * dim); + for (size_t row = 0; row < database_size; ++row) { + database[row * dim] = DataType(static_cast(row % 11) - 5.0f); + database[row * dim + 1] = DataType(static_cast((row * 3) % 13) - 6.0f); + database[row * dim + 2] = DataType(static_cast((row * 5) % 17) - 8.0f); + } + + std::vector result_ids; + std::vector result_distances; + by_idx_distance( + query.data(), + database.data(), + selected_ids.data(), + dim, + selected_ids.size(), + [](size_t) -> std::optional { return true; }, + [&](float distance, int64_t id) { + result_ids.push_back(id); + result_distances.push_back(distance); + }); + + REQUIRE(result_ids == std::vector{0, 1, 2, 3, 4}); + REQUIRE(result_distances.size() == selected_ids.size()); + for (size_t i = 0; i < selected_ids.size(); ++i) { + const auto id = selected_ids[i]; + const float expected = direct_distance(query.data(), database.data() + id * dim, dim); + REQUIRE_THAT(result_distances[i], Catch::Matchers::WithinAbs(expected, 0.001f)); + } +} + +struct RecordingDistanceComputer : faiss::DistanceComputer { + void + set_query(const float*) override { + } + + float + operator()(faiss::idx_t id) override { + return static_cast(id); + } + + float + symmetric_dis(faiss::idx_t, faiss::idx_t) override { + return 0.0f; + } +}; + +} // namespace + TEST_CASE("Test Distance Compute", "[distance]") { std::mt19937 rng; std::uniform_int_distribution<> distrib(1, 100000); @@ -122,3 +189,141 @@ TEST_CASE("Test Distance Compute", "[distance]") { } } } + +TEST_CASE("Test typed by-idx distance callbacks use positions", "[distance][typed][by_idx]") { + using namespace faiss::cppcontrib::knowhere; + + SECTION("INT8 inner product") { + CheckByIdxDistancePositions<::knowhere::int8>( + [](auto... args) { int8_vec_inner_products_ny_by_idx_if(args...); }, + int8_vec_inner_product); + } + + SECTION("INT8 L2") { + CheckByIdxDistancePositions<::knowhere::int8>( + [](auto... args) { int8_vec_L2sqr_ny_by_idx_if(args...); }, + int8_vec_L2sqr); + } + + SECTION("FP16 inner product") { + CheckByIdxDistancePositions<::knowhere::fp16>( + [](auto... args) { fp16_vec_inner_products_ny_by_idx_if(args...); }, + fp16_vec_inner_product); + } + + SECTION("FP16 L2") { + CheckByIdxDistancePositions<::knowhere::fp16>( + [](auto... args) { fp16_vec_L2sqr_ny_by_idx_if(args...); }, + fp16_vec_L2sqr); + } + + SECTION("BF16 inner product") { + CheckByIdxDistancePositions<::knowhere::bf16>( + [](auto... args) { bf16_vec_inner_products_ny_by_idx_if(args...); }, + bf16_vec_inner_product); + } + + SECTION("BF16 L2") { + CheckByIdxDistancePositions<::knowhere::bf16>( + [](auto... args) { bf16_vec_L2sqr_ny_by_idx_if(args...); }, + bf16_vec_L2sqr); + } +} + +TEST_CASE("Test generic by-idx distance callbacks use positions", "[distance][by_idx]") { + const std::array ids = {5, 9, 17}; + RecordingDistanceComputer distance_computer; + std::vector callback_positions; + std::vector distances; + + faiss::cppcontrib::knowhere::distance_compute_by_idx_if( + ids.data(), + ids.size(), + &distance_computer, + [](size_t) -> std::optional { return true; }, + [&](float distance, faiss::idx_t position) { + distances.push_back(distance); + callback_positions.push_back(position); + }); + + REQUIRE(callback_positions == std::vector{0, 1, 2}); + REQUIRE(distances == std::vector{5.0f, 9.0f, 17.0f}); +} + +TEST_CASE("Test typed selector results use database IDs", "[distance][typed][selector]") { + using namespace faiss::cppcontrib::knowhere; + constexpr size_t dim = 3; + constexpr size_t database_size = 18; + const std::array selected_ids = {5, 9, 17}; + const std::array<::knowhere::int8, dim> query = {2, -3, 5}; + std::vector<::knowhere::int8> database(database_size * dim); + for (size_t row = 0; row < database_size; ++row) { + database[row * dim] = static_cast<::knowhere::int8>(row % 7 + 1); + database[row * dim + 1] = static_cast<::knowhere::int8>(row % 5 + 2); + database[row * dim + 2] = static_cast<::knowhere::int8>(row % 3 + 3); + } + faiss::IDSelectorArray selector(selected_ids.size(), selected_ids.data()); + + auto check_results = [&](const std::array& distances, + const std::array& result_ids, + auto expected_distance) { + auto sorted_result_ids = result_ids; + std::sort(sorted_result_ids.begin(), sorted_result_ids.end()); + REQUIRE(sorted_result_ids == selected_ids); + for (size_t i = 0; i < result_ids.size(); ++i) { + REQUIRE_THAT(distances[i], Catch::Matchers::WithinAbs(expected_distance(result_ids[i]), 0.001f)); + } + }; + + SECTION("inner product") { + std::array distances; + std::array result_ids; + knn_inner_product_typed( + query.data(), database.data(), dim, 1, database_size, 3, distances.data(), result_ids.data(), &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_inner_product(query.data(), database.data() + id * dim, dim); + }); + } + + SECTION("L2") { + std::array distances; + std::array result_ids; + knn_L2sqr_typed(query.data(), + database.data(), + dim, + 1, + database_size, + 3, + distances.data(), + result_ids.data(), + nullptr, + &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_L2sqr(query.data(), database.data() + id * dim, dim); + }); + } + + SECTION("cosine") { + std::vector y_inv_norms(database_size); + for (size_t id = 0; id < database_size; ++id) { + y_inv_norms[id] = 1.0f / sqrtf(int8_vec_norm_L2sqr(database.data() + id * dim, dim)); + } + const float x_inv_norm = 1.0f / sqrtf(int8_vec_norm_L2sqr(query.data(), dim)); + std::array distances; + std::array result_ids; + knn_cosine_typed(query.data(), + database.data(), + y_inv_norms.data(), + dim, + 1, + database_size, + 3, + distances.data(), + result_ids.data(), + &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_inner_product(query.data(), database.data() + id * dim, dim) * x_inv_norm * + y_inv_norms[id]; + }); + } +} diff --git a/tests/ut/test_utils.cc b/tests/ut/test_utils.cc index e14053ed3..858ef1e76 100644 --- a/tests/ut/test_utils.cc +++ b/tests/ut/test_utils.cc @@ -21,6 +21,7 @@ #include "knowhere/heap.h" #include "knowhere/utils.h" #include "knowhere/version.h" +#include "roaring/roaring.h" #include "utils.h" namespace { @@ -113,6 +114,69 @@ TEST_CASE("Test Bitset Generation", "[utils]") { } } +TEST_CASE("Test BitsetView Roaring", "[utils]") { + std::vector dense(2, 0); + for (auto id : {1, 3, 8, 13}) { + dense[id >> 3] |= 0x1 << (id & 0x7); + } + + auto* roaring = roaring_bitmap_create(); + for (auto id : {1, 3, 8, 13}) { + roaring_bitmap_add(roaring, id); + } + + knowhere::BitsetView dense_view(dense.data(), 16); + knowhere::BitsetView roaring_view(roaring, 16); + REQUIRE(roaring_view.is_roaring()); + REQUIRE(roaring_view.count() == 0); + REQUIRE(roaring_view.get_filtered_out_num_() == dense_view.get_filtered_out_num_()); + REQUIRE(roaring_view.filter_ratio() == dense_view.filter_ratio()); + REQUIRE(roaring_view.get_first_valid_index() == dense_view.get_first_valid_index()); + for (size_t i = 0; i < 16; ++i) { + REQUIRE(roaring_view.test(i) == dense_view.test(i)); + } + REQUIRE(roaring_view.ToDense() == dense); + + roaring_view.set_id_offset(2); + dense_view.set_id_offset(2); + for (size_t i = 0; i < 14; ++i) { + REQUIRE(roaring_view.test(i) == dense_view.test(i)); + } + + std::vector out_ids{0, 1, 3, 7, 8, 13}; + knowhere::BitsetView roaring_out_ids_view(roaring, 16); + knowhere::BitsetView dense_out_ids_view(dense.data(), 16); + roaring_out_ids_view.set_out_ids(out_ids.data(), out_ids.size()); + dense_out_ids_view.set_out_ids(out_ids.data(), out_ids.size()); + REQUIRE(roaring_out_ids_view.count() == dense_out_ids_view.count()); + for (size_t i = 0; i < out_ids.size(); ++i) { + REQUIRE(roaring_out_ids_view.test(i) == dense_out_ids_view.test(i)); + } + + roaring_bitmap_free(roaring); +} + +TEST_CASE("Test BitsetView Frozen Roaring", "[utils]") { + auto* roaring = roaring_bitmap_create(); + for (auto id : {0, 2, 4, 10}) { + roaring_bitmap_add(roaring, id); + } + const auto frozen_size = roaring_bitmap_frozen_size_in_bytes(roaring); + std::vector frozen(frozen_size); + roaring_bitmap_frozen_serialize(roaring, frozen.data()); + + auto frozen_view = knowhere::BitsetView::FromFrozenRoaring(frozen.data(), frozen.size(), 12); + REQUIRE(frozen_view.is_roaring()); + REQUIRE(frozen_view.get_filtered_out_num_() == 4); + REQUIRE(frozen_view.get_first_valid_index() == 1); + REQUIRE(frozen_view.test(0)); + REQUIRE(!frozen_view.test(1)); + REQUIRE(frozen_view.test(10)); + REQUIRE(frozen_view.test(12)); + + roaring_bitmap_free(roaring); +} + namespace { constexpr size_t kHeapSize = 10; constexpr size_t kElementCount = 10000; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h index 082128c8f..c5e580eef 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h @@ -1270,4 +1270,3 @@ void int8_vec_L2sqr_ny_by_idx_if( } } } //namespace faiss - diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp index f300de47e..af2a1a6a7 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include @@ -30,6 +32,61 @@ namespace faiss::cppcontrib::knowhere { namespace { +bool make_roaring_valid_ids( + const ::knowhere::BitsetView& bitset_view, + size_t ny, + std::vector& ids) { + ids.clear(); + if (!bitset_view.can_iterate_roaring_without_mapping()) { + return false; + } + + const uint64_t range_start = bitset_view.id_offset(); + const uint64_t range_end = range_start + ny; + std::unique_ptr range( + roaring_bitmap_from_range(range_start, range_end, 1), + roaring_bitmap_free); + if (range == nullptr) { + return false; + } + std::unique_ptr selected( + roaring_bitmap_andnot(range.get(), bitset_view.roaring()), + roaring_bitmap_free); + if (selected == nullptr) { + return false; + } + + const auto nselected = roaring_bitmap_get_cardinality(selected.get()); + ids.resize(nselected); + std::vector selected_ids(nselected); + roaring_bitmap_to_uint32_array(selected.get(), selected_ids.data()); + for (size_t i = 0; i < nselected; ++i) { + ids[i] = static_cast(selected_ids[i] - range_start); + } + return true; +} + +template +void dispatch_typed_with_roaring_fast_path( + const IDSelector* sel, + size_t ny, + Fn&& fn) { + if (const auto* sel_bs = + dynamic_cast(sel)) { + std::vector roaring_ids; + if (make_roaring_valid_ids(sel_bs->bitset_view, ny, roaring_ids)) { + IDSelectorArray sel_array(roaring_ids.size(), roaring_ids.data()); + fn(sel_array); + } else { + fn(*sel_bs); + } + } else if (sel == nullptr) { + fn(IDSelectorAll()); + } else { + fn(*sel); + } +} + template void exhaustive_inner_product_impl_typed( const DataType* __restrict x, @@ -52,17 +109,20 @@ void exhaustive_inner_product_impl_typed( resi.add_result(ip, j); }; if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, &selector](const float distance, + const idx_t position) { + resi.add_result(distance, selector.ids[position]); + }; if constexpr (std::is_same_v) { fp16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -103,17 +163,20 @@ void exhaustive_L2sqr_seq_impl_typed( resi.add_result(ip, j); }; if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, &selector](const float distance, + const idx_t position) { + resi.add_result(distance, selector.ids[position]); + }; if constexpr (std::is_same_v) { fp16_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -170,17 +233,34 @@ void exhaustive_cosine_seq_impl_typed( resi.begin(i); // the lambda that applies a filtered element if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, + &selector, + x_inv_norm, + y, + y_inv_norms, + d, + norm_computer](const float ip, + const idx_t position) { + const auto id = selector.ids[position]; + const float y_inv_norm = (y_inv_norms != nullptr) + ? y_inv_norms[id] + : 1.0f / + std::max( + sqrtf(norm_computer(y + id * d, d)), + FLT_MIN); + + resi.add_result(ip * (x_inv_norm * y_inv_norm), id); + }; if constexpr (std::is_same_v) { fp16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -221,26 +301,18 @@ void knn_inner_product_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed( + x, y, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed( + x, y, d, nx, ny, res, selector); + }); } if (imin != 0) { @@ -266,15 +338,9 @@ void all_inner_product_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -288,15 +354,9 @@ void all_inner_product_distances_typed( float* output, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, output); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, selector); + }); return; } @@ -322,26 +382,18 @@ void knn_L2sqr_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed( + x, y, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed( + x, y, d, nx, ny, res, selector); + }); } if (imin != 0) { for (size_t i = 0; i < nx * k; i++) { @@ -367,14 +419,9 @@ void all_L2sqr_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -389,14 +436,9 @@ void all_L2sqr_distances_typed( const float* y_norms, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, output); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, selector); + }); return; } @@ -422,30 +464,18 @@ void knn_cosine_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_cosine_seq_impl_typed( + x, y, y_inv_norms, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_cosine_seq_impl_typed( + x, y, y_inv_norms, d, nx, ny, res, selector); + }); } if (imin != 0) { for (size_t i = 0; i < nx * k; i++) { @@ -471,16 +501,10 @@ void all_cosine_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, res, *sel); - } + x, y, y_inv_norms, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -495,16 +519,10 @@ void all_cosine_distances_typed( float* distances, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, distances); - if (const auto* sel_bs = - dynamic_cast(sel)) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, res, *sel); - } + x, y, y_inv_norms, d, nx, ny, res, selector); + }); return; } @@ -523,14 +541,9 @@ void range_search_L2sqr_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, selector); + }); return; } @@ -545,15 +558,9 @@ void range_search_inner_product_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, selector); + }); return; } @@ -570,16 +577,10 @@ void range_search_cosine_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, resh, *sel); - } + x, y, y_inv_norms, d, nx, ny, resh, selector); + }); return; } @@ -962,4 +963,3 @@ template void faiss::cppcontrib::knowhere::range_search_cosine_typed<::knowhere: const faiss::IDSelector*); } -