diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md index c1bf840160..0a16fed4f1 100644 --- a/cpp/CHANGELOG.md +++ b/cpp/CHANGELOG.md @@ -35,6 +35,7 @@ Please mark all change in change log and use the ticket from JIRA. - MS-626 - Refactor DataObj to support cache any type data ## New Feature +- MS-627 - Integrate new index: IVFSQHybrid ## Task - MS-554 - Change license to Apache 2.0 diff --git a/cpp/src/core/cmake/ThirdPartyPackagesCore.cmake b/cpp/src/core/cmake/ThirdPartyPackagesCore.cmake index 63d991a1e5..5691d4f9dc 100644 --- a/cpp/src/core/cmake/ThirdPartyPackagesCore.cmake +++ b/cpp/src/core/cmake/ThirdPartyPackagesCore.cmake @@ -235,8 +235,10 @@ else() message(STATUS "FAISS URL = ${FAISS_SOURCE_URL}") endif() # set(FAISS_MD5 "a589663865a8558205533c8ac414278c") -# set(FAISS_MD5 "57da9c4f599cc8fa4260488b1c96e1cc") # commit-id 6dbdf75987c34a2c853bd172ea0d384feea8358c -set(FAISS_MD5 "21deb1c708490ca40ecb899122c01403") # commit-id 643e48f479637fd947e7b93fa4ca72b38ecc9a39 +# set(FAISS_MD5 "57da9c4f599cc8fa4260488b1c96e1cc") # commit-id 6dbdf75987c34a2c853bd172ea0d384feea8358c branch-0.2.0 +# set(FAISS_MD5 "21deb1c708490ca40ecb899122c01403") # commit-id 643e48f479637fd947e7b93fa4ca72b38ecc9a39 branch-0.2.0 +# set(FAISS_MD5 "072db398351cca6e88f52d743bbb9fa0") # commit-id 3a2344d04744166af41ef1a74449d68a315bfe17 branch-0.2.1 +set(FAISS_MD5 "5af237d77947ee632f169bcb36feee2b") # commit-id 2c8affd0da60354e4322fa4c0224519e7912b9c4 branch-0.2.1 if(DEFINED ENV{KNOWHERE_ARROW_URL}) set(ARROW_SOURCE_URL "$ENV{KNOWHERE_ARROW_URL}") diff --git a/cpp/src/core/knowhere/CMakeLists.txt b/cpp/src/core/knowhere/CMakeLists.txt index 4d596e243a..cb0d5895d1 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 d538c0dea1..a5e8f90f34 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp @@ -126,32 +126,22 @@ 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, float* distances, int64_t* labels, const Config& cfg) { std::lock_guard lk(mutex_); // TODO(linxj): gpu index support GenParams - if (auto device_index = std::static_pointer_cast(index_)) { + if (auto device_index = std::dynamic_pointer_cast(index_)) { auto search_cfg = std::dynamic_pointer_cast(cfg); device_index->setNumProbes(search_cfg->nprobe); { - // TODO(linxj): allocate mem + // TODO(linxj): allocate gpu mem ResScope rs(res_, gpu_id_); device_index->search(n, (float*)data, k, distances, labels); } + } else { + KNOWHERE_THROW_MSG("Not a GpuIndexIVF type."); } } @@ -159,12 +149,16 @@ VectorIndexPtr GPUIVF::CopyGpuToCpu(const Config& config) { std::lock_guard lk(mutex_); - faiss::Index* device_index = index_.get(); - faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_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); - std::shared_ptr new_index; - new_index.reset(host_index); - return std::make_shared(new_index); + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); + } else { + return std::make_shared(index_); + } } VectorIndexPtr 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 8edd0fe30e..fa9a206c48 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h @@ -78,10 +78,6 @@ class GPUIVF : public IVF, public GPUIndex { VectorIndexPtr Clone() final; - // TODO(linxj): Deprecated - virtual IVFIndexPtr - Copy_index_gpu_to_cpu(); - protected: void search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp index 428710d355..510ab46bd6 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp @@ -146,7 +146,7 @@ IVF::set_index_model(IndexModelPtr model) { std::shared_ptr IVF::GenParams(const Config& config) { - auto params = std::make_shared(); + auto params = std::make_shared(); auto search_cfg = std::dynamic_pointer_cast(config); params->nprobe = search_cfg->nprobe; diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp new file mode 100644 index 0000000000..60b1770fc3 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp @@ -0,0 +1,208 @@ +// +// 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 "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "faiss/AutoTune.h" +#include "faiss/gpu/GpuAutoTune.h" +#include "faiss/gpu/GpuIndexIVF.h" +#include "knowhere/adapter/VectorAdapter.h" +#include "knowhere/common/Exception.h" + +namespace knowhere { + +IndexModelPtr +IVFSQHybrid::Train(const DatasetPtr& dataset, const Config& config) { + auto build_cfg = std::dynamic_pointer_cast(config); + if (build_cfg != nullptr) { + build_cfg->CheckValid(); // throw exception + } + gpu_id_ = build_cfg->gpu_id; + + GETTENSOR(dataset) + + std::stringstream index_type; + index_type << "IVF" << build_cfg->nlist << "," + << "SQ8Hybrid"; + auto build_index = faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(build_cfg->metric_type)); + + auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (temp_resource != nullptr) { + ResScope rs(temp_resource, gpu_id_, true); + auto device_index = faiss::gpu::index_cpu_to_gpu(temp_resource->faiss_res.get(), gpu_id_, build_index); + device_index->train(rows, (float*)p_data); + + std::shared_ptr host_index = nullptr; + host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index)); + + delete device_index; + delete build_index; + + return std::make_shared(host_index); + } else { + KNOWHERE_THROW_MSG("Build IVFSQHybrid can't get gpu resource"); + } +} + +VectorIndexPtr +IVFSQHybrid::CopyGpuToCpu(const Config& config) { + std::lock_guard lk(mutex_); + + 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); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); + } else { + // TODO(linxj): why? jinhai + return std::make_shared(index_); + } +} + +VectorIndexPtr +IVFSQHybrid::CopyCpuToGpu(const int64_t& device_id, const Config& config) { + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + faiss::IndexComposition index_composition; + index_composition.index = index_.get(); + index_composition.quantizer = nullptr; + index_composition.mode = 0; // copy all + + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, &index_composition, &option); + + 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"); + } +} + +void +IVFSQHybrid::LoadImpl(const BinarySet& index_binary) { + FaissBaseIndex::LoadImpl(index_binary); // load on cpu +} + +void +IVFSQHybrid::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, + const Config& cfg) { + 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); + } +} + +QuantizerPtr +IVFSQHybrid::LoadQuantizer(const Config& conf) { + auto quantizer_conf = std::dynamic_pointer_cast(conf); + if (quantizer_conf != nullptr) { + if (quantizer_conf->mode != 1) { + KNOWHERE_THROW_MSG("mode only support 1 in this func"); + } + } + gpu_id_ = quantizer_conf->gpu_id; + + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { + ResScope rs(res, gpu_id_, false); + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + auto index_composition = new faiss::IndexComposition; + index_composition->index = index_.get(); + index_composition->quantizer = nullptr; + 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; + + 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"); + } +} + +void +IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) { + auto ivf_quantizer = std::dynamic_pointer_cast(q); + if (ivf_quantizer == nullptr) { + KNOWHERE_THROW_MSG("Quantizer type error"); + } + + 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; + } +} + +void +IVFSQHybrid::UnsetQuantizer() { + auto* ivf_index = dynamic_cast(index_.get()); + if (ivf_index == nullptr) { + KNOWHERE_THROW_MSG("Index type error"); + } + + 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"); + } +} + +} // namespace knowhere diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h new file mode 100644 index 0000000000..1ec67760ff --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h @@ -0,0 +1,82 @@ +// 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. + +#pragma once + +#include +#include + +#include "IndexGPUIVFSQ.h" +#include "Quantizer.h" + +namespace knowhere { + +struct FaissIVFQuantizer : public Quantizer { + faiss::gpu::GpuIndexFlat* quantizer = nullptr; +}; +using FaissIVFQuantizerPtr = std::shared_ptr; + +class IVFSQHybrid : public GPUIVFSQ { + public: + explicit IVFSQHybrid(const int& device_id) : GPUIVFSQ(device_id) { + 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) { + gpu_mode = true; + } + + public: + QuantizerPtr + LoadQuantizer(const Config& conf); + + void + SetQuantizer(const QuantizerPtr& q); + + void + UnsetQuantizer(); + + void + LoadData(const knowhere::QuantizerPtr& q, const Config& conf); + + IndexModelPtr + Train(const DatasetPtr& dataset, const Config& config) override; + + VectorIndexPtr + CopyGpuToCpu(const Config& config) override; + + VectorIndexPtr + CopyCpuToGpu(const int64_t& device_id, const Config& config) override; + + protected: + void + search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; + + void + LoadImpl(const BinarySet& index_binary) override; + + protected: + bool gpu_mode = false; +}; + +} // namespace knowhere diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/Quantizer.h b/cpp/src/core/knowhere/knowhere/index/vector_index/Quantizer.h new file mode 100644 index 0000000000..ea74e97c82 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/Quantizer.h @@ -0,0 +1,35 @@ +// 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. + +#pragma once + +#include +#include "knowhere/common/Config.h" + +namespace knowhere { + +struct Quantizer { + virtual ~Quantizer() = default; +}; +using QuantizerPtr = std::shared_ptr; + +struct QuantizerCfg : Cfg { + uint64_t mode = -1; // 0: all data, 1: copy quantizer, 2: copy data +}; +using QuantizerConfig = std::shared_ptr; + +} // namespace knowhere 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 16191b8183..5ff2bfc2e3 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/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" namespace knowhere { namespace cloner { @@ -37,6 +38,10 @@ 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/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h index bf570ca4eb..73f959baa0 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 @@ -96,6 +96,11 @@ class ResScope { 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) : device_id(device_id), move(false), own(true) { diff --git a/cpp/src/core/thirdparty/versions.txt b/cpp/src/core/thirdparty/versions.txt index 59a496f5e5..9ee845f1e3 100644 --- a/cpp/src/core/thirdparty/versions.txt +++ b/cpp/src/core/thirdparty/versions.txt @@ -3,4 +3,4 @@ BOOST_VERSION=1.70.0 GTEST_VERSION=1.8.1 LAPACK_VERSION=v3.8.0 OPENBLAS_VERSION=v0.3.6 -FAISS_VERSION=branch-0.2.0 \ No newline at end of file +FAISS_VERSION=branch-0.2.1 \ No newline at end of file diff --git a/cpp/src/core/unittest/CMakeLists.txt b/cpp/src/core/unittest/CMakeLists.txt index 6b5775c2cd..0a52a2ed83 100644 --- a/cpp/src/core/unittest/CMakeLists.txt +++ b/cpp/src/core/unittest/CMakeLists.txt @@ -42,6 +42,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/unittest/test_ivf.cpp b/cpp/src/core/unittest/test_ivf.cpp index 17eb888ddc..3fbc5e0141 100644 --- a/cpp/src/core/unittest/test_ivf.cpp +++ b/cpp/src/core/unittest/test_ivf.cpp @@ -24,7 +24,6 @@ #include #include -#include "knowhere/adapter/Structure.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" @@ -33,6 +32,7 @@ #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/utils.h" @@ -67,6 +67,8 @@ 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); } } @@ -74,6 +76,7 @@ enum class ParameterType { ivf, ivfpq, ivfsq, + ivfsqhybrid, nsg, }; @@ -107,7 +110,7 @@ class ParamGenerator { tempconf->nbits = 8; tempconf->metric_type = kn::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; @@ -139,6 +142,16 @@ 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 kn::cloner::CopyCpuToGpu(index_, device_id, kn::Config()); + } + return index_; + } + protected: std::string index_type; kn::Config conf; @@ -151,7 +164,8 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest, // std::make_tuple("IVFPQ", ParameterType::ivfpq), // std::make_tuple("GPUIVFPQ", ParameterType::ivfpq), std::make_tuple("IVFSQ", ParameterType::ivfsq), - std::make_tuple("GPUIVFSQ", ParameterType::ivfsq))); + std::make_tuple("GPUIVFSQ", ParameterType::ivfsq), + std::make_tuple("IVFSQHybrid", ParameterType::ivfsqhybrid))); void AssertAnns(const kn::DatasetPtr& result, const int& nq, const int& k) { @@ -191,11 +205,67 @@ TEST_P(IVFTest, ivf_basic) { index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + + auto new_idx = ChooseTodo(); + auto result = new_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); // 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); + } + + { + 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; } // @@ -248,7 +318,8 @@ TEST_P(IVFTest, ivf_serialize) { index_->set_index_model(model); index_->Add(base_dataset, conf); - auto result = index_->Search(query_dataset, conf); + auto new_idx = ChooseTodo(); + auto result = new_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); } @@ -272,7 +343,8 @@ TEST_P(IVFTest, ivf_serialize) { index_->Load(binaryset); EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + auto new_idx = ChooseTodo(); + auto result = new_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); } } @@ -288,7 +360,8 @@ TEST_P(IVFTest, clone_test) { index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + auto new_idx = ChooseTodo(); + auto result = new_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); // PrintResult(result, nq, k); @@ -301,31 +374,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(); - }, - kn::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({ @@ -369,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; @@ -385,7 +462,8 @@ TEST_P(IVFTest, seal_test) { index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + auto new_idx = ChooseTodo(); + auto result = new_idx->Search(query_dataset, conf); AssertAnns(result, nq, conf->k); auto cpu_idx = kn::cloner::CopyGpuToCpu(index_, kn::Config()); @@ -506,8 +584,8 @@ TEST_F(GPURESTEST, gpuivfsq) { 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 result = index_->Search(query_dataset, conf); + // AssertAnns(result, nq, k); auto cpu_idx = kn::cloner::CopyGpuToCpu(index_, kn::Config()); cpu_idx->Seal(); @@ -579,8 +657,8 @@ TEST_F(GPURESTEST, copyandsearch) { 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 result = index_->Search(query_dataset, conf); + // AssertAnns(result, nq, k); auto cpu_idx = kn::cloner::CopyGpuToCpu(index_, kn::Config()); cpu_idx->Seal(); diff --git a/cpp/src/server/DBWrapper.cpp b/cpp/src/server/DBWrapper.cpp index 306863f8ae..bb3bd012ab 100644 --- a/cpp/src/server/DBWrapper.cpp +++ b/cpp/src/server/DBWrapper.cpp @@ -24,6 +24,7 @@ #include #include +#include #include #include diff --git a/cpp/src/utils/ValidationUtil.cpp b/cpp/src/utils/ValidationUtil.cpp index 4345ebb704..b982a31f5e 100644 --- a/cpp/src/utils/ValidationUtil.cpp +++ b/cpp/src/utils/ValidationUtil.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include diff --git a/cpp/src/wrapper/VecImpl.cpp b/cpp/src/wrapper/VecImpl.cpp index a6023052db..1ed20c8029 100644 --- a/cpp/src/wrapper/VecImpl.cpp +++ b/cpp/src/wrapper/VecImpl.cpp @@ -20,6 +20,7 @@ #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #include "utils/Log.h" @@ -239,7 +240,7 @@ IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, co 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); @@ -264,5 +265,75 @@ IVFMixIndex::Load(const 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(const 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()); + } + return Status::OK(); +} + +Status +IVFHybridIndex::UnsetQuantizer() { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + new_idx->UnsetQuantizer(); + } 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(); +} + +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 } // namespace milvus diff --git a/cpp/src/wrapper/VecImpl.h b/cpp/src/wrapper/VecImpl.h index 2be822084f..f35a6ac4cd 100644 --- a/cpp/src/wrapper/VecImpl.h +++ b/cpp/src/wrapper/VecImpl.h @@ -91,6 +91,21 @@ class IVFMixIndex : public VecIndexImpl { Load(const knowhere::BinarySet& index_binary) override; }; +class IVFHybridIndex : public IVFMixIndex { + public: + knowhere::QuantizerPtr + LoadQuantizer(const Config& conf) override; + + Status + SetQuantizer(const knowhere::QuantizerPtr& q) override; + + Status + UnsetQuantizer() override; + + Status + LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override; +}; + class BFIndex : public VecIndexImpl { public: explicit BFIndex(std::shared_ptr index) diff --git a/cpp/src/wrapper/VecIndex.cpp b/cpp/src/wrapper/VecIndex.cpp index 90a58f588c..041b3a8df5 100644 --- a/cpp/src/wrapper/VecIndex.cpp +++ b/cpp/src/wrapper/VecIndex.cpp @@ -25,6 +25,7 @@ #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/IndexKDT.h" #include "knowhere/index/vector_index/IndexNSG.h" #include "utils/Log.h" @@ -142,7 +143,11 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { index = std::make_shared(gpu_device); break; } - case IndexType::NSG_MIX: { + 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 63d4ef6903..34d3d2f761 100644 --- a/cpp/src/wrapper/VecIndex.h +++ b/cpp/src/wrapper/VecIndex.h @@ -23,6 +23,7 @@ #include "cache/DataObj.h" #include "knowhere/common/BinarySet.h" #include "knowhere/common/Config.h" +#include "knowhere/index/vector_index/Quantizer.h" #include "utils/Status.h" namespace milvus { @@ -42,6 +43,7 @@ enum class IndexType { FAISS_IVFSQ8_MIX, FAISS_IVFSQ8_CPU, FAISS_IVFSQ8_GPU, + FAISS_IVFSQ8_HYBRID, // only support build on gpu. NSG_MIX, }; @@ -90,6 +92,29 @@ class VecIndex : public cache::DataObj { virtual Status Load(const knowhere::BinarySet& index_binary) = 0; + + // TODO(linxj): refactor later + //////////////// + virtual knowhere::QuantizerPtr + LoadQuantizer(const Config& conf) { + return nullptr; + } + + virtual Status + LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { + return Status::OK(); + } + + virtual Status + SetQuantizer(const knowhere::QuantizerPtr& q) { + return Status::OK(); + } + + virtual Status + UnsetQuantizer() { + return Status::OK(); + } + //////////////// }; extern Status