Merge branch 'fix_MS_453' into 'branch-0.4.0'

MS-453 GPU search error when nprobe set more than 1024

See merge request megasearch/milvus!497

Former-commit-id: cd6d47dd9705a751dcf757de29f7c49664c3e65d

cpp/CHANGELOG.md

@@ -28,6 +28,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - 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

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