提交 502f1a95 编写于 作者: X xj.lin

MS-27 support gpu config


Former-commit-id: 08749b66413000571d733a28303eed3944220a9b
上级 84396d05
...@@ -17,6 +17,7 @@ Please mark all change in change log and use the ticket from JIRA. ...@@ -17,6 +17,7 @@ Please mark all change in change log and use the ticket from JIRA.
- MS-26 - cmake. Add thirdparty packages - MS-26 - cmake. Add thirdparty packages
- MS-31 - cmake: add prometheus - MS-31 - cmake: add prometheus
- MS-33 - cmake: add -j4 to make third party packages build faster - 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 ### Task
......
...@@ -58,10 +58,10 @@ endif() ...@@ -58,10 +58,10 @@ endif()
if(CMAKE_BUILD_TYPE STREQUAL "Release") if(CMAKE_BUILD_TYPE STREQUAL "Release")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
if (GPU_VERSION STREQUAL "ON") # if (GPU_VERSION STREQUAL "ON")
set(ENABLE_LICENSE "ON") # set(ENABLE_LICENSE "ON")
add_definitions("-DENABLE_LICENSE") # add_definitions("-DENABLE_LICENSE")
endif () # endif ()
else() else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
endif() endif()
......
...@@ -6,14 +6,19 @@ ...@@ -6,14 +6,19 @@
#include "mutex" #include "mutex"
#ifdef GPU_VERSION #ifdef GPU_VERSION
#include <faiss/gpu/StandardGpuResources.h> #include <faiss/gpu/StandardGpuResources.h>
#include "faiss/gpu/GpuIndexIVFFlat.h" #include <faiss/gpu/GpuIndexIVFFlat.h>
#include "faiss/gpu/GpuAutoTune.h" #include <faiss/gpu/GpuAutoTune.h>
#endif #endif
#include "faiss/IndexFlat.h"
#include <faiss/IndexFlat.h>
#include <easylogging++.h>
#include "server/ServerConfig.h"
#include "IndexBuilder.h" #include "IndexBuilder.h"
...@@ -21,6 +26,31 @@ namespace zilliz { ...@@ -21,6 +26,31 @@ namespace zilliz {
namespace vecwise { namespace vecwise {
namespace engine { 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; using std::vector;
static std::mutex gpu_resource; static std::mutex gpu_resource;
...@@ -44,7 +74,7 @@ Index_ptr IndexBuilder::build_all(const long &nb, ...@@ -44,7 +74,7 @@ Index_ptr IndexBuilder::build_all(const long &nb,
std::lock_guard<std::mutex> lk(gpu_resource); std::lock_guard<std::mutex> lk(gpu_resource);
faiss::gpu::StandardGpuResources res; 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) { if (!device_index->is_trained) {
nt == 0 || xt == nullptr ? device_index->train(nb, xb) nt == 0 || xt == nullptr ? device_index->train(nb, xb)
: device_index->train(nt, xt); : device_index->train(nt, xt);
......
...@@ -32,14 +32,14 @@ class IndexBuilder { ...@@ -32,14 +32,14 @@ class IndexBuilder {
const long &nt = 0, const long &nt = 0,
const std::vector<float> &xt = std::vector<float>()); const std::vector<float> &xt = std::vector<float>());
void train(const long &nt, //void train(const long &nt,
const std::vector<float> &xt); // const std::vector<float> &xt);
//
Index_ptr add(const long &nb, //Index_ptr add(const long &nb,
const std::vector<float> &xb, // const std::vector<float> &xb,
const std::vector<long> &ids); // const std::vector<long> &ids);
//
void set_build_option(const Operand_ptr &opd); //void set_build_option(const Operand_ptr &opd);
protected: protected:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册