diff --git a/thirdparty/faiss/faiss/IndexBinaryIVF.cpp b/thirdparty/faiss/faiss/IndexBinaryIVF.cpp index 622350f67..e2322ed0c 100644 --- a/thirdparty/faiss/faiss/IndexBinaryIVF.cpp +++ b/thirdparty/faiss/faiss/IndexBinaryIVF.cpp @@ -431,7 +431,7 @@ struct IVFBinaryScannerJaccard: BinaryInvertedListScanner { const idx_t *ids, int32_t *simi, idx_t *idxi, size_t k, - const BitsetView bitset = nullptr) const override { + const BitsetView bitset) const override { using C = CMax; float* psimi = (float*)simi; size_t nup = 0; @@ -601,7 +601,7 @@ void search_knn_binary_dis_heap( idx_t* labels, bool store_pairs, const IVFSearchParameters* params, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { idx_t nprobe = params ? params->nprobe : ivf.nprobe; nprobe = std::min((idx_t)ivf.nlist, nprobe); idx_t max_codes = params ? params->max_codes : ivf.max_codes; diff --git a/thirdparty/faiss/faiss/IndexPQ.h b/thirdparty/faiss/faiss/IndexPQ.h index 332dbbfb1..5e65c00a6 100644 --- a/thirdparty/faiss/faiss/IndexPQ.h +++ b/thirdparty/faiss/faiss/IndexPQ.h @@ -46,7 +46,7 @@ struct IndexPQ : IndexFlatCodes { idx_t k, float* distances, idx_t* labels, - const BitsetView bitset) const override; + const BitsetView bitset = nullptr) const override; /* The standalone codec interface */ void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override; diff --git a/thirdparty/faiss/faiss/utils/BinaryDistance.h b/thirdparty/faiss/faiss/utils/BinaryDistance.h index 651a73305..c1b11d388 100644 --- a/thirdparty/faiss/faiss/utils/BinaryDistance.h +++ b/thirdparty/faiss/faiss/utils/BinaryDistance.h @@ -153,7 +153,7 @@ namespace faiss { size_t ncodes, std::vector& result, size_t buffer_size, - const BitsetView bitset); + const BitsetView bitset = nullptr); extern template void binary_range_search, int>( @@ -166,7 +166,7 @@ namespace faiss { size_t ncodes, std::vector& result, size_t buffer_size, - const BitsetView bitset); + const BitsetView bitset = nullptr); extern template void binary_range_search, float>( @@ -179,7 +179,7 @@ namespace faiss { size_t ncodes, std::vector& result, size_t buffer_size, - const BitsetView bitset); + const BitsetView bitset = nullptr); extern template void binary_range_search, bool>( @@ -192,7 +192,7 @@ namespace faiss { size_t ncodes, std::vector& result, size_t buffer_size, - const BitsetView bitset); + const BitsetView bitset = nullptr); } // namespace faiss diff --git a/thirdparty/faiss/faiss/utils/distances.cpp b/thirdparty/faiss/faiss/utils/distances.cpp index 1d34d2e8e..decc76f28 100644 --- a/thirdparty/faiss/faiss/utils/distances.cpp +++ b/thirdparty/faiss/faiss/utils/distances.cpp @@ -109,7 +109,7 @@ void exhaustive_parallel_on_nx( size_t ny, ResultHandler& res, fvec_func_ptr dis_compute_func, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { using SingleResultHandler = typename ResultHandler::SingleResultHandler; #pragma omp parallel { @@ -141,7 +141,7 @@ void exhaustive_parallel_on_ny( size_t ny, ResultHandler& res, fvec_func_ptr dis_compute_func, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { using SingleResultHandler = typename ResultHandler::SingleResultHandler; size_t k = res.k; size_t thread_max_num = omp_get_max_threads(); @@ -206,7 +206,7 @@ void exhaustive_L2sqr_IP_seq( size_t ny, ResultHandler& res, fvec_func_ptr dis_compute_func, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { size_t thread_max_num = omp_get_max_threads(); if (ny > parallel_policy_threshold || (nx < thread_max_num / 2 && ny >= thread_max_num * 32)) { exhaustive_parallel_on_ny(x, y, d, nx, ny, res, dis_compute_func, bitset); @@ -224,7 +224,7 @@ void exhaustive_inner_product_blas( size_t nx, size_t ny, ResultHandler& res, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { // BLAS does not like empty matrices if (nx == 0 || ny == 0) return; @@ -369,7 +369,7 @@ static void knn_jaccard_blas( size_t ny, ResultHandler& res, const DistanceCorrection& corr, - const BitsetView bitset = nullptr) { + const BitsetView bitset) { // BLAS does not like empty matrices if (nx == 0 || ny == 0) return; diff --git a/unittest/test_binaryidmap.cpp b/unittest/test_binaryidmap.cpp index ea6334f9b..9e1ce3a66 100644 --- a/unittest/test_binaryidmap.cpp +++ b/unittest/test_binaryidmap.cpp @@ -22,6 +22,10 @@ using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; +typedef int (*binary_int_func_ptr)(const int64_t*, const int64_t*); +typedef float (*binary_float_func_ptr)(const int64_t*, const int64_t*); +typedef bool (*binary_bool_func_ptr)(const int64_t*, const int64_t*); + class BinaryIDMAPTest : public DataGen, public TestWithParam { protected: void @@ -33,6 +37,89 @@ class BinaryIDMAPTest : public DataGen, public TestWithParam { void TearDown() override{}; + // for hamming + void + RunRangeSearchBF( + std::vector>& golden_result, + std::vector& golden_cnt, + binary_int_func_ptr func, + float radius) { + for (auto i = 0; i < nq; ++i) { + const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; + const int64_t* pb = reinterpret_cast(xb_bin.data()); + for (auto j = 0; j < nb; ++j) { + auto dist = func(pq, pb + j); + if (dist < radius) { + golden_result[i][j] = true; + golden_cnt[i]++; + } + } + } + }; + + // for jaccard & tanimoto + void + RunRangeSearchBF( + std::vector>& golden_result, + std::vector& golden_cnt, + binary_float_func_ptr func, + float radius) { + for (auto i = 0; i < nq; ++i) { + const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; + const int64_t* pb = reinterpret_cast(xb_bin.data()); + for (auto j = 0; j < nb; ++j) { + auto dist = func(pq, pb + j); + if (dist < radius) { + golden_result[i][j] = true; + golden_cnt[i]++; + } + } + } + }; + + // for superstructure & substructure + void + RunRangeSearchBF( + std::vector>& golden_result, + std::vector& golden_cnt, + binary_bool_func_ptr func, + float radius) { + for (auto i = 0; i < nq; ++i) { + const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; + const int64_t* pb = reinterpret_cast(xb_bin.data()); + for (auto j = 0; j < nb; ++j) { + auto dist = func(pq, pb + j); + if (dist > radius) { + golden_result[i][j] = true; + golden_cnt[i]++; + } + } + } + }; + + void + CheckRangeSearchResult( + std::vector>& golden_result, + std::vector& golden_cnt, + std::vector& results) { + for (auto i = 0; i < nq; ++i) { + int correct_cnt = 0; + for (auto& res_space : results[i]) { + auto qnr = res_space->buffer_size * res_space->buffers.size() - res_space->buffer_size + res_space->wp; + for (auto j = 0; j < qnr; ++j) { + auto bno = j / res_space->buffer_size; + auto pos = j % res_space->buffer_size; + auto idx = res_space->buffers[bno].ids[pos]; + ASSERT_EQ(golden_result[i][idx], true); + if (golden_result[i][idx]) { + correct_cnt++; + } + } + } + ASSERT_EQ(correct_cnt, golden_cnt[i]); + } + } + protected: knowhere::BinaryIDMAPPtr index_ = nullptr; }; @@ -93,8 +180,10 @@ TEST_P(BinaryIDMAPTest, binaryidmap_basic) { TEST_P(BinaryIDMAPTest, binaryidmap_serialize) { auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { - FileIOWriter writer(filename); - writer(static_cast(bin->data.get()), bin->size); + { + FileIOWriter writer(filename); + writer(static_cast(bin->data.get()), bin->size); + } FileIOReader reader(filename); reader(ret, bin->size); @@ -107,33 +196,31 @@ TEST_P(BinaryIDMAPTest, binaryidmap_serialize) { {knowhere::Metric::TYPE, MetricType}, }; - { - // serialize index - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); - auto re_result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(re_result, nq, k); - // PrintResult(re_result, nq, k); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto binaryset = index_->Serialize(conf); - auto bin = binaryset.GetByName("BinaryIVF"); - - std::string filename = temp_path("/tmp/bianryidmap_test_serialize.bin"); - auto load_data = new uint8_t[bin->size]; - serialize(filename, bin, load_data); - - binaryset.clear(); - std::shared_ptr data(load_data); - binaryset.Append("BinaryIVF", data, bin->size); - - index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(result, nq, k); - // PrintResult(result, nq, k); - } + // serialize index + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); + auto result1 = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result1, nq, k); + // PrintResult(result1, nq, k); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto binaryset = index_->Serialize(conf); + auto bin = binaryset.GetByName("BinaryIVF"); + + std::string filename = temp_path("/tmp/bianryidmap_test_serialize.bin"); + auto load_data = new uint8_t[bin->size]; + serialize(filename, bin, load_data); + + binaryset.clear(); + std::shared_ptr data(load_data); + binaryset.Append("BinaryIVF", data, bin->size); + + index_->Load(binaryset); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto result2 = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result2, nq, k); + // PrintResult(result2, nq, k); } TEST_P(BinaryIDMAPTest, binaryidmap_slice) { @@ -145,24 +232,22 @@ TEST_P(BinaryIDMAPTest, binaryidmap_slice) { {knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, knowhere::index_file_slice_size}, }; - { - // serialize index - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); - auto re_result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(re_result, nq, k); - // PrintResult(re_result, nq, k); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto binaryset = index_->Serialize(conf); - - index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(result, nq, k); - // PrintResult(result, nq, k); - } + // serialize index + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); + auto result1 = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result1, nq, k); + // PrintResult(result1, nq, k); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto binaryset = index_->Serialize(conf); + + index_->Load(binaryset); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto result2 = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result2, nq, k); + // PrintResult(result2, nq, k); } TEST_P(BinaryIDMAPTest, binaryidmap_range_search) { @@ -174,50 +259,22 @@ TEST_P(BinaryIDMAPTest, binaryidmap_range_search) { {knowhere::Metric::TYPE, MetricType}, }; - std::vector> idmap(nq, std::vector(nb, false)); - std::vector bf_cnt(nq, 0); + std::vector> golden_result(nq, std::vector(nb, false)); + std::vector golden_cnt(nq, 0); + // hamming int hamming_radius = 10; auto hamming_dis = [](const int64_t* pa, const int64_t* pb) -> int { return __builtin_popcountl((*pa) ^ (*pb)); }; - auto bruteforce_ham = [&]() { - // std::cout << "show hamming bruteforce ans:" << std::endl; - for (auto i = 0; i < nq; ++i) { - const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; - const int64_t* pb = reinterpret_cast(xb_bin.data()); - for (auto j = 0; j < nb; ++j) { - auto dist = hamming_dis(pq, pb + j); - if (dist < hamming_radius) { - idmap[i][j] = true; - bf_cnt[i]++; - // std::cout << "i = " << i << ", j = " << j << std::endl; - } - } - } - }; + // jaccard float jaccard_radius = 0.4; auto jaccard_dis = [](const int64_t* pa, const int64_t* pb) -> float { auto intersection = __builtin_popcountl((*pa) & (*pb)); auto Union = __builtin_popcountl((*pa) | (*pb)); return 1.0 - (double)intersection / (double)Union; }; - auto bruteforce_jcd = [&]() { - // std::cout << "show jaccard bruteforce ans:" << std::endl; - for (auto i = 0; i < nq; ++i) { - const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; - const int64_t* pb = reinterpret_cast(xb_bin.data()); - for (auto j = 0; j < nb; ++j) { - auto dist = jaccard_dis(pq, pb + j); - if (dist < jaccard_radius) { - idmap[i][j] = true; - bf_cnt[i]++; - // std::cout << "query_id = " << i << ", ans_j = " << j << ", dis_j = " << dist - // << std::endl; - } - } - } - }; + // tanimoto float tanimoto_radius = 0.2; auto tanimoto_dis = [](const int64_t* pa, const int64_t* pb) -> float { auto intersection = __builtin_popcountl((*pa) & (*pb)); @@ -225,145 +282,61 @@ TEST_P(BinaryIDMAPTest, binaryidmap_range_search) { auto jcd = 1.0 - (double)intersection / (double)Union; return (-log2(1 - jcd)); }; - auto bruteforce_tnm = [&]() { - // std::cout << "show tanimoto bruteforce ans:" << std::endl; - for (auto i = 0; i < nq; ++i) { - const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; - const int64_t* pb = reinterpret_cast(xb_bin.data()); - for (auto j = 0; j < nb; ++j) { - auto dist = tanimoto_dis(pq, pb + j); - if (dist < tanimoto_radius) { - idmap[i][j] = true; - bf_cnt[i]++; - // std::cout << "query_id = " << i << ", ans_j = " << j << ", dis_j = " << dist - // << std::endl; - } - } - } - }; + // superstructure bool superstructure_radius = false; auto superstructure_dis = [](const int64_t* pa, const int64_t* pb) -> bool { return ((*pa) & (*pb)) == (*pb); }; - auto bruteforce_sup = [&]() { - for (auto i = 0; i < nq; ++i) { - const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; - const int64_t* pb = reinterpret_cast(xb_bin.data()); - for (auto j = 0; j < nb; ++j) { - auto dist = superstructure_dis(pq, pb + j); - if (dist > superstructure_radius) { - idmap[i][j] = true; - bf_cnt[i]++; - } - } - } - }; + // substructure int substructure_radius = false; auto substructure_dis = [](const int64_t* pa, const int64_t* pb) -> bool { return ((*pa) & (*pb)) == (*pa); }; - auto bruteforce_sub = [&]() { - for (auto i = 0; i < nq; ++i) { - const int64_t* pq = reinterpret_cast(xq_bin.data()) + i; - const int64_t* pb = reinterpret_cast(xb_bin.data()); - for (auto j = 0; j < nb; ++j) { - auto dist = substructure_dis(pq, pb + j); - if (dist > substructure_radius) { - idmap[i][j] = true; - bf_cnt[i]++; - } - } - } - }; - auto mt = conf[knowhere::Metric::TYPE].get(); - // std::cout << "current metric_type = " << mt << std::endl; - float set_radius = radius; - if ("HAMMING" == mt) { - set_radius = hamming_radius; - bruteforce_ham(); - // for (auto i = 0;i < nq; ++ i) - // std::cout << bf_cnt[i] << " "; - // std::cout << std::endl; - } else if ("JACCARD" == mt) { - set_radius = jaccard_radius; - bruteforce_jcd(); - // for (auto i = 0;i < nq; ++ i) - // std::cout << bf_cnt[i] << " "; - // std::cout << std::endl; - } else if ("TANIMOTO" == mt) { - set_radius = tanimoto_radius; - bruteforce_tnm(); - // for (auto i = 0;i < nq; ++ i) - // std::cout << bf_cnt[i] << " "; - // std::cout << std::endl; - } else if ("SUPERSTRUCTURE" == mt) { - set_radius = superstructure_radius; - bruteforce_sup(); - // for (auto i = 0;i < nq; ++ i) - // std::cout << bf_cnt[i] << " "; - // std::cout << std::endl; - } else if ("SUBSTRUCTURE" == mt) { - set_radius = substructure_radius; - bruteforce_sub(); - // for (auto i = 0;i < nq; ++ i) - // std::cout << bf_cnt[i] << " "; - // std::cout << std::endl; + auto metric_type = conf[knowhere::Metric::TYPE].get(); + float curr_radius = radius; + if (metric_type == "HAMMING") { + curr_radius = hamming_radius; + RunRangeSearchBF(golden_result, golden_cnt, hamming_dis, hamming_radius); + } else if (metric_type == "JACCARD") { + curr_radius = jaccard_radius; + RunRangeSearchBF(golden_result, golden_cnt, jaccard_dis, jaccard_radius); + } else if (metric_type == "TANIMOTO") { + curr_radius = tanimoto_radius; + RunRangeSearchBF(golden_result, golden_cnt, tanimoto_dis, tanimoto_radius); + } else if (metric_type == "SUPERSTRUCTURE") { + curr_radius = superstructure_radius; + RunRangeSearchBF(golden_result, golden_cnt, superstructure_dis, superstructure_radius); + } else if (metric_type == "SUBSTRUCTURE") { + curr_radius = substructure_radius; + RunRangeSearchBF(golden_result, golden_cnt, substructure_dis, substructure_radius); } else { std::cout << "unsupport type of metric type" << std::endl; } - conf[knowhere::IndexParams::range_search_radius] = set_radius; - // std::cout << "current radius = " << conf[knowhere::IndexParams::range_search_radius].get() << - // std::endl; - auto compare_res = [&](std::vector& results) { - // std::cout << "show faiss ans:" << std::endl; - for (auto i = 0; i < nq; ++i) { - int correct_cnt = 0; - for (auto& res_space : results[i]) { - auto qnr = res_space->buffer_size * res_space->buffers.size() - res_space->buffer_size + res_space->wp; - for (auto j = 0; j < qnr; ++j) { - auto bno = j / res_space->buffer_size; - auto pos = j % res_space->buffer_size; - // std::cout << "query_id = " << i << ", ans_j = " << - // res_space->buffers[bno].ids[pos] << ", dis_j = " << - // res_space->buffers[bno].dis[pos] << std::endl; - ASSERT_EQ(idmap[i][res_space->buffers[bno].ids[pos]], true); - if (idmap[i][res_space->buffers[bno].ids[pos]]) { - correct_cnt++; - } - } - } - ASSERT_EQ(correct_cnt, bf_cnt[i]); - } - }; + conf[knowhere::IndexParams::range_search_radius] = curr_radius; - { - // serialize index - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); + // serialize index + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); - std::vector results; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq_bin.data() + i * dim / 8); - results.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + std::vector results1; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq_bin.data() + i * dim / 8); + results1.push_back(index_->QueryByDistance(qd, conf, nullptr)); + } + CheckRangeSearchResult(golden_result, golden_cnt, results1); - compare_res(results); - // - auto binaryset = index_->Serialize(conf); - index_->Load(binaryset); + auto binaryset = index_->Serialize(conf); + index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - { - std::vector rresults; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq_bin.data() + i * dim / 8); - rresults.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); - compare_res(rresults); - } + std::vector results2; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq_bin.data() + i * dim / 8); + results2.push_back(index_->QueryByDistance(qd, conf, nullptr)); } + CheckRangeSearchResult(golden_result, golden_cnt, results1); } diff --git a/unittest/test_binaryivf.cpp b/unittest/test_binaryivf.cpp index f5f2e9e95..c2e2d7214 100644 --- a/unittest/test_binaryivf.cpp +++ b/unittest/test_binaryivf.cpp @@ -95,8 +95,10 @@ TEST_P(BinaryIVFTest, binaryivf_basic) { TEST_P(BinaryIVFTest, binaryivf_serialize) { auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { - FileIOWriter writer(filename); - writer(static_cast(bin->data.get()), bin->size); + { + FileIOWriter writer(filename); + writer(static_cast(bin->data.get()), bin->size); + } FileIOReader reader(filename); reader(ret, bin->size); diff --git a/unittest/test_idmap.cpp b/unittest/test_idmap.cpp index a7ac4d7c6..bc183a410 100644 --- a/unittest/test_idmap.cpp +++ b/unittest/test_idmap.cpp @@ -13,33 +13,35 @@ #include #include -#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/index/IndexType.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #ifdef KNOWHERE_GPU_VERSION #include #include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif -#include "Helper.h" +#include "knowhere/utils/distances_simd.h" #include "unittest/utils.h" using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; +typedef float (*fvec_func_ptr)(const float*, const float*, size_t); + class IDMAPTest : public DataGen, public TestWithParam { protected: void SetUp() override { + Init_with_default(); #ifdef KNOWHERE_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif index_mode_ = GetParam(); - Init_with_default(); index_ = std::make_shared(); } @@ -50,25 +52,73 @@ class IDMAPTest : public DataGen, public TestWithParam { #endif } + void + RunRangeSearchBF( + std::vector>& golden_result, + std::vector& golden_cnt, + fvec_func_ptr func, + float radius, + bool check_small) { + for (auto i = 0; i < nq; ++i) { + const float* pq = xq.data() + i * dim; + for (auto j = 0; j < nb; ++j) { + const float* pb = xb.data() + j * dim; + auto dist = func(pq, pb, dim); + if ((check_small && dist < radius) || (!check_small && dist > radius)) { + golden_result[i][j] = true; + golden_cnt[i]++; + } + } + } + }; + + void + CheckRangeSearchResult( + std::vector>& golden_result, + std::vector& golden_cnt, + std::vector& results) { + for (auto i = 0; i < nq; ++i) { + int correct_cnt = 0; + for (auto& res_space : results[i]) { + auto qnr = + res_space->buffer_size * res_space->buffers.size() - res_space->buffer_size + res_space->wp; + for (auto j = 0; j < qnr; ++j) { + auto bno = j / res_space->buffer_size; + auto pos = j % res_space->buffer_size; + auto idx = res_space->buffers[bno].ids[pos]; + ASSERT_EQ(golden_result[i][idx], true); + if (golden_result[i][idx]) { + correct_cnt++; + } + } + } + ASSERT_EQ(correct_cnt, golden_cnt[i]); + } + } + protected: knowhere::IDMAPPtr index_ = nullptr; knowhere::IndexMode index_mode_; }; -INSTANTIATE_TEST_CASE_P(IDMAPParameters, - IDMAPTest, - Values( +INSTANTIATE_TEST_CASE_P( + IDMAPParameters, + IDMAPTest, + Values( #ifdef KNOWHERE_GPU_VERSION - knowhere::IndexMode::MODE_GPU, + knowhere::IndexMode::MODE_GPU, #endif - knowhere::IndexMode::MODE_CPU)); + knowhere::IndexMode::MODE_CPU) + ); TEST_P(IDMAPTest, idmap_basic) { ASSERT_TRUE(!xb.empty()); - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::meta::TOPK, k}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} + }; // null faiss index { @@ -84,21 +134,21 @@ TEST_P(IDMAPTest, idmap_basic) { ASSERT_TRUE(index_->GetRawVectors() != nullptr); auto result = index_->Query(query_dataset, conf, nullptr); AssertAnns(result, nq, k); - // PrintResult(result, nq, k); + // PrintResult(result, nq, k); - if (index_mode_ == knowhere::IndexMode::MODE_GPU) { #ifdef KNOWHERE_GPU_VERSION + if (index_mode_ == knowhere::IndexMode::MODE_GPU) { // cpu to gpu index_ = std::dynamic_pointer_cast(index_->CopyCpuToGpu(DEVICEID, conf)); -#endif } +#endif auto binaryset = index_->Serialize(conf); auto new_index = std::make_shared(); new_index->Load(binaryset); auto result2 = new_index->Query(query_dataset, conf, nullptr); AssertAnns(result2, nq, k); - // PrintResult(re_result, nq, k); + // PrintResult(re_result, nq, k); #if 0 auto result3 = new_index->QueryById(id_dataset, conf); @@ -127,333 +177,207 @@ TEST_P(IDMAPTest, idmap_basic) { TEST_P(IDMAPTest, idmap_serialize) { auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { - FileIOWriter writer(filename); - writer(static_cast(bin->data.get()), bin->size); + { + FileIOWriter writer(filename); + writer(static_cast(bin->data.get()), bin->size); + } FileIOReader reader(filename); reader(ret, bin->size); }; - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::meta::TOPK, k}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} + }; - { - // serialize index - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); - if (index_mode_ == knowhere::IndexMode::MODE_GPU) { #ifdef KNOWHERE_GPU_VERSION - // cpu to gpu - index_ = std::dynamic_pointer_cast(index_->CopyCpuToGpu(DEVICEID, conf)); + if (index_mode_ == knowhere::IndexMode::MODE_GPU) { + // cpu to gpu + index_ = std::dynamic_pointer_cast(index_->CopyCpuToGpu(DEVICEID, conf)); + } #endif - } - auto re_result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(re_result, nq, k); - // PrintResult(re_result, nq, k); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto binaryset = index_->Serialize(conf); - auto bin = binaryset.GetByName("IVF"); - - std::string filename = temp_path("/tmp/idmap_test_serialize.bin"); - auto load_data = new uint8_t[bin->size]; - serialize(filename, bin, load_data); - - binaryset.clear(); - std::shared_ptr data(load_data); - binaryset.Append("IVF", data, bin->size); - - index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(result, nq, k); - // PrintResult(result, nq, k); - } + auto re_result = index_->Query(query_dataset, conf, nullptr); + AssertAnns(re_result, nq, k); + // PrintResult(re_result, nq, k); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto binaryset = index_->Serialize(conf); + auto bin = binaryset.GetByName("IVF"); + + std::string filename = temp_path("/tmp/idmap_test_serialize.bin"); + auto load_data = new uint8_t[bin->size]; + serialize(filename, bin, load_data); + + binaryset.clear(); + std::shared_ptr data(load_data); + binaryset.Append("IVF", data, bin->size); + + index_->Load(binaryset); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result, nq, k); + // PrintResult(result, nq, k); } TEST_P(IDMAPTest, idmap_slice) { - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, knowhere::index_file_slice_size}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::meta::TOPK, k}, + {knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, knowhere::index_file_slice_size}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} + }; - { - // serialize index - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); - if (index_mode_ == knowhere::IndexMode::MODE_GPU) { #ifdef KNOWHERE_GPU_VERSION - // cpu to gpu - index_ = std::dynamic_pointer_cast(index_->CopyCpuToGpu(DEVICEID, conf)); + if (index_mode_ == knowhere::IndexMode::MODE_GPU) { + // cpu to gpu + index_ = std::dynamic_pointer_cast(index_->CopyCpuToGpu(DEVICEID, conf)); + } #endif - } - auto re_result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(re_result, nq, k); - // PrintResult(re_result, nq, k); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto binaryset = index_->Serialize(conf); - - index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Query(query_dataset, conf, nullptr); - AssertAnns(result, nq, k); - // PrintResult(result, nq, k); - } + auto re_result = index_->Query(query_dataset, conf, nullptr); + AssertAnns(re_result, nq, k); + // PrintResult(re_result, nq, k); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto binaryset = index_->Serialize(conf); + + index_->Load(binaryset); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf, nullptr); + AssertAnns(result, nq, k); + // PrintResult(result, nq, k); } TEST_P(IDMAPTest, idmap_range_search_l2) { - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::IndexParams::range_search_radius, radius}, - {knowhere::IndexParams::range_search_buffer_size, buffer_size}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; - - auto l2dis = [](const float* pa, const float* pb, size_t dim) -> float { - float ret = 0; - for (auto i = 0; i < dim; ++i) { - auto dif = (pa[i] - pb[i]); - ret += dif * dif; - } - return ret; - }; - - std::vector> idmap(nq, std::vector(nb, false)); - std::vector bf_cnt(nq, 0); - - auto bruteforce = [&]() { - auto rds = radius * radius; - for (auto i = 0; i < nq; ++i) { - const float* pq = xq.data() + i * dim; - for (auto j = 0; j < nb; ++j) { - const float* pb = xb.data() + j * dim; - auto dist = l2dis(pq, pb, dim); - if (dist < rds) { - idmap[i][j] = true; - bf_cnt[i]++; - } - } - } + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::IndexParams::range_search_radius, radius}, + {knowhere::IndexParams::range_search_buffer_size, buffer_size}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} }; - bruteforce(); - - auto compare_res = [&](std::vector& results) { - { // compare the result - // std::cout << "size of result: " << results.size() << std::endl; - for (auto i = 0; i < nq; ++i) { - int correct_cnt = 0; - // std::cout << "query id = " << i << ", result[i].size = " << results[i].size() << - // std::endl; - for (auto& res_space : results[i]) { - // std::cout << "buffer size = " << res_space->buffer_size << ", wp = " << - // res_space->wp << ", size of buffers = " << res_space->buffers.size() << - // std::endl; - auto qnr = - res_space->buffer_size * res_space->buffers.size() - res_space->buffer_size + res_space->wp; - // std::cout << "qnr = " << qnr << std::endl; - for (auto j = 0; j < qnr; ++j) { - auto bno = j / res_space->buffer_size; - auto pos = j % res_space->buffer_size; - ASSERT_EQ(idmap[i][res_space->buffers[bno].ids[pos]], true); - if (idmap[i][res_space->buffers[bno].ids[pos]]) { - correct_cnt++; - } - } - } - ASSERT_EQ(correct_cnt, bf_cnt[i]); - } - } - }; - - { - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); + std::vector> golden_result(nq, std::vector(nb, false)); + std::vector golden_cnt(nq, 0); + RunRangeSearchBF(golden_result, golden_cnt, faiss::fvec_L2sqr_ref, radius * radius, true); - std::vector results; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); - results.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); - compare_res(results); + std::vector results1; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); + results1.push_back(index_->QueryByDistance(qd, conf, nullptr)); + } + CheckRangeSearchResult(golden_result, golden_cnt, results1); - auto binaryset = index_->Serialize(conf); - index_->Load(binaryset); + auto binaryset = index_->Serialize(conf); + index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - { // query again and compare the result - std::vector rresults; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); - rresults.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); - compare_res(rresults); - } + // query again and compare the result + std::vector results2; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); + results2.push_back(index_->QueryByDistance(qd, conf, nullptr)); } + CheckRangeSearchResult(golden_result, golden_cnt, results2); } TEST_P(IDMAPTest, idmap_range_search_ip) { - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::IndexParams::range_search_radius, radius}, - {knowhere::IndexParams::range_search_buffer_size, buffer_size}, - {knowhere::Metric::TYPE, knowhere::Metric::IP}}; - - auto ipdis = [](const float* pa, const float* pb, size_t dim) -> float { - float ret = 0; - for (auto i = 0; i < dim; ++i) { - ret += pa[i] * pb[i]; - } - return ret; + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::IndexParams::range_search_radius, radius}, + {knowhere::IndexParams::range_search_buffer_size, buffer_size}, + {knowhere::Metric::TYPE, knowhere::Metric::IP} }; - std::vector> idmap(nq, std::vector(nb, false)); - std::vector bf_cnt(nq, 0); - - auto bruteforce = [&]() { - for (auto i = 0; i < nq; ++i) { - const float* pq = xq.data() + i * dim; - for (auto j = 0; j < nb; ++j) { - const float* pb = xb.data() + j * dim; - auto dist = ipdis(pq, pb, dim); - if (dist > radius) { - idmap[i][j] = true; - bf_cnt[i]++; - } - } - } - }; + std::vector> golden_result(nq, std::vector(nb, false)); + std::vector golden_cnt(nq, 0); + RunRangeSearchBF(golden_result, golden_cnt, faiss::fvec_inner_product_ref, radius, false); - bruteforce(); - - auto compare_res = [&](std::vector& results) { - { // compare the result - for (auto i = 0; i < nq; ++i) { - int correct_cnt = 0; - for (auto& res_space : results[i]) { - auto qnr = - res_space->buffer_size * res_space->buffers.size() - res_space->buffer_size + res_space->wp; - for (auto j = 0; j < qnr; ++j) { - auto bno = j / res_space->buffer_size; - auto pos = j % res_space->buffer_size; - ASSERT_EQ(idmap[i][res_space->buffers[bno].ids[pos]], true); - if (idmap[i][res_space->buffers[bno].ids[pos]]) { - correct_cnt++; - } - } - } - ASSERT_EQ(correct_cnt, bf_cnt[i]); - } - } - }; - - { - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); - - std::vector results; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); - results.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); - compare_res(results); + std::vector results1; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); + results1.push_back(index_->QueryByDistance(qd, conf, nullptr)); + } + CheckRangeSearchResult(golden_result, golden_cnt, results1); - auto binaryset = index_->Serialize(conf); - index_->Load(binaryset); + auto binaryset = index_->Serialize(conf); + index_->Load(binaryset); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dim(), dim); - { // query again and compare the result - std::vector rresults; - for (auto i = 0; i < nq; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); - rresults.push_back(index_->QueryByDistance(qd, conf, nullptr)); - } + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dim(), dim); - compare_res(rresults); - } + // query again and compare the result + std::vector results2; + for (auto i = 0; i < nq; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq.data() + i * dim); + results2.push_back(index_->QueryByDistance(qd, conf, nullptr)); } + CheckRangeSearchResult(golden_result, golden_cnt, results2); } TEST_P(IDMAPTest, idmap_dynamic_result_set) { - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::IndexParams::range_search_radius, radius}, - {knowhere::IndexParams::range_search_buffer_size, buffer_size}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; - - auto l2dis = [](const float* pa, const float* pb, size_t dim) -> float { - float ret = 0; - for (auto i = 0; i < dim; ++i) { - auto dif = (pa[i] - pb[i]); - ret += dif * dif; - } - return ret; - }; - - std::vector> idmap(nq, std::vector(nb, false)); - std::vector bf_cnt(nq, 0); - - auto bruteforce = [&]() { - auto rds = radius * radius; - for (auto i = 0; i < nq; ++i) { - const float* pq = xq.data() + i * dim; - for (auto j = 0; j < nb; ++j) { - const float* pb = xb.data() + j * dim; - auto dist = l2dis(pq, pb, dim); - if (dist < rds) { - idmap[i][j] = true; - bf_cnt[i]++; - } - } - } + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::IndexParams::range_search_radius, radius}, + {knowhere::IndexParams::range_search_buffer_size, buffer_size}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} }; - bruteforce(); + std::vector> golden_result(nq, std::vector(nb, false)); + std::vector golden_cnt(nq, 0); + RunRangeSearchBF(golden_result, golden_cnt, faiss::fvec_L2sqr_ref, radius * radius, true); - auto check_rst = [&](knowhere::DynamicResultSet& rst, knowhere::ResultSetPostProcessType rspt) { - { // compare the result - for (auto i = 0; i < rst.count - 1; ++i) { - if (rspt == knowhere::ResultSetPostProcessType::SortAsc) - ASSERT_LE(rst.distances.get() + i, rst.distances.get() + i + 1); - else if (rspt == knowhere::ResultSetPostProcessType::SortDesc) - ASSERT_GE(rst.distances.get() + i, rst.distances.get() + i + 1); - } + auto check_rst = [&](knowhere::DynamicResultSet& rst, knowhere::ResultSetPostProcessType type) { + for (auto i = 0; i < rst.count - 1; ++i) { + if (type == knowhere::ResultSetPostProcessType::SortAsc) + ASSERT_LE(rst.distances.get() + i, rst.distances.get() + i + 1); + else if (type == knowhere::ResultSetPostProcessType::SortDesc) + ASSERT_GE(rst.distances.get() + i, rst.distances.get() + i + 1); } }; - { - knowhere::DynamicResultCollector collector; - index_->Train(base_dataset, conf); - index_->AddWithoutIds(base_dataset, knowhere::Config()); + knowhere::DynamicResultCollector collector; + index_->Train(base_dataset, conf); + index_->AddWithoutIds(base_dataset, knowhere::Config()); - for (auto i = 0; i < 3; ++i) { - auto qd = knowhere::GenDataset(1, dim, xq.data()); - collector.Append(index_->QueryByDistance(qd, conf, nullptr)); - } + for (auto i = 0; i < 3; ++i) { + auto qd = knowhere::GenDataset(1, dim, xq.data()); + collector.Append(index_->QueryByDistance(qd, conf, nullptr)); + } - auto rst = collector.Merge(1000, knowhere::ResultSetPostProcessType::SortAsc); - ASSERT_LE(rst.count, 1000); + auto rst = collector.Merge(1000, knowhere::ResultSetPostProcessType::SortAsc); + ASSERT_LE(rst.count, 1000); - check_rst(rst, knowhere::ResultSetPostProcessType::SortAsc); - } + check_rst(rst, knowhere::ResultSetPostProcessType::SortAsc); } #ifdef KNOWHERE_GPU_VERSION TEST_P(IDMAPTest, idmap_copy) { ASSERT_TRUE(!xb.empty()); - knowhere::Config conf{{knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + knowhere::Config conf{ + {knowhere::meta::DIM, dim}, + {knowhere::meta::TOPK, k}, + {knowhere::Metric::TYPE, knowhere::Metric::L2} + }; index_->Train(base_dataset, conf); index_->AddWithoutIds(base_dataset, conf); @@ -464,45 +388,41 @@ TEST_P(IDMAPTest, idmap_copy) { AssertAnns(result, nq, k); // PrintResult(result, nq, k); - { - // clone - // auto clone_index = index_->Clone(); - // auto clone_result = clone_index->Search(query_dataset, conf); - // AssertAnns(clone_result, nq, k); - } - - { - // cpu to gpu - ASSERT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(index_, -1, conf)); - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); - auto clone_result = clone_index->Query(query_dataset, conf, nullptr); - - AssertAnns(clone_result, nq, k); - ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawVectors(); }, - knowhere::KnowhereException); - ASSERT_ANY_THROW(clone_index->Serialize(conf)); - - auto binary = clone_index->Serialize(conf); - clone_index->Load(binary); - auto new_result = clone_index->Query(query_dataset, conf, nullptr); - AssertAnns(new_result, nq, k); - - // auto clone_gpu_idx = clone_index->Clone(); - // auto clone_gpu_res = clone_gpu_idx->Search(query_dataset, conf); - // AssertAnns(clone_gpu_res, nq, k); - - // gpu to cpu - auto host_index = knowhere::cloner::CopyGpuToCpu(clone_index, conf); - auto host_result = host_index->Query(query_dataset, conf, nullptr); - AssertAnns(host_result, nq, k); - ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawVectors() != nullptr); - - // gpu to gpu - auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); - auto new_device_index = - std::static_pointer_cast(device_index)->CopyGpuToGpu(DEVICEID, conf); - auto device_result = new_device_index->Query(query_dataset, conf, nullptr); - AssertAnns(device_result, nq, k); - } + // clone + // auto clone_index = index_->Clone(); + // auto clone_result = clone_index->Search(query_dataset, conf); + // AssertAnns(clone_result, nq, k); + + // cpu to gpu + ASSERT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(index_, -1, conf)); + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); + auto clone_result = clone_index->Query(query_dataset, conf, nullptr); + + AssertAnns(clone_result, nq, k); + ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawVectors(); }, + knowhere::KnowhereException); + ASSERT_ANY_THROW(clone_index->Serialize(conf)); + + auto binary = clone_index->Serialize(conf); + clone_index->Load(binary); + auto new_result = clone_index->Query(query_dataset, conf, nullptr); + AssertAnns(new_result, nq, k); + + // auto clone_gpu_idx = clone_index->Clone(); + // auto clone_gpu_res = clone_gpu_idx->Search(query_dataset, conf); + // AssertAnns(clone_gpu_res, nq, k); + + // gpu to cpu + auto host_index = knowhere::cloner::CopyGpuToCpu(clone_index, conf); + auto host_result = host_index->Query(query_dataset, conf, nullptr); + AssertAnns(host_result, nq, k); + ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawVectors() != nullptr); + + // gpu to gpu + auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); + auto new_device_index = + std::static_pointer_cast(device_index)->CopyGpuToGpu(DEVICEID, conf); + auto device_result = new_device_index->Query(query_dataset, conf, nullptr); + AssertAnns(device_result, nq, k); } #endif diff --git a/unittest/test_ivf.cpp b/unittest/test_ivf.cpp index b122100dc..069b7eae6 100644 --- a/unittest/test_ivf.cpp +++ b/unittest/test_ivf.cpp @@ -178,8 +178,10 @@ TEST_P(IVFTest, ivf_basic_gpu) { TEST_P(IVFTest, ivf_serialize) { auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { - FileIOWriter writer(filename); - writer(static_cast(bin->data.get()), bin->size); + { + FileIOWriter writer(filename); + writer(static_cast(bin->data.get()), bin->size); + } FileIOReader reader(filename); reader(ret, bin->size);