diff --git a/cpp/src/core/knowhere/CMakeLists.txt b/cpp/src/core/knowhere/CMakeLists.txt index 4d596e243ab4ec5d501e66028eee51a54d472d49..cb0d5895d103c127aa126aad50a84823d152cb5d 100644 --- a/cpp/src/core/knowhere/CMakeLists.txt +++ b/cpp/src/core/knowhere/CMakeLists.txt @@ -50,6 +50,7 @@ set(index_srcs knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp knowhere/index/vector_index/IndexIVFSQ.cpp knowhere/index/vector_index/IndexGPUIVFSQ.cpp + knowhere/index/vector_index/IndexIVFSQHybrid.cpp knowhere/index/vector_index/IndexIVFPQ.cpp knowhere/index/vector_index/IndexGPUIVFPQ.cpp knowhere/index/vector_index/FaissBaseIndex.cpp diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp index e2154eaca5de2500795bfc0560b2d6f5052784ce..667f8ec4f15a0ef037f3cb17663468fe6db43028 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp @@ -127,17 +127,6 @@ void GPUIVF::LoadImpl(const BinarySet &index_binary) { } } -IVFIndexPtr GPUIVF::Copy_index_gpu_to_cpu() { - 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); -} - void GPUIVF::search_impl(int64_t n, const float *data, int64_t k, diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h index cd04fe4c5433f677b2c5e3444d7b9b17119e9331..82695a7e4a34ffcc55e0507b5d03b5d389f23cd1 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h @@ -75,9 +75,6 @@ public: VectorIndexPtr Clone() final; - // TODO(linxj): Deprecated - virtual IVFIndexPtr Copy_index_gpu_to_cpu(); - protected: void search_impl(int64_t n, diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index d8ccdd1087bd78a2999b16187c558f7f5f4717a4..13d4523705e6e2c82622484d039f9b3b0783d300 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -22,6 +22,7 @@ #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "Cloner.h" @@ -38,6 +39,10 @@ VectorIndexPtr CopyGpuToCpu(const VectorIndexPtr &index, const Config &config) { } VectorIndexPtr CopyCpuToGpu(const VectorIndexPtr &index, const int64_t &device_id, const Config &config) { + if (auto device_index = std::dynamic_pointer_cast(index)) { + return device_index->CopyCpuToGpu(device_id, config); + } + if (auto device_index = std::dynamic_pointer_cast(index)) { return device_index->CopyGpuToGpu(device_id, config); } diff --git a/cpp/src/core/test/test_ivf.cpp b/cpp/src/core/test/test_ivf.cpp index a08d28d7d885532e08d1edbf6ae0d6912a935ffa..38b28483ea95c1af6c1263baa5d1fb8050a177e6 100644 --- a/cpp/src/core/test/test_ivf.cpp +++ b/cpp/src/core/test/test_ivf.cpp @@ -25,6 +25,7 @@ #include #include +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" #include "knowhere/adapter/Structure.h" @@ -209,6 +210,37 @@ 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); +// AssertAnns(result, nq, conf->k); +// //PrintResult(result, nq, k); +//} + //TEST_P(IVFTest, gpu_to_cpu) { // if (index_type.find("GPU") == std::string::npos) { return; } // diff --git a/cpp/src/wrapper/VecImpl.cpp b/cpp/src/wrapper/VecImpl.cpp index 65d4c60796dee44d03e665b62a2a46373edd7255..0bbe97cbbed1b265911146b241938ba36fc26ddf 100644 --- a/cpp/src/wrapper/VecImpl.cpp +++ b/cpp/src/wrapper/VecImpl.cpp @@ -19,6 +19,7 @@ #include "wrapper/VecImpl.h" #include "utils/Log.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/helpers/Cloner.h" @@ -251,7 +252,7 @@ IVFMixIndex::BuildAll(const int64_t &nb, index_->set_index_model(model); index_->Add(dataset, cfg); - if (auto device_index = std::dynamic_pointer_cast(index_)) { + if (auto device_index = std::dynamic_pointer_cast(index_)) { auto host_index = device_index->CopyGpuToCpu(Config()); index_ = host_index; type = ConvertToCpuIndexType(type); @@ -276,6 +277,33 @@ IVFMixIndex::Load(const zilliz::knowhere::BinarySet &index_binary) { return Status::OK(); } +knowhere::QuantizerPtr IVFHybridIndex::LoadQuantizer(const Config& conf) { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)){ + return new_idx->LoadQuantizer(conf); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + } +} + +Status IVFHybridIndex::SetQuantizer(knowhere::QuantizerPtr q) { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + new_idx->SetQuantizer(q); + } 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()); + } +} + } // namespace engine } // namespace milvus } // namespace zilliz diff --git a/cpp/src/wrapper/VecImpl.h b/cpp/src/wrapper/VecImpl.h index a91b24562b9a607f8aed312ddc25b58a4c5651e3..7a131f7e07fe28c03600dd104a400ed1ab25e7d1 100644 --- a/cpp/src/wrapper/VecImpl.h +++ b/cpp/src/wrapper/VecImpl.h @@ -101,6 +101,13 @@ class IVFMixIndex : public VecIndexImpl { Load(const zilliz::knowhere::BinarySet &index_binary) override; }; +class IVFHybridIndex : public IVFMixIndex { + public: + knowhere::QuantizerPtr LoadQuantizer(const Config& conf) override; + + Status SetQuantizer(knowhere::QuantizerPtr q) override; +}; + class BFIndex : public VecIndexImpl { public: explicit BFIndex(std::shared_ptr index) : VecIndexImpl(std::move(index), diff --git a/cpp/src/wrapper/VecIndex.cpp b/cpp/src/wrapper/VecIndex.cpp index 734f66338f1d0cbfb540ff0a48b544be89e45c48..27bba6149e016c5bf6550d2efbde841c93d6b881 100644 --- a/cpp/src/wrapper/VecIndex.cpp +++ b/cpp/src/wrapper/VecIndex.cpp @@ -25,6 +25,7 @@ #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexKDT.h" #include "knowhere/index/vector_index/IndexNSG.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/common/Exception.h" #include "VecImpl.h" #include "utils/Log.h" @@ -137,6 +138,10 @@ GetVecIndexFactory(const IndexType &type, const Config &cfg) { index = std::make_shared(gpu_device); break; } + case IndexType::FAISS_IVFSQ8_HYBRID: { + index = std::make_shared(gpu_device); + break; + } case IndexType::NSG_MIX: { // TODO(linxj): bug. index = std::make_shared(gpu_device); break; diff --git a/cpp/src/wrapper/VecIndex.h b/cpp/src/wrapper/VecIndex.h index e4fe387260f4b94278d2a7b9f531db96b23a021f..8a9c91377bf7c47bafea18992def56d2d0e849ef 100644 --- a/cpp/src/wrapper/VecIndex.h +++ b/cpp/src/wrapper/VecIndex.h @@ -24,6 +24,7 @@ #include "utils/Status.h" #include "knowhere/common/Config.h" #include "knowhere/common/BinarySet.h" +#include "knowhere/index/vector_index/Quantizer.h" namespace zilliz { namespace milvus { @@ -43,6 +44,7 @@ enum class IndexType { FAISS_IVFSQ8_MIX, FAISS_IVFSQ8_CPU, FAISS_IVFSQ8_GPU, + FAISS_IVFSQ8_HYBRID, // only support build on gpu. NSG_MIX, }; @@ -100,6 +102,14 @@ class VecIndex { virtual Status Load(const zilliz::knowhere::BinarySet &index_binary) = 0; + + // TODO(linxj): refactor later + virtual knowhere::QuantizerPtr + LoadQuantizer(const Config& conf) {} + + // TODO(linxj): refactor later + virtual Status + SetQuantizer(knowhere::QuantizerPtr q) {} }; extern Status