From 4b91faa50c68873c4721f8ba3f4a4c4b8eb437ce Mon Sep 17 00:00:00 2001 From: "xiaojun.lin" Date: Thu, 19 Sep 2019 21:21:38 +0800 Subject: [PATCH] MS-565 split class version 1 Former-commit-id: 1a60045de13e055d4b0a7e9da389b7e377c141af --- cpp/src/core/knowhere/CMakeLists.txt | 12 +- .../knowhere/adapter/SptagAdapter.cpp | 2 +- .../index/vector_index/FaissBaseIndex.cpp | 80 +++++++++ .../index/vector_index/FaissBaseIndex.h | 52 ++++++ .../index/vector_index/IndexGPUIVF.cpp | 83 +--------- .../knowhere/index/vector_index/IndexGPUIVF.h | 110 ++++++------- .../index/vector_index/IndexGPUIVFPQ.cpp | 68 ++++++++ .../index/vector_index/IndexGPUIVFPQ.h | 47 ++++++ .../index/vector_index/IndexGPUIVFSQ.cpp | 74 +++++++++ .../index/vector_index/IndexGPUIVFSQ.h | 43 +++++ .../index/vector_index/IndexIDMAP.cpp | 1 + .../knowhere/index/vector_index/IndexIDMAP.h | 6 +- .../knowhere/index/vector_index/IndexIVF.cpp | 153 +----------------- .../knowhere/index/vector_index/IndexIVF.h | 148 +++++++---------- .../index/vector_index/IndexIVFPQ.cpp | 60 +++++++ .../knowhere/index/vector_index/IndexIVFPQ.h | 47 ++++++ .../index/vector_index/IndexIVFSQ.cpp | 70 ++++++++ .../knowhere/index/vector_index/IndexIVFSQ.h | 45 ++++++ .../knowhere/index/vector_index/IndexKDT.cpp | 4 +- .../{utils => helpers}/Cloner.cpp | 4 +- .../vector_index/{utils => helpers}/Cloner.h | 0 .../{utils => helpers}/Definitions.h | 0 .../FaissGpuResourceMgr.cpp | 0 .../{utils => helpers}/FaissGpuResourceMgr.h | 0 .../index/vector_index/helpers/FaissIO.cpp | 65 ++++++++ .../index/vector_index/helpers/FaissIO.h | 48 ++++++ .../{utils => helpers}/KDTParameterMgr.cpp | 0 .../{utils => helpers}/KDTParameterMgr.h | 0 .../knowhere/index/vector_index/nsg/NSGIO.h | 1 + cpp/src/wrapper/KnowhereResource.cpp | 2 +- cpp/src/wrapper/vec_impl.cpp | 2 +- cpp/src/wrapper/vec_index.cpp | 6 +- cpp/unittest/db/utils.cpp | 4 +- cpp/unittest/wrapper/wrapper_test.cpp | 3 +- 34 files changed, 845 insertions(+), 395 deletions(-) create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.h create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.h create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.h rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/Cloner.cpp (92%) rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/Cloner.h (100%) rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/Definitions.h (100%) rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/FaissGpuResourceMgr.cpp (100%) rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/FaissGpuResourceMgr.h (100%) create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp create mode 100644 cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.h rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/KDTParameterMgr.cpp (100%) rename cpp/src/core/knowhere/knowhere/index/vector_index/{utils => helpers}/KDTParameterMgr.h (100%) diff --git a/cpp/src/core/knowhere/CMakeLists.txt b/cpp/src/core/knowhere/CMakeLists.txt index 728e198f..76750f4f 100644 --- a/cpp/src/core/knowhere/CMakeLists.txt +++ b/cpp/src/core/knowhere/CMakeLists.txt @@ -42,13 +42,19 @@ set(index_srcs knowhere/index/vector_index/IndexIDMAP.cpp knowhere/index/vector_index/IndexIVF.cpp knowhere/index/vector_index/IndexGPUIVF.cpp - knowhere/index/vector_index/utils/KDTParameterMgr.cpp + knowhere/index/vector_index/helpers/KDTParameterMgr.cpp knowhere/index/vector_index/IndexNSG.cpp knowhere/index/vector_index/nsg/NSG.cpp knowhere/index/vector_index/nsg/NSGIO.cpp knowhere/index/vector_index/nsg/NSGHelper.cpp - knowhere/index/vector_index/utils/Cloner.cpp - knowhere/index/vector_index/utils/FaissGpuResourceMgr.cpp + knowhere/index/vector_index/helpers/Cloner.cpp + knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp + knowhere/index/vector_index/IndexIVFSQ.cpp + knowhere/index/vector_index/IndexGPUIVFSQ.cpp + knowhere/index/vector_index/IndexIVFPQ.cpp + knowhere/index/vector_index/IndexGPUIVFPQ.cpp + knowhere/index/vector_index/FaissBaseIndex.cpp + knowhere/index/vector_index/helpers/FaissIO.cpp ) set(depend_libs diff --git a/cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp b/cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp index 3c5d466d..ee5a4a01 100644 --- a/cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp +++ b/cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -#include "knowhere/index/vector_index/utils/Definitions.h" +#include "knowhere/index/vector_index/helpers/Definitions.h" #include "SptagAdapter.h" #include "Structure.h" diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp new file mode 100644 index 00000000..a3ab7956 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp @@ -0,0 +1,80 @@ +// 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 +#include + +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/helpers/FaissIO.h" +#include "FaissBaseIndex.h" + + +namespace zilliz { +namespace knowhere { + +FaissBaseIndex::FaissBaseIndex(std::shared_ptr index) : index_(std::move(index)) {} + +BinarySet FaissBaseIndex::SerializeImpl() { + try { + faiss::Index *index = index_.get(); + + SealImpl(); + + MemoryIOWriter writer; + faiss::write_index(index, &writer); + auto data = std::make_shared(); + data.reset(writer.data_); + + BinarySet res_set; + // TODO(linxj): use virtual func Name() instead of raw string. + res_set.Append("IVF", data, writer.rp); + return res_set; + } catch (std::exception &e) { + KNOWHERE_THROW_MSG(e.what()); + } +} + +void FaissBaseIndex::LoadImpl(const BinarySet &index_binary) { + auto binary = index_binary.GetByName("IVF"); + + MemoryIOReader reader; + reader.total = binary->size; + reader.data_ = binary->data.get(); + + faiss::Index *index = faiss::read_index(&reader); + + index_.reset(index); +} + +void FaissBaseIndex::SealImpl() { +// TODO(linxj): enable +//#ifdef ZILLIZ_FAISS + faiss::Index *index = index_.get(); + auto idx = dynamic_cast(index); + if (idx != nullptr) { + idx->to_readonly(); + } + //else { + // KNOHWERE_ERROR_MSG("Seal failed"); + //} +//#endif +} + +} // knowhere +} // zilliz + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.h b/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.h new file mode 100644 index 00000000..29edbbd6 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/FaissBaseIndex.h @@ -0,0 +1,52 @@ +// 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 "knowhere/common/BinarySet.h" + + +namespace zilliz { +namespace knowhere { + +class FaissBaseIndex { + protected: + explicit FaissBaseIndex(std::shared_ptr index); + + virtual BinarySet + SerializeImpl(); + + virtual void + LoadImpl(const BinarySet &index_binary); + + virtual void + SealImpl(); + + protected: + std::shared_ptr index_ = nullptr; +}; + +} // knowhere +} // zilliz + + + 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 df260831..a77702ac 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp @@ -20,18 +20,17 @@ #include #include #include -#include #include #include #include #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/utils/Cloner.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" #include "knowhere/adapter/VectorAdapter.h" #include "IndexGPUIVF.h" +#include "knowhere/index/vector_index/helpers/FaissIO.h" -#include namespace zilliz { namespace knowhere { @@ -187,84 +186,6 @@ void GPUIVF::Add(const DatasetPtr &dataset, const Config &config) { } } -IndexModelPtr GPUIVFPQ::Train(const DatasetPtr &dataset, const Config &config) { - auto nlist = config["nlist"].as(); - auto M = config["M"].as(); // number of subquantizers(subvectors) - auto nbits = config["nbits"].as();// number of bit per subvector index - auto gpu_num = config.get_with_default("gpu_id", gpu_id_); - auto metric_type = config["metric_type"].as_string() == "L2" ? - faiss::METRIC_L2 : faiss::METRIC_L2; // IP not support. - - GETTENSOR(dataset) - - // TODO(linxj): set device here. - // TODO(linxj): set gpu resource here. - faiss::gpu::StandardGpuResources res; - faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, nlist, M, nbits, metric_type); - 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); -} - -std::shared_ptr GPUIVFPQ::GenParams(const Config &config) { - auto params = std::make_shared(); - params->nprobe = config.get_with_default("nprobe", size_t(1)); - //params->scan_table_threshold = 0; - //params->polysemous_ht = 0; - //params->max_codes = 0; - - return params; -} - -VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config &config) { - KNOWHERE_THROW_MSG("not support yet"); -} - -IndexModelPtr GPUIVFSQ::Train(const DatasetPtr &dataset, const Config &config) { - auto nlist = config["nlist"].as(); - auto nbits = config["nbits"].as(); // TODO(linxj): gpu only support SQ4 SQ8 SQ16 - gpu_id_ = config.get_with_default("gpu_id", gpu_id_); - auto metric_type = config["metric_type"].as_string() == "L2" ? - faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; - - GETTENSOR(dataset) - - std::stringstream index_type; - index_type << "IVF" << nlist << "," << "SQ" << nbits; - auto build_index = faiss::index_factory(dim, index_type.str().c_str(), 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 IVFSQ can't get gpu resource"); - } -} - -VectorIndexPtr GPUIVFSQ::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); - - std::shared_ptr new_index; - new_index.reset(host_index); - return std::make_shared(new_index); -} - void GPUIndex::SetGpuDevice(const int &gpu_id) { gpu_id_ = gpu_id; } 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 f972799a..cd04fe4c 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h @@ -20,79 +20,79 @@ #include "IndexIVF.h" -#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" namespace zilliz { namespace knowhere { class GPUIndex { - public: - explicit GPUIndex(const int &device_id) : gpu_id_(device_id) {} - GPUIndex(const int& device_id, const ResPtr& resource): gpu_id_(device_id), res_(resource){} +public: + explicit GPUIndex(const int &device_id) : gpu_id_(device_id) {} - virtual VectorIndexPtr CopyGpuToCpu(const Config &config) = 0; - virtual VectorIndexPtr CopyGpuToGpu(const int64_t &device_id, const Config &config) = 0; + GPUIndex(const int& device_id, const ResPtr& resource): gpu_id_(device_id), res_(resource) {} - void SetGpuDevice(const int &gpu_id); - const int64_t &GetGpuDevice(); + virtual VectorIndexPtr + CopyGpuToCpu(const Config &config) = 0; - protected: - int64_t gpu_id_; - ResWPtr res_; + virtual VectorIndexPtr + CopyGpuToGpu(const int64_t &device_id, const Config &config) = 0; + + void + SetGpuDevice(const int &gpu_id); + + const int64_t & + GetGpuDevice(); + +protected: + int64_t gpu_id_; + ResWPtr res_; }; class GPUIVF : public IVF, public GPUIndex { - public: - explicit GPUIVF(const int &device_id) : IVF(), GPUIndex(device_id) {} - explicit GPUIVF(std::shared_ptr index, const int64_t &device_id, ResPtr &resource) - : IVF(std::move(index)), GPUIndex(device_id, resource) {}; - IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; - void Add(const DatasetPtr &dataset, const Config &config) override; - void set_index_model(IndexModelPtr model) override; - //DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override; - VectorIndexPtr CopyGpuToCpu(const Config &config) override; - VectorIndexPtr CopyGpuToGpu(const int64_t &device_id, const Config &config) override; - 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; - BinarySet SerializeImpl() override; - void LoadImpl(const BinarySet &index_binary) override; -}; +public: + explicit GPUIVF(const int &device_id) : IVF(), GPUIndex(device_id) {} -class GPUIVFSQ : public GPUIVF { - public: - explicit GPUIVFSQ(const int &device_id) : GPUIVF(device_id) {} - explicit GPUIVFSQ(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; + explicit GPUIVF(std::shared_ptr index, const int64_t &device_id, ResPtr &resource) + : IVF(std::move(index)), GPUIndex(device_id, resource) {}; - public: - VectorIndexPtr CopyGpuToCpu(const Config &config) override; -}; + IndexModelPtr + Train(const DatasetPtr &dataset, const Config &config) override; -class GPUIVFPQ : public GPUIVF { - public: - explicit GPUIVFPQ(const int &device_id) : GPUIVF(device_id) {} - IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; + void + Add(const DatasetPtr &dataset, const Config &config) override; - public: - VectorIndexPtr CopyGpuToCpu(const Config &config) override; + void + set_index_model(IndexModelPtr model) override; - protected: - // TODO(linxj): remove GenParams. - std::shared_ptr GenParams(const Config &config) override; -}; + //DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override; + VectorIndexPtr + CopyGpuToCpu(const Config &config) override; + VectorIndexPtr + CopyGpuToGpu(const int64_t &device_id, const Config &config) override; + + 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; + + BinarySet + SerializeImpl() override; + + void + LoadImpl(const BinarySet &index_binary) override; +}; } // namespace knowhere } // namespace zilliz diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp new file mode 100644 index 00000000..437d8f89 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp @@ -0,0 +1,68 @@ +// 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 +#include +#include + +#include "IndexGPUIVFPQ.h" +#include "knowhere/common/Exception.h" +#include "knowhere/adapter/VectorAdapter.h" + + +namespace zilliz { +namespace knowhere { + +IndexModelPtr GPUIVFPQ::Train(const DatasetPtr &dataset, const Config &config) { + auto nlist = config["nlist"].as(); + auto M = config["M"].as(); // number of subquantizers(subvectors) + auto nbits = config["nbits"].as();// number of bit per subvector index + auto gpu_num = config.get_with_default("gpu_id", gpu_id_); + auto metric_type = config["metric_type"].as_string() == "L2" ? + faiss::METRIC_L2 : faiss::METRIC_L2; // IP not support. + + GETTENSOR(dataset) + + // TODO(linxj): set device here. + // TODO(linxj): set gpu resource here. + faiss::gpu::StandardGpuResources res; + faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, nlist, M, nbits, metric_type); + 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); +} + +std::shared_ptr GPUIVFPQ::GenParams(const Config &config) { + auto params = std::make_shared(); + params->nprobe = config.get_with_default("nprobe", size_t(1)); + //params->scan_table_threshold = 0; + //params->polysemous_ht = 0; + //params->max_codes = 0; + + return params; +} + +VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config &config) { + KNOWHERE_THROW_MSG("not support yet"); +} + +} // knowhere +} // zilliz diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h new file mode 100644 index 00000000..c407ee0c --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h @@ -0,0 +1,47 @@ +// 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 "IndexGPUIVF.h" + +namespace zilliz { +namespace knowhere { + +class GPUIVFPQ : public GPUIVF { +public: + explicit GPUIVFPQ(const int &device_id) : GPUIVF(device_id) {} + + IndexModelPtr + Train(const DatasetPtr &dataset, const Config &config) override; + +public: + VectorIndexPtr + CopyGpuToCpu(const Config &config) override; + +protected: + // TODO(linxj): remove GenParams. + std::shared_ptr + GenParams(const Config &config) override; +}; + +} // knowhere +} // zilliz + + + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp new file mode 100644 index 00000000..2504dd45 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp @@ -0,0 +1,74 @@ +// 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 + +#include "knowhere/adapter/VectorAdapter.h" +#include "knowhere/common/Exception.h" +#include "IndexGPUIVFSQ.h" +#include "IndexIVFSQ.h" + + +namespace zilliz { +namespace knowhere { + + IndexModelPtr GPUIVFSQ::Train(const DatasetPtr &dataset, const Config &config) { + auto nlist = config["nlist"].as(); + auto nbits = config["nbits"].as(); // TODO(linxj): gpu only support SQ4 SQ8 SQ16 + gpu_id_ = config.get_with_default("gpu_id", gpu_id_); + auto metric_type = config["metric_type"].as_string() == "L2" ? + faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; + + GETTENSOR(dataset) + + std::stringstream index_type; + index_type << "IVF" << nlist << "," << "SQ" << nbits; + auto build_index = faiss::index_factory(dim, index_type.str().c_str(), 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 IVFSQ can't get gpu resource"); + } + } + + VectorIndexPtr GPUIVFSQ::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); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); + } + +} // knowhere +} // zilliz + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h new file mode 100644 index 00000000..67790796 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h @@ -0,0 +1,43 @@ +// 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 "IndexGPUIVF.h" + + +namespace zilliz { +namespace knowhere { + +class GPUIVFSQ : public GPUIVF { +public: + explicit GPUIVFSQ(const int &device_id) : GPUIVF(device_id) {} + + explicit GPUIVFSQ(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; + + VectorIndexPtr + CopyGpuToCpu(const Config &config) override; +}; + +} // knowhere +} // zilliz + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp index c81e69f7..a53a73da 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp @@ -25,6 +25,7 @@ #include "knowhere/common/Exception.h" #include "knowhere/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/helpers/FaissIO.h" #include "IndexIDMAP.h" diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.h index 16d5ece5..106759fa 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIDMAP.h @@ -25,10 +25,10 @@ namespace zilliz { namespace knowhere { -class IDMAP : public VectorIndex, public BasicIndex { +class IDMAP : public VectorIndex, public FaissBaseIndex { public: - IDMAP() : BasicIndex(nullptr) {}; - explicit IDMAP(std::shared_ptr index) : BasicIndex(std::move(index)) {}; + IDMAP() : FaissBaseIndex(nullptr) {}; + explicit IDMAP(std::shared_ptr index) : FaissBaseIndex(std::move(index)) {}; BinarySet Serialize() override; void Load(const BinarySet &index_binary) override; void Train(const Config &config); 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 aaaafdad..a7adfe92 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp @@ -238,7 +238,7 @@ void IVF::Seal() { } -IVFIndexModel::IVFIndexModel(std::shared_ptr index) : BasicIndex(std::move(index)) {} +IVFIndexModel::IVFIndexModel(std::shared_ptr index) : FaissBaseIndex(std::move(index)) {} BinarySet IVFIndexModel::Serialize() { if (!index_ || !index_->is_trained) { @@ -257,159 +257,8 @@ void IVFIndexModel::SealImpl() { // do nothing } -IndexModelPtr IVFSQ::Train(const DatasetPtr &dataset, const Config &config) { - auto nlist = config["nlist"].as(); - auto nbits = config["nbits"].as(); // TODO(linxj): only support SQ4 SQ6 SQ8 SQ16 - auto metric_type = config["metric_type"].as_string() == "L2" ? - faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; - - GETTENSOR(dataset) - - std::stringstream index_type; - index_type << "IVF" << nlist << "," << "SQ" << nbits; - auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type); - build_index->train(rows, (float *) p_data); - - std::shared_ptr ret_index; - ret_index.reset(build_index); - return std::make_shared(ret_index); -} - -VectorIndexPtr IVFSQ::Clone_impl(const std::shared_ptr &index) { - return std::make_shared(index); -} - -VectorIndexPtr IVFSQ::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; - - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get(), &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"); - } -} - -IndexModelPtr IVFPQ::Train(const DatasetPtr &dataset, const Config &config) { - auto nlist = config["nlist"].as(); - auto M = config["M"].as(); // number of subquantizers(subvector) - auto nbits = config["nbits"].as();// number of bit per subvector index - auto metric_type = config["metric_type"].as_string() == "L2" ? - faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; - - GETTENSOR(dataset) - - faiss::Index *coarse_quantizer = new faiss::IndexFlat(dim, metric_type); - auto index = std::make_shared(coarse_quantizer, dim, nlist, M, nbits); - index->train(rows, (float *) p_data); - - return std::make_shared(index); -} -std::shared_ptr IVFPQ::GenParams(const Config &config) { - auto params = std::make_shared(); - params->nprobe = config.get_with_default("nprobe", size_t(1)); - //params->scan_table_threshold = 0; - //params->polysemous_ht = 0; - //params->max_codes = 0; - - return params; -} - -VectorIndexPtr IVFPQ::Clone_impl(const std::shared_ptr &index) { - return std::make_shared(index); -} - -BasicIndex::BasicIndex(std::shared_ptr index) : index_(std::move(index)) {} - -BinarySet BasicIndex::SerializeImpl() { - try { - faiss::Index *index = index_.get(); - - SealImpl(); - - MemoryIOWriter writer; - faiss::write_index(index, &writer); - auto data = std::make_shared(); - data.reset(writer.data_); - - BinarySet res_set; - // TODO(linxj): use virtual func Name() instead of raw string. - res_set.Append("IVF", data, writer.rp); - return res_set; - } catch (std::exception &e) { - KNOWHERE_THROW_MSG(e.what()); - } -} - -void BasicIndex::LoadImpl(const BinarySet &index_binary) { - auto binary = index_binary.GetByName("IVF"); - - MemoryIOReader reader; - reader.total = binary->size; - reader.data_ = binary->data.get(); - faiss::Index *index = faiss::read_index(&reader); - - index_.reset(index); -} - -void BasicIndex::SealImpl() { -// TODO(linxj): enable -//#ifdef ZILLIZ_FAISS - faiss::Index *index = index_.get(); - auto idx = dynamic_cast(index); - if (idx != nullptr) { - idx->to_readonly(); - } - //else { - // KNOHWERE_ERROR_MSG("Seal failed"); - //} -//#endif -} - -// TODO(linxj): Get From Config File -static size_t magic_num = 2; -size_t MemoryIOWriter::operator()(const void *ptr, size_t size, size_t nitems) { - auto total_need = size * nitems + rp; - - if (!data_) { // data == nullptr - total = total_need * magic_num; - rp = size * nitems; - data_ = new uint8_t[total]; - memcpy((void *) (data_), ptr, rp); - } - - if (total_need > total) { - total = total_need * magic_num; - auto new_data = new uint8_t[total]; - memcpy((void *) new_data, (void *) data_, rp); - delete data_; - data_ = new_data; - - memcpy((void *) (data_ + rp), ptr, size * nitems); - rp = total_need; - } else { - memcpy((void *) (data_ + rp), ptr, size * nitems); - rp = total_need; - } - - return nitems; -} - -size_t MemoryIOReader::operator()(void *ptr, size_t size, size_t nitems) { - if (rp >= total) return 0; - size_t nremain = (total - rp) / size; - if (nremain < nitems) nitems = nremain; - memcpy(ptr, (void *) (data_ + rp), size * nitems); - rp += size * nitems; - return nitems; -} } diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.h index 79f41a02..72f98a5a 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.h @@ -21,136 +21,106 @@ #include #include -#include -#include -#include - #include "VectorIndex.h" +#include "FaissBaseIndex.h" +#include "faiss/IndexIVF.h" namespace zilliz { namespace knowhere { -class BasicIndex { - protected: - explicit BasicIndex(std::shared_ptr index); - virtual BinarySet SerializeImpl(); - virtual void LoadImpl(const BinarySet &index_binary); - virtual void - SealImpl(); - - protected: - std::shared_ptr index_ = nullptr; -}; - using Graph = std::vector>; -class IVF : public VectorIndex, protected BasicIndex { +class IVF : public VectorIndex, protected FaissBaseIndex { public: - IVF() : BasicIndex(nullptr) {}; - explicit IVF(std::shared_ptr index) : BasicIndex(std::move(index)) {} - VectorIndexPtr Clone() override;; - IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; - void set_index_model(IndexModelPtr model) override; - void Add(const DatasetPtr &dataset, const Config &config) override; - void AddWithoutIds(const DatasetPtr &dataset, const Config &config); - DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override; - void GenGraph(const int64_t &k, Graph &graph, const DatasetPtr &dataset, const Config &config); - BinarySet Serialize() override; - void Load(const BinarySet &index_binary) override; - int64_t Count() override; - int64_t Dimension() override; - - void - Seal() override; - - virtual VectorIndexPtr CopyCpuToGpu(const int64_t &device_id, const Config &config); + IVF() : FaissBaseIndex(nullptr) {}; + explicit IVF(std::shared_ptr index) : FaissBaseIndex(std::move(index)) {} - protected: - virtual std::shared_ptr GenParams(const Config &config); + VectorIndexPtr + Clone() override;; - virtual VectorIndexPtr Clone_impl(const std::shared_ptr &index); + IndexModelPtr + Train(const DatasetPtr &dataset, const Config &config) override; - virtual void search_impl(int64_t n, - const float *data, - int64_t k, - float *distances, - int64_t *labels, - const Config &cfg); + void + set_index_model(IndexModelPtr model) override; - protected: - std::mutex mutex_; -}; + void + Add(const DatasetPtr &dataset, const Config &config) override; -using IVFIndexPtr = std::shared_ptr; + void + AddWithoutIds(const DatasetPtr &dataset, const Config &config); -class IVFSQ : public IVF { - public: - explicit IVFSQ(std::shared_ptr index) : IVF(std::move(index)) {} - IVFSQ() = default; - IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; - VectorIndexPtr CopyCpuToGpu(const int64_t &device_id, const Config &config) override; - protected: - VectorIndexPtr Clone_impl(const std::shared_ptr &index) override; -}; + DatasetPtr + Search(const DatasetPtr &dataset, const Config &config) override; -class IVFPQ : public IVF { - public: - explicit IVFPQ(std::shared_ptr index) : IVF(std::move(index)) {} - IVFPQ() = default; - IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override; - protected: - std::shared_ptr GenParams(const Config &config) override; - VectorIndexPtr Clone_impl(const std::shared_ptr &index) override; -}; + void + GenGraph(const int64_t &k, Graph &graph, const DatasetPtr &dataset, const Config &config); + BinarySet + Serialize() override; -//class OPQIVFPQ : public IVFPQ { -// public: -// PreprocessorPtr BuildPreprocessor(const Dataset &dataset, const Config &config) override; -//}; + void + Load(const BinarySet &index_binary) override; + int64_t + Count() override; -class GPUIVF; + int64_t + Dimension() override; + void + Seal() override; -struct MemoryIOWriter : public faiss::IOWriter { - uint8_t *data_ = nullptr; - size_t total = 0; - size_t rp = 0; + virtual VectorIndexPtr + CopyCpuToGpu(const int64_t &device_id, const Config &config); - size_t operator()(const void *ptr, size_t size, size_t nitems) override; -}; + protected: + virtual std::shared_ptr + GenParams(const Config &config); -struct MemoryIOReader : public faiss::IOReader { - uint8_t *data_; - size_t rp = 0; - size_t total = 0; + virtual VectorIndexPtr + Clone_impl(const std::shared_ptr &index); - size_t operator()(void *ptr, size_t size, size_t nitems) override; + virtual void + search_impl(int64_t n, + const float *data, + int64_t k, + float *distances, + int64_t *labels, + const Config &cfg); + protected: + std::mutex mutex_; }; +using IVFIndexPtr = std::shared_ptr; -class IVFIndexModel : public IndexModel, public BasicIndex { +class GPUIVF; +class IVFIndexModel : public IndexModel, public FaissBaseIndex { friend IVF; friend GPUIVF; public: explicit IVFIndexModel(std::shared_ptr index); - IVFIndexModel() : BasicIndex(nullptr) {}; - BinarySet Serialize() override; + + IVFIndexModel() : FaissBaseIndex(nullptr) {}; + + BinarySet + Serialize() override; + + void + Load(const BinarySet &binary) override; + protected: - void SealImpl() override; - public: - void Load(const BinarySet &binary) override; + void + SealImpl() override; protected: std::mutex mutex_; }; - using IVFIndexModelPtr = std::shared_ptr; } diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp new file mode 100644 index 00000000..eca63a5f --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp @@ -0,0 +1,60 @@ +// 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 +#include + +#include "IndexIVFPQ.h" +#include "knowhere/common/Exception.h" +#include "knowhere/adapter/VectorAdapter.h" + +namespace zilliz { +namespace knowhere { + +IndexModelPtr IVFPQ::Train(const DatasetPtr &dataset, const Config &config) { + auto nlist = config["nlist"].as(); + auto M = config["M"].as(); // number of subquantizers(subvector) + auto nbits = config["nbits"].as();// number of bit per subvector index + auto metric_type = config["metric_type"].as_string() == "L2" ? + faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; + + GETTENSOR(dataset) + + faiss::Index *coarse_quantizer = new faiss::IndexFlat(dim, metric_type); + auto index = std::make_shared(coarse_quantizer, dim, nlist, M, nbits); + index->train(rows, (float *) p_data); + + return std::make_shared(index); +} + +std::shared_ptr IVFPQ::GenParams(const Config &config) { + auto params = std::make_shared(); + params->nprobe = config.get_with_default("nprobe", size_t(1)); + //params->scan_table_threshold = 0; + //params->polysemous_ht = 0; + //params->max_codes = 0; + + return params; +} + +VectorIndexPtr IVFPQ::Clone_impl(const std::shared_ptr &index) { + return std::make_shared(index); +} + +} // knowhere +} // zilliz diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.h new file mode 100644 index 00000000..427f76fe --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFPQ.h @@ -0,0 +1,47 @@ +// 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 "IndexIVF.h" + +namespace zilliz { +namespace knowhere { + +class IVFPQ : public IVF { +public: + explicit IVFPQ(std::shared_ptr index) : IVF(std::move(index)) {} + + IVFPQ() = default; + + IndexModelPtr + Train(const DatasetPtr &dataset, const Config &config) override; + +protected: + std::shared_ptr + GenParams(const Config &config) override; + + VectorIndexPtr + Clone_impl(const std::shared_ptr &index) override; +}; + +} // knowhere +} // zilliz + + + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp new file mode 100644 index 00000000..00f191ff --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp @@ -0,0 +1,70 @@ +// 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 + +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#include "knowhere/adapter/VectorAdapter.h" +#include "IndexIVFSQ.h" +#include "IndexGPUIVFSQ.h" + + +namespace zilliz { +namespace knowhere { + +IndexModelPtr IVFSQ::Train(const DatasetPtr &dataset, const Config &config) { + auto nlist = config["nlist"].as(); + auto nbits = config["nbits"].as(); // TODO(linxj): only support SQ4 SQ6 SQ8 SQ16 + auto metric_type = config["metric_type"].as_string() == "L2" ? + faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; + + GETTENSOR(dataset) + + std::stringstream index_type; + index_type << "IVF" << nlist << "," << "SQ" << nbits; + auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type); + build_index->train(rows, (float *) p_data); + + std::shared_ptr ret_index; + ret_index.reset(build_index); + return std::make_shared(ret_index); +} + +VectorIndexPtr IVFSQ::Clone_impl(const std::shared_ptr &index) { + return std::make_shared(index); +} + +VectorIndexPtr IVFSQ::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; + + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get(), &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"); + } +} + +} // knowhere +} // zilliz diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.h b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.h new file mode 100644 index 00000000..5ee771f9 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQ.h @@ -0,0 +1,45 @@ +// 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 "IndexIVF.h" + +namespace zilliz { +namespace knowhere { + +class IVFSQ : public IVF { +public: + explicit IVFSQ(std::shared_ptr index) : IVF(std::move(index)) {} + + IVFSQ() = default; + + IndexModelPtr + Train(const DatasetPtr &dataset, const Config &config) override; + + VectorIndexPtr + CopyCpuToGpu(const int64_t &device_id, const Config &config) override; + +protected: + VectorIndexPtr + Clone_impl(const std::shared_ptr &index) override; +}; + +} // knowhere +} // zilliz + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexKDT.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexKDT.cpp index 65084695..5e049b72 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/IndexKDT.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/IndexKDT.cpp @@ -25,9 +25,9 @@ #undef mkdir #include "IndexKDT.h" -#include "knowhere/index/vector_index/utils/Definitions.h" +#include "knowhere/index/vector_index/helpers/Definitions.h" //#include "knowhere/index/preprocessor/normalize.h" -#include "knowhere/index/vector_index/utils/KDTParameterMgr.h" +#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h" #include "knowhere/adapter/SptagAdapter.h" #include "knowhere/common/Exception.h" diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/Cloner.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp similarity index 92% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/Cloner.cpp rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index 02e9c292..d8ccdd10 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/Cloner.cpp +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -18,6 +18,8 @@ #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/index/vector_index/IndexIDMAP.h" #include "Cloner.h" @@ -43,7 +45,7 @@ VectorIndexPtr CopyCpuToGpu(const VectorIndexPtr &index, const int64_t &device_i 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 tranfer to gpu"); + KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu"); } 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/cpp/src/core/knowhere/knowhere/index/vector_index/utils/Cloner.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.h similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/Cloner.h rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.h diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/Definitions.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Definitions.h similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/Definitions.h rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Definitions.h diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/FaissGpuResourceMgr.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/FaissGpuResourceMgr.cpp rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/FaissGpuResourceMgr.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/FaissGpuResourceMgr.h rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp new file mode 100644 index 00000000..48167fe9 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp @@ -0,0 +1,65 @@ +// 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 + +#include "FaissIO.h" + +namespace zilliz { +namespace knowhere { + +// TODO(linxj): Get From Config File +static size_t magic_num = 2; +size_t MemoryIOWriter::operator()(const void *ptr, size_t size, size_t nitems) { + auto total_need = size * nitems + rp; + + if (!data_) { // data == nullptr + total = total_need * magic_num; + rp = size * nitems; + data_ = new uint8_t[total]; + memcpy((void *) (data_), ptr, rp); + } + + if (total_need > total) { + total = total_need * magic_num; + auto new_data = new uint8_t[total]; + memcpy((void *) new_data, (void *) data_, rp); + delete data_; + data_ = new_data; + + memcpy((void *) (data_ + rp), ptr, size * nitems); + rp = total_need; + } else { + memcpy((void *) (data_ + rp), ptr, size * nitems); + rp = total_need; + } + + return nitems; +} + +size_t MemoryIOReader::operator()(void *ptr, size_t size, size_t nitems) { + if (rp >= total) return 0; + size_t nremain = (total - rp) / size; + if (nremain < nitems) nitems = nremain; + memcpy(ptr, (void *) (data_ + rp), size * nitems); + rp += size * nitems; + return nitems; +} + +} // knowhere +} // zilliz diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.h new file mode 100644 index 00000000..cc089ec3 --- /dev/null +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissIO.h @@ -0,0 +1,48 @@ +// 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 + +namespace zilliz { +namespace knowhere { + +struct MemoryIOWriter : public faiss::IOWriter { + uint8_t *data_ = nullptr; + size_t total = 0; + size_t rp = 0; + + size_t + operator()(const void *ptr, size_t size, size_t nitems) override; +}; + +struct MemoryIOReader : public faiss::IOReader { + uint8_t *data_; + size_t rp = 0; + size_t total = 0; + + size_t + operator()(void *ptr, size_t size, size_t nitems) override; +}; + +} // knowhere +} // zilliz + + + diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/KDTParameterMgr.cpp b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/KDTParameterMgr.cpp rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/utils/KDTParameterMgr.h b/cpp/src/core/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h similarity index 100% rename from cpp/src/core/knowhere/knowhere/index/vector_index/utils/KDTParameterMgr.h rename to cpp/src/core/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h diff --git a/cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSGIO.h b/cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSGIO.h index e34eb992..3d6786c6 100644 --- a/cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSGIO.h +++ b/cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSGIO.h @@ -18,6 +18,7 @@ #pragma once +#include "knowhere/index/vector_index/helpers/FaissIO.h" #include "NSG.h" #include "knowhere/index/vector_index/IndexIVF.h" diff --git a/cpp/src/wrapper/KnowhereResource.cpp b/cpp/src/wrapper/KnowhereResource.cpp index 02633447..007858de 100644 --- a/cpp/src/wrapper/KnowhereResource.cpp +++ b/cpp/src/wrapper/KnowhereResource.cpp @@ -17,7 +17,7 @@ #include "KnowhereResource.h" -#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #include "server/ServerConfig.h" #include diff --git a/cpp/src/wrapper/vec_impl.cpp b/cpp/src/wrapper/vec_impl.cpp index cb26b422..accb612c 100644 --- a/cpp/src/wrapper/vec_impl.cpp +++ b/cpp/src/wrapper/vec_impl.cpp @@ -20,7 +20,7 @@ #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/utils/Cloner.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" #include "vec_impl.h" #include "data_transfer.h" diff --git a/cpp/src/wrapper/vec_index.cpp b/cpp/src/wrapper/vec_index.cpp index 213dc24f..7ed176d9 100644 --- a/cpp/src/wrapper/vec_index.cpp +++ b/cpp/src/wrapper/vec_index.cpp @@ -16,8 +16,12 @@ // under the License. #include "knowhere/index/vector_index/IndexIVF.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexKDT.h" #include "knowhere/index/vector_index/IndexNSG.h" #include "knowhere/common/Exception.h" diff --git a/cpp/unittest/db/utils.cpp b/cpp/unittest/db/utils.cpp index 19277d1e..e9950a85 100644 --- a/cpp/unittest/db/utils.cpp +++ b/cpp/unittest/db/utils.cpp @@ -25,9 +25,7 @@ #include "cache/CpuCacheMgr.h" #include "db/DBFactory.h" #include "db/Options.h" -#include "server/ServerConfig.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" INITIALIZE_EASYLOGGINGPP diff --git a/cpp/unittest/wrapper/wrapper_test.cpp b/cpp/unittest/wrapper/wrapper_test.cpp index fbac6b13..83896fed 100644 --- a/cpp/unittest/wrapper/wrapper_test.cpp +++ b/cpp/unittest/wrapper/wrapper_test.cpp @@ -20,8 +20,7 @@ #include "utils/easylogging++.h" #include "src/wrapper/vec_index.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #include "utils.h" -- GitLab