From 77b28a39819dbaecd6dbef0eee613482b302f7f7 Mon Sep 17 00:00:00 2001 From: "xiaojun.lin" Date: Thu, 10 Oct 2019 14:28:00 +0800 Subject: [PATCH] enable hybrid IVFSQ8 Former-commit-id: 9bbf835e63818ef6cc5f792d0ab11247794d0c3a --- .../index/vector_index/IndexIVFSQHybrid.cpp | 51 +++++-- .../index/vector_index/IndexIVFSQHybrid.h | 14 +- .../helpers/FaissGpuResourceMgr.h | 5 + cpp/src/core/test/CMakeLists.txt | 1 + cpp/src/core/test/test_ivf.cpp | 130 +++++++++++------- cpp/src/wrapper/VecImpl.cpp | 21 +++ cpp/src/wrapper/VecImpl.h | 3 + cpp/src/wrapper/VecIndex.h | 8 +- 8 files changed, 172 insertions(+), 61 deletions(-) diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp index b2110df9..c36a5d07 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp @@ -63,7 +63,7 @@ VectorIndexPtr IVFSQHybrid::CopyGpuToCpu(const Config &config) { std::lock_guard lk(mutex_); - if (auto device_idx = std::dynamic_pointer_cast(index_)) { + if (auto device_idx = std::dynamic_pointer_cast(index_)) { faiss::Index *device_index = index_.get(); faiss::Index *host_index = faiss::gpu::index_gpu_to_cpu(device_index); @@ -113,6 +113,7 @@ IVFSQHybrid::search_impl(int64_t n, if (gpu_mode) { GPUIVF::search_impl(n, data, k, distances, labels, cfg); } else { + ResScope rs(res_, gpu_id_); IVF::search_impl(n, data, k, distances, labels, cfg); } } @@ -121,7 +122,9 @@ QuantizerPtr IVFSQHybrid::LoadQuantizer(const Config &conf) { auto quantizer_conf = std::dynamic_pointer_cast(conf); if (quantizer_conf != nullptr) { - quantizer_conf->CheckValid(); // throw exception + if(quantizer_conf->mode != 1) { + KNOWHERE_THROW_MSG("mode only support 1 in this func"); + } } gpu_id_ = quantizer_conf->gpu_id; @@ -133,13 +136,14 @@ IVFSQHybrid::LoadQuantizer(const Config &conf) { auto index_composition = new faiss::IndexComposition; index_composition->index = index_.get(); index_composition->quantizer = nullptr; - index_composition->mode = quantizer_conf->mode; // 1 or 2 + index_composition->mode = quantizer_conf->mode; // only 1 auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index_composition, &option); delete gpu_index; - std::shared_ptr q; - q->quantizer = index_composition; + auto q = std::make_shared(); + q->quantizer = index_composition->quantizer; + res_ = res; return q; } else { KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); @@ -153,15 +157,13 @@ IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) { KNOWHERE_THROW_MSG("Quantizer type error"); } - if (ivf_quantizer->quantizer->mode == 2) gpu_mode = true; // all in gpu - faiss::IndexIVF *ivf_index = dynamic_cast(index_.get()); faiss::gpu::GpuIndexFlat *is_gpu_flat_index = dynamic_cast(ivf_index->quantizer); if (is_gpu_flat_index == nullptr) { delete ivf_index->quantizer; - ivf_index->quantizer = ivf_quantizer->quantizer->quantizer; + ivf_index->quantizer = ivf_quantizer->quantizer; } } @@ -175,5 +177,38 @@ IVFSQHybrid::UnsetQuantizer() { ivf_index->quantizer = nullptr; } +void +IVFSQHybrid::LoadData(const knowhere::QuantizerPtr &q, const Config &conf) { + auto quantizer_conf = std::dynamic_pointer_cast(conf); + if (quantizer_conf != nullptr) { + if(quantizer_conf->mode != 2) { + KNOWHERE_THROW_MSG("mode only support 2 in this func"); + } + } + if (quantizer_conf->gpu_id != gpu_id_) { + KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card"); + } + + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { + ResScope rs(res, gpu_id_, false); + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + auto ivf_quantizer = std::dynamic_pointer_cast(q); + if (ivf_quantizer == nullptr) KNOWHERE_THROW_MSG("quantizer type not faissivfquantizer"); + + auto index_composition = new faiss::IndexComposition; + index_composition->index = index_.get(); + index_composition->quantizer = ivf_quantizer->quantizer; + index_composition->mode = quantizer_conf->mode; // only 2 + + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index_composition, &option); + index_.reset(gpu_index); + gpu_mode = true; // all in gpu + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); + } +} + } } diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h index f2f7c200..687eeda7 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h @@ -29,15 +29,20 @@ namespace zilliz { namespace knowhere { struct FaissIVFQuantizer : public Quantizer { - faiss::IndexComposition *quantizer = nullptr; + faiss::gpu::GpuIndexFlat *quantizer = nullptr; }; using FaissIVFQuantizerPtr = std::shared_ptr; class IVFSQHybrid : public GPUIVFSQ { public: - explicit IVFSQHybrid(const int &device_id) : GPUIVFSQ(device_id) {} + explicit IVFSQHybrid(const int &device_id) : GPUIVFSQ(device_id) { + gpu_mode = false; + } - explicit IVFSQHybrid(std::shared_ptr index) : GPUIVFSQ(-1) {gpu_mode = false;} + explicit IVFSQHybrid(std::shared_ptr index) : GPUIVFSQ(-1) { + index_ = index; + gpu_mode = false; + } explicit IVFSQHybrid(std::shared_ptr index, const int64_t &device_id, ResPtr &resource) : GPUIVFSQ(index, device_id, resource) { @@ -54,6 +59,9 @@ class IVFSQHybrid : public GPUIVFSQ { void UnsetQuantizer(); + void + LoadData(const knowhere::QuantizerPtr &q, const Config& conf); + IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h index 3b4786a2..c1585e1b 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h @@ -100,6 +100,11 @@ public: Lock(); } + ResScope(ResWPtr &res, const int64_t& device_id, const bool& isown) + : resource(res), device_id(device_id), move(true), own(isown) { + Lock(); + } + // specif for search // get the ownership of gpuresource and gpu ResScope(ResWPtr &res, const int64_t &device_id) diff --git a/cpp/src/core/test/CMakeLists.txt b/cpp/src/core/test/CMakeLists.txt index cc6ea5bd..0cf83de2 100644 --- a/cpp/src/core/test/CMakeLists.txt +++ b/cpp/src/core/test/CMakeLists.txt @@ -41,6 +41,7 @@ set(ivf_srcs ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp + ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp ) diff --git a/cpp/src/core/test/test_ivf.cpp b/cpp/src/core/test/test_ivf.cpp index 38b28483..f1d19565 100644 --- a/cpp/src/core/test/test_ivf.cpp +++ b/cpp/src/core/test/test_ivf.cpp @@ -65,6 +65,8 @@ IVFIndexPtr IndexFactory(const std::string &type) { return std::make_shared(); } else if (type == "GPUIVFSQ") { return std::make_shared(device_id); + } else if (type == "IVFSQHybrid") { + return std::make_shared(device_id); } } @@ -72,6 +74,7 @@ enum class ParameterType { ivf, ivfpq, ivfsq, + ivfsqhybrid, nsg, }; @@ -105,7 +108,7 @@ class ParamGenerator { tempconf->metric_type = METRICTYPE::L2; return tempconf; } - else if (type == ParameterType::ivfsq) { + else if (type == ParameterType::ivfsq || type == ParameterType::ivfsqhybrid) { auto tempconf = std::make_shared(); tempconf->d = DIM; tempconf->gpu_id = device_id; @@ -160,6 +163,7 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest, // std::make_tuple("GPUIVFPQ", ParameterType::ivfpq), std::make_tuple("IVFSQ", ParameterType::ivfsq), std::make_tuple("GPUIVFSQ", ParameterType::ivfsq) + std::make_tuple("IVFSQHybrid", ParameterType::ivfsqhybrid) ) ); @@ -210,36 +214,60 @@ TEST_P(IVFTest, ivf_basic) { //PrintResult(result, nq, k); } -//TEST_P(IVFTest, hybrid) { -// 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 iss_idx = std::make_shared(device_id); -// -// auto binaryset = index_->Serialize(); -// iss_idx->Load(binaryset); -// -// auto quantizer_conf = std::make_shared(); -// quantizer_conf->mode = 1; -// quantizer_conf->gpu_id = 1; -// auto q = iss_idx->LoadQuantizer(quantizer_conf); -// iss_idx->SetQuantizer(q); -// auto result = iss_idx->Search(query_dataset, conf); +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); -// //PrintResult(result, nq, 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); + } + + { + 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; + hybrid_2_idx->LoadData(q, quantizer_conf); + + auto result = hybrid_2_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; } @@ -350,29 +378,35 @@ TEST_P(IVFTest, clone_test) { } }; +// { +// // clone in place +// std::vector support_idx_vec{"IVF", "GPUIVF", "IVFPQ", "IVFSQ", "GPUIVFSQ"}; +// 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 = index_->Clone(); +// auto clone_result = clone_index->Search(query_dataset, conf); +// //AssertAnns(result, nq, conf->k); +// AssertEqual(result, clone_result); +// std::cout << "inplace clone [" << index_type << "] success" << std::endl; +// }); +// } else { +// EXPECT_THROW({ +// std::cout << "inplace clone [" << index_type << "] failed" << std::endl; +// auto clone_index = index_->Clone(); +// }, KnowhereException); +// } +// } + { - // clone in place - std::vector support_idx_vec{"IVF", "GPUIVF", "IVFPQ", "IVFSQ", "GPUIVFSQ"}; - 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 = index_->Clone(); - auto clone_result = clone_index->Search(query_dataset, conf); - //AssertAnns(result, nq, conf->k); - AssertEqual(result, clone_result); - std::cout << "inplace clone [" << index_type << "] success" << std::endl; - }); - } else { - EXPECT_THROW({ - std::cout << "inplace clone [" << index_type << "] failed" << std::endl; - auto clone_index = index_->Clone(); - }, KnowhereException); + if (index_type == "IVFSQHybrid") { + return; } } { // copy from gpu to cpu - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ"}; + 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()) { EXPECT_NO_THROW({ @@ -412,7 +446,7 @@ TEST_P(IVFTest, clone_test) { TEST_P(IVFTest, seal_test) { //FaissGpuResourceMgr::GetInstance().InitDevice(device_id); - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ"}; + 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()) { return; diff --git a/cpp/src/wrapper/VecImpl.cpp b/cpp/src/wrapper/VecImpl.cpp index 4530757f..17b9ec61 100644 --- a/cpp/src/wrapper/VecImpl.cpp +++ b/cpp/src/wrapper/VecImpl.cpp @@ -304,6 +304,7 @@ IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) { WRAPPER_LOG_ERROR << e.what(); return Status(KNOWHERE_ERROR, e.what()); } + return Status::OK(); } Status @@ -323,6 +324,26 @@ IVFHybridIndex::UnsetQuantizer() { WRAPPER_LOG_ERROR << e.what(); return Status(KNOWHERE_ERROR, e.what()); } + return Status::OK(); +} + +Status IVFHybridIndex::LoadData(const knowhere::QuantizerPtr &q, const Config &conf) { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + new_idx->LoadData(q, conf); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + return Status(KNOWHERE_ERROR, "not support"); + } + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); } } // namespace engine diff --git a/cpp/src/wrapper/VecImpl.h b/cpp/src/wrapper/VecImpl.h index ba0934cf..c5e2484d 100644 --- a/cpp/src/wrapper/VecImpl.h +++ b/cpp/src/wrapper/VecImpl.h @@ -111,6 +111,9 @@ class IVFHybridIndex : public IVFMixIndex { Status UnsetQuantizer() override; + + Status + LoadData(const knowhere::QuantizerPtr &q, const Config &conf) override; }; class BFIndex : public VecIndexImpl { diff --git a/cpp/src/wrapper/VecIndex.h b/cpp/src/wrapper/VecIndex.h index 24352b8d..e45f3d05 100644 --- a/cpp/src/wrapper/VecIndex.h +++ b/cpp/src/wrapper/VecIndex.h @@ -104,15 +104,19 @@ class VecIndex { Load(const zilliz::knowhere::BinarySet &index_binary) = 0; // TODO(linxj): refactor later + //////////////// virtual knowhere::QuantizerPtr - LoadQuantizer(const Config& conf) { return Status::OK(); } + LoadQuantizer(const Config& conf) { return nullptr; } + + virtual Status + LoadData(const knowhere::QuantizerPtr &q, const Config &conf) { return Status::OK(); } - // TODO(linxj): refactor later virtual Status SetQuantizer(const knowhere::QuantizerPtr& q) { return Status::OK(); } virtual Status UnsetQuantizer() { return Status::OK(); } + //////////////// }; extern Status -- GitLab