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

IVF_SQ8 and IVF_PQ cannot be built on multiple GPUs (#5624) 

* IVF_SQ8 and IVF_PQ cannot be built on multiple GPUs

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>

* index created by std::make_shared

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
pull/5633/head
shengjun.li 2021-06-05 17:19:34 +08:00 committed by GitHub
parent 7dac20c35c
commit 8ead67ef95
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 29 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
internal/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp
View File

@ -38,14 +38,12 @@ GPUIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
if (gpu_res != nullptr) {
ResScope rs(gpu_res, gpu_id_, true);
faiss::gpu::GpuIndexIVFFlatConfig idx_config;
idx_config.device = gpu_id_;
idx_config.device = static_cast<int32_t>(gpu_id_);
int32_t nlist = config[IndexParams::nlist];
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
auto device_index =
new faiss::gpu::GpuIndexIVFFlat(gpu_res->faiss_res.get(), dim, nlist, metric_type, idx_config);
device_index->train(rows, reinterpret_cast<const float*>(p_data));
index_.reset(device_index);
index_ = std::make_shared<faiss::gpu::GpuIndexIVFFlat>(gpu_res->faiss_res.get(), dim, nlist, metric_type,
idx_config);
index_->train(rows, reinterpret_cast<const float*>(p_data));
res_ = gpu_res;
} else {
KNOWHERE_THROW_MSG("Build IVF can't get gpu resource");

12
internal/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp
View File

@ -31,11 +31,15 @@ GPUIVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
if (gpu_res != nullptr) {
ResScope rs(gpu_res, gpu_id_, true);
auto device_index = new faiss::gpu::GpuIndexIVFPQ(
gpu_res->faiss_res.get(), dim, config[IndexParams::nlist].get<int64_t>(), config[IndexParams::m],
config[IndexParams::nbits], GetMetricType(config[Metric::TYPE].get<std::string>()));
faiss::gpu::GpuIndexIVFPQConfig idx_config;
idx_config.device = static_cast<int32_t>(gpu_id_);
int32_t nlist = config[IndexParams::nlist];
int32_t m = config[IndexParams::m];
int32_t nbits = config[IndexParams::nbits];
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
index_ = std::make_shared<faiss::gpu::GpuIndexIVFPQ>(gpu_res->faiss_res.get(), dim, nlist, m, nbits,
metric_type, idx_config);
device_index->train(rows, reinterpret_cast<const float*>(p_data));
index_.reset(device_index);
res_ = gpu_res;
} else {
KNOWHERE_THROW_MSG("Build IVFPQ can't get gpu resource");

12
internal/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp
View File

@ -33,11 +33,13 @@ GPUIVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
if (gpu_res != nullptr) {
ResScope rs(gpu_res, gpu_id_, true);
auto device_index = new faiss::gpu::GpuIndexIVFScalarQuantizer(
gpu_res->faiss_res.get(), dim, config[IndexParams::nlist].get<int64_t>(), faiss::QuantizerType::QT_8bit,
GetMetricType(config[Metric::TYPE].get<std::string>()));
device_index->train(rows, reinterpret_cast<const float*>(p_data));
index_.reset(device_index);
faiss::gpu::GpuIndexIVFScalarQuantizerConfig idx_config;
idx_config.device = static_cast<int32_t>(gpu_id_);
int32_t nlist = config[IndexParams::nlist];
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
index_ = std::make_shared<faiss::gpu::GpuIndexIVFScalarQuantizer>(
gpu_res->faiss_res.get(), dim, nlist, faiss::QuantizerType::QT_8bit, metric_type, true, idx_config);
index_->train(rows, (float*)p_data);
res_ = gpu_res;
} else {
KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");

21
internal/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp
View File

@ -33,27 +33,22 @@ IVFSQHybrid::Train(const DatasetPtr& dataset_ptr, const Config& config) {
GET_TENSOR_DATA_DIM(dataset_ptr)
gpu_id_ = config[knowhere::meta::DEVICEID];
std::stringstream index_type;
index_type << "IVF" << config[IndexParams::nlist] << ","
<< "SQ8Hybrid";
auto build_index =
faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(config[Metric::TYPE].get<std::string>()));
auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
if (gpu_res != nullptr) {
ResScope rs(gpu_res, gpu_id_, true);
auto device_index = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, build_index);
device_index->train(rows, reinterpret_cast<const float*>(p_data));
index_.reset(device_index);
faiss::gpu::GpuIndexIVFSQHybridConfig idx_config;
idx_config.device = static_cast<int32_t>(gpu_id_);
int32_t nlist = config[IndexParams::nlist];
faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
index_ = std::make_shared<faiss::gpu::GpuIndexIVFSQHybrid>(
gpu_res->faiss_res.get(), dim, nlist, faiss::QuantizerType::QT_8bit, metric_type, true, idx_config);
index_->train(rows, reinterpret_cast<const float*>(p_data));
res_ = gpu_res;
gpu_mode_ = 2;
index_mode_ = IndexMode::MODE_GPU;
} else {
delete build_index;
KNOWHERE_THROW_MSG("Build IVFSQHybrid can't get gpu resource");
}
delete build_index;
}
VecIndexPtr

2
internal/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h
View File

@ -38,12 +38,14 @@ class IVFSQHybrid : public GPUIVFSQ {
explicit IVFSQHybrid(const int& device_id) : GPUIVFSQ(device_id) {
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H;
gpu_mode_ = 0;
index_mode_ = IndexMode::MODE_CPU;
}
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {
index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H;
index_ = index;
gpu_mode_ = 0;
index_mode_ = IndexMode::MODE_CPU;
}
explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index, const int64_t device_id, ResPtr& resource)