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