From 8c412093caecefbc005dbbf39e1e3436083fd722 Mon Sep 17 00:00:00 2001 From: "zirui.chen" Date: Thu, 14 Nov 2019 11:13:35 +0800 Subject: [PATCH] Support new Index type IVFPQ #127 --- CHANGELOG.md | 1 + core/src/db/engine/ExecutionEngine.h | 4 +- core/src/db/engine/ExecutionEngineImpl.cpp | 4 + .../index/vector_index/IndexGPUIVFPQ.cpp | 33 ++++-- .../index/vector_index/IndexGPUIVFPQ.h | 5 + .../index/vector_index/IndexIVFPQ.cpp | 18 +++ .../knowhere/index/vector_index/IndexIVFPQ.h | 3 + .../index/vector_index/helpers/Cloner.cpp | 2 +- core/src/index/unittest/test_ivf.cpp | 4 +- core/src/wrapper/ConfAdapter.cpp | 46 ++++++-- core/src/wrapper/ConfAdapter.h | 3 + core/src/wrapper/ConfAdapterMgr.cpp | 1 + core/src/wrapper/VecIndex.cpp | 13 ++- core/src/wrapper/VecIndex.h | 2 + core/unittest/wrapper/test_wrapper.cpp | 98 ++++++++-------- core/unittest/wrapper/utils.h | 106 ++++++++++-------- 16 files changed, 226 insertions(+), 117 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d46ed6070..56def33b59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ Please mark all change in change log and use the ticket from JIRA. ## Feature - \#12 - Pure CPU version for Milvus +- #127 - Support new Index type IVFPQ ## Improvement diff --git a/core/src/db/engine/ExecutionEngine.h b/core/src/db/engine/ExecutionEngine.h index 51c77eb78e..86a014cf66 100644 --- a/core/src/db/engine/ExecutionEngine.h +++ b/core/src/db/engine/ExecutionEngine.h @@ -26,6 +26,7 @@ namespace milvus { namespace engine { +// TODO(linxj): replace with VecIndex::IndexType enum class EngineType { INVALID = 0, FAISS_IDMAP = 1, @@ -33,7 +34,8 @@ enum class EngineType { FAISS_IVFSQ8, NSG_MIX, FAISS_IVFSQ8H, - MAX_VALUE = FAISS_IVFSQ8H, + FAISS_PQ, + MAX_VALUE = FAISS_PQ, }; enum class MetricType { diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index 19c699bda7..3579a86291 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -116,6 +116,10 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) { index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID); break; } + case EngineType::FAISS_PQ: { + index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX); + break; + } default: { ENGINE_LOG_ERROR << "Unsupported index type"; return nullptr; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp index b027539c37..cbd4f4f09c 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp @@ -39,17 +39,19 @@ GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) { GETTENSOR(dataset) - // TODO(linxj): set device here. - // TODO(linxj): set gpu resource here. - faiss::gpu::StandardGpuResources res; - faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, build_cfg->nlist, build_cfg->m, build_cfg->nbits, - GetMetricType(build_cfg->metric_type)); // IP not support - device_index.train(rows, (float*)p_data); - - std::shared_ptr host_index = nullptr; - host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index)); - - return std::make_shared(host_index); + auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (temp_resource != nullptr) { + ResScope rs(temp_resource, gpu_id_, true); + auto device_index = new faiss::gpu::GpuIndexIVFPQ(temp_resource->faiss_res.get(), dim, build_cfg->nlist, + build_cfg->m, build_cfg->nbits, + GetMetricType(build_cfg->metric_type)); // IP not support + device_index->train(rows, (float*)p_data); + std::shared_ptr host_index = nullptr; + host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index)); + return std::make_shared(host_index); + } else { + KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource"); + } } std::shared_ptr @@ -66,7 +68,14 @@ GPUIVFPQ::GenParams(const Config& config) { VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config& config) { - KNOWHERE_THROW_MSG("not support yet"); + std::lock_guard lk(mutex_); + + faiss::Index* device_index = index_.get(); + faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); } } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h index 13ea1075ca..f699445d65 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h @@ -18,6 +18,7 @@ #pragma once #include +#include #include "IndexGPUIVF.h" @@ -28,6 +29,10 @@ class GPUIVFPQ : public GPUIVF { explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) { } + GPUIVFPQ(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) + : GPUIVF(std::move(index), device_id, resource) { + } + IndexModelPtr Train(const DatasetPtr& dataset, const Config& config) override; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp index 03acbf31d7..e807b8c421 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp @@ -17,11 +17,15 @@ #include #include +#include + #include #include #include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" namespace knowhere { @@ -60,4 +64,18 @@ IVFPQ::Clone_impl(const std::shared_ptr& index) { return std::make_shared(index); } +VectorIndexPtr +IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); + + std::shared_ptr device_index; + device_index.reset(gpu_index); + return std::make_shared(device_index, device_id, res); + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); + } +} + } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h index 69aaa5090b..fc50c68389 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h @@ -34,6 +34,9 @@ class IVFPQ : public IVF { IndexModelPtr Train(const DatasetPtr& dataset, const Config& config) override; + VectorIndexPtr + CopyCpuToGpu(const int64_t& device_id, const Config& config) override; + protected: std::shared_ptr GenParams(const Config& config) override; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index 5ff2bfc2e3..363e629eca 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -49,7 +49,7 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config if (auto cpu_index = std::dynamic_pointer_cast(index)) { return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { - KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu"); + return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp index ae1034b9e0..fa294adc0e 100644 --- a/core/src/index/unittest/test_ivf.cpp +++ b/core/src/index/unittest/test_ivf.cpp @@ -211,7 +211,7 @@ TEST_P(IVFTest, clone_test) { { // copy from gpu to cpu - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"}; + std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "GPUIVFPQ", "IVFSQHybrid"}; auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ @@ -236,7 +236,7 @@ TEST_P(IVFTest, clone_test) { { // copy to gpu - std::vector support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ"}; + std::vector support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ", "IVFPQ", "GPUIVFPQ"}; auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp index 2dcf6bab7e..90d8c6523b 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/wrapper/ConfAdapter.cpp @@ -76,7 +76,7 @@ IVFConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) { if (size <= TYPICAL_COUNT / 16384 + 1) { // handle less row count, avoid nlist set to 0 return 1; - } else if (int(size / TYPICAL_COUNT) * nlist == 0) { + } else if (int(size / TYPICAL_COUNT) * nlist <= 0) { // calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT return int(size / TYPICAL_COUNT * 16384); } @@ -87,7 +87,11 @@ knowhere::Config IVFConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { auto conf = std::make_shared(); conf->k = metaconf.k; - conf->nprobe = metaconf.nprobe; + + if (metaconf.nprobe <= 0) + conf->nprobe = 16; // hardcode here + else + conf->nprobe = metaconf.nprobe; switch (type) { case IndexType::FAISS_IVFFLAT_GPU: @@ -123,11 +127,33 @@ IVFPQConfAdapter::Match(const TempMetaConf& metaconf) { conf->metric_type = metaconf.metric_type; conf->gpu_id = conf->gpu_id; conf->nbits = 8; - conf->m = 8; + + if (!(conf->d % 4)) + conf->m = conf->d / 4; // compression radio = 16 + else if (!(conf->d % 2)) + conf->m = conf->d / 2; // compression radio = 8 + else if (!(conf->d % 3)) + conf->m = conf->d / 3; // compression radio = 12 + else + conf->m = conf->d; // same as SQ8, compression radio = 4 + MatchBase(conf); return conf; } +knowhere::Config +IVFPQConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { + auto conf = std::make_shared(); + conf->k = metaconf.k; + + if (metaconf.nprobe <= 0) + conf->nprobe = 16; // hardcode here + else + conf->nprobe = metaconf.nprobe; + + return conf; +} + knowhere::Config NSGConfAdapter::Match(const TempMetaConf& metaconf) { auto conf = std::make_shared(); @@ -136,13 +162,14 @@ NSGConfAdapter::Match(const TempMetaConf& metaconf) { conf->metric_type = metaconf.metric_type; conf->gpu_id = conf->gpu_id; + double factor = metaconf.size / TYPICAL_COUNT; auto scale_factor = round(metaconf.dim / 128.0); scale_factor = scale_factor >= 4 ? 4 : scale_factor; - conf->nprobe = 6 + 10 * scale_factor; - conf->knng = 100 + 100 * scale_factor; - conf->search_length = 40 + 5 * scale_factor; - conf->out_degree = 50 + 5 * scale_factor; - conf->candidate_pool_size = 200 + 100 * scale_factor; + conf->nprobe = conf->nlist > 10000 ? conf->nlist * 0.02 : conf->nlist * 0.1; + conf->knng = (100 + 100 * scale_factor) * factor; + conf->search_length = (40 + 5 * scale_factor) * factor; + conf->out_degree = (50 + 5 * scale_factor) * factor; + conf->candidate_pool_size = (200 + 100 * scale_factor) * factor; MatchBase(conf); // WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist @@ -156,6 +183,9 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) auto conf = std::make_shared(); conf->k = metaconf.k; conf->search_length = metaconf.search_length; + if (metaconf.search_length == TEMPMETA_DEFAULT_VALUE) { + conf->search_length = 30; // TODO(linxj): hardcode here. + } return conf; } diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/wrapper/ConfAdapter.h index 4c8e528a2d..4723afb097 100644 --- a/core/src/wrapper/ConfAdapter.h +++ b/core/src/wrapper/ConfAdapter.h @@ -79,6 +79,9 @@ class IVFPQConfAdapter : public IVFConfAdapter { public: knowhere::Config Match(const TempMetaConf& metaconf) override; + + knowhere::Config + MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override; }; class NSGConfAdapter : public IVFConfAdapter { diff --git a/core/src/wrapper/ConfAdapterMgr.cpp b/core/src/wrapper/ConfAdapterMgr.cpp index b329588cf2..d0eba04529 100644 --- a/core/src/wrapper/ConfAdapterMgr.cpp +++ b/core/src/wrapper/ConfAdapterMgr.cpp @@ -53,6 +53,7 @@ AdapterMgr::RegisterAdapter() { REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_CPU, ivfpq_cpu); REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_GPU, ivfpq_gpu); + REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix); REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix); } diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp index 2323783c3d..9f9e71876d 100644 --- a/core/src/wrapper/VecIndex.cpp +++ b/core/src/wrapper/VecIndex.cpp @@ -168,6 +168,10 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { index = std::make_shared(gpu_device); break; } + case IndexType::FAISS_IVFPQ_MIX: { + index = std::make_shared(gpu_device); + return std::make_shared(index, IndexType::FAISS_IVFPQ_MIX); + } default: { return nullptr; } } return std::make_shared(index, type); @@ -276,6 +280,10 @@ ConvertToCpuIndexType(const IndexType& type) { case IndexType::FAISS_IVFSQ8_MIX: { return IndexType::FAISS_IVFSQ8_CPU; } + case IndexType::FAISS_IVFPQ_GPU: + case IndexType::FAISS_IVFPQ_MIX: { + return IndexType::FAISS_IVFPQ_CPU; + } default: { return type; } } } @@ -291,9 +299,12 @@ ConvertToGpuIndexType(const IndexType& type) { case IndexType::FAISS_IVFSQ8_CPU: { return IndexType::FAISS_IVFSQ8_GPU; } + case IndexType::FAISS_IVFPQ_MIX: + case IndexType::FAISS_IVFPQ_CPU: { + return IndexType::FAISS_IVFPQ_GPU; + } default: { return type; } } } - } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/VecIndex.h b/core/src/wrapper/VecIndex.h index cd9ba86951..efe01a25d7 100644 --- a/core/src/wrapper/VecIndex.h +++ b/core/src/wrapper/VecIndex.h @@ -33,6 +33,7 @@ namespace engine { using Config = knowhere::Config; +// TODO(linxj): replace with string, Do refactor serialization enum class IndexType { INVALID = 0, FAISS_IDMAP = 1, @@ -47,6 +48,7 @@ enum class IndexType { FAISS_IVFSQ8_GPU, FAISS_IVFSQ8_HYBRID, // only support build on gpu. NSG_MIX, + FAISS_IVFPQ_MIX, }; class VecIndex; diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index 097cb69075..5b0b01c00c 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -17,11 +17,7 @@ #include "external/easyloggingpp/easylogging++.h" #include "wrapper/VecIndex.h" - -#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" -#endif - #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "wrapper/utils.h" @@ -29,61 +25,73 @@ INITIALIZE_EASYLOGGINGPP -using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; +using ::testing::Combine; class KnowhereWrapperTest - : public DataGenBase, - public TestWithParam<::std::tuple> { + : public DataGenBase, + public TestWithParam<::std::tuple> { protected: - void - SetUp() override { -#ifdef MILVUS_GPU_VERSION + void SetUp() override { knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); -#endif std::string generator_type; std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); + milvus::engine::TempMetaConf tempconf; + tempconf.metric_type = knowhere::METRICTYPE::L2; + tempconf.gpu_id = DEVICEID; + tempconf.size = nb; + tempconf.dim = dim; + tempconf.k = k; + index_ = GetVecIndexFactory(index_type); - conf = ParamGenerator::GetInstance().Gen(index_type); - conf->k = k; - conf->d = dim; - conf->gpu_id = DEVICEID; + conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); + searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); + +// conf->k = k; +// conf->d = dim; +// conf->gpu_id = DEVICEID; } - void - TearDown() override { -#ifdef MILVUS_GPU_VERSION + void TearDown() override { knowhere::FaissGpuResourceMgr::GetInstance().Free(); -#endif } protected: milvus::engine::IndexType index_type; milvus::engine::VecIndexPtr index_ = nullptr; knowhere::Config conf; + knowhere::Config searchconf; }; -INSTANTIATE_TEST_CASE_P( - WrapperParam, KnowhereWrapperTest, - Values( -//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] -#ifdef MILVUS_GPU_VERSION - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 100000, 10, 10), - // std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, - // 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), -// std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), -#endif - // std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 100000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10))); +INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, + Values( + //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, + "Default", + 64, + 100000, + 10, + 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, + "Default", + 64, + 100000, + 10, + 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", DIM, 10000, 10, 10), +// std::make_tuple(milvus::engine::IndexType::NSG_MIX, "Default", DIM, 100000, 10, 10), +// std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", DIM, NB, 10, 10) + ) +); TEST_P(KnowhereWrapperTest, BASE_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -93,12 +101,10 @@ TEST_P(KnowhereWrapperTest, BASE_TEST) { std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } -#ifdef MILVUS_GPU_VERSION - TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -107,13 +113,13 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); { 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); + dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); } AssertResult(res_ids, res_dis); } @@ -125,12 +131,11 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { 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); + dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); } AssertResult(res_ids, res_dis); } } -#endif TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -139,7 +144,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); { @@ -152,7 +157,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } @@ -166,13 +171,12 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } } #include "wrapper/ConfAdapter.h" - TEST(whatever, test_config) { milvus::engine::TempMetaConf conf; auto nsg_conf = std::make_shared(); diff --git a/core/unittest/wrapper/utils.h b/core/unittest/wrapper/utils.h index c16b8fc930..917c130273 100644 --- a/core/unittest/wrapper/utils.h +++ b/core/unittest/wrapper/utils.h @@ -15,18 +15,21 @@ // specific language governing permissions and limitations // under the License. + #pragma once #include -#include -#include -#include #include #include +#include +#include +#include +#include -#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "wrapper/VecIndex.h" #include "wrapper/utils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +#include "wrapper/ConfAdapterMgr.h" class DataGenBase; @@ -40,29 +43,31 @@ constexpr int64_t PINMEM = 1024 * 1024 * 200; constexpr int64_t TEMPMEM = 1024 * 1024 * 300; constexpr int64_t RESNUM = 2; -static const char* CONFIG_PATH = "/tmp/milvus_test"; -static const char* CONFIG_FILE = "/server_config.yaml"; +static const char *CONFIG_PATH = "/tmp/milvus_test"; +static const char *CONFIG_FILE = "/server_config.yaml"; class KnowhereTest : public ::testing::Test { protected: - void - SetUp() override; - void - TearDown() override; + void SetUp() override; + void TearDown() override; }; 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, 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); + 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); + void AssertResult(const std::vector& ids, const std::vector& dis); int dim = DIM; int nb = NB; @@ -79,14 +84,22 @@ class DataGenBase { class ParamGenerator { public: - static ParamGenerator& - GetInstance() { + static ParamGenerator& GetInstance() { static ParamGenerator instance; return instance; } - knowhere::Config - Gen(const milvus::engine::IndexType& type) { + knowhere::Config GenSearchConf(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) { + auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); + return adapter->MatchSearch(conf, type); + } + + knowhere::Config GenBuild(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) { + auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); + return adapter->Match(conf); + } + + knowhere::Config Gen(const milvus::engine::IndexType& type) { switch (type) { case milvus::engine::IndexType::FAISS_IDMAP: { auto tempconf = std::make_shared(); @@ -113,34 +126,37 @@ class ParamGenerator { 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; - // } + case milvus::engine::IndexType::FAISS_IVFPQ_CPU: + case milvus::engine::IndexType::FAISS_IVFPQ_GPU: + case milvus::engine::IndexType::FAISS_IVFPQ_MIX: { + 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 SanityCheck : public DataGenBase { + +//class SanityCheck : public DataGenBase { // public: // void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, long *ids, // const int &k, long *gt_ids, float *gt_dis) override; //}; +