milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enable hybrid IVFSQ8 

Former-commit-id: 9bbf835e63818ef6cc5f792d0ab11247794d0c3a
pull/191/head
xiaojun.lin 2019-10-10 14:28:00 +08:00
parent d4babbcae1
commit 77b28a3981
8 changed files with 172 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
View File

@ -63,7 +63,7 @@ VectorIndexPtr
IVFSQHybrid::CopyGpuToCpu(const Config &config) {
std::lock_guard<std::mutex> lk(mutex_);
if (auto device_idx = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {
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);
@ -113,6 +113,7 @@ IVFSQHybrid::search_impl(int64_t n,
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);
}
}
@ -121,7 +122,9 @@ QuantizerPtr
IVFSQHybrid::LoadQuantizer(const Config &conf) {
auto quantizer_conf = std::dynamic_pointer_cast<QuantizerCfg>(conf);
if (quantizer_conf != nullptr) {
quantizer_conf->CheckValid(); // throw exception
if(quantizer_conf->mode != 1) {
KNOWHERE_THROW_MSG("mode only support 1 in this func");
}
}
gpu_id_ = quantizer_conf->gpu_id;
@ -133,13 +136,14 @@ IVFSQHybrid::LoadQuantizer(const Config &conf) {
auto index_composition = new faiss::IndexComposition;
index_composition->index = index_.get();
index_composition->quantizer = nullptr;
index_composition->mode = quantizer_conf->mode; // 1 or 2
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;
std::shared_ptr<FaissIVFQuantizer> q;
q->quantizer = index_composition;
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");
@ -153,15 +157,13 @@ IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) {
KNOWHERE_THROW_MSG("Quantizer type error");
}
if (ivf_quantizer->quantizer->mode == 2) gpu_mode = true; // all in gpu
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->quantizer;
ivf_index->quantizer = ivf_quantizer->quantizer;
}
}
@ -175,5 +177,38 @@ IVFSQHybrid::UnsetQuantizer() {
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");
}
}
}
}

14
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h
View File

@ -29,15 +29,20 @@ namespace zilliz {
namespace knowhere {
struct FaissIVFQuantizer : public Quantizer {
faiss::IndexComposition *quantizer = nullptr;
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) {}
explicit IVFSQHybrid(const int &device_id) : GPUIVFSQ(device_id) {
gpu_mode = false;
}
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {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) {
@ -54,6 +59,9 @@ class IVFSQHybrid : public GPUIVFSQ {
void
UnsetQuantizer();
void
LoadData(const knowhere::QuantizerPtr &q, const Config& conf);
IndexModelPtr
Train(const DatasetPtr &dataset, const Config &config) override;

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

@ -100,6 +100,11 @@ public:
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)

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

@ -41,6 +41,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
)

130
cpp/src/core/test/test_ivf.cpp
View File

@ -65,6 +65,8 @@ IVFIndexPtr IndexFactory(const std::string &type) {
return std::make_shared<IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<GPUIVFSQ>(device_id);
} else if (type == "IVFSQHybrid") {
return std::make_shared<IVFSQHybrid>(device_id);
}
}
@ -72,6 +74,7 @@ enum class ParameterType {
ivf,
ivfpq,
ivfsq,
ivfsqhybrid,
nsg,
};
@ -105,7 +108,7 @@ class ParamGenerator {
tempconf->metric_type = METRICTYPE::L2;
return tempconf;
}
else if (type == ParameterType::ivfsq) {
else if (type == ParameterType::ivfsq || type == ParameterType::ivfsqhybrid) {
auto tempconf = std::make_shared<IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
@ -160,6 +163,7 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
// std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std::make_tuple("IVFSQ", ParameterType::ivfsq),
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)
std::make_tuple("IVFSQHybrid", ParameterType::ivfsqhybrid)
)
);
@ -210,36 +214,60 @@ TEST_P(IVFTest, ivf_basic) {
//PrintResult(result, nq, k);
}
//TEST_P(IVFTest, hybrid) {
// assert(!xb.empty());
//
// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
// index_->set_preprocessor(preprocessor);
//
// auto model = index_->Train(base_dataset, conf);
// index_->set_index_model(model);
// index_->Add(base_dataset, conf);
// EXPECT_EQ(index_->Count(), nb);
// EXPECT_EQ(index_->Dimension(), dim);
//
//// auto new_idx = ChooseTodo();
//// auto result = new_idx->Search(query_dataset, conf);
//// AssertAnns(result, nq, conf->k);
//
// auto iss_idx = std::make_shared<IVFSQHybrid>(device_id);
//
// auto binaryset = index_->Serialize();
// iss_idx->Load(binaryset);
//
// auto quantizer_conf = std::make_shared<QuantizerCfg>();
// quantizer_conf->mode = 1;
// quantizer_conf->gpu_id = 1;
// auto q = iss_idx->LoadQuantizer(quantizer_conf);
// iss_idx->SetQuantizer(q);
// auto result = iss_idx->Search(query_dataset, conf);
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);
// //PrintResult(result, nq, k);
//}
{
auto hybrid_1_idx = std::make_shared<IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_1_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<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<IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_2_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<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; }
@ -350,29 +378,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();
}, 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({
@ -412,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;

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

@ -304,6 +304,7 @@ IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
Status
@ -323,6 +324,26 @@ IVFHybridIndex::UnsetQuantizer() {
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

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

@ -111,6 +111,9 @@ class IVFHybridIndex : public IVFMixIndex {
Status
UnsetQuantizer() override;
Status
LoadData(const knowhere::QuantizerPtr &q, const Config &conf) override;
};
class BFIndex : public VecIndexImpl {

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

@ -104,15 +104,19 @@ class VecIndex {
Load(const zilliz::knowhere::BinarySet &index_binary) = 0;
// TODO(linxj): refactor later
////////////////
virtual knowhere::QuantizerPtr
LoadQuantizer(const Config& conf) { return Status::OK(); }
LoadQuantizer(const Config& conf) { return nullptr; }
virtual Status
LoadData(const knowhere::QuantizerPtr &q, const Config &conf) { return Status::OK(); }
// TODO(linxj): refactor later
virtual Status
SetQuantizer(const knowhere::QuantizerPtr& q) { return Status::OK(); }
virtual Status
UnsetQuantizer() { return Status::OK(); }
////////////////
};
extern Status