Merge remote-tracking branch 'upstream/branch-0.4.0' into cmake

Former-commit-id: 89e1b3ae3c22c02d40eed555c4d3dcd9d76faa7a

cpp/CHANGELOG.md

@@ -27,6 +27,8 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-470 - Drop index success, which table not created
 - MS-471 - code coverage run failed
 - MS-492 - Drop index failed if index have been created with index_type: FLAT
+- MS-493 - Knowhere unittest crash
+- MS-453 - GPU search error when nprobe set more than 1024
 
 ## Improvement
 - MS-327 - Clean code for milvus
@@ -85,6 +87,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-455 - Distribute tasks by minimal cost in scheduler
 - MS-460 - Put transport speed as weight when choosing neighbour to execute task
 - MS-459 - Add cache for pick function in tasktable
+- MS-476 - Improve search performance
 - MS-482 - Change search stream transport to unary in grpc
 - MS-487 - Define metric type in CreateTable
 - MS-488 - Improve code format in scheduler

cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h

@@ -9,7 +9,7 @@ namespace zilliz {
 namespace knowhere {
 
 struct Resource {
-    Resource(std::shared_ptr<faiss::gpu::StandardGpuResources> &r): faiss_res(r) {
+    explicit Resource(std::shared_ptr<faiss::gpu::StandardGpuResources> &r): faiss_res(r) {
         static int64_t global_id = 0;
         id = global_id++;
     }
@@ -32,6 +32,11 @@ class FaissGpuResourceMgr {
     static FaissGpuResourceMgr &
     GetInstance();
 
+    // Free gpu resource, avoid cudaGetDevice error when deallocate.
+    // this func should be invoke before main return
+    void
+    Free();
+
     void
     AllocateTempMem(ResPtr &resource, const int64_t& device_id, const int64_t& size);
 

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

@@ -282,7 +282,7 @@ void FaissGpuResourceMgr::InitResource() {
     for(auto& device : devices_params_) {
         auto& resource_vec = idle_[device.first];
 
-        for (int i = 0; i < device.second.resource_num; ++i) {
+        for (int64_t i = 0; i < device.second.resource_num; ++i) {
             auto res = std::make_shared<faiss::gpu::StandardGpuResources>();
 
             // TODO(linxj): enable set pinned memory
@@ -351,6 +351,18 @@ void FaissGpuResourceMgr::MoveToIdle(const int64_t &device_id, const ResPtr &res
     idle_[device_id].insert(it, res);
 }
 
+void FaissGpuResourceMgr::Free() {
+    for (auto &item : in_use_) {
+        auto& res_vec = item.second;
+        res_vec.clear();
+    }
+    for (auto &item : idle_) {
+        auto& res_vec = item.second;
+        res_vec.clear();
+    }
+    is_init = false;
+}
+
 void GPUIndex::SetGpuDevice(const int &gpu_id) {
     gpu_id_ = gpu_id;
 }

cpp/src/core/test/test_idmap.cpp

@@ -26,6 +26,11 @@ class IDMAPTest : public DataGen, public ::testing::Test {
         Init_with_default();
         index_ = std::make_shared<IDMAP>();
     }
+
+    void TearDown() override {
+        FaissGpuResourceMgr::GetInstance().Free();
+    }
+
  protected:
     IDMAPPtr index_ = nullptr;
 };

cpp/src/core/test/test_ivf.cpp

@@ -7,13 +7,11 @@
 #include <gtest/gtest.h>
 
 #include <iostream>
-#include <sstream>
 #include <thread>
 
 #include <faiss/AutoTune.h>
 #include <faiss/gpu/GpuAutoTune.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
-#include <faiss/gpu/GpuClonerOptions.h>
 
 #include "knowhere/index/vector_index/gpu_ivf.h"
 #include "knowhere/index/vector_index/ivf.h"
@@ -58,6 +56,9 @@ class IVFTest
         index_ = IndexFactory(index_type);
         FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024*1024*200, 1024*1024*300, 2);
     }
+    void TearDown() override {
+        FaissGpuResourceMgr::GetInstance().Free();
+    }
 
  protected:
     std::string index_type;
@@ -369,6 +370,7 @@ class GPURESTEST
     void TearDown() override {
         delete ids;
         delete dis;
+        FaissGpuResourceMgr::GetInstance().Free();
     }
 
  protected:

cpp/src/server/Server.cpp

@@ -19,6 +19,7 @@
 #include <unistd.h>
 #include <string.h>
 #include <src/scheduler/SchedInst.h>
+#include "knowhere/index/vector_index/gpu_ivf.h"
 
 #include "metrics/Metrics.h"
 #include "DBWrapper.h"
@@ -232,6 +233,7 @@ Server::StopService() {
     grpc::GrpcMilvusServer::StopService();
     DBWrapper::GetInstance().StopService();
     engine::StopSchedulerService();
+    knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource.
 }
 
 }

cpp/src/wrapper/knowhere/vec_impl.cpp

@@ -72,8 +72,11 @@ server::KnowhereError VecIndexImpl::Search(const long &nq, const float *xq, floa
         auto k = cfg["k"].as<int>();
         auto dataset = GenDataset(nq, dim, xq);
 
-        Config search_cfg;
-        auto res = index_->Search(dataset, cfg);
+        Config search_cfg = cfg;
+
+        ParameterValidation(type, search_cfg);
+
+        auto res = index_->Search(dataset, search_cfg);
         auto ids_array = res->array()[0];
         auto dis_array = res->array()[1];
 

cpp/src/wrapper/knowhere/vec_index.cpp

@@ -71,7 +71,7 @@ size_t FileIOWriter::operator()(void *ptr, size_t size) {
 }
 
 
-VecIndexPtr GetVecIndexFactory(const IndexType &type, const Config& cfg) {
+VecIndexPtr GetVecIndexFactory(const IndexType &type, const Config &cfg) {
     std::shared_ptr<zilliz::knowhere::VectorIndex> index;
     auto gpu_device = cfg.get_with_default("gpu_id", 0);
     switch (type) {
@@ -235,6 +235,31 @@ void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Co
     }
 }
 
+#if CUDA_VERSION > 9000
+#define GPU_MAX_NRPOBE 2048
+#else
+#define GPU_MAX_NRPOBE 1024
+#endif
+
+void ParameterValidation(const IndexType &type, Config &cfg) {
+    switch (type) {
+        case IndexType::FAISS_IVFSQ8_GPU:
+        case IndexType::FAISS_IVFFLAT_GPU:
+        case IndexType::FAISS_IVFPQ_GPU: {
+            if (cfg.get_with_default("nprobe", 0) != 0) {
+                auto nprobe = cfg["nprobe"].as<int>();
+                if (nprobe > GPU_MAX_NRPOBE) {
+                    WRAPPER_LOG_WARNING << "When search with GPU, nprobe shoud be no more than " << GPU_MAX_NRPOBE << ", but you passed " << nprobe
+                                      << ". Search with " << GPU_MAX_NRPOBE << " instead";
+                    cfg.insert_or_assign("nprobe", GPU_MAX_NRPOBE);
+                }
+            }
+            break;
+        }
+        default:break;
+    }
+}
+
 IndexType ConvertToCpuIndexType(const IndexType &type) {
     // TODO(linxj): add IDMAP
     switch (type) {

cpp/src/wrapper/knowhere/vec_index.h

@@ -14,6 +14,8 @@
 #include "knowhere/common/config.h"
 #include "knowhere/common/binary_set.h"
 
+#include "cuda.h"
+
 
 namespace zilliz {
 namespace milvus {
@@ -90,6 +92,8 @@ extern VecIndexPtr LoadVecIndex(const IndexType &index_type, const zilliz::knowh
 
 extern void AutoGenParams(const IndexType& type, const long& size, Config& cfg);
 
+extern void ParameterValidation(const IndexType& type, Config& cfg);
+
 extern IndexType ConvertToCpuIndexType(const IndexType& type);
 extern IndexType ConvertToGpuIndexType(const IndexType& type);
 

cpp/unittest/knowhere/knowhere_test.cpp

@@ -40,6 +40,9 @@ class KnowhereWrapperTest
 
         index_ = GetVecIndexFactory(index_type);
     }
+    void TearDown() override {
+        zilliz::knowhere::FaissGpuResourceMgr::GetInstance().Free();
+    }
 
     void AssertResult(const std::vector<long> &ids, const std::vector<float> &dis) {
         EXPECT_EQ(ids.size(), nq * k);