From 8540b3567acb4124cdb874d0928e03118c1dccb3 Mon Sep 17 00:00:00 2001 From: wxyu Date: Sun, 13 Oct 2019 13:20:11 +0800 Subject: [PATCH] SQ8H in GPU part2 Former-commit-id: 4c48987574ed24bf8a543d97520eb3a6b554fca5 --- .../index/vector_index/IndexIVFSQHybrid.cpp | 8 +++++--- .../index/vector_index/IndexIVFSQHybrid.h | 3 +-- cpp/src/core/unittest/test_ivf.cpp | 4 ++-- cpp/src/db/engine/ExecutionEngineImpl.cpp | 15 ++++++++++----- cpp/src/wrapper/VecImpl.cpp | 9 +++------ cpp/src/wrapper/VecImpl.h | 2 +- cpp/src/wrapper/VecIndex.h | 4 ++-- 7 files changed, 24 insertions(+), 21 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 5f916f33..c6c92913 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp @@ -180,7 +180,7 @@ IVFSQHybrid::UnsetQuantizer() { ivf_index->quantizer = nullptr; } -void +VectorIndexPtr IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { auto quantizer_conf = std::dynamic_pointer_cast(conf); if (quantizer_conf != nullptr) { @@ -207,8 +207,10 @@ IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { 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 = 2; // all in gpu + std::shared_ptr new_idx; + new_idx.reset(gpu_index); + auto sq_idx = std::make_shared(new_idx, gpu_id_, res); + return sq_idx; } 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 4273a412..d0c58baa 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h @@ -60,8 +60,7 @@ class IVFSQHybrid : public GPUIVFSQ { void UnsetQuantizer(); - // todo(xiaojun): return void => VecIndex - void + VectorIndexPtr LoadData(const knowhere::QuantizerPtr& q, const Config& conf); IndexModelPtr diff --git a/cpp/src/core/unittest/test_ivf.cpp b/cpp/src/core/unittest/test_ivf.cpp index c6faea91..c5066e96 100644 --- a/cpp/src/core/unittest/test_ivf.cpp +++ b/cpp/src/core/unittest/test_ivf.cpp @@ -253,9 +253,9 @@ TEST_P(IVFTest, hybrid) { 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 gpu_idx = hybrid_2_idx->LoadData(q, quantizer_conf); - auto result = hybrid_2_idx->Search(query_dataset, conf); + auto result = gpu_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); PrintResult(result, nq, k); } diff --git a/cpp/src/db/engine/ExecutionEngineImpl.cpp b/cpp/src/db/engine/ExecutionEngineImpl.cpp index 6b8bb622..c70a5c3b 100644 --- a/cpp/src/db/engine/ExecutionEngineImpl.cpp +++ b/cpp/src/db/engine/ExecutionEngineImpl.cpp @@ -256,11 +256,14 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) { conf->gpu_id = device_id; if (quantizer) { + std::cout << "cache hit" << std::endl; // cache hit conf->mode = 2; - index_->SetQuantizer(quantizer->Data()); - index_->LoadData(quantizer->Data(), conf); + auto new_index = index_->LoadData(quantizer->Data(), conf); + index_ = new_index; } else { + std::cout << "cache miss" << std::endl; + // cache hit // cache miss if (index_ == nullptr) { ENGINE_LOG_ERROR << "ExecutionEngineImpl: index is null, failed to copy to gpu"; @@ -268,9 +271,9 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) { } conf->mode = 1; auto q = index_->LoadQuantizer(conf); - index_->SetQuantizer(q); conf->mode = 2; - index_->LoadData(q, conf); + auto new_index = index_->LoadData(q, conf); + index_ = new_index; // cache auto cached_quantizer = std::make_shared(q); @@ -445,7 +448,9 @@ ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, int64_t npr auto status = index_->Search(n, data, distances, labels, conf); - HybridUnset(); + if (hybrid) { + HybridUnset(); + } if (!status.ok()) { ENGINE_LOG_ERROR << "Search error"; diff --git a/cpp/src/wrapper/VecImpl.cpp b/cpp/src/wrapper/VecImpl.cpp index 1ed20c80..3ff79690 100644 --- a/cpp/src/wrapper/VecImpl.cpp +++ b/cpp/src/wrapper/VecImpl.cpp @@ -315,24 +315,21 @@ IVFHybridIndex::UnsetQuantizer() { return Status::OK(); } -Status +VecIndexPtr 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); + return std::make_shared(new_idx->LoadData(q, conf), type); } 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(); + return nullptr; } } // namespace engine diff --git a/cpp/src/wrapper/VecImpl.h b/cpp/src/wrapper/VecImpl.h index fd9bb79c..1f5ca296 100644 --- a/cpp/src/wrapper/VecImpl.h +++ b/cpp/src/wrapper/VecImpl.h @@ -106,7 +106,7 @@ class IVFHybridIndex : public IVFMixIndex { Status UnsetQuantizer() override; - Status + VecIndexPtr LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override; }; diff --git a/cpp/src/wrapper/VecIndex.h b/cpp/src/wrapper/VecIndex.h index f5fdd494..55981ef5 100644 --- a/cpp/src/wrapper/VecIndex.h +++ b/cpp/src/wrapper/VecIndex.h @@ -103,9 +103,9 @@ class VecIndex : public cache::DataObj { return nullptr; } - virtual Status + virtual VecIndexPtr LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { - return Status::OK(); + return nullptr; } virtual Status -- GitLab