// 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 "utils/easylogging++.h" #include "src/wrapper/VecIndex.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "wrapper/utils.h" #include INITIALIZE_EASYLOGGINGPP namespace { namespace ms = zilliz::milvus::engine; namespace kw = zilliz::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; } kw::Config Gen(const ms::IndexType &type) { switch (type) { case ms::IndexType::FAISS_IDMAP: { auto tempconf = std::make_shared(); tempconf->metric_type = zilliz::knowhere::METRICTYPE::L2; return tempconf; } case ms::IndexType::FAISS_IVFFLAT_CPU: case ms::IndexType::FAISS_IVFFLAT_GPU: case ms::IndexType::FAISS_IVFFLAT_MIX: { auto tempconf = std::make_shared(); tempconf->nlist = 100; tempconf->nprobe = 16; tempconf->metric_type = zilliz::knowhere::METRICTYPE::L2; return tempconf; } case ms::IndexType::FAISS_IVFSQ8_CPU: case ms::IndexType::FAISS_IVFSQ8_GPU: case ms::IndexType::FAISS_IVFSQ8_MIX: { auto tempconf = std::make_shared(); tempconf->nlist = 100; tempconf->nprobe = 16; tempconf->nbits = 8; tempconf->metric_type = zilliz::knowhere::METRICTYPE::L2; return tempconf; } case ms::IndexType::FAISS_IVFPQ_CPU: case ms::IndexType::FAISS_IVFPQ_GPU: { auto tempconf = std::make_shared(); tempconf->nlist = 100; tempconf->nprobe = 16; tempconf->nbits = 8; tempconf->m = 8; tempconf->metric_type = zilliz::knowhere::METRICTYPE::L2; return tempconf; } case ms::IndexType::NSG_MIX: { auto tempconf = std::make_shared(); 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 = zilliz::knowhere::METRICTYPE::L2; return tempconf; } } } }; class KnowhereWrapperTest : public TestWithParam<::std::tuple> { protected: void SetUp() override { zilliz::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); std::string generator_type; std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); auto generator = std::make_shared(); generator->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; } void TearDown() override { zilliz::knowhere::FaissGpuResourceMgr::GetInstance().Free(); } void AssertResult(const std::vector &ids, const std::vector &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: ms::IndexType index_type; kw::Config conf; int dim = DIM; int nb = NB; int nq = 10; int k = 10; std::vector xb; std::vector xq; std::vector ids; ms::VecIndexPtr index_ = nullptr; // Ground Truth std::vector gt_ids; std::vector gt_dis; }; INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, Values( //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] std::make_tuple(ms::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 100000, 10, 10), std::make_tuple(ms::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), std::make_tuple(ms::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 100000, 10, 10), std::make_tuple(ms::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10), std::make_tuple(ms::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), std::make_tuple(ms::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), // std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), // std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), std::make_tuple(ms::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10) ) ); TEST_P(KnowhereWrapperTest, BASE_TEST) { EXPECT_EQ(index_->GetType(), index_type); auto elems = nq * k; std::vector res_ids(elems); std::vector res_dis(elems); 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); } TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { EXPECT_EQ(index_->GetType(), index_type); auto elems = nq * k; std::vector res_ids(elems); std::vector res_dis(elems); 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); { auto dev_idx = index_->CopyToGpu(DEVICE_ID); for (int i = 0; i < 10; ++i) { dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); } AssertResult(res_ids, res_dis); } { std::string file_location = "/tmp/knowhere_gpu_file"; write_index(index_, file_location); auto new_index = ms::read_index(file_location); auto dev_idx = new_index->CopyToGpu(DEVICE_ID); for (int i = 0; i < 10; ++i) { dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); } AssertResult(res_ids, res_dis); } } //TEST_P(KnowhereWrapperTest, TO_CPU_TEST) { // // dev //} TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { EXPECT_EQ(index_->GetType(), index_type); auto elems = nq * k; std::vector res_ids(elems); std::vector res_dis(elems); 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); { auto binary = index_->Serialize(); auto type = index_->GetType(); auto new_index = GetVecIndexFactory(type); new_index->Load(binary); EXPECT_EQ(new_index->Dimension(), index_->Dimension()); EXPECT_EQ(new_index->Count(), index_->Count()); std::vector res_ids(elems); std::vector res_dis(elems); new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); AssertResult(res_ids, res_dis); } { std::string file_location = "/tmp/knowhere"; write_index(index_, file_location); auto new_index = ms::read_index(file_location); EXPECT_EQ(new_index->GetType(), ConvertToCpuIndexType(index_type)); EXPECT_EQ(new_index->Dimension(), index_->Dimension()); EXPECT_EQ(new_index->Count(), index_->Count()); std::vector res_ids(elems); std::vector res_dis(elems); new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); AssertResult(res_ids, res_dis); } } // TODO(linxj): add exception test //TEST_P(KnowhereWrapperTest, exception_test) { //}