admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-29 23:15:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'branch-0.5.1' into 'branch-0.5.0'

Browse Source 
MS-627 Integrate new index: IVFSQHybrid

See merge request megasearch/milvus!683

Former-commit-id: 7f176330c1b32e4fa161e5c11a747266017c89b2
This commit is contained in:
jinhai 2019-10-10 15:35:49 +08:00
commit 88906b1aec
20 changed files with 587 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
cpp/CHANGELOG.md
View File

@ -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

6
cpp/src/core/cmake/ThirdPartyPackagesCore.cmake
View File

@ -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}")

1
cpp/src/core/knowhere/CMakeLists.txt
View File

@ -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

32
cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
View File

@ -126,32 +126,22 @@ GPUIVF::LoadImpl(const BinarySet& index_binary) {
}
}
IVFIndexPtr
GPUIVF::Copy_index_gpu_to_cpu() {
std::lock_guard<std::mutex> lk(mutex_);
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IVF>(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<std::mutex> lk(mutex_);
// TODO(linxj): gpu index support GenParams
if (auto device_index = std::static_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {
if (auto device_index = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {
auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(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<std::mutex> 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<faiss::gpu::GpuIndexIVF>(index_)) {
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IVF>(new_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IVF>(new_index);
} else {
return std::make_shared<IVF>(index_);
}
}
VectorIndexPtr

4
cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.h
View File

@ -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;

2
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVF.cpp
View File

@ -146,7 +146,7 @@ IVF::set_index_model(IndexModelPtr model) {
std::shared_ptr<faiss::IVFSearchParameters>
IVF::GenParams(const Config& config) {
auto params = std::make_shared<faiss::IVFPQSearchParameters>();
auto params = std::make_shared<faiss::IVFSearchParameters>();
auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(config);
params->nprobe = search_cfg->nprobe;

208
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp Normal file
View File

@ -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<IVFSQCfg>(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<faiss::Index> host_index = nullptr;
host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));
delete device_index;
delete build_index;
return std::make_shared<IVFIndexModel>(host_index);
} else {
KNOWHERE_THROW_MSG("Build IVFSQHybrid can't get gpu resource");
}
}
VectorIndexPtr
IVFSQHybrid::CopyGpuToCpu(const Config& config) {
std::lock_guard<std::mutex> lk(mutex_);
if (auto device_idx = std::dynamic_pointer_cast<faiss::IndexIVF>(index_)) {
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IVFSQHybrid>(new_index);
} else {
// TODO(linxj): why? jinhai
return std::make_shared<IVFSQHybrid>(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<faiss::Index> device_index;
device_index.reset(gpu_index);
return std::make_shared<IVFSQHybrid>(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<QuantizerCfg>(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<FaissIVFQuantizer>();
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<FaissIVFQuantizer>(q);
if (ivf_quantizer == nullptr) {
KNOWHERE_THROW_MSG("Quantizer type error");
}
faiss::IndexIVF* ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
faiss::gpu::GpuIndexFlat* is_gpu_flat_index = dynamic_cast<faiss::gpu::GpuIndexFlat*>(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<faiss::IndexIVF*>(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<QuantizerCfg>(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<FaissIVFQuantizer>(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

82
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h Normal file
View File

@ -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 <faiss/index_io.h>
#include <memory>
#include "IndexGPUIVFSQ.h"
#include "Quantizer.h"
namespace knowhere {
struct FaissIVFQuantizer : public Quantizer {
faiss::gpu::GpuIndexFlat* quantizer = nullptr;
};
using FaissIVFQuantizerPtr = std::shared_ptr<FaissIVFQuantizer>;
class IVFSQHybrid : public GPUIVFSQ {
public:
explicit IVFSQHybrid(const int& device_id) : GPUIVFSQ(device_id) {
gpu_mode = false;
}
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {
index_ = index;
gpu_mode = false;
}
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> 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

35
cpp/src/core/knowhere/knowhere/index/vector_index/Quantizer.h Normal file
View File

@ -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 <memory>
#include "knowhere/common/Config.h"
namespace knowhere {
struct Quantizer {
virtual ~Quantizer() = default;
};
using QuantizerPtr = std::shared_ptr<Quantizer>;
struct QuantizerCfg : Cfg {
uint64_t mode = -1; // 0: all data, 1: copy quantizer, 2: copy data
};
using QuantizerConfig = std::shared_ptr<QuantizerCfg>;
} // namespace knowhere

5
cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
View File

@ -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<IVFSQHybrid>(index)) {
return device_index->CopyCpuToGpu(device_id, config);
}
if (auto device_index = std::dynamic_pointer_cast<GPUIndex>(index)) {
return device_index->CopyGpuToGpu(device_id, config);
}

5
cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
View File

@ -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) {

2
cpp/src/core/thirdparty/versions.txt vendored
View File

@ -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
FAISS_VERSION=branch-0.2.1

1
cpp/src/core/unittest/CMakeLists.txt
View File

@ -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
)

142
cpp/src/core/unittest/test_ivf.cpp
View File

@ -24,7 +24,6 @@
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexIVFFlat.h>
#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<kn::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<kn::GPUIVFSQ>(device_id);
} else if (type == "IVFSQHybrid") {
return std::make_shared<kn::IVFSQHybrid>(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<kn::IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
@ -139,6 +142,16 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
kn::FaissGpuResourceMgr::GetInstance().Free();
}
kn::VectorIndexPtr
ChooseTodo() {
std::vector<std::string> 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<kn::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_1_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<kn::QuantizerCfg>();
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<kn::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_2_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<kn::QuantizerCfg>();
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<std::string> 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<std::string> 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<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ"};
std::vector<std::string> 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<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ"};
std::vector<std::string> 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();

1
cpp/src/server/DBWrapper.cpp
View File

@ -24,6 +24,7 @@
#include <faiss/utils.h>
#include <omp.h>
#include <cmath>
#include <string>
#include <vector>

1
cpp/src/utils/ValidationUtil.cpp
View File

@ -22,6 +22,7 @@
#include <arpa/inet.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cmath>
#include <regex>
#include <string>

73
cpp/src/wrapper/VecImpl.cpp
View File

@ -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<knowhere::GPUIVF>(index_)) {
if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIndex>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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

15
cpp/src/wrapper/VecImpl.h
View File

@ -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<knowhere::VectorIndex> index)

7
cpp/src/wrapper/VecIndex.cpp
View File

@ -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<knowhere::GPUIVFSQ>(gpu_device);
break;
}
case IndexType::NSG_MIX: {
case IndexType::FAISS_IVFSQ8_HYBRID: {
index = std::make_shared<knowhere::IVFSQHybrid>(gpu_device);
break;
}
case IndexType::NSG_MIX: { // TODO(linxj): bug.
index = std::make_shared<knowhere::NSG>(gpu_device);
break;
}

25
cpp/src/wrapper/VecIndex.h
View File

@ -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