diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp index 34c81991c93ae50ed2dd989804b8dab2bbd95bd2..fe5bf0990a4724fc77555745c51851f002524b1f 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp @@ -189,6 +189,8 @@ IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { if (quantizer_conf->mode != 2) { KNOWHERE_THROW_MSG("mode only support 2 in this func"); } + } else { + KNOWHERE_THROW_MSG("conf error"); } // if (quantizer_conf->gpu_id != gpu_id_) { // KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card"); diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp index d74c6bc562e1b79555fef20f5d7e2f5ad24dc5d0..837629c6eb77c2614c7457d56567d9505f83cc25 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp @@ -63,7 +63,7 @@ FaissGpuResourceMgr::InitResource() { mutex_cache_.emplace(device_id, std::make_unique()); - // std::cout << "Device Id: " << device_id << std::endl; + // std::cout << "Device Id: " << DEVICEID << std::endl; auto& device_param = device.second; auto& bq = idle_map_[device_id]; @@ -119,7 +119,7 @@ void FaissGpuResourceMgr::Dump() { for (auto& item : idle_map_) { auto& bq = item.second; - std::cout << "device_id: " << item.first << ", resource count:" << bq.Size(); + std::cout << "DEVICEID: " << item.first << ", resource count:" << bq.Size(); } } diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt index 1812bc25030f8bb5fa437405520fd6f4eaced5bf..8a5e089486507ce3c2c04b7bb5929f9e14a89926 100644 --- a/core/src/index/unittest/CMakeLists.txt +++ b/core/src/index/unittest/CMakeLists.txt @@ -73,9 +73,17 @@ target_link_libraries(test_kdt SPTAGLibStatic ${depend_libs} ${unittest_libs} ${basic_libs}) +add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs}) +target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs}) + +add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs}) +target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs}) + install(TARGETS test_ivf DESTINATION unittest) install(TARGETS test_idmap DESTINATION unittest) install(TARGETS test_kdt DESTINATION unittest) +install(TARGETS test_gpuresource DESTINATION unittest) +install(TARGETS test_customized_index DESTINATION unittest) #add_subdirectory(faiss_ori) add_subdirectory(test_nsg) diff --git a/core/src/index/unittest/Helper.h b/core/src/index/unittest/Helper.h new file mode 100644 index 0000000000000000000000000000000000000000..d11a484c0334c1f399e189794fedbfe426b7f549 --- /dev/null +++ b/core/src/index/unittest/Helper.h @@ -0,0 +1,120 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" + +constexpr int DEVICEID = 0; +constexpr int64_t DIM = 128; +constexpr int64_t NB = 10000; +constexpr int64_t NQ = 10; +constexpr int64_t K = 10; +constexpr int64_t PINMEM = 1024 * 1024 * 200; +constexpr int64_t TEMPMEM = 1024 * 1024 * 300; +constexpr int64_t RESNUM = 2; + +knowhere::IVFIndexPtr +IndexFactory(const std::string& type) { + if (type == "IVF") { + return std::make_shared(); + } else if (type == "IVFPQ") { + return std::make_shared(); + } else if (type == "GPUIVF") { + return std::make_shared(DEVICEID); + } else if (type == "GPUIVFPQ") { + return std::make_shared(DEVICEID); + } else if (type == "IVFSQ") { + return std::make_shared(); + } else if (type == "GPUIVFSQ") { + return std::make_shared(DEVICEID); + } else if (type == "IVFSQHybrid") { + return std::make_shared(DEVICEID); + } +} + +enum class ParameterType { + ivf, + ivfpq, + ivfsq, +}; + +class ParamGenerator { + public: + static ParamGenerator& + GetInstance() { + static ParamGenerator instance; + return instance; + } + + knowhere::Config + Gen(const ParameterType& type) { + if (type == ParameterType::ivf) { + auto tempconf = std::make_shared(); + tempconf->d = DIM; + tempconf->gpu_id = DEVICEID; + tempconf->nlist = 100; + tempconf->nprobe = 4; + tempconf->k = K; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } else if (type == ParameterType::ivfpq) { + auto tempconf = std::make_shared(); + tempconf->d = DIM; + tempconf->gpu_id = DEVICEID; + tempconf->nlist = 100; + tempconf->nprobe = 4; + tempconf->k = K; + tempconf->m = 4; + tempconf->nbits = 8; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } else if (type == ParameterType::ivfsq) { + auto tempconf = std::make_shared(); + tempconf->d = DIM; + tempconf->gpu_id = DEVICEID; + tempconf->nlist = 100; + tempconf->nprobe = 4; + tempconf->k = K; + tempconf->nbits = 8; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + } +}; + +#include + +class TestGpuIndexBase : public ::testing::Test { + protected: + void + SetUp() override { + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); + } + + void + TearDown() override { + knowhere::FaissGpuResourceMgr::GetInstance().Free(); + } +}; diff --git a/core/src/index/unittest/test_customized_index.cpp b/core/src/index/unittest/test_customized_index.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1e0b1d932dd0e02cd86e87af1cf4c3b0f9154878 --- /dev/null +++ b/core/src/index/unittest/test_customized_index.cpp @@ -0,0 +1,122 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 "unittest/Helper.h" +#include "unittest/utils.h" + +class SingleIndexTest : public DataGen, public TestGpuIndexBase { + protected: + void + SetUp() override { + TestGpuIndexBase::SetUp(); + Generate(DIM, NB, NQ); + k = K; + } + + void + TearDown() override { + TestGpuIndexBase::TearDown(); + } + + protected: + std::string index_type; + knowhere::IVFIndexPtr index_ = nullptr; +}; + +#ifdef CUSTOMIZATION +TEST_F(SingleIndexTest, IVFSQHybrid) { + assert(!xb.empty()); + + index_type = "IVFSQHybrid"; + index_ = IndexFactory(index_type); + auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq); + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + + auto model = index_->Train(base_dataset, conf); + index_->set_index_model(model); + index_->Add(base_dataset, conf); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dimension(), dim); + + auto binaryset = index_->Serialize(); + { + // copy cpu to gpu + auto cpu_idx = std::make_shared(DEVICEID); + cpu_idx->Load(binaryset); + + { + for (int i = 0; i < 3; ++i) { + auto gpu_idx = cpu_idx->CopyCpuToGpu(DEVICEID, conf); + auto result = gpu_idx->Search(query_dataset, conf); + AssertAnns(result, nq, conf->k); + // PrintResult(result, nq, k); + } + } + } + + { + // quantization already in gpu, only copy data + auto cpu_idx = std::make_shared(DEVICEID); + cpu_idx->Load(binaryset); + + auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf); + auto gpu_idx = pair.first; + auto quantization = pair.second; + + auto result = gpu_idx->Search(query_dataset, conf); + AssertAnns(result, nq, conf->k); + // PrintResult(result, nq, k); + + auto quantizer_conf = std::make_shared(); + quantizer_conf->mode = 2; // only copy data + quantizer_conf->gpu_id = DEVICEID; + for (int i = 0; i < 2; ++i) { + auto hybrid_idx = std::make_shared(DEVICEID); + hybrid_idx->Load(binaryset); + + auto new_idx = hybrid_idx->LoadData(quantization, quantizer_conf); + auto result = new_idx->Search(query_dataset, conf); + AssertAnns(result, nq, conf->k); + // PrintResult(result, nq, k); + } + } + + { + // quantization already in gpu, only set quantization + auto cpu_idx = std::make_shared(DEVICEID); + cpu_idx->Load(binaryset); + + auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf); + auto quantization = pair.second; + + for (int i = 0; i < 2; ++i) { + auto hybrid_idx = std::make_shared(DEVICEID); + hybrid_idx->Load(binaryset); + + hybrid_idx->SetQuantizer(quantization); + auto result = hybrid_idx->Search(query_dataset, conf); + AssertAnns(result, nq, conf->k); + // PrintResult(result, nq, k); + hybrid_idx->UnsetQuantizer(); + } + } +} + +#endif diff --git a/core/src/index/unittest/test_gpuresource.cpp b/core/src/index/unittest/test_gpuresource.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e3485d26ba9ecda5afe812c9179d5947882af236 --- /dev/null +++ b/core/src/index/unittest/test_gpuresource.cpp @@ -0,0 +1,309 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 "knowhere/common/Exception.h" +#include "knowhere/common/Timer.h" +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" + +#include "unittest/Helper.h" +#include "unittest/utils.h" + +class GPURESTEST : public DataGen, public TestGpuIndexBase { + protected: + void + SetUp() override { + TestGpuIndexBase::SetUp(); + Generate(DIM, NB, NQ); + + k = K; + elems = nq * k; + ids = (int64_t*)malloc(sizeof(int64_t) * elems); + dis = (float*)malloc(sizeof(float) * elems); + } + + void + TearDown() override { + delete ids; + delete dis; + TestGpuIndexBase::TearDown(); + } + + protected: + std::string index_type; + knowhere::IVFIndexPtr index_ = nullptr; + + int64_t* ids = nullptr; + float* dis = nullptr; + int64_t elems = 0; +}; + +TEST_F(GPURESTEST, copyandsearch) { + // search and copy at the same time + printf("==================\n"); + + index_type = "GPUIVF"; + index_ = IndexFactory(index_type); + + auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf); + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(base_dataset, conf); + index_->set_index_model(model); + index_->Add(base_dataset, conf); + auto result = index_->Search(query_dataset, conf); + AssertAnns(result, nq, k); + + auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + cpu_idx->Seal(); + auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + + constexpr int64_t search_count = 50; + constexpr int64_t load_count = 15; + auto search_func = [&] { + // TimeRecorder tc("search&load"); + for (int i = 0; i < search_count; ++i) { + search_idx->Search(query_dataset, conf); + // if (i > search_count - 6 || i == 0) + // tc.RecordSection("search once"); + } + // tc.ElapseFromBegin("search finish"); + }; + auto load_func = [&] { + // TimeRecorder tc("search&load"); + for (int i = 0; i < load_count; ++i) { + knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + // if (i > load_count -5 || i < 5) + // tc.RecordSection("Copy to gpu"); + } + // tc.ElapseFromBegin("load finish"); + }; + + knowhere::TimeRecorder tc("Basic"); + knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + tc.RecordSection("Copy to gpu once"); + search_idx->Search(query_dataset, conf); + tc.RecordSection("Search once"); + search_func(); + tc.RecordSection("Search total cost"); + load_func(); + tc.RecordSection("Copy total cost"); + + std::thread search_thread(search_func); + std::thread load_thread(load_func); + search_thread.join(); + load_thread.join(); + tc.RecordSection("Copy&Search total"); +} + +TEST_F(GPURESTEST, trainandsearch) { + index_type = "GPUIVF"; + index_ = IndexFactory(index_type); + + auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf); + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(base_dataset, conf); + auto new_index = IndexFactory(index_type); + new_index->set_index_model(model); + new_index->Add(base_dataset, conf); + auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config()); + cpu_idx->Seal(); + auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + + constexpr int train_count = 5; + constexpr int search_count = 200; + auto train_stage = [&] { + for (int i = 0; i < train_count; ++i) { + auto model = index_->Train(base_dataset, conf); + auto test_idx = IndexFactory(index_type); + test_idx->set_index_model(model); + test_idx->Add(base_dataset, conf); + } + }; + auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) { + for (int i = 0; i < search_count; ++i) { + auto result = search_idx->Search(query_dataset, conf); + AssertAnns(result, nq, k); + } + }; + + // TimeRecorder tc("record"); + // train_stage(); + // tc.RecordSection("train cost"); + // search_stage(search_idx); + // tc.RecordSection("search cost"); + + { + // search and build parallel + std::thread search_thread(search_stage, std::ref(search_idx)); + std::thread train_thread(train_stage); + train_thread.join(); + search_thread.join(); + } + { + // build parallel + std::thread train_1(train_stage); + std::thread train_2(train_stage); + train_1.join(); + train_2.join(); + } + { + // search parallel + auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + std::thread search_1(search_stage, std::ref(search_idx)); + std::thread search_2(search_stage, std::ref(search_idx_2)); + search_1.join(); + search_2.join(); + } +} + +#ifdef CompareToOriFaiss +TEST_F(GPURESTEST, gpu_ivf_resource_test) { + assert(!xb.empty()); + + { + index_ = std::make_shared(-1); + ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), -1); + std::dynamic_pointer_cast(index_)->SetGpuDevice(DEVICEID); + ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), DEVICEID); + + auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq); + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(base_dataset, conf); + index_->set_index_model(model); + index_->Add(base_dataset, conf); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dimension(), dim); + + // knowhere::TimeRecorder tc("knowere GPUIVF"); + for (int i = 0; i < search_count; ++i) { + index_->Search(query_dataset, conf); + if (i > search_count - 6 || i < 5) + // tc.RecordSection("search once"); + } + // tc.ElapseFromBegin("search all"); + } + knowhere::FaissGpuResourceMgr::GetInstance().Dump(); + + // { + // // ori faiss IVF-Search + // faiss::gpu::StandardGpuResources res; + // faiss::gpu::GpuIndexIVFFlatConfig idx_config; + // idx_config.device = DEVICEID; + // faiss::gpu::GpuIndexIVFFlat device_index(&res, dim, 1638, faiss::METRIC_L2, idx_config); + // device_index.train(nb, xb.data()); + // device_index.add(nb, xb.data()); + // + // knowhere::TimeRecorder tc("ori IVF"); + // for (int i = 0; i < search_count; ++i) { + // device_index.search(nq, xq.data(), k, dis, ids); + // if (i > search_count - 6 || i < 5) + // tc.RecordSection("search once"); + // } + // tc.ElapseFromBegin("search all"); + // } +} + +TEST_F(GPURESTEST, gpuivfsq) { + { + // knowhere gpu ivfsq + index_type = "GPUIVFSQ"; + index_ = IndexFactory(index_type); + + auto conf = std::make_shared(); + conf->nlist = 1638; + conf->d = dim; + conf->gpu_id = DEVICEID; + conf->metric_type = knowhere::METRICTYPE::L2; + conf->k = k; + conf->nbits = 8; + conf->nprobe = 1; + + auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(base_dataset, conf); + index_->set_index_model(model); + index_->Add(base_dataset, conf); + // auto result = index_->Search(query_dataset, conf); + // AssertAnns(result, nq, k); + + auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + cpu_idx->Seal(); + + knowhere::TimeRecorder tc("knowhere GPUSQ8"); + auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + tc.RecordSection("Copy to gpu"); + for (int i = 0; i < search_count; ++i) { + search_idx->Search(query_dataset, conf); + if (i > search_count - 6 || i < 5) + tc.RecordSection("search once"); + } + tc.ElapseFromBegin("search all"); + } + + { + // Ori gpuivfsq Test + const char* index_description = "IVF1638,SQ8"; + faiss::Index* ori_index = faiss::index_factory(dim, index_description, faiss::METRIC_L2); + + faiss::gpu::StandardGpuResources res; + auto device_index = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, ori_index); + device_index->train(nb, xb.data()); + device_index->add(nb, xb.data()); + + auto cpu_index = faiss::gpu::index_gpu_to_cpu(device_index); + auto idx = dynamic_cast(cpu_index); + if (idx != nullptr) { + idx->to_readonly(); + } + delete device_index; + delete ori_index; + + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + knowhere::TimeRecorder tc("ori GPUSQ8"); + faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, cpu_index, &option); + tc.RecordSection("Copy to gpu"); + for (int i = 0; i < search_count; ++i) { + search_idx->search(nq, xq.data(), k, dis, ids); + if (i > search_count - 6 || i < 5) + tc.RecordSection("search once"); + } + tc.ElapseFromBegin("search all"); + delete cpu_index; + delete search_idx; + } +} +#endif diff --git a/core/src/index/unittest/test_idmap.cpp b/core/src/index/unittest/test_idmap.cpp index 077f7328f74da35d533f82f544068e18d89747ff..d1ff3ee0465dcc3dc7d6a02b79980ab4bbc5a805 100644 --- a/core/src/index/unittest/test_idmap.cpp +++ b/core/src/index/unittest/test_idmap.cpp @@ -23,54 +23,28 @@ #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/helpers/Cloner.h" +#include "Helper.h" #include "unittest/utils.h" -static int device_id = 0; -class IDMAPTest : public DataGen, public ::testing::Test { +class IDMAPTest : public DataGen, public TestGpuIndexBase { protected: void SetUp() override { - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); + TestGpuIndexBase::SetUp(); + Init_with_default(); index_ = std::make_shared(); } void TearDown() override { - knowhere::FaissGpuResourceMgr::GetInstance().Free(); + TestGpuIndexBase::TearDown(); } protected: knowhere::IDMAPPtr index_ = nullptr; }; -void -AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(i, *(ids->data()->GetValues(1, i * k))); - } -} - -void -PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - auto dists = result->array()[1]; - - std::stringstream ss_id; - std::stringstream ss_dist; - for (auto i = 0; i < 10; i++) { - for (auto j = 0; j < k; ++j) { - ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - } - ss_id << std::endl; - ss_dist << std::endl; - } - std::cout << "id\n" << ss_id.str() << std::endl; - std::cout << "dist\n" << ss_dist.str() << std::endl; -} - TEST_F(IDMAPTest, idmap_basic) { ASSERT_TRUE(!xb.empty()); @@ -87,7 +61,7 @@ TEST_F(IDMAPTest, idmap_basic) { ASSERT_TRUE(index_->GetRawIds() != nullptr); auto result = index_->Search(query_dataset, conf); AssertAnns(result, nq, k); - PrintResult(result, nq, k); + // PrintResult(result, nq, k); index_->Seal(); auto binaryset = index_->Serialize(); @@ -95,7 +69,7 @@ TEST_F(IDMAPTest, idmap_basic) { new_index->Load(binaryset); auto re_result = index_->Search(query_dataset, conf); AssertAnns(re_result, nq, k); - PrintResult(re_result, nq, k); + // PrintResult(re_result, nq, k); } TEST_F(IDMAPTest, idmap_serialize) { @@ -118,7 +92,7 @@ TEST_F(IDMAPTest, idmap_serialize) { index_->Add(base_dataset, knowhere::Config()); auto re_result = index_->Search(query_dataset, conf); AssertAnns(re_result, nq, k); - PrintResult(re_result, nq, k); + // PrintResult(re_result, nq, k); EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); auto binaryset = index_->Serialize(); @@ -138,7 +112,7 @@ TEST_F(IDMAPTest, idmap_serialize) { EXPECT_EQ(index_->Dimension(), dim); auto result = index_->Search(query_dataset, conf); AssertAnns(result, nq, k); - PrintResult(result, nq, k); + // PrintResult(result, nq, k); } } @@ -169,7 +143,7 @@ TEST_F(IDMAPTest, copy_test) { { // cpu to gpu - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf); + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); auto clone_result = clone_index->Search(query_dataset, conf); AssertAnns(clone_result, nq, k); ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawVectors(); }, @@ -194,9 +168,9 @@ TEST_F(IDMAPTest, copy_test) { ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawIds() != nullptr); // gpu to gpu - auto device_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf); + auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); auto new_device_index = - std::static_pointer_cast(device_index)->CopyGpuToGpu(device_id, conf); + std::static_pointer_cast(device_index)->CopyGpuToGpu(DEVICEID, conf); auto device_result = new_device_index->Search(query_dataset, conf); AssertAnns(device_result, nq, k); } diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp index 4aaddb108b6e454e7d0376283e091b14f28be0ce..fae27b0dd3ec2125133bf205c9dbd4ec4c05effa 100644 --- a/core/src/index/unittest/test_ivf.cpp +++ b/core/src/index/unittest/test_ivf.cpp @@ -35,99 +35,25 @@ #include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/helpers/Cloner.h" +#include "unittest/Helper.h" #include "unittest/utils.h" using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; -constexpr int device_id = 0; -constexpr int64_t DIM = 128; -constexpr int64_t NB = 1000000 / 100; -constexpr int64_t NQ = 10; -constexpr int64_t K = 10; - -knowhere::IVFIndexPtr -IndexFactory(const std::string& type) { - if (type == "IVF") { - return std::make_shared(); - } else if (type == "IVFPQ") { - return std::make_shared(); - } else if (type == "GPUIVF") { - return std::make_shared(device_id); - } else if (type == "GPUIVFPQ") { - return std::make_shared(device_id); - } else if (type == "IVFSQ") { - return std::make_shared(); - } else if (type == "GPUIVFSQ") { - return std::make_shared(device_id); - } else if (type == "IVFSQHybrid") { - return std::make_shared(device_id); - } -} - -enum class ParameterType { - ivf, - ivfpq, - ivfsq, - nsg, -}; - -class ParamGenerator { - public: - static ParamGenerator& - GetInstance() { - static ParamGenerator instance; - return instance; - } - - knowhere::Config - Gen(const ParameterType& type) { - if (type == ParameterType::ivf) { - auto tempconf = std::make_shared(); - tempconf->d = DIM; - tempconf->gpu_id = device_id; - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->k = K; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } else if (type == ParameterType::ivfpq) { - auto tempconf = std::make_shared(); - tempconf->d = DIM; - tempconf->gpu_id = device_id; - tempconf->nlist = 25; - tempconf->nprobe = 4; - tempconf->k = K; - tempconf->m = 4; - tempconf->nbits = 8; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } else if (type == ParameterType::ivfsq) { - auto tempconf = std::make_shared(); - tempconf->d = DIM; - tempconf->gpu_id = device_id; - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->k = K; - tempconf->nbits = 8; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - } -}; - class IVFTest : public DataGen, public TestWithParam<::std::tuple> { protected: void SetUp() override { + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); + ParameterType parameter_type; std::tie(index_type, parameter_type) = GetParam(); // Init_with_default(); Generate(DIM, NB, NQ); index_ = IndexFactory(index_type); conf = ParamGenerator::GetInstance().Gen(parameter_type); - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 600, 2); } void @@ -140,7 +66,7 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple gpu_idx{"GPUIVFSQ"}; auto finder = std::find(gpu_idx.cbegin(), gpu_idx.cend(), index_type); if (finder != gpu_idx.cend()) { - return knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config()); + return knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); } return index_; } @@ -162,33 +88,6 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest, #endif std::make_tuple("GPUIVFSQ", ParameterType::ivfsq))); -void -AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(i, *(ids->data()->GetValues(1, i * k))); - } -} - -void -PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - auto dists = result->array()[1]; - - std::stringstream ss_id; - std::stringstream ss_dist; - for (auto i = 0; i < 10; i++) { - for (auto j = 0; j < k; ++j) { - ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - } - ss_id << std::endl; - ss_dist << std::endl; - } - std::cout << "id\n" << ss_id.str() << std::endl; - std::cout << "dist\n" << ss_dist.str() << std::endl; -} - TEST_P(IVFTest, ivf_basic) { assert(!xb.empty()); @@ -207,85 +106,6 @@ TEST_P(IVFTest, ivf_basic) { // PrintResult(result, nq, k); } -TEST_P(IVFTest, hybrid) { - if (index_type != "IVFSQHybrid") { - return; - } - assert(!xb.empty()); - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - - // auto new_idx = ChooseTodo(); - // auto result = new_idx->Search(query_dataset, conf); - // AssertAnns(result, nq, conf->k); - - { - auto hybrid_1_idx = std::make_shared(device_id); - - auto binaryset = index_->Serialize(); - hybrid_1_idx->Load(binaryset); - - auto quantizer_conf = std::make_shared(); - quantizer_conf->mode = 1; - quantizer_conf->gpu_id = device_id; - auto q = hybrid_1_idx->LoadQuantizer(quantizer_conf); - hybrid_1_idx->SetQuantizer(q); - auto result = hybrid_1_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf->k); - PrintResult(result, nq, k); - hybrid_1_idx->UnsetQuantizer(); - } - - { - auto hybrid_2_idx = std::make_shared(device_id); - - auto binaryset = index_->Serialize(); - hybrid_2_idx->Load(binaryset); - - auto quantizer_conf = std::make_shared(); - quantizer_conf->mode = 1; - quantizer_conf->gpu_id = device_id; - auto q = hybrid_2_idx->LoadQuantizer(quantizer_conf); - quantizer_conf->mode = 2; - auto gpu_idx = hybrid_2_idx->LoadData(q, quantizer_conf); - - auto result = gpu_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf->k); - PrintResult(result, nq, k); - } -} - -// TEST_P(IVFTest, gpu_to_cpu) { -// if (index_type.find("GPU") == std::string::npos) { return; } -// -// // else -// assert(!xb.empty()); -// -// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); -// index_->set_preprocessor(preprocessor); -// -// auto model = index_->Train(base_dataset, conf); -// index_->set_index_model(model); -// index_->Add(base_dataset, conf); -// EXPECT_EQ(index_->Count(), nb); -// EXPECT_EQ(index_->Dimension(), dim); -// auto result = index_->Search(query_dataset, conf); -// AssertAnns(result, nq, k); -// -// if (auto device_index = std::dynamic_pointer_cast(index_)) { -// auto host_index = device_index->Copy_index_gpu_to_cpu(); -// auto result = host_index->Search(query_dataset, conf); -// AssertAnns(result, nq, k); -// } -//} - TEST_P(IVFTest, ivf_serialize) { auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { FileIOWriter writer(filename); @@ -423,7 +243,7 @@ TEST_P(IVFTest, clone_test) { auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config()); + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); auto clone_result = clone_index->Search(query_dataset, conf); AssertEqual(result, clone_result); std::cout << "clone C <=> G [" << index_type << "] success" << std::endl; @@ -432,7 +252,7 @@ TEST_P(IVFTest, clone_test) { EXPECT_THROW( { std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl; - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config()); + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); }, knowhere::KnowhereException); } @@ -440,9 +260,7 @@ TEST_P(IVFTest, clone_test) { } #ifdef CUSTOMIZATION -TEST_P(IVFTest, seal_test) { - // FaissGpuResourceMgr::GetInstance().InitDevice(device_id); - +TEST_P(IVFTest, gpu_seal_test) { std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"}; auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder == support_idx_vec.cend()) { @@ -466,309 +284,13 @@ TEST_P(IVFTest, seal_test) { auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); knowhere::TimeRecorder tc("CopyToGpu"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); + knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); auto without_seal = tc.RecordSection("Without seal"); cpu_idx->Seal(); tc.RecordSection("seal cost"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); + knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); auto with_seal = tc.RecordSection("With seal"); ASSERT_GE(without_seal, with_seal); } -#endif - -class GPURESTEST : public DataGen, public ::testing::Test { - protected: - void - SetUp() override { - Generate(128, 1000000, 1000); - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); - - k = 100; - elems = nq * k; - ids = (int64_t*)malloc(sizeof(int64_t) * elems); - dis = (float*)malloc(sizeof(float) * elems); - } - - void - TearDown() override { - delete ids; - delete dis; - knowhere::FaissGpuResourceMgr::GetInstance().Free(); - } - - protected: - std::string index_type; - knowhere::IVFIndexPtr index_ = nullptr; - - int64_t* ids = nullptr; - float* dis = nullptr; - int64_t elems = 0; -}; - -const int search_count = 18; -const int load_count = 3; - -TEST_F(GPURESTEST, gpu_ivf_resource_test) { - assert(!xb.empty()); - - { - index_ = std::make_shared(-1); - ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), -1); - std::dynamic_pointer_cast(index_)->SetGpuDevice(device_id); - ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), device_id); - - auto conf = std::make_shared(); - conf->nlist = 1638; - conf->d = dim; - conf->gpu_id = device_id; - conf->metric_type = knowhere::METRICTYPE::L2; - conf->k = k; - conf->nprobe = 1; - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - - knowhere::TimeRecorder tc("knowere GPUIVF"); - for (int i = 0; i < search_count; ++i) { - index_->Search(query_dataset, conf); - if (i > search_count - 6 || i < 5) - tc.RecordSection("search once"); - } - tc.ElapseFromBegin("search all"); - } - knowhere::FaissGpuResourceMgr::GetInstance().Dump(); - - { - // IVF-Search - faiss::gpu::StandardGpuResources res; - faiss::gpu::GpuIndexIVFFlatConfig idx_config; - idx_config.device = device_id; - faiss::gpu::GpuIndexIVFFlat device_index(&res, dim, 1638, faiss::METRIC_L2, idx_config); - device_index.train(nb, xb.data()); - device_index.add(nb, xb.data()); - - knowhere::TimeRecorder tc("ori IVF"); - for (int i = 0; i < search_count; ++i) { - device_index.search(nq, xq.data(), k, dis, ids); - if (i > search_count - 6 || i < 5) - tc.RecordSection("search once"); - } - tc.ElapseFromBegin("search all"); - } -} -#ifdef CUSTOMIZATION -TEST_F(GPURESTEST, gpuivfsq) { - { - // knowhere gpu ivfsq - index_type = "GPUIVFSQ"; - index_ = IndexFactory(index_type); - - auto conf = std::make_shared(); - conf->nlist = 1638; - conf->d = dim; - conf->gpu_id = device_id; - conf->metric_type = knowhere::METRICTYPE::L2; - conf->k = k; - conf->nbits = 8; - conf->nprobe = 1; - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - // auto result = index_->Search(query_dataset, conf); - // AssertAnns(result, nq, k); - - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - cpu_idx->Seal(); - - knowhere::TimeRecorder tc("knowhere GPUSQ8"); - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - tc.RecordSection("Copy to gpu"); - for (int i = 0; i < search_count; ++i) { - search_idx->Search(query_dataset, conf); - if (i > search_count - 6 || i < 5) - tc.RecordSection("search once"); - } - tc.ElapseFromBegin("search all"); - } - - { - // Ori gpuivfsq Test - const char* index_description = "IVF1638,SQ8"; - faiss::Index* ori_index = faiss::index_factory(dim, index_description, faiss::METRIC_L2); - - faiss::gpu::StandardGpuResources res; - auto device_index = faiss::gpu::index_cpu_to_gpu(&res, device_id, ori_index); - device_index->train(nb, xb.data()); - device_index->add(nb, xb.data()); - - auto cpu_index = faiss::gpu::index_gpu_to_cpu(device_index); - auto idx = dynamic_cast(cpu_index); - if (idx != nullptr) { - idx->to_readonly(); - } - delete device_index; - delete ori_index; - - faiss::gpu::GpuClonerOptions option; - option.allInGpu = true; - - knowhere::TimeRecorder tc("ori GPUSQ8"); - faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, device_id, cpu_index, &option); - tc.RecordSection("Copy to gpu"); - for (int i = 0; i < search_count; ++i) { - search_idx->search(nq, xq.data(), k, dis, ids); - if (i > search_count - 6 || i < 5) - tc.RecordSection("search once"); - } - tc.ElapseFromBegin("search all"); - delete cpu_index; - delete search_idx; - } -} #endif - -TEST_F(GPURESTEST, copyandsearch) { - // search and copy at the same time - printf("==================\n"); - - index_type = "GPUIVF"; - index_ = IndexFactory(index_type); - - auto conf = std::make_shared(); - conf->nlist = 1638; - conf->d = dim; - conf->gpu_id = device_id; - conf->metric_type = knowhere::METRICTYPE::L2; - conf->k = k; - conf->nbits = 8; - conf->nprobe = 1; - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - // auto result = index_->Search(query_dataset, conf); - // AssertAnns(result, nq, k); - - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - cpu_idx->Seal(); - - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - - auto search_func = [&] { - // TimeRecorder tc("search&load"); - for (int i = 0; i < search_count; ++i) { - search_idx->Search(query_dataset, conf); - // if (i > search_count - 6 || i == 0) - // tc.RecordSection("search once"); - } - // tc.ElapseFromBegin("search finish"); - }; - auto load_func = [&] { - // TimeRecorder tc("search&load"); - for (int i = 0; i < load_count; ++i) { - knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - // if (i > load_count -5 || i < 5) - // tc.RecordSection("Copy to gpu"); - } - // tc.ElapseFromBegin("load finish"); - }; - - knowhere::TimeRecorder tc("basic"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - tc.RecordSection("Copy to gpu once"); - search_idx->Search(query_dataset, conf); - tc.RecordSection("search once"); - search_func(); - tc.RecordSection("only search total"); - load_func(); - tc.RecordSection("only copy total"); - - std::thread search_thread(search_func); - std::thread load_thread(load_func); - search_thread.join(); - load_thread.join(); - tc.RecordSection("Copy&search total"); -} - -TEST_F(GPURESTEST, TrainAndSearch) { - index_type = "GPUIVF"; - index_ = IndexFactory(index_type); - - auto conf = std::make_shared(); - conf->nlist = 1638; - conf->d = dim; - conf->gpu_id = device_id; - conf->metric_type = knowhere::METRICTYPE::L2; - conf->k = k; - conf->nbits = 8; - conf->nprobe = 1; - - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - auto new_index = IndexFactory(index_type); - new_index->set_index_model(model); - new_index->Add(base_dataset, conf); - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config()); - cpu_idx->Seal(); - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - - constexpr int train_count = 1; - constexpr int search_count = 5000; - auto train_stage = [&] { - for (int i = 0; i < train_count; ++i) { - auto model = index_->Train(base_dataset, conf); - auto test_idx = IndexFactory(index_type); - test_idx->set_index_model(model); - test_idx->Add(base_dataset, conf); - } - }; - auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) { - for (int i = 0; i < search_count; ++i) { - auto result = search_idx->Search(query_dataset, conf); - AssertAnns(result, nq, k); - } - }; - - // TimeRecorder tc("record"); - // train_stage(); - // tc.RecordSection("train cost"); - // search_stage(search_idx); - // tc.RecordSection("search cost"); - - { - // search and build parallel - std::thread search_thread(search_stage, std::ref(search_idx)); - std::thread train_thread(train_stage); - train_thread.join(); - search_thread.join(); - } - { - // build parallel - std::thread train_1(train_stage); - std::thread train_2(train_stage); - train_1.join(); - train_2.join(); - } - { - // search parallel - auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config()); - std::thread search_1(search_stage, std::ref(search_idx)); - std::thread search_2(search_stage, std::ref(search_idx_2)); - search_1.join(); - search_2.join(); - } -} - -// TODO(lxj): Add exception test diff --git a/core/src/index/unittest/test_kdt.cpp b/core/src/index/unittest/test_kdt.cpp index f9058cc4d272fccab8e97a73fc0bc4ef5221ad0d..54008818758356e5bfa2bf7800a910c6ff7561b6 100644 --- a/core/src/index/unittest/test_kdt.cpp +++ b/core/src/index/unittest/test_kdt.cpp @@ -52,33 +52,6 @@ class KDTTest : public DataGen, public ::testing::Test { std::shared_ptr index_ = nullptr; }; -void -AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(i, *(ids->data()->GetValues(1, i * k))); - } -} - -void -PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - auto dists = result->array()[1]; - - std::stringstream ss_id; - std::stringstream ss_dist; - for (auto i = 0; i < 10; i++) { - for (auto j = 0; j < k; ++j) { - ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - } - ss_id << std::endl; - ss_dist << std::endl; - } - std::cout << "id\n" << ss_id.str() << std::endl; - std::cout << "dist\n" << ss_dist.str() << std::endl; -} - // TODO(lxj): add test about count() and dimension() TEST_F(KDTTest, kdt_basic) { assert(!xb.empty()); diff --git a/core/src/index/unittest/test_nsg/test_nsg.cpp b/core/src/index/unittest/test_nsg/test_nsg.cpp index 14d899a488ef7d031a48e8c1c59be431a8a98615..11b9becce49e371616677e7de50401c4cb4ef571 100644 --- a/core/src/index/unittest/test_nsg/test_nsg.cpp +++ b/core/src/index/unittest/test_nsg/test_nsg.cpp @@ -30,19 +30,19 @@ using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; -constexpr int64_t DEVICE_ID = 0; +constexpr int64_t DEVICEID = 0; class NSGInterfaceTest : public DataGen, public ::testing::Test { protected: void SetUp() override { // Init_with_default(); - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2); + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2); Generate(256, 1000000 / 100, 1); index_ = std::make_shared(); auto tmp_conf = std::make_shared(); - tmp_conf->gpu_id = DEVICE_ID; + tmp_conf->gpu_id = DEVICEID; tmp_conf->knng = 20; tmp_conf->nprobe = 8; tmp_conf->nlist = 163; @@ -69,14 +69,6 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test { knowhere::Config search_conf; }; -void -AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->array()[0]; - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(i, *(ids->data()->GetValues(1, i * k))); - } -} - TEST_F(NSGInterfaceTest, basic_test) { assert(!xb.empty()); diff --git a/core/src/index/unittest/utils.cpp b/core/src/index/unittest/utils.cpp index cdfc56b1cb672a7f22424caacfdb9d7721c76130..d4a59bafbb2550373177e5c0fe7a20224e161959 100644 --- a/core/src/index/unittest/utils.cpp +++ b/core/src/index/unittest/utils.cpp @@ -17,6 +17,7 @@ #include "unittest/utils.h" +#include #include #include #include @@ -147,3 +148,30 @@ generate_query_dataset(int64_t nb, int64_t dim, float* xb) { auto dataset = std::make_shared(std::move(tensors), tensor_schema); return dataset; } + +void +AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { + auto ids = result->array()[0]; + for (auto i = 0; i < nq; i++) { + EXPECT_EQ(i, *(ids->data()->GetValues(1, i * k))); + } +} + +void +PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { + auto ids = result->array()[0]; + auto dists = result->array()[1]; + + std::stringstream ss_id; + std::stringstream ss_dist; + for (auto i = 0; i < 10; i++) { + for (auto j = 0; j < k; ++j) { + ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; + ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; + } + ss_id << std::endl; + ss_dist << std::endl; + } + std::cout << "id\n" << ss_id.str() << std::endl; + std::cout << "dist\n" << ss_dist.str() << std::endl; +} diff --git a/core/src/index/unittest/utils.h b/core/src/index/unittest/utils.h index acc3e89183972190955639b70c19da3b8c15272c..b39cf9ea14d48e005241e84c62d2fbd294ed8190 100644 --- a/core/src/index/unittest/utils.h +++ b/core/src/index/unittest/utils.h @@ -68,6 +68,12 @@ generate_dataset(int64_t nb, int64_t dim, float* xb, int64_t* ids); knowhere::DatasetPtr generate_query_dataset(int64_t nb, int64_t dim, float* xb); +void +AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k); + +void +PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k); + struct FileIOWriter { std::fstream fs; std::string name; diff --git a/core/unittest/wrapper/test_hybrid_index.cpp b/core/unittest/wrapper/test_hybrid_index.cpp new file mode 100644 index 0000000000000000000000000000000000000000..757d5b20988e56c6da4e1359868290f7fe78eae2 --- /dev/null +++ b/core/unittest/wrapper/test_hybrid_index.cpp @@ -0,0 +1,133 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 "wrapper/VecIndex.h" +#include "wrapper/utils.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" + +#include +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" + +using ::testing::TestWithParam; +using ::testing::Values; +using ::testing::Combine; + +class KnowhereHybrid + : public DataGenBase, public ::testing::Test { + protected: + void SetUp() override { + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); + + dim = 128; + nb = 10000; + nq = 100; + k = 100; + GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); + } + + void TearDown() override { + knowhere::FaissGpuResourceMgr::GetInstance().Free(); + } + + protected: + milvus::engine::IndexType index_type; + milvus::engine::VecIndexPtr index_ = nullptr; + knowhere::Config conf; +}; + +#ifdef CUSTOMIZATION +TEST_F(KnowhereHybrid, test_interface) { + assert(!xb.empty()); + + index_type = milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID; + index_ = GetVecIndexFactory(index_type); + conf = ParamGenerator::GetInstance().Gen(index_type); + + auto elems = nq * k; + std::vector res_ids(elems); + std::vector res_dis(elems); + + conf->gpu_id = DEVICEID; + conf->d = dim; + conf->k = k; + index_->BuildAll(nb, xb.data(), ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + AssertResult(res_ids, res_dis); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->Dimension(), dim); + + auto binaryset = index_->Serialize(); + { + // cpu -> gpu + auto cpu_idx = GetVecIndexFactory(index_type); + cpu_idx->Load(binaryset); + { + for (int i = 0; i < 2; ++i) { + auto gpu_idx = cpu_idx->CopyToGpu(DEVICEID, conf); + gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + AssertResult(res_ids, res_dis); + } + } + } + + { + // quantization already in gpu, only copy data + auto cpu_idx = GetVecIndexFactory(index_type); + cpu_idx->Load(binaryset); + + auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf); + auto gpu_idx = pair.first; + auto quantization = pair.second; + + gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + AssertResult(res_ids, res_dis); + + auto quantizer_conf = std::make_shared(); + quantizer_conf->mode = 2; + quantizer_conf->gpu_id = DEVICEID; + for (int i = 0; i < 2; ++i) { + auto hybrid_idx = GetVecIndexFactory(index_type); + hybrid_idx->Load(binaryset); + + hybrid_idx->LoadData(quantization, quantizer_conf); + hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + AssertResult(res_ids, res_dis); + } + } + + { + // quantization already in gpu, only set quantization + auto cpu_idx = GetVecIndexFactory(index_type); + cpu_idx->Load(binaryset); + + auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf); + auto quantization = pair.second; + + for (int i = 0; i < 2; ++i) { + auto hybrid_idx = GetVecIndexFactory(index_type); + hybrid_idx->Load(binaryset); + + hybrid_idx->SetQuantizer(quantization); + hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + AssertResult(res_ids, res_dis); + hybrid_idx->UnsetQuantizer(); + } + } +} + +#endif diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index 7accef649c8bc3986b0333acd01647f6c04a6bc0..e16a4e2bb243c5d12ed1f9be62d4ac3aed95e218 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -25,150 +25,36 @@ INITIALIZE_EASYLOGGINGPP -namespace { - -namespace ms = milvus::engine; -namespace kw = knowhere; - -} // namespace - using ::testing::TestWithParam; using ::testing::Values; using ::testing::Combine; -constexpr int64_t DIM = 128; -constexpr int64_t NB = 100000; -constexpr int64_t DEVICE_ID = 0; - -class ParamGenerator { - public: - static ParamGenerator& GetInstance() { - static ParamGenerator instance; - return instance; - } - - knowhere::Config Gen(const milvus::engine::IndexType& type) { - switch (type) { - case milvus::engine::IndexType::FAISS_IDMAP: { - auto tempconf = std::make_shared(); - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - case milvus::engine::IndexType::FAISS_IVFFLAT_CPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_GPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: { - auto tempconf = std::make_shared(); - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - case milvus::engine::IndexType::FAISS_IVFSQ8_CPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_GPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: { - auto tempconf = std::make_shared(); - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->nbits = 8; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - case milvus::engine::IndexType::FAISS_IVFPQ_CPU: - case milvus::engine::IndexType::FAISS_IVFPQ_GPU: { - auto tempconf = std::make_shared(); - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->nbits = 8; - tempconf->m = 8; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - case milvus::engine::IndexType::NSG_MIX: { - auto tempconf = std::make_shared(); - tempconf->nlist = 100; - tempconf->nprobe = 16; - tempconf->search_length = 8; - tempconf->knng = 200; - tempconf->search_length = 40; // TODO(linxj): be 20 when search - tempconf->out_degree = 60; - tempconf->candidate_pool_size = 200; - tempconf->metric_type = knowhere::METRICTYPE::L2; - return tempconf; - } - } - } -}; - class KnowhereWrapperTest - : public TestWithParam<::std::tuple> { + : public DataGenBase, + public TestWithParam<::std::tuple> { protected: void SetUp() override { - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, - 1024 * 1024 * 200, - 1024 * 1024 * 300, - 2); + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); std::string generator_type; std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); - - auto generator = std::make_shared(); - generator->GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); + GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); index_ = GetVecIndexFactory(index_type); - conf = ParamGenerator::GetInstance().Gen(index_type); conf->k = k; conf->d = dim; - conf->gpu_id = DEVICE_ID; + conf->gpu_id = DEVICEID; } void TearDown() override { knowhere::FaissGpuResourceMgr::GetInstance().Free(); } - void AssertResult(const std::vector& ids, const std::vector& dis) { - EXPECT_EQ(ids.size(), nq * k); - EXPECT_EQ(dis.size(), nq * k); - - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(ids[i * k], gt_ids[i * k]); - //EXPECT_EQ(dis[i * k], gt_dis[i * k]); - } - - int match = 0; - for (int i = 0; i < nq; ++i) { - for (int j = 0; j < k; ++j) { - for (int l = 0; l < k; ++l) { - if (ids[i * nq + j] == gt_ids[i * nq + l]) match++; - } - } - } - - auto precision = float(match) / (nq * k); - EXPECT_GT(precision, 0.5); - std::cout << std::endl << "Precision: " << precision - << ", match: " << match - << ", total: " << nq * k - << std::endl; - } - protected: milvus::engine::IndexType index_type; - knowhere::Config conf; - - int dim = DIM; - int nb = NB; - int nq = 10; - int k = 10; - std::vector xb; - std::vector xq; - std::vector ids; - milvus::engine::VecIndexPtr index_ = nullptr; - - // Ground Truth - std::vector gt_ids; - std::vector gt_dis; + knowhere::Config conf; }; INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, @@ -220,7 +106,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { AssertResult(res_ids, res_dis); { - auto dev_idx = index_->CopyToGpu(DEVICE_ID); + auto dev_idx = index_->CopyToGpu(DEVICEID); for (int i = 0; i < 10; ++i) { dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); } @@ -232,7 +118,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { write_index(index_, file_location); auto new_index = milvus::engine::read_index(file_location); - auto dev_idx = new_index->CopyToGpu(DEVICE_ID); + auto dev_idx = new_index->CopyToGpu(DEVICEID); for (int i = 0; i < 10; ++i) { dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); } @@ -240,10 +126,6 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { } } -//TEST_P(KnowhereWrapperTest, TO_CPU_TEST) { -// // dev -//} - TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -282,8 +164,3 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { AssertResult(res_ids, res_dis); } } - -// TODO(linxj): add exception test -//TEST_P(KnowhereWrapperTest, exception_test) { -//} - diff --git a/core/unittest/wrapper/utils.cpp b/core/unittest/wrapper/utils.cpp index f2bb83b482b5d8a90d0a8e92135be496e6a83912..445b7a2de6f158dc1cd34aace4a096c383401130 100644 --- a/core/unittest/wrapper/utils.cpp +++ b/core/unittest/wrapper/utils.cpp @@ -16,6 +16,7 @@ // under the License. +#include #include #include "wrapper/utils.h" @@ -59,3 +60,30 @@ DataGenBase::GenData(const int &dim, gt_dis.resize(nq * k); GenData(dim, nb, nq, xb.data(), xq.data(), ids.data(), k, gt_ids.data(), gt_dis.data()); } + +void +DataGenBase::AssertResult(const std::vector& ids, const std::vector& dis) { + EXPECT_EQ(ids.size(), nq * k); + EXPECT_EQ(dis.size(), nq * k); + + for (auto i = 0; i < nq; i++) { + EXPECT_EQ(ids[i * k], gt_ids[i * k]); + //EXPECT_EQ(dis[i * k], gt_dis[i * k]); + } + + int match = 0; + for (int i = 0; i < nq; ++i) { + for (int j = 0; j < k; ++j) { + for (int l = 0; l < k; ++l) { + if (ids[i * nq + j] == gt_ids[i * nq + l]) match++; + } + } + } + + auto precision = float(match) / (nq * k); + EXPECT_GT(precision, 0.5); + std::cout << std::endl << "Precision: " << precision + << ", match: " << match + << ", total: " << nq * k + << std::endl; +} diff --git a/core/unittest/wrapper/utils.h b/core/unittest/wrapper/utils.h index ff4ce9c23a52c038a298a60b7b06cb671927821a..b895b599e026a5cf77717f485cdb0744ddfa4617 100644 --- a/core/unittest/wrapper/utils.h +++ b/core/unittest/wrapper/utils.h @@ -24,24 +24,110 @@ #include #include +#include "wrapper/VecIndex.h" +#include "wrapper/utils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" + class DataGenBase; using DataGenPtr = std::shared_ptr; +constexpr int64_t DIM = 128; +constexpr int64_t NB = 100000; +constexpr int64_t NQ = 10; +constexpr int64_t DEVICEID = 0; +constexpr int64_t PINMEM = 1024 * 1024 * 200; +constexpr int64_t TEMPMEM = 1024 * 1024 * 300; +constexpr int64_t RESNUM = 2; + class DataGenBase { public: - virtual void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, int64_t *ids, - const int &k, int64_t *gt_ids, float *gt_dis); - - virtual void GenData(const int &dim, - const int &nb, - const int &nq, - std::vector &xb, - std::vector &xq, - std::vector &ids, - const int &k, - std::vector >_ids, - std::vector >_dis); + virtual void GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids, + const int& k, int64_t* gt_ids, float* gt_dis); + + virtual void GenData(const int& dim, + const int& nb, + const int& nq, + std::vector& xb, + std::vector& xq, + std::vector& ids, + const int& k, + std::vector& gt_ids, + std::vector& gt_dis); + + void AssertResult(const std::vector& ids, const std::vector& dis); + + int dim = DIM; + int nb = NB; + int nq = NQ; + int k = 10; + std::vector xb; + std::vector xq; + std::vector ids; + + // Ground Truth + std::vector gt_ids; + std::vector gt_dis; +}; + +class ParamGenerator { + public: + static ParamGenerator& GetInstance() { + static ParamGenerator instance; + return instance; + } + + knowhere::Config Gen(const milvus::engine::IndexType& type) { + switch (type) { + case milvus::engine::IndexType::FAISS_IDMAP: { + auto tempconf = std::make_shared(); + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + case milvus::engine::IndexType::FAISS_IVFFLAT_CPU: + case milvus::engine::IndexType::FAISS_IVFFLAT_GPU: + case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: { + auto tempconf = std::make_shared(); + tempconf->nlist = 100; + tempconf->nprobe = 16; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + case milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID: + case milvus::engine::IndexType::FAISS_IVFSQ8_CPU: + case milvus::engine::IndexType::FAISS_IVFSQ8_GPU: + case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: { + auto tempconf = std::make_shared(); + tempconf->nlist = 100; + tempconf->nprobe = 16; + tempconf->nbits = 8; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + case milvus::engine::IndexType::FAISS_IVFPQ_CPU: + case milvus::engine::IndexType::FAISS_IVFPQ_GPU: { + auto tempconf = std::make_shared(); + tempconf->nlist = 100; + tempconf->nprobe = 16; + tempconf->nbits = 8; + tempconf->m = 8; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + case milvus::engine::IndexType::NSG_MIX: { + auto tempconf = std::make_shared(); + tempconf->nlist = 100; + tempconf->nprobe = 16; + tempconf->search_length = 8; + tempconf->knng = 200; + tempconf->search_length = 40; // TODO(linxj): be 20 when search + tempconf->out_degree = 60; + tempconf->candidate_pool_size = 200; + tempconf->metric_type = knowhere::METRICTYPE::L2; + return tempconf; + } + } + } };