mirror of https://github.com/milvus-io/milvus.git
MS-27 support gpu config
Former-commit-id: 08749b66413000571d733a28303eed3944220a9bpull/191/head
parent
84396d0521
commit
502f1a956c
cpp
|
@ -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()
|
||||||
|
|
||||||
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 @@
|
||||||
|
|
||||||
#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 {
|
||||||
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,
|
||||||
|
|
||||||
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 {
|
||||||
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:
|
||||||
|
|
Loading…
Reference in New Issue