MS-648 add ut (472679ba) · Commits · SUMMER2020 / students / proj-2018085

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

+2 −0

Original line number	Diff line number	Diff line
		@@ -189,6 +189,8 @@ IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) {
		if (quantizer_conf->mode != 2) {
		KNOWHERE_THROW_MSG("mode only support 2 in this func");
		}
		} else {
		KNOWHERE_THROW_MSG("conf error");
		}
		// if (quantizer_conf->gpu_id != gpu_id_) {
		// KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card");

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

+2 −2

Original line number	Diff line number	Diff line
		@@ -63,7 +63,7 @@ FaissGpuResourceMgr::InitResource() {

		mutex_cache_.emplace(device_id, std::make_unique<std::mutex>());

		// std::cout << "Device Id: " << device_id << std::endl;
		// std::cout << "Device Id: " << DEVICEID << std::endl;
		auto& device_param = device.second;
		auto& bq = idle_map_[device_id];

		@@ -119,7 +119,7 @@ void
		FaissGpuResourceMgr::Dump() {
		for (auto& item : idle_map_) {
		auto& bq = item.second;
		std::cout << "device_id: " << item.first << ", resource count:" << bq.Size();
		std::cout << "DEVICEID: " << item.first << ", resource count:" << bq.Size();
		}
		}

core/src/index/unittest/CMakeLists.txt

+8 −0

Original line number	Diff line number	Diff line
		@@ -73,9 +73,17 @@ target_link_libraries(test_kdt
		SPTAGLibStatic
		${depend_libs} ${unittest_libs} ${basic_libs})

		add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})
		target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})

		add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})
		target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})

		install(TARGETS test_ivf DESTINATION unittest)
		install(TARGETS test_idmap DESTINATION unittest)
		install(TARGETS test_kdt DESTINATION unittest)
		install(TARGETS test_gpuresource DESTINATION unittest)
		install(TARGETS test_customized_index DESTINATION unittest)

		#add_subdirectory(faiss_ori)
		add_subdirectory(test_nsg)

core/src/index/unittest/test_customized_index.cpp

0 → 100644

+108 −0

Original line number	Diff line number	Diff line
		//
		// Created by link on 2019/10/17.
		//
		#include <gtest/gtest.h>

		#include "unittest/Helper.h"
		#include "unittest/utils.h"

		class SingleIndexTest : public DataGen, public TestGpuIndexBase {
		protected:
		void
		SetUp() override {
		TestGpuIndexBase::SetUp();
		Generate(DIM, NB, NQ);
		k = K;
		}

		void
		TearDown() override {
		TestGpuIndexBase::TearDown();
		}

		protected:
		std::string index_type;
		knowhere::IVFIndexPtr index_ = nullptr;
		};

		#ifdef CUSTOMIZATION
		TEST_F(SingleIndexTest, IVFSQHybrid) {
		assert(!xb.empty());

		index_type = "IVFSQHybrid";
		index_ = IndexFactory(index_type);
		auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq);
		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 binaryset = index_->Serialize();
		{
		// copy cpu to gpu
		auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
		cpu_idx->Load(binaryset);

		{
		for (int i = 0; i < 3; ++i) {
		auto gpu_idx = cpu_idx->CopyCpuToGpu(DEVICEID, conf);
		auto result = gpu_idx->Search(query_dataset, conf);
		AssertAnns(result, nq, conf->k);
		// PrintResult(result, nq, k);
		}
		}
		}

		{
		// quantization already in gpu, only copy data
		auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
		cpu_idx->Load(binaryset);

		auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);
		auto gpu_idx = pair.first;
		auto quantization = pair.second;

		auto result = gpu_idx->Search(query_dataset, conf);
		AssertAnns(result, nq, conf->k);
		// PrintResult(result, nq, k);

		auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
		quantizer_conf->mode = 2; // only copy data
		quantizer_conf->gpu_id = DEVICEID;
		for (int i = 0; i < 2; ++i) {
		auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
		hybrid_idx->Load(binaryset);

		auto new_idx = hybrid_idx->LoadData(quantization, quantizer_conf);
		auto result = new_idx->Search(query_dataset, conf);
		AssertAnns(result, nq, conf->k);
		// PrintResult(result, nq, k);
		}
		}

		{
		// quantization already in gpu, only set quantization
		auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
		cpu_idx->Load(binaryset);

		auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);
		auto quantization = pair.second;

		for (int i = 0; i < 2; ++i) {
		auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
		hybrid_idx->Load(binaryset);

		hybrid_idx->SetQuantizer(quantization);
		auto result = hybrid_idx->Search(query_dataset, conf);
		AssertAnns(result, nq, conf->k);
		// PrintResult(result, nq, k);
		hybrid_idx->UnsetQuantizer();
		}
		}
		}

		#endif

core/src/index/unittest/test_gpuresource.cpp

0 → 100644

+311 −0

Original line number	Diff line number	Diff line
		// 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 <gtest/gtest.h>

		#include <iostream>
		#include <thread>

		#include <faiss/AutoTune.h>
		#include <faiss/gpu/GpuAutoTune.h>
		#include <faiss/gpu/GpuIndexIVFFlat.h>

		#include "knowhere/common/Exception.h"
		#include "knowhere/common/Timer.h"
		#include "knowhere/index/vector_index/IndexGPUIVF.h"
		#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
		#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
		#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"
		#include "unittest/Helper.h"

		class GPURESTEST : public DataGen, public TestGpuIndexBase {
		protected:
		void
		SetUp() override {
		TestGpuIndexBase::SetUp();
		Generate(DIM, NB, NQ);

		k = K;
		elems = nq * k;
		ids = (int64_t)malloc(sizeof(int64_t) elems);
		dis = (float)malloc(sizeof(float) elems);
		}

		void
		TearDown() override {
		delete ids;
		delete dis;
		TestGpuIndexBase::TearDown();
		}

		protected:
		std::string index_type;
		knowhere::IVFIndexPtr index_ = nullptr;

		int64_t* ids = nullptr;
		float* dis = nullptr;
		int64_t elems = 0;
		};

		TEST_F(GPURESTEST, copyandsearch) {
		// search and copy at the same time
		printf("==================\n");

		index_type = "GPUIVF";
		index_ = IndexFactory(index_type);

		auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf);
		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);
		auto result = index_->Search(query_dataset, conf);
		AssertAnns(result, nq, k);

		auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
		cpu_idx->Seal();
		auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());

		constexpr int64_t search_count = 50;
		constexpr int64_t load_count = 15;
		auto search_func = [&] {
		// TimeRecorder tc("search&load");
		for (int i = 0; i < search_count; ++i) {
		search_idx->Search(query_dataset, conf);
		// if (i > search_count - 6 \|\| i == 0)
		// tc.RecordSection("search once");
		}
		// tc.ElapseFromBegin("search finish");
		};
		auto load_func = [&] {
		// TimeRecorder tc("search&load");
		for (int i = 0; i < load_count; ++i) {
		knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
		// if (i > load_count -5 \|\| i < 5)
		// tc.RecordSection("Copy to gpu");
		}
		// tc.ElapseFromBegin("load finish");
		};

		knowhere::TimeRecorder tc("Basic");
		knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
		tc.RecordSection("Copy to gpu once");
		search_idx->Search(query_dataset, conf);
		tc.RecordSection("Search once");
		search_func();
		tc.RecordSection("Search total cost");
		load_func();
		tc.RecordSection("Copy total cost");

		std::thread search_thread(search_func);
		std::thread load_thread(load_func);
		search_thread.join();
		load_thread.join();
		tc.RecordSection("Copy&Search total");
		}

		TEST_F(GPURESTEST, trainandsearch) {
		index_type = "GPUIVF";
		index_ = IndexFactory(index_type);

		auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf);
		auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
		index_->set_preprocessor(preprocessor);
		auto model = index_->Train(base_dataset, conf);
		auto new_index = IndexFactory(index_type);
		new_index->set_index_model(model);
		new_index->Add(base_dataset, conf);
		auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config());
		cpu_idx->Seal();
		auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());

		constexpr int train_count = 5;
		constexpr int search_count = 200;
		auto train_stage = [&] {
		for (int i = 0; i < train_count; ++i) {
		auto model = index_->Train(base_dataset, conf);
		auto test_idx = IndexFactory(index_type);
		test_idx->set_index_model(model);
		test_idx->Add(base_dataset, conf);
		}
		};
		auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) {
		for (int i = 0; i < search_count; ++i) {
		auto result = search_idx->Search(query_dataset, conf);
		AssertAnns(result, nq, k);
		}
		};

		// TimeRecorder tc("record");
		// train_stage();
		// tc.RecordSection("train cost");
		// search_stage(search_idx);
		// tc.RecordSection("search cost");

		{
		// search and build parallel
		std::thread search_thread(search_stage, std::ref(search_idx));
		std::thread train_thread(train_stage);
		train_thread.join();
		search_thread.join();
		}
		{
		// build parallel
		std::thread train_1(train_stage);
		std::thread train_2(train_stage);
		train_1.join();
		train_2.join();
		}
		{
		// search parallel
		auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
		std::thread search_1(search_stage, std::ref(search_idx));
		std::thread search_2(search_stage, std::ref(search_idx_2));
		search_1.join();
		search_2.join();
		}
		}

		#ifdef CompareToOriFaiss
		TEST_F(GPURESTEST, gpu_ivf_resource_test) {
		assert(!xb.empty());

		{
		index_ = std::make_shared<knowhere::GPUIVF>(-1);
		ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), -1);
		std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->SetGpuDevice(DEVICEID);
		ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), DEVICEID);

		auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq);
		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);

		// knowhere::TimeRecorder tc("knowere GPUIVF");
		for (int i = 0; i < search_count; ++i) {
		index_->Search(query_dataset, conf);
		if (i > search_count - 6 \|\| i < 5)
		// tc.RecordSection("search once");
		}
		// tc.ElapseFromBegin("search all");
		}
		knowhere::FaissGpuResourceMgr::GetInstance().Dump();

		// {
		// // ori faiss IVF-Search
		// faiss::gpu::StandardGpuResources res;
		// faiss::gpu::GpuIndexIVFFlatConfig idx_config;
		// idx_config.device = DEVICEID;
		// faiss::gpu::GpuIndexIVFFlat device_index(&res, dim, 1638, faiss::METRIC_L2, idx_config);
		// device_index.train(nb, xb.data());
		// device_index.add(nb, xb.data());
		//
		// knowhere::TimeRecorder tc("ori IVF");
		// for (int i = 0; i < search_count; ++i) {
		// device_index.search(nq, xq.data(), k, dis, ids);
		// if (i > search_count - 6 \|\| i < 5)
		// tc.RecordSection("search once");
		// }
		// tc.ElapseFromBegin("search all");
		// }
		}

		TEST_F(GPURESTEST, gpuivfsq) {
		{
		// knowhere gpu ivfsq
		index_type = "GPUIVFSQ";
		index_ = IndexFactory(index_type);

		auto conf = std::make_shared<knowhere::IVFSQCfg>();
		conf->nlist = 1638;
		conf->d = dim;
		conf->gpu_id = DEVICEID;
		conf->metric_type = knowhere::METRICTYPE::L2;
		conf->k = k;
		conf->nbits = 8;
		conf->nprobe = 1;

		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);
		// auto result = index_->Search(query_dataset, conf);
		// AssertAnns(result, nq, k);

		auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
		cpu_idx->Seal();

		knowhere::TimeRecorder tc("knowhere GPUSQ8");
		auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
		tc.RecordSection("Copy to gpu");
		for (int i = 0; i < search_count; ++i) {
		search_idx->Search(query_dataset, conf);
		if (i > search_count - 6 \|\| i < 5)
		tc.RecordSection("search once");
		}
		tc.ElapseFromBegin("search all");
		}

		{
		// Ori gpuivfsq Test
		const char* index_description = "IVF1638,SQ8";
		faiss::Index* ori_index = faiss::index_factory(dim, index_description, faiss::METRIC_L2);

		faiss::gpu::StandardGpuResources res;
		auto device_index = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, ori_index);
		device_index->train(nb, xb.data());
		device_index->add(nb, xb.data());

		auto cpu_index = faiss::gpu::index_gpu_to_cpu(device_index);
		auto idx = dynamic_cast<faiss::IndexIVF*>(cpu_index);
		if (idx != nullptr) {
		idx->to_readonly();
		}
		delete device_index;
		delete ori_index;

		faiss::gpu::GpuClonerOptions option;
		option.allInGpu = true;

		knowhere::TimeRecorder tc("ori GPUSQ8");
		faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, cpu_index, &option);
		tc.RecordSection("Copy to gpu");
		for (int i = 0; i < search_count; ++i) {
		search_idx->search(nq, xq.data(), k, dis, ids);
		if (i > search_count - 6 \|\| i < 5)
		tc.RecordSection("search once");
		}
		tc.ElapseFromBegin("search all");
		delete cpu_index;
		delete search_idx;
		}
		}
		#endif