diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 55da3c82a04457a5e1747b2aaf1405347d67d972..08898aacdfec90582eeb25efeaa03222295ef684 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -34,7 +34,7 @@ else() endif() message(STATUS "Build type = ${BUILD_TYPE}") -#add_definitions(-DNEW_SCHEDULER) +add_definitions(-DNEW_SCHEDULER) project(milvus VERSION "${MILVUS_VERSION}") project(milvus_engine LANGUAGES CUDA CXX) diff --git a/cpp/conf/server_config.template b/cpp/conf/server_config.template index 0e28737d994ce35391be4fb1367bde08e9f38464..a112bd4c4084a147d8208c8e62c8a5d5dea6a268 100644 --- a/cpp/conf/server_config.template +++ b/cpp/conf/server_config.template @@ -66,24 +66,24 @@ resource_config: enable_loader: true enable_executor: true - gtx1060: - type: GPU - memory: 6 - device_id: 0 - enable_loader: true - enable_executor: true +# gtx1060: +# type: GPU +# memory: 6 +# device_id: 0 +# enable_loader: true +# enable_executor: true - gtx1660: - type: GPU - memory: 6 - device_id: 1 - enable_loader: true - enable_executor: true +# gtx1660: +# type: GPU +# memory: 6 +# device_id: 1 +# enable_loader: false +# enable_executor: false # connection list, length: 0~N # format: -${resource_name}===${resource_name} connections: - ssda===cpu - - cpu===gtx1060 - - cpu===gtx1660 +# - cpu===gtx1060 +# - cpu===gtx1660 diff --git a/cpp/src/scheduler/SchedInst.cpp b/cpp/src/scheduler/SchedInst.cpp index 972c2e260658f2ba56a6e868ec19777b7f92ff39..42352eed5e38f8d3790e3a9e0d1401ce35ef7ecc 100644 --- a/cpp/src/scheduler/SchedInst.cpp +++ b/cpp/src/scheduler/SchedInst.cpp @@ -7,6 +7,7 @@ #include "SchedInst.h" #include "server/ServerConfig.h" #include "ResourceFactory.h" +#include "knowhere/index/vector_index/gpu_ivf.h" namespace zilliz { namespace milvus { @@ -36,8 +37,12 @@ SchedServInit() { device_id, enable_loader, enable_executor)); + + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id); } + knowhere::FaissGpuResourceMgr::GetInstance().InitResource(); + auto default_connection = Connection("default_connection", 500.0); auto connections = config.GetSequence(server::CONFIG_RESOURCE_CONNECTIONS); for (auto &conn : connections) { diff --git a/cpp/src/wrapper/knowhere/vec_impl.cpp b/cpp/src/wrapper/knowhere/vec_impl.cpp index b0fb4c07990b8d106a79a235a4bed819fc032b51..0b9855c639f973acac5e5187560abba8195760e4 100644 --- a/cpp/src/wrapper/knowhere/vec_impl.cpp +++ b/cpp/src/wrapper/knowhere/vec_impl.cpp @@ -241,8 +241,9 @@ server::KnowhereError IVFMixIndex::BuildAll(const long &nb, index_->Add(dataset, cfg); if (auto device_index = std::dynamic_pointer_cast(index_)) { - auto host_index = device_index->Copy_index_gpu_to_cpu(); + auto host_index = device_index->CopyGpuToCpu(Config()); index_ = host_index; + type = TransferToCpuIndexType(type); } else { WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed"; } diff --git a/cpp/src/wrapper/knowhere/vec_index.cpp b/cpp/src/wrapper/knowhere/vec_index.cpp index c31c5261dabedf1054d0e8279b33fe4f3dcefbd7..9ac0d8b3ad8e14d925081cf91a540ac56a5640f0 100644 --- a/cpp/src/wrapper/knowhere/vec_index.cpp +++ b/cpp/src/wrapper/knowhere/vec_index.cpp @@ -106,6 +106,10 @@ VecIndexPtr GetVecIndexFactory(const IndexType &type) { index = std::make_shared(0); return std::make_shared(index, IndexType::FAISS_IVFSQ8_MIX); } + case IndexType::FAISS_IVFSQ8: { + index = std::make_shared(); + break; + } case IndexType::NSG_MIX: { // TODO(linxj): bug. index = std::make_shared(0); break; @@ -194,10 +198,10 @@ server::KnowhereError write_index(VecIndexPtr index, const std::string &location // TODO(linxj): redo here. void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Config &cfg) { auto nlist = cfg.get_with_default("nlist", 0); - if (size <= TYPICAL_COUNT/16384 + 1) { + if (size <= TYPICAL_COUNT / 16384 + 1) { //handle less row count, avoid nlist set to 0 cfg["nlist"] = 1; - } else if (int(size/TYPICAL_COUNT) * nlist == 0) { + } else if (int(size / TYPICAL_COUNT) * nlist == 0) { //calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT cfg["nlist"] = int(size / TYPICAL_COUNT * 16384); } @@ -225,6 +229,20 @@ void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Co } } +IndexType TransferToCpuIndexType(const IndexType &type) { + switch (type) { + case IndexType::FAISS_IVFFLAT_MIX: { + return IndexType::FAISS_IVFFLAT_CPU; + } + case IndexType::FAISS_IVFSQ8_MIX: { + return IndexType::FAISS_IVFSQ8; + } + default: { + return IndexType::INVALID; + } + } +} + } } } diff --git a/cpp/src/wrapper/knowhere/vec_index.h b/cpp/src/wrapper/knowhere/vec_index.h index 19f0c6d36051cb801bd39344591bcfb17bc7da96..1c45ce89fc78defcbb9390b3c8afb4446b2870b2 100644 --- a/cpp/src/wrapper/knowhere/vec_index.h +++ b/cpp/src/wrapper/knowhere/vec_index.h @@ -32,6 +32,7 @@ enum class IndexType { FAISS_IVFPQ_GPU, SPTAG_KDT_RNT_CPU, FAISS_IVFSQ8_MIX, + FAISS_IVFSQ8, NSG_MIX, }; @@ -88,6 +89,8 @@ extern VecIndexPtr LoadVecIndex(const IndexType &index_type, const zilliz::knowh extern void AutoGenParams(const IndexType& type, const long& size, Config& cfg); +extern IndexType TransferToCpuIndexType(const IndexType& type); + } } } diff --git a/cpp/unittest/knowhere/knowhere_test.cpp b/cpp/unittest/knowhere/knowhere_test.cpp index abe7c84d4e74ca1b609987d81976b3809a33fd99..1fd23250dea70c94633f2ef7f0af1b208ef8f648 100644 --- a/cpp/unittest/knowhere/knowhere_test.cpp +++ b/cpp/unittest/knowhere/knowhere_test.cpp @@ -8,13 +8,14 @@ #include #include +#include "knowhere/index/vector_index/gpu_ivf.h" #include "utils.h" INITIALIZE_EASYLOGGINGPP using namespace zilliz::milvus::engine; -using namespace zilliz::knowhere; +//using namespace zilliz::knowhere; using ::testing::TestWithParam; using ::testing::Values; @@ -66,8 +67,8 @@ class KnowhereWrapperTest Config train_cfg; Config search_cfg; - int dim = 64; - int nb = 10000; + int dim = 512; + int nb = 1000000; int nq = 10; int k = 10; std::vector xb; @@ -94,27 +95,27 @@ INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, // Config::object{{"nlist", 100}, {"dim", 64}}, // Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 40}} //), - std::make_tuple(IndexType::FAISS_IVFFLAT_MIX, "Default", - 64, 100000, 10, 10, - Config::object{{"nlist", 1000}, {"dim", 64}, {"metric_type", "L2"}}, - Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}} - ), - std::make_tuple(IndexType::FAISS_IDMAP, "Default", - 64, 100000, 10, 10, - Config::object{{"dim", 64}, {"metric_type", "L2"}}, - Config::object{{"dim", 64}, {"k", 10}} - ), +// std::make_tuple(IndexType::FAISS_IVFFLAT_MIX, "Default", +// 64, 100000, 10, 10, +// Config::object{{"nlist", 1000}, {"dim", 64}, {"metric_type", "L2"}}, +// Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}} +// ), +// std::make_tuple(IndexType::FAISS_IDMAP, "Default", +// 64, 100000, 10, 10, +// Config::object{{"dim", 64}, {"metric_type", "L2"}}, +// Config::object{{"dim", 64}, {"k", 10}} +// ), std::make_tuple(IndexType::FAISS_IVFSQ8_MIX, "Default", - 64, 100000, 10, 10, - Config::object{{"dim", 64}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}}, - Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}} - ), - std::make_tuple(IndexType::NSG_MIX, "Default", - 128, 250000, 10, 10, - Config::object{{"dim", 128}, {"nlist", 8192}, {"nprobe", 16}, {"metric_type", "L2"}, - {"knng", 200}, {"search_length", 40}, {"out_degree", 60}, {"candidate_pool_size", 200}}, - Config::object{{"k", 10}, {"search_length", 20}} + 512, 1000000, 10, 10, + Config::object{{"dim", 512}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}}, + Config::object{{"dim", 512}, {"k", 10}, {"nprobe", 5}} ) +// std::make_tuple(IndexType::NSG_MIX, "Default", +// 128, 250000, 10, 10, +// Config::object{{"dim", 128}, {"nlist", 8192}, {"nprobe", 16}, {"metric_type", "L2"}, +// {"knng", 200}, {"search_length", 40}, {"out_degree", 60}, {"candidate_pool_size", 200}}, +// Config::object{{"k", 10}, {"search_length", 20}} +// ) //std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", // 64, 10000, 10, 10, // Config::object{{"TPTNumber", 1}, {"dim", 64}}, @@ -135,6 +136,34 @@ TEST_P(KnowhereWrapperTest, base_test) { AssertResult(res_ids, res_dis); } +TEST_P(KnowhereWrapperTest, to_gpu_test) { + EXPECT_EQ(index_->GetType(), index_type); + + zilliz::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0); + zilliz::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(1); + + auto elems = nq * k; + std::vector res_ids(elems); + std::vector res_dis(elems); + + index_->BuildAll(nb, xb.data(), ids.data(), train_cfg); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); + AssertResult(res_ids, res_dis); + { + index_->CopyToGpu(1); + } + + std::string file_location = "/tmp/whatever"; + write_index(index_, file_location); + auto new_index = read_index(file_location); + + auto dev_idx = new_index->CopyToGpu(1); + for (int i = 0; i < 10000; ++i) { + dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); + } + AssertResult(res_ids, res_dis); +} + TEST_P(KnowhereWrapperTest, serialize) { EXPECT_EQ(index_->GetType(), index_type);