Partial change to prepare for GPU index type support (#22591)

Signed-off-by: Yudong Cai <yudong.cai@zilliz.com>
pull/22775/head
Cai Yudong 2023-03-14 23:21:56 +08:00 committed by GitHub
parent b57e476089
commit ab3cbdfc61
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 68 additions and 4 deletions

View File

@ -35,6 +35,13 @@ milvus: build-cpp print-build-info
GO111MODULE=on $(GO) build -ldflags="-r $${RPATH} -X '$(OBJPREFIX).BuildTags=$(BUILD_TAGS)' -X '$(OBJPREFIX).BuildTime=$(BUILD_TIME)' -X '$(OBJPREFIX).GitCommit=$(GIT_COMMIT)' -X '$(OBJPREFIX).GoVersion=$(GO_VERSION)'" \
${AARCH64_FLAG} -o $(INSTALL_PATH)/milvus $(PWD)/cmd/main.go 1>/dev/null
milvus-gpu: build-cpp-gpu print-build-info
@echo "Building Milvus-gpu ..."
@source $(PWD)/scripts/setenv.sh && \
mkdir -p $(INSTALL_PATH) && go env -w CGO_ENABLED="1" && \
GO111MODULE=on $(GO) build -ldflags="-r $${RPATH} -X '$(OBJPREFIX).BuildTags=$(BUILD_TAGS)' -X '$(OBJPREFIX).BuildTime=$(BUILD_TIME)' -X '$(OBJPREFIX).GitCommit=$(GIT_COMMIT)' -X '$(OBJPREFIX).GoVersion=$(GO_VERSION)'" \
${AARCH64_FLAG} -o $(INSTALL_PATH)/milvus-gpu $(PWD)/cmd/main.go 1>/dev/null
get-build-deps:
@(env bash $(PWD)/scripts/install_deps.sh)
@ -155,6 +162,10 @@ build-cpp: download-milvus-proto
@echo "Building Milvus cpp library ..."
@(env bash $(PWD)/scripts/core_build.sh -t ${mode} -f "$(CUSTOM_THIRDPARTY_PATH)" -n ${disk_index})
build-cpp-gpu: download-milvus-proto
@echo "Building Milvus cpp gpu library ..."
@(env bash $(PWD)/scripts/core_build.sh -t ${mode} -g -f "$(CUSTOM_THIRDPARTY_PATH)" -n ${disk_index})
build-cpp-embd: download-milvus-proto
@echo "Building **Embedded** Milvus cpp library ..."
@(env bash $(PWD)/scripts/core_build.sh -b -t ${mode} -f "$(CUSTOM_THIRDPARTY_PATH)" -n ${disk_index})

View File

@ -25,6 +25,10 @@ add_definitions(-DELPP_THREAD_SAFE)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
message( STATUS "Building using CMake version: ${CMAKE_VERSION}" )
if ( MILVUS_GPU_VERSION )
add_definitions(-DMILVUS_GPU_VERSION)
endif ()
project(core)
include(CheckCXXCompilerFlag)
if ( APPLE )

View File

@ -77,6 +77,13 @@ KnowhereSetSimdType(const char* value) {
}
}
void
KnowhereInitGPU(const int32_t gpu_id, const int32_t res_num) {
#ifdef MILVUS_GPU_VERSION
knowhere::KnowhereConfig::InitGPUResource(gpu_id, res_num);
#endif
}
void
KnowhereInitThreadPool(const uint32_t num_threads) {
knowhere::ThreadPool::InitGlobalThreadPool(num_threads);

View File

@ -25,6 +25,9 @@ KnowhereInitImpl(const char*);
std::string
KnowhereSetSimdType(const char*);
void
KnowhereInitGPU(const int32_t, const int32_t);
void
KnowhereInitThreadPool(const uint32_t);

View File

@ -27,3 +27,8 @@ IndexBuilderSetSimdType(const char* value) {
ret[real_type.length()] = 0;
return ret;
}
void
IndexBuilderInitGPU(const int32_t gpu_id, const int32_t res_num) {
milvus::config::KnowhereInitGPU(gpu_id, res_num);
}

View File

@ -22,6 +22,9 @@ IndexBuilderInit(const char*);
char*
IndexBuilderSetSimdType(const char*);
void
IndexBuilderInitGPU(const int32_t, const int32_t);
#ifdef __cplusplus
};
#endif

View File

@ -58,4 +58,9 @@ SegcoreSetSimdType(const char* value) {
return ret;
}
extern "C" void
SegcoreInitGPU(const int32_t gpu_id, const int32_t res_num) {
milvus::config::KnowhereInitGPU(gpu_id, res_num);
}
} // namespace milvus::segcore

View File

@ -31,6 +31,9 @@ SegcoreSetNprobe(const int64_t);
char*
SegcoreSetSimdType(const char*);
void
SegcoreInitGPU(const int32_t, const int32_t);
void
SegcoreSetThreadPoolNum(const uint32_t num_threads);

View File

@ -12,7 +12,7 @@
#-------------------------------------------------------------------------------
# Using default c and cxx compiler in our build tree
# Thirdpart cxx and c flags
add_compile_options(-O3 -fPIC -Wno-error -fopenmp)
append_flags(CMAKE_CXX_FLAGS FLAGS "-O3 -fPIC -Wno-error -fopenmp")
if (NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD)
set(EP_LOG_OPTIONS LOG_CONFIGURE 1 LOG_BUILD 1 LOG_INSTALL 1 LOG_DOWNLOAD 1)

View File

@ -22,6 +22,10 @@ else ()
set(WITH_DISKANN OFF CACHE BOOL "" FORCE )
endif ()
if ( MILVUS_GPU_VERSION STREQUAL "ON" )
set(USE_CUDA ON CACHE BOOL "" FORCE )
endif ()
set( CMAKE_PREFIX_PATH ${CONAN_BOOST_ROOT} )
FetchContent_Declare(
knowhere

View File

@ -32,9 +32,6 @@
constexpr int64_t DIM = 16;
constexpr int64_t NQ = 10;
constexpr int64_t K = 4;
#ifdef MILVUS_GPU_VERSION
int DEVICEID = 0;
#endif
namespace indexcgo = milvus::proto::indexcgo;
namespace schemapb = milvus::proto::schema;

View File

@ -156,6 +156,11 @@ func (i *IndexNode) initKnowhere() {
cCPUNum := C.int(hardware.GetCPUNum())
C.InitCpuNum(cCPUNum)
// init GPU resource
cGpuId := C.int32_t(0)
cResNum := C.int32_t(1)
C.IndexBuilderInitGPU(cGpuId, cResNum)
initcore.InitLocalStorageConfig(Params)
}

View File

@ -232,6 +232,11 @@ func (node *QueryNode) InitSegcore() {
cCPUNum := C.int(hardware.GetCPUNum())
C.InitCpuNum(cCPUNum)
// init GPU resource
cGpuId := C.int32_t(0)
cResNum := C.int32_t(1)
C.SegcoreInitGPU(cGpuId, cResNum)
initcore.InitLocalStorageConfig(Params)
mmapDirPath := paramtable.Get().QueryNodeCfg.MmapDirPath.GetValue()

View File

@ -46,6 +46,12 @@ func (mgr *ConfAdapterMgrImpl) GetAdapter(indexType string) (ConfAdapter, error)
}
func (mgr *ConfAdapterMgrImpl) registerConfAdapter() {
mgr.adapters[IndexFaissGPUIDMap] = newBaseConfAdapter()
mgr.adapters[IndexFaissGPUIvfFlat] = newIVFConfAdapter()
mgr.adapters[IndexFaissGPUIvfPQ] = newIVFPQConfAdapter()
mgr.adapters[IndexFaissGPUIvfSQ8] = newIVFSQConfAdapter()
mgr.adapters[IndexRaftIvfFlat] = newIVFConfAdapter()
mgr.adapters[IndexRaftIvfPQ] = newIVFPQConfAdapter()
mgr.adapters[IndexFaissIDMap] = newBaseConfAdapter()
mgr.adapters[IndexFaissIvfFlat] = newIVFConfAdapter()
mgr.adapters[IndexFaissIvfPQ] = newIVFPQConfAdapter()

View File

@ -16,6 +16,12 @@ type IndexType = string
// IndexType definitions
const (
IndexFaissGPUIDMap IndexType = "GPU_FLAT" // no index is built.
IndexFaissGPUIvfFlat IndexType = "GPU_IVF_FLAT"
IndexFaissGPUIvfPQ IndexType = "GPU_IVF_PQ"
IndexFaissGPUIvfSQ8 IndexType = "GPU_IVF_SQ8"
IndexRaftIvfFlat IndexType = "RAFT_IVF_FLAT"
IndexRaftIvfPQ IndexType = "RAFT_IVF_PQ"
IndexFaissIDMap IndexType = "FLAT" // no index is built.
IndexFaissIvfFlat IndexType = "IVF_FLAT"
IndexFaissIvfPQ IndexType = "IVF_PQ"