enable hybrid IVFSQ8 (edde9e80) · Commits · SUMMER2020 / students / proj-2018085

cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp

+43 −8

Original line number	Diff line number	Diff line
		@@ -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");
		}
		}

		}
		}

cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h

+11 −3

Original line number	Diff line number	Diff line
		@@ -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;

cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h

+5 −0

Original line number	Diff line number	Diff line
		@@ -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)

cpp/src/core/test/CMakeLists.txt

+1 −0

Original line number	Diff line number	Diff line
		@@ -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
		)

cpp/src/core/test/test_ivf.cpp

+82 −48

Original line number	Diff line number	Diff line
		@@ -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;