admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-01-02 00:45:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'knowhere_ut' into '0.5.0'

Browse Source 
MS-648 Knowhere ut

See merge request megasearch/milvus!757

Former-commit-id: 7729a6a45c4e3a91b76d85b3bfe7f26374115060
This commit is contained in:
jinhai 2019-10-17 16:50:55 +08:00
commit de86414d3c
16 changed files with 888 additions and 708 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

@ -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();
}
}

8
core/src/index/unittest/CMakeLists.txt
View File

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

120
core/src/index/unittest/Helper.h Normal file
View File

@ -0,0 +1,120 @@
// 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 <memory>
#include <string>
#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"
constexpr int DEVICEID = 0;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 10000;
constexpr int64_t NQ = 10;
constexpr int64_t K = 10;
constexpr int64_t PINMEM = 1024 * 1024 * 200;
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
constexpr int64_t RESNUM = 2;
knowhere::IVFIndexPtr
IndexFactory(const std::string& type) {
if (type == "IVF") {
return std::make_shared<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(DEVICEID);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(DEVICEID);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(DEVICEID);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
}
}
enum class ParameterType {
ivf,
ivfpq,
ivfsq,
};
class ParamGenerator {
public:
static ParamGenerator&
GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config
Gen(const ParameterType& type) {
if (type == ParameterType::ivf) {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfpq) {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->m = 4;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfsq) {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
};
#include <gtest/gtest.h>
class TestGpuIndexBase : public ::testing::Test {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
}
void
TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
};

122
core/src/index/unittest/test_customized_index.cpp Normal file
View File

@ -0,0 +1,122 @@
// 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 "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

309
core/src/index/unittest/test_gpuresource.cpp Normal file
View File

@ -0,0 +1,309 @@
// 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/Helper.h"
#include "unittest/utils.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

50
core/src/index/unittest/test_idmap.cpp
View File

@ -23,54 +23,28 @@
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "Helper.h"
#include "unittest/utils.h"
static int device_id = 0;
class IDMAPTest : public DataGen, public ::testing::Test {
class IDMAPTest : public DataGen, public TestGpuIndexBase {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2);
TestGpuIndexBase::SetUp();
Init_with_default();
index_ = std::make_shared<knowhere::IDMAP>();
}
void
TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
TestGpuIndexBase::TearDown();
}
protected:
knowhere::IDMAPPtr index_ = nullptr;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
TEST_F(IDMAPTest, idmap_basic) {
ASSERT_TRUE(!xb.empty());
@ -87,7 +61,7 @@ TEST_F(IDMAPTest, idmap_basic) {
ASSERT_TRUE(index_->GetRawIds() != nullptr);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
// PrintResult(result, nq, k);
index_->Seal();
auto binaryset = index_->Serialize();
@ -95,7 +69,7 @@ TEST_F(IDMAPTest, idmap_basic) {
new_index->Load(binaryset);
auto re_result = index_->Search(query_dataset, conf);
AssertAnns(re_result, nq, k);
PrintResult(re_result, nq, k);
// PrintResult(re_result, nq, k);
}
TEST_F(IDMAPTest, idmap_serialize) {
@ -118,7 +92,7 @@ TEST_F(IDMAPTest, idmap_serialize) {
index_->Add(base_dataset, knowhere::Config());
auto re_result = index_->Search(query_dataset, conf);
AssertAnns(re_result, nq, k);
PrintResult(re_result, nq, k);
// PrintResult(re_result, nq, k);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
auto binaryset = index_->Serialize();
@ -138,7 +112,7 @@ TEST_F(IDMAPTest, idmap_serialize) {
EXPECT_EQ(index_->Dimension(), dim);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
// PrintResult(result, nq, k);
}
}
@ -169,7 +143,7 @@ TEST_F(IDMAPTest, copy_test) {
{
// cpu to gpu
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf);
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf);
auto clone_result = clone_index->Search(query_dataset, conf);
AssertAnns(clone_result, nq, k);
ASSERT_THROW({ std::static_pointer_cast<knowhere::GPUIDMAP>(clone_index)->GetRawVectors(); },
@ -194,9 +168,9 @@ TEST_F(IDMAPTest, copy_test) {
ASSERT_TRUE(std::static_pointer_cast<knowhere::IDMAP>(host_index)->GetRawIds() != nullptr);
// gpu to gpu
auto device_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf);
auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf);
auto new_device_index =
std::static_pointer_cast<knowhere::GPUIDMAP>(device_index)->CopyGpuToGpu(device_id, conf);
std::static_pointer_cast<knowhere::GPUIDMAP>(device_index)->CopyGpuToGpu(DEVICEID, conf);
auto device_result = new_device_index->Search(query_dataset, conf);
AssertAnns(device_result, nq, k);
}

498
core/src/index/unittest/test_ivf.cpp
View File

@ -35,99 +35,25 @@
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "unittest/Helper.h"
#include "unittest/utils.h"
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
constexpr int device_id = 0;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 1000000 / 100;
constexpr int64_t NQ = 10;
constexpr int64_t K = 10;
knowhere::IVFIndexPtr
IndexFactory(const std::string& type) {
if (type == "IVF") {
return std::make_shared<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(device_id);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(device_id);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(device_id);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(device_id);
}
}
enum class ParameterType {
ivf,
ivfpq,
ivfsq,
nsg,
};
class ParamGenerator {
public:
static ParamGenerator&
GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config
Gen(const ParameterType& type) {
if (type == ParameterType::ivf) {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->k = K;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfpq) {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 25;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->m = 4;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfsq) {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->k = K;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
};
class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, ParameterType>> {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
ParameterType parameter_type;
std::tie(index_type, parameter_type) = GetParam();
// Init_with_default();
Generate(DIM, NB, NQ);
index_ = IndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(parameter_type);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
}
void
@ -140,7 +66,7 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
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 knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
return knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
}
return index_;
}
@ -162,33 +88,6 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
#endif
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)));
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
TEST_P(IVFTest, ivf_basic) {
assert(!xb.empty());
@ -207,85 +106,6 @@ TEST_P(IVFTest, ivf_basic) {
// 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<knowhere::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_1_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<knowhere::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);
hybrid_1_idx->UnsetQuantizer();
}
{
auto hybrid_2_idx = std::make_shared<knowhere::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_2_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 1;
quantizer_conf->gpu_id = device_id;
auto q = hybrid_2_idx->LoadQuantizer(quantizer_conf);
quantizer_conf->mode = 2;
auto gpu_idx = hybrid_2_idx->LoadData(q, quantizer_conf);
auto result = gpu_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; }
//
// // else
// 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 result = index_->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
//
// if (auto device_index = std::dynamic_pointer_cast<GPUIVF>(index_)) {
// auto host_index = device_index->Copy_index_gpu_to_cpu();
// auto result = host_index->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
// }
//}
TEST_P(IVFTest, ivf_serialize) {
auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) {
FileIOWriter writer(filename);
@ -423,7 +243,7 @@ TEST_P(IVFTest, clone_test) {
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 = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
auto clone_result = clone_index->Search(query_dataset, conf);
AssertEqual(result, clone_result);
std::cout << "clone C <=> G [" << index_type << "] success" << std::endl;
@ -432,7 +252,7 @@ TEST_P(IVFTest, clone_test) {
EXPECT_THROW(
{
std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl;
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
},
knowhere::KnowhereException);
}
@ -440,9 +260,7 @@ TEST_P(IVFTest, clone_test) {
}
#ifdef CUSTOMIZATION
TEST_P(IVFTest, seal_test) {
// FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
TEST_P(IVFTest, gpu_seal_test) {
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()) {
@ -466,309 +284,13 @@ TEST_P(IVFTest, seal_test) {
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
knowhere::TimeRecorder tc("CopyToGpu");
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
auto without_seal = tc.RecordSection("Without seal");
cpu_idx->Seal();
tc.RecordSection("seal cost");
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
auto with_seal = tc.RecordSection("With seal");
ASSERT_GE(without_seal, with_seal);
}
#endif
class GPURESTEST : public DataGen, public ::testing::Test {
protected:
void
SetUp() override {
Generate(128, 1000000, 1000);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2);
k = 100;
elems = nq * k;
ids = (int64_t*)malloc(sizeof(int64_t) * elems);
dis = (float*)malloc(sizeof(float) * elems);
}
void
TearDown() override {
delete ids;
delete dis;
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
protected:
std::string index_type;
knowhere::IVFIndexPtr index_ = nullptr;
int64_t* ids = nullptr;
float* dis = nullptr;
int64_t elems = 0;
};
const int search_count = 18;
const int load_count = 3;
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(device_id);
ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), device_id);
auto conf = std::make_shared<knowhere::IVFCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
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);
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();
{
// IVF-Search
faiss::gpu::StandardGpuResources res;
faiss::gpu::GpuIndexIVFFlatConfig idx_config;
idx_config.device = device_id;
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");
}
}
#ifdef CUSTOMIZATION
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 = device_id;
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, device_id, 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, device_id, 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, device_id, 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
TEST_F(GPURESTEST, copyandsearch) {
// search and copy at the same time
printf("==================\n");
index_type = "GPUIVF";
index_ = IndexFactory(index_type);
auto conf = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
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();
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
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, device_id, 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, device_id, knowhere::Config());
tc.RecordSection("Copy to gpu once");
search_idx->Search(query_dataset, conf);
tc.RecordSection("search once");
search_func();
tc.RecordSection("only search total");
load_func();
tc.RecordSection("only copy total");
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 = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
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);
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, device_id, knowhere::Config());
constexpr int train_count = 1;
constexpr int search_count = 5000;
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, device_id, 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();
}
}
// TODO(lxj): Add exception test

27
core/src/index/unittest/test_kdt.cpp
View File

@ -52,33 +52,6 @@ class KDTTest : public DataGen, public ::testing::Test {
std::shared_ptr<knowhere::CPUKDTRNG> index_ = nullptr;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
// TODO(lxj): add test about count() and dimension()
TEST_F(KDTTest, kdt_basic) {
assert(!xb.empty());

14
core/src/index/unittest/test_nsg/test_nsg.cpp
View File

@ -30,19 +30,19 @@ using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
constexpr int64_t DEVICE_ID = 0;
constexpr int64_t DEVICEID = 0;
class NSGInterfaceTest : public DataGen, public ::testing::Test {
protected:
void
SetUp() override {
// Init_with_default();
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
Generate(256, 1000000 / 100, 1);
index_ = std::make_shared<knowhere::NSG>();
auto tmp_conf = std::make_shared<knowhere::NSGCfg>();
tmp_conf->gpu_id = DEVICE_ID;
tmp_conf->gpu_id = DEVICEID;
tmp_conf->knng = 20;
tmp_conf->nprobe = 8;
tmp_conf->nlist = 163;
@ -69,14 +69,6 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test {
knowhere::Config search_conf;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
TEST_F(NSGInterfaceTest, basic_test) {
assert(!xb.empty());

28
core/src/index/unittest/utils.cpp
View File

@ -17,6 +17,7 @@
#include "unittest/utils.h"
#include <gtest/gtest.h>
#include <memory>
#include <string>
#include <utility>
@ -147,3 +148,30 @@ generate_query_dataset(int64_t nb, int64_t dim, float* xb) {
auto dataset = std::make_shared<knowhere::Dataset>(std::move(tensors), tensor_schema);
return dataset;
}
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}

6
core/src/index/unittest/utils.h
View File

@ -68,6 +68,12 @@ generate_dataset(int64_t nb, int64_t dim, float* xb, int64_t* ids);
knowhere::DatasetPtr
generate_query_dataset(int64_t nb, int64_t dim, float* xb);
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k);
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k);
struct FileIOWriter {
std::fstream fs;
std::string name;

133
core/unittest/wrapper/test_hybrid_index.cpp Normal file
View File

@ -0,0 +1,133 @@
// 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 "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include <gtest/gtest.h>
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
class KnowhereHybrid
: public DataGenBase, public ::testing::Test {
protected:
void SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
dim = 128;
nb = 10000;
nq = 100;
k = 100;
GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
}
void TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
protected:
milvus::engine::IndexType index_type;
milvus::engine::VecIndexPtr index_ = nullptr;
knowhere::Config conf;
};
#ifdef CUSTOMIZATION
TEST_F(KnowhereHybrid, test_interface) {
assert(!xb.empty());
index_type = milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID;
index_ = GetVecIndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(index_type);
auto elems = nq * k;
std::vector<int64_t> res_ids(elems);
std::vector<float> res_dis(elems);
conf->gpu_id = DEVICEID;
conf->d = dim;
conf->k = k;
index_->BuildAll(nb, xb.data(), ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
auto binaryset = index_->Serialize();
{
// cpu -> gpu
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
{
for (int i = 0; i < 2; ++i) {
auto gpu_idx = cpu_idx->CopyToGpu(DEVICEID, conf);
gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
}
}
}
{
// quantization already in gpu, only copy data
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf);
auto gpu_idx = pair.first;
auto quantization = pair.second;
gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 2;
quantizer_conf->gpu_id = DEVICEID;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = GetVecIndexFactory(index_type);
hybrid_idx->Load(binaryset);
hybrid_idx->LoadData(quantization, quantizer_conf);
hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
}
}
{
// quantization already in gpu, only set quantization
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf);
auto quantization = pair.second;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = GetVecIndexFactory(index_type);
hybrid_idx->Load(binaryset);
hybrid_idx->SetQuantizer(quantization);
hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
hybrid_idx->UnsetQuantizer();
}
}
}
#endif

139
core/unittest/wrapper/test_wrapper.cpp
View File

@ -25,150 +25,36 @@
INITIALIZE_EASYLOGGINGPP
namespace {
namespace ms = milvus::engine;
namespace kw = knowhere;
} // namespace
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 100000;
constexpr int64_t DEVICE_ID = 0;
class ParamGenerator {
public:
static ParamGenerator& GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config Gen(const milvus::engine::IndexType& type) {
switch (type) {
case milvus::engine::IndexType::FAISS_IDMAP: {
auto tempconf = std::make_shared<knowhere::Cfg>();
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFFLAT_CPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_GPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFSQ8_CPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_GPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->m = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::NSG_MIX: {
auto tempconf = std::make_shared<knowhere::NSGCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->search_length = 8;
tempconf->knng = 200;
tempconf->search_length = 40; // TODO(linxj): be 20 when search
tempconf->out_degree = 60;
tempconf->candidate_pool_size = 200;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
}
};
class KnowhereWrapperTest
: public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
protected:
void SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID,
1024 * 1024 * 200,
1024 * 1024 * 300,
2);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
std::string generator_type;
std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam();
auto generator = std::make_shared<DataGenBase>();
generator->GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
index_ = GetVecIndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(index_type);
conf->k = k;
conf->d = dim;
conf->gpu_id = DEVICE_ID;
conf->gpu_id = DEVICEID;
}
void TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis) {
EXPECT_EQ(ids.size(), nq * k);
EXPECT_EQ(dis.size(), nq * k);
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(ids[i * k], gt_ids[i * k]);
//EXPECT_EQ(dis[i * k], gt_dis[i * k]);
}
int match = 0;
for (int i = 0; i < nq; ++i) {
for (int j = 0; j < k; ++j) {
for (int l = 0; l < k; ++l) {
if (ids[i * nq + j] == gt_ids[i * nq + l]) match++;
}
}
}
auto precision = float(match) / (nq * k);
EXPECT_GT(precision, 0.5);
std::cout << std::endl << "Precision: " << precision
<< ", match: " << match
<< ", total: " << nq * k
<< std::endl;
}
protected:
milvus::engine::IndexType index_type;
knowhere::Config conf;
int dim = DIM;
int nb = NB;
int nq = 10;
int k = 10;
std::vector<float> xb;
std::vector<float> xq;
std::vector<int64_t> ids;
milvus::engine::VecIndexPtr index_ = nullptr;
// Ground Truth
std::vector<int64_t> gt_ids;
std::vector<float> gt_dis;
knowhere::Config conf;
};
INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
@ -220,7 +106,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
AssertResult(res_ids, res_dis);
{
auto dev_idx = index_->CopyToGpu(DEVICE_ID);
auto dev_idx = index_->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
}
@ -232,7 +118,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
write_index(index_, file_location);
auto new_index = milvus::engine::read_index(file_location);
auto dev_idx = new_index->CopyToGpu(DEVICE_ID);
auto dev_idx = new_index->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
}
@ -240,10 +126,6 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
}
}
//TEST_P(KnowhereWrapperTest, TO_CPU_TEST) {
// // dev
//}
TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
EXPECT_EQ(index_->GetType(), index_type);
@ -282,8 +164,3 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
AssertResult(res_ids, res_dis);
}
}
// TODO(linxj): add exception test
//TEST_P(KnowhereWrapperTest, exception_test) {
//}

28
core/unittest/wrapper/utils.cpp
View File

@ -16,6 +16,7 @@
// under the License.
#include <gtest/gtest.h>
#include <faiss/IndexFlat.h>
#include "wrapper/utils.h"
@ -59,3 +60,30 @@ DataGenBase::GenData(const int &dim,
gt_dis.resize(nq * k);
GenData(dim, nb, nq, xb.data(), xq.data(), ids.data(), k, gt_ids.data(), gt_dis.data());
}
void
DataGenBase::AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis) {
EXPECT_EQ(ids.size(), nq * k);
EXPECT_EQ(dis.size(), nq * k);
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(ids[i * k], gt_ids[i * k]);
//EXPECT_EQ(dis[i * k], gt_dis[i * k]);
}
int match = 0;
for (int i = 0; i < nq; ++i) {
for (int j = 0; j < k; ++j) {
for (int l = 0; l < k; ++l) {
if (ids[i * nq + j] == gt_ids[i * nq + l]) match++;
}
}
}
auto precision = float(match) / (nq * k);
EXPECT_GT(precision, 0.5);
std::cout << std::endl << "Precision: " << precision
<< ", match: " << match
<< ", total: " << nq * k
<< std::endl;
}

108
core/unittest/wrapper/utils.h
View File

@ -24,24 +24,110 @@
#include <cstdio>
#include <fstream>
#include "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
class DataGenBase;
using DataGenPtr = std::shared_ptr<DataGenBase>;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 100000;
constexpr int64_t NQ = 10;
constexpr int64_t DEVICEID = 0;
constexpr int64_t PINMEM = 1024 * 1024 * 200;
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
constexpr int64_t RESNUM = 2;
class DataGenBase {
public:
virtual void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, int64_t *ids,
const int &k, int64_t *gt_ids, float *gt_dis);
virtual void GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids,
const int& k, int64_t* gt_ids, float* gt_dis);
virtual void GenData(const int &dim,
const int &nb,
const int &nq,
std::vector<float> &xb,
std::vector<float> &xq,
std::vector<int64_t> &ids,
const int &k,
std::vector<int64_t> &gt_ids,
std::vector<float> &gt_dis);
virtual void GenData(const int& dim,
const int& nb,
const int& nq,
std::vector<float>& xb,
std::vector<float>& xq,
std::vector<int64_t>& ids,
const int& k,
std::vector<int64_t>& gt_ids,
std::vector<float>& gt_dis);
void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
int dim = DIM;
int nb = NB;
int nq = NQ;
int k = 10;
std::vector<float> xb;
std::vector<float> xq;
std::vector<int64_t> ids;
// Ground Truth
std::vector<int64_t> gt_ids;
std::vector<float> gt_dis;
};
class ParamGenerator {
public:
static ParamGenerator& GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config Gen(const milvus::engine::IndexType& type) {
switch (type) {
case milvus::engine::IndexType::FAISS_IDMAP: {
auto tempconf = std::make_shared<knowhere::Cfg>();
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFFLAT_CPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_GPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID:
case milvus::engine::IndexType::FAISS_IVFSQ8_CPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_GPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->m = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::NSG_MIX: {
auto tempconf = std::make_shared<knowhere::NSGCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->search_length = 8;
tempconf->knng = 200;
tempconf->search_length = 40; // TODO(linxj): be 20 when search
tempconf->out_degree = 60;
tempconf->candidate_pool_size = 200;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
}
};