From 502f1a956c0a064d3aacb4c2bc7534b5ccd373cc Mon Sep 17 00:00:00 2001
From: "xj.lin" <xiaojun.lin@zilliz.com>
Date: Fri, 31 May 2019 17:37:16 +0800
Subject: [PATCH] MS-27 support gpu config

Former-commit-id: 08749b66413000571d733a28303eed3944220a9b
---
 cpp/CHANGELOG.md                 |  1 +
 cpp/CMakeLists.txt               |  8 +++----
 cpp/src/wrapper/IndexBuilder.cpp | 38 ++++++++++++++++++++++++++++----
 cpp/src/wrapper/IndexBuilder.h   | 16 +++++++-------
 4 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md
index 7ecac566..c52c9e35 100644
--- a/cpp/CHANGELOG.md
+++ b/cpp/CHANGELOG.md
@@ -17,6 +17,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-26 - cmake. Add thirdparty packages
 - MS-31 - cmake: add prometheus
 - MS-33 - cmake: add -j4 to make third party packages build faster
+- MS-27 - support gpu config and disable license build config in cmake
 
 ### Task
 
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index f4e32948..58ba3aea 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -58,10 +58,10 @@ endif()
 
 if(CMAKE_BUILD_TYPE STREQUAL "Release")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
-    if (GPU_VERSION STREQUAL "ON")
-        set(ENABLE_LICENSE "ON")
-        add_definitions("-DENABLE_LICENSE")
-    endif ()
+    # if (GPU_VERSION STREQUAL "ON")
+    #     set(ENABLE_LICENSE "ON")
+    #     add_definitions("-DENABLE_LICENSE")
+    # endif ()
 else()
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
 endif()
diff --git a/cpp/src/wrapper/IndexBuilder.cpp b/cpp/src/wrapper/IndexBuilder.cpp
index 240cf95b..9470a33b 100644
--- a/cpp/src/wrapper/IndexBuilder.cpp
+++ b/cpp/src/wrapper/IndexBuilder.cpp
@@ -6,14 +6,19 @@
 
 #include "mutex"
 
+
 #ifdef GPU_VERSION
 #include <faiss/gpu/StandardGpuResources.h>
-#include "faiss/gpu/GpuIndexIVFFlat.h"
-#include "faiss/gpu/GpuAutoTune.h"
+#include <faiss/gpu/GpuIndexIVFFlat.h>
+#include <faiss/gpu/GpuAutoTune.h>
 #endif
 
-#include "faiss/IndexFlat.h"
 
+#include <faiss/IndexFlat.h>
+#include <easylogging++.h>
+
+
+#include "server/ServerConfig.h"
 #include "IndexBuilder.h"
 
 
@@ -21,6 +26,31 @@ namespace zilliz {
 namespace vecwise {
 namespace engine {
 
+class GpuResources {
+ public:
+    static GpuResources &GetInstance() {
+        static GpuResources instance;
+        return instance;
+    }
+
+    void SelectGpu() {
+        using namespace zilliz::vecwise::server;
+        ServerConfig &config = ServerConfig::GetInstance();
+        ConfigNode server_config = config.GetConfig(CONFIG_SERVER);
+        gpu_num = server_config.GetInt32Value("gpu_index", 0);
+    }
+
+    int32_t GetGpu() {
+        return gpu_num;
+    }
+
+ private:
+    GpuResources() : gpu_num(0) { SelectGpu(); }
+
+ private:
+    int32_t gpu_num;
+};
+
 using std::vector;
 
 static std::mutex gpu_resource;
@@ -44,7 +74,7 @@ Index_ptr IndexBuilder::build_all(const long &nb,
 
         std::lock_guard<std::mutex> lk(gpu_resource);
         faiss::gpu::StandardGpuResources res;
-        auto device_index = faiss::gpu::index_cpu_to_gpu(&res, 0, ori_index);
+        auto device_index = faiss::gpu::index_cpu_to_gpu(&res, GpuResources::GetInstance().GetGpu(), ori_index);
         if (!device_index->is_trained) {
             nt == 0 || xt == nullptr ? device_index->train(nb, xb)
                                      : device_index->train(nt, xt);
diff --git a/cpp/src/wrapper/IndexBuilder.h b/cpp/src/wrapper/IndexBuilder.h
index e4819326..996e3970 100644
--- a/cpp/src/wrapper/IndexBuilder.h
+++ b/cpp/src/wrapper/IndexBuilder.h
@@ -32,14 +32,14 @@ class IndexBuilder {
                                 const long &nt = 0,
                                 const std::vector<float> &xt = std::vector<float>());
 
-    void train(const long &nt,
-               const std::vector<float> &xt);
-
-    Index_ptr add(const long &nb,
-                  const std::vector<float> &xb,
-                  const std::vector<long> &ids);
-
-    void set_build_option(const Operand_ptr &opd);
+    //void train(const long &nt,
+    //           const std::vector<float> &xt);
+    //
+    //Index_ptr add(const long &nb,
+    //              const std::vector<float> &xb,
+    //              const std::vector<long> &ids);
+    //
+    //void set_build_option(const Operand_ptr &opd);
 
 
  protected:
-- 
GitLab