From 4ca6b59b912770d2832ffd81fc0ca9a2ef157900 Mon Sep 17 00:00:00 2001 From: youny626 Date: Tue, 5 Nov 2019 10:23:59 +0800 Subject: [PATCH] add CPU version --- core/CMakeLists.txt | 242 +++++---- core/build.sh | 18 +- core/cmake/DefineOptions.cmake | 5 + core/conf/server_config.template | 13 +- core/src/CMakeLists.txt | 42 +- core/src/config.h | 3 + core/src/config.h.in | 3 + core/src/db/DB.h | 7 +- core/src/db/DBImpl.cpp | 18 +- core/src/db/DBImpl.h | 9 +- core/src/db/Types.h | 8 +- core/src/db/engine/ExecutionEngineImpl.cpp | 38 +- core/src/index/CMakeLists.txt | 47 +- core/src/index/cmake/DefineOptionsCore.cmake | 11 +- .../index/cmake/ThirdPartyPackagesCore.cmake | 322 ++++++----- core/src/index/knowhere/CMakeLists.txt | 57 +- .../index/vector_index/IndexGPUIDMAP.cpp | 113 ++++ .../index/vector_index/IndexGPUIDMAP.h | 47 ++ .../index/vector_index/IndexIDMAP.cpp | 385 ++++++-------- .../knowhere/index/vector_index/IndexIDMAP.h | 29 - .../knowhere/index/vector_index/IndexIVF.cpp | 24 +- .../index/vector_index/IndexIVFSQ.cpp | 17 +- .../knowhere/index/vector_index/IndexNSG.cpp | 8 +- core/src/index/unittest/CMakeLists.txt | 86 +-- core/src/index/unittest/Helper.h | 32 +- .../unittest/faiss_benchmark/CMakeLists.txt | 64 ++- .../index/unittest/faiss_ori/CMakeLists.txt | 63 ++- core/src/index/unittest/test_idmap.cpp | 6 +- core/src/index/unittest/test_ivf.cpp | 86 +-- .../index/unittest/test_nsg/CMakeLists.txt | 3 +- core/src/index/unittest/test_nsg/test_nsg.cpp | 6 + core/src/main.cpp | 33 +- core/src/metrics/SystemInfo.cpp | 25 +- core/src/scheduler/JobMgr.cpp | 31 +- core/src/scheduler/SchedInst.cpp | 13 +- core/src/scheduler/SchedInst.h | 1 + core/src/scheduler/TaskCreator.cpp | 11 +- core/src/scheduler/Utils.cpp | 4 + .../scheduler/action/PushTaskToNeighbour.cpp | 102 ++-- core/src/scheduler/job/BuildIndexJob.cpp | 5 +- core/src/scheduler/job/SearchJob.cpp | 16 +- core/src/scheduler/job/SearchJob.h | 22 +- core/src/scheduler/optimizer/OnlyGPUPass.cpp | 2 +- core/src/scheduler/task/SearchTask.cpp | 115 ++-- core/src/scheduler/task/SearchTask.h | 5 +- core/src/sdk/grpc/ClientProxy.cpp | 2 +- core/src/server/Config.cpp | 75 +-- core/src/server/Config.h | 4 +- core/src/server/Server.cpp | 2 +- core/src/server/grpc_impl/GrpcRequestTask.cpp | 28 +- core/src/utils/ValidationUtil.cpp | 10 + core/src/wrapper/ConfAdapter.cpp | 2 +- core/src/wrapper/KnowhereResource.cpp | 9 + core/src/wrapper/VecImpl.cpp | 501 +++++++----------- core/src/wrapper/VecImpl.h | 35 -- core/src/wrapper/VecIndex.cpp | 47 +- .../WrapperException.cpp} | 16 +- core/src/wrapper/WrapperException.h | 36 ++ core/src/wrapper/gpu/GPUVecImpl.cpp | 164 ++++++ core/src/wrapper/gpu/GPUVecImpl.h | 66 +++ core/thirdparty/versions.txt | 1 + core/ubuntu_build_deps.sh | 6 + core/unittest/CMakeLists.txt | 20 +- core/unittest/db/CMakeLists.txt | 2 +- core/unittest/db/test_db.cpp | 47 +- core/unittest/db/test_db_mysql.cpp | 22 +- core/unittest/db/test_mem.cpp | 24 +- core/unittest/db/test_search.cpp | 168 +++--- core/unittest/db/utils.cpp | 13 +- core/unittest/metrics/test_metrics.cpp | 5 +- core/unittest/scheduler/CMakeLists.txt | 2 +- core/unittest/server/CMakeLists.txt | 2 +- core/unittest/server/test_config.cpp | 120 ----- core/unittest/server/test_rpc.cpp | 2 +- core/unittest/server/test_util.cpp | 2 + core/unittest/server/utils.cpp | 8 +- core/unittest/wrapper/CMakeLists.txt | 8 +- core/unittest/wrapper/test_wrapper.cpp | 61 ++- core/unittest/wrapper/utils.cpp | 6 + 79 files changed, 2103 insertions(+), 1610 deletions(-) create mode 100644 core/src/config.h create mode 100644 core/src/config.h.in create mode 100644 core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp create mode 100644 core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h rename core/src/{version.h.macro => wrapper/WrapperException.cpp} (74%) create mode 100644 core/src/wrapper/WrapperException.h create mode 100644 core/src/wrapper/gpu/GPUVecImpl.cpp create mode 100644 core/src/wrapper/gpu/GPUVecImpl.h diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt index 402e65fb10..a59e80a6e8 100644 --- a/core/CMakeLists.txt +++ b/core/CMakeLists.txt @@ -23,9 +23,9 @@ message(STATUS "Building using CMake version: ${CMAKE_VERSION}") set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_CURRENT_SOURCE_DIR}/cmake") -MACRO (GET_CURRENT_TIME CURRENT_TIME) +MACRO(GET_CURRENT_TIME CURRENT_TIME) execute_process(COMMAND "date" +"%Y-%m-%d %H:%M.%S" OUTPUT_VARIABLE ${CURRENT_TIME}) -ENDMACRO (GET_CURRENT_TIME) +ENDMACRO(GET_CURRENT_TIME) GET_CURRENT_TIME(BUILD_TIME) string(REGEX REPLACE "\n" "" BUILD_TIME ${BUILD_TIME}) @@ -42,23 +42,20 @@ GET_GIT_BRANCH_NAME(GIT_BRANCH_NAME) message(STATUS "GIT_BRANCH_NAME = ${GIT_BRANCH_NAME}") if(NOT GIT_BRANCH_NAME STREQUAL "") string(REGEX REPLACE "\n" "" GIT_BRANCH_NAME ${GIT_BRANCH_NAME}) -endif() +endif () -set(MILVUS_VERSION "0.5.1") +set(MILVUS_VERSION "${GIT_BRANCH_NAME}") string(REGEX MATCH "[0-9]+\\.[0-9]+\\.[0-9]" MILVUS_VERSION "${MILVUS_VERSION}") -find_package(ClangTools) -set(BUILD_SUPPORT_DIR "${CMAKE_SOURCE_DIR}/build-support") - -if(CMAKE_BUILD_TYPE STREQUAL "Release") +if (CMAKE_BUILD_TYPE STREQUAL "Release") set(BUILD_TYPE "Release") -else() +else () set(BUILD_TYPE "Debug") -endif() +endif () message(STATUS "Build type = ${BUILD_TYPE}") project(milvus VERSION "${MILVUS_VERSION}") -project(milvus_engine LANGUAGES CUDA CXX) +project(milvus_engine LANGUAGES CXX) unset(CMAKE_EXPORT_COMPILE_COMMANDS CACHE) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) @@ -67,15 +64,15 @@ set(MILVUS_VERSION_MAJOR "${milvus_VERSION_MAJOR}") set(MILVUS_VERSION_MINOR "${milvus_VERSION_MINOR}") set(MILVUS_VERSION_PATCH "${milvus_VERSION_PATCH}") -if(MILVUS_VERSION_MAJOR STREQUAL "" +if (MILVUS_VERSION_MAJOR STREQUAL "" OR MILVUS_VERSION_MINOR STREQUAL "" OR MILVUS_VERSION_PATCH STREQUAL "") message(WARNING "Failed to determine Milvus version from git branch name") set(MILVUS_VERSION "0.5.0") -endif() +endif () message(STATUS "Build version = ${MILVUS_VERSION}") -configure_file(${CMAKE_CURRENT_SOURCE_DIR}/src/version.h.macro ${CMAKE_CURRENT_SOURCE_DIR}/src/version.h) +configure_file(${CMAKE_CURRENT_SOURCE_DIR}/src/config.h.in ${CMAKE_CURRENT_SOURCE_DIR}/src/config.h @ONLY) message(STATUS "Milvus version: " "${MILVUS_VERSION_MAJOR}.${MILVUS_VERSION_MINOR}.${MILVUS_VERSION_PATCH} " @@ -84,46 +81,33 @@ message(STATUS "Milvus version: " set(CMAKE_CXX_STANDARD 14) set(CMAKE_CXX_STANDARD_REQUIRED on) -if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)") +if (CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)") message(STATUS "Building milvus_engine on x86 architecture") set(MILVUS_BUILD_ARCH x86_64) -elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)") +elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)") message(STATUS "Building milvus_engine on ppc architecture") set(MILVUS_BUILD_ARCH ppc64le) -else() +else () message(WARNING "Unknown processor type") message(WARNING "CMAKE_SYSTEM_PROCESSOR=${CMAKE_SYSTEM_PROCESSOR}") set(MILVUS_BUILD_ARCH unknown) -endif() - -find_package (Python COMPONENTS Interpreter Development) - -find_package(CUDA) -set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda") - -if(CMAKE_BUILD_TYPE STREQUAL "Release") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp") - set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3") -else() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp") - set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g") -endif() +endif () # Ensure that a default make is set -if("${MAKE}" STREQUAL "") - if(NOT MSVC) +if ("${MAKE}" STREQUAL "") + if (NOT MSVC) find_program(MAKE make) - endif() -endif() + endif () +endif () find_path(MYSQL_INCLUDE_DIR - NAMES "mysql.h" - PATH_SUFFIXES "mysql") + NAMES "mysql.h" + PATH_SUFFIXES "mysql") if (${MYSQL_INCLUDE_DIR} STREQUAL "MYSQL_INCLUDE_DIR-NOTFOUND") message(FATAL_ERROR "Could not found MySQL include directory") -else() +else () include_directories(${MYSQL_INCLUDE_DIR}) -endif() +endif () set(MILVUS_SOURCE_DIR ${PROJECT_SOURCE_DIR}) set(MILVUS_BINARY_DIR ${PROJECT_BINARY_DIR}) @@ -134,26 +118,50 @@ include(DefineOptions) include(BuildUtils) include(ThirdPartyPackages) -config_summary() +set(MILVUS_GPU_VERSION false) +if (MILVUS_CPU_VERSION) + message(STATUS "Building Milvus CPU version") + add_compile_definitions("MILVUS_CPU_VERSION") +else () + message(STATUS "Building Milvus GPU version") + set(MILVUS_GPU_VERSION true) + add_compile_definitions("MILVUS_GPU_VERSION") + enable_language(CUDA) + find_package(CUDA 10 REQUIRED) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda") +endif () + +if (CMAKE_BUILD_TYPE STREQUAL "Release") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp") + if (MILVUS_GPU_VERSION) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3") + endif () +else () + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp") + if (MILVUS_GPU_VERSION) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g") + endif () +endif () if (CUSTOMIZATION) add_definitions(-DCUSTOMIZATION) endif (CUSTOMIZATION) +config_summary() add_subdirectory(src) if (BUILD_UNIT_TEST STREQUAL "ON") if (BUILD_COVERAGE STREQUAL "ON") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage") - endif() + endif () add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unittest) -endif() +endif () add_custom_target(Clean-All COMMAND ${CMAKE_BUILD_TOOL} clean) -if("${MILVUS_DB_PATH}" STREQUAL "") +if ("${MILVUS_DB_PATH}" STREQUAL "") set(MILVUS_DB_PATH "/tmp/milvus") -endif() +endif () configure_file(${CMAKE_CURRENT_SOURCE_DIR}/conf/server_config.template ${CMAKE_CURRENT_SOURCE_DIR}/conf/server_config.yaml) configure_file(${CMAKE_CURRENT_SOURCE_DIR}/conf/log_config.template ${CMAKE_CURRENT_SOURCE_DIR}/conf/log_config.conf) @@ -169,19 +177,22 @@ install(FILES DESTINATION conf) +find_package(Python COMPONENTS Interpreter Development) +find_package(ClangTools) +set(BUILD_SUPPORT_DIR "${CMAKE_SOURCE_DIR}/build-support") # # "make lint" target # -if(NOT MILVUS_VERBOSE_LINT) - set(MILVUS_LINT_QUIET "--quiet") -endif() +if (NOT MILVUS_VERBOSE_LINT) + set(MILVUS_LINT_QUIET "--quiet") +endif () -if(NOT LINT_EXCLUSIONS_FILE) - # source files matching a glob from a line in this file - # will be excluded from linting (cpplint, clang-tidy, clang-format) - set(LINT_EXCLUSIONS_FILE ${BUILD_SUPPORT_DIR}/lint_exclusions.txt) -endif() +if (NOT LINT_EXCLUSIONS_FILE) + # source files matching a glob from a line in this file + # will be excluded from linting (cpplint, clang-tidy, clang-format) + set(LINT_EXCLUSIONS_FILE ${BUILD_SUPPORT_DIR}/lint_exclusions.txt) +endif () find_program(CPPLINT_BIN NAMES cpplint cpplint.py HINTS ${BUILD_SUPPORT_DIR}) message(STATUS "Found cpplint executable at ${CPPLINT_BIN}") @@ -190,77 +201,76 @@ message(STATUS "Found cpplint executable at ${CPPLINT_BIN}") # "make lint" targets # add_custom_target(lint - ${PYTHON_EXECUTABLE} - ${BUILD_SUPPORT_DIR}/run_cpplint.py - --cpplint_binary - ${CPPLINT_BIN} - --exclude_globs - ${LINT_EXCLUSIONS_FILE} - --source_dir - ${CMAKE_CURRENT_SOURCE_DIR} - ${MILVUS_LINT_QUIET}) + ${PYTHON_EXECUTABLE} + ${BUILD_SUPPORT_DIR}/run_cpplint.py + --cpplint_binary + ${CPPLINT_BIN} + --exclude_globs + ${LINT_EXCLUSIONS_FILE} + --source_dir + ${CMAKE_CURRENT_SOURCE_DIR} + ${MILVUS_LINT_QUIET}) # # "make clang-format" and "make check-clang-format" targets # -if(${CLANG_FORMAT_FOUND}) - # runs clang format and updates files in place. - add_custom_target(clang-format - ${PYTHON_EXECUTABLE} - ${BUILD_SUPPORT_DIR}/run_clang_format.py - --clang_format_binary - ${CLANG_FORMAT_BIN} - --exclude_globs - ${LINT_EXCLUSIONS_FILE} - --source_dir - ${CMAKE_CURRENT_SOURCE_DIR}/src - --fix - ${MILVUS_LINT_QUIET}) +if (${CLANG_FORMAT_FOUND}) + # runs clang format and updates files in place. + add_custom_target(clang-format + ${PYTHON_EXECUTABLE} + ${BUILD_SUPPORT_DIR}/run_clang_format.py + --clang_format_binary + ${CLANG_FORMAT_BIN} + --exclude_globs + ${LINT_EXCLUSIONS_FILE} + --source_dir + ${CMAKE_CURRENT_SOURCE_DIR}/src + --fix + ${MILVUS_LINT_QUIET}) - # runs clang format and exits with a non-zero exit code if any files need to be reformatted - add_custom_target(check-clang-format - ${PYTHON_EXECUTABLE} - ${BUILD_SUPPORT_DIR}/run_clang_format.py - --clang_format_binary - ${CLANG_FORMAT_BIN} - --exclude_globs - ${LINT_EXCLUSIONS_FILE} - --source_dir - ${CMAKE_CURRENT_SOURCE_DIR}/src - ${MILVUS_LINT_QUIET}) -endif() + # runs clang format and exits with a non-zero exit code if any files need to be reformatted + add_custom_target(check-clang-format + ${PYTHON_EXECUTABLE} + ${BUILD_SUPPORT_DIR}/run_clang_format.py + --clang_format_binary + ${CLANG_FORMAT_BIN} + --exclude_globs + ${LINT_EXCLUSIONS_FILE} + --source_dir + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${MILVUS_LINT_QUIET}) +endif () # # "make clang-tidy" and "make check-clang-tidy" targets # -if(${CLANG_TIDY_FOUND}) - # runs clang-tidy and attempts to fix any warning automatically - add_custom_target(clang-tidy - ${PYTHON_EXECUTABLE} - ${BUILD_SUPPORT_DIR}/run_clang_tidy.py - --clang_tidy_binary - ${CLANG_TIDY_BIN} - --exclude_globs - ${LINT_EXCLUSIONS_FILE} - --compile_commands - ${CMAKE_BINARY_DIR}/compile_commands.json - --source_dir - ${CMAKE_CURRENT_SOURCE_DIR}/src - --fix - ${MILVUS_LINT_QUIET}) - - # runs clang-tidy and exits with a non-zero exit code if any errors are found. - add_custom_target(check-clang-tidy - ${PYTHON_EXECUTABLE} - ${BUILD_SUPPORT_DIR}/run_clang_tidy.py - --clang_tidy_binary - ${CLANG_TIDY_BIN} - --exclude_globs - ${LINT_EXCLUSIONS_FILE} - --compile_commands - ${CMAKE_BINARY_DIR}/compile_commands.json - --source_dir - ${CMAKE_CURRENT_SOURCE_DIR}/src - ${MILVUS_LINT_QUIET}) -endif() +if (${CLANG_TIDY_FOUND}) + # runs clang-tidy and attempts to fix any warning automatically + add_custom_target(clang-tidy + ${PYTHON_EXECUTABLE} + ${BUILD_SUPPORT_DIR}/run_clang_tidy.py + --clang_tidy_binary + ${CLANG_TIDY_BIN} + --exclude_globs + ${LINT_EXCLUSIONS_FILE} + --compile_commands + ${CMAKE_BINARY_DIR}/compile_commands.json + --source_dir + ${CMAKE_CURRENT_SOURCE_DIR}/src + --fix + ${MILVUS_LINT_QUIET}) + # runs clang-tidy and exits with a non-zero exit code if any errors are found. + add_custom_target(check-clang-tidy + ${PYTHON_EXECUTABLE} + ${BUILD_SUPPORT_DIR}/run_clang_tidy.py + --clang_tidy_binary + ${CLANG_TIDY_BIN} + --exclude_globs + ${LINT_EXCLUSIONS_FILE} + --compile_commands + ${CMAKE_BINARY_DIR}/compile_commands.json + --source_dir + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${MILVUS_LINT_QUIET}) +endif () \ No newline at end of file diff --git a/core/build.sh b/core/build.sh index a662fe497a..819278b94a 100755 --- a/core/build.sh +++ b/core/build.sh @@ -12,6 +12,8 @@ USE_JFROG_CACHE="OFF" RUN_CPPLINT="OFF" CUSTOMIZATION="OFF" # default use ori faiss CUDA_COMPILER=/usr/local/cuda/bin/nvcc +CPU_VERSION="OFF" +WITH_MKL="OFF" CUSTOMIZED_FAISS_URL="${FAISS_URL:-NONE}" wget -q --method HEAD ${CUSTOMIZED_FAISS_URL} @@ -21,7 +23,7 @@ else CUSTOMIZATION="OFF" fi -while getopts "p:d:t:ulrcgjhx" arg +while getopts "p:d:t:ulrcgjhxzm" arg do case $arg in p) @@ -58,6 +60,12 @@ do x) CUSTOMIZATION="OFF" # force use ori faiss ;; + z) + CPU_VERSION="ON" + ;; + m) + WITH_MKL="ON" + ;; h) # help echo " @@ -71,10 +79,12 @@ parameter: -c: code coverage(default: OFF) -g: profiling(default: OFF) -j: use jfrog cache build directory(default: OFF) +-z: build pure CPU version(default: OFF) +-m: build with MKL(default: OFF) -h: help usage: -./build.sh -p \${INSTALL_PREFIX} -t \${BUILD_TYPE} [-u] [-l] [-r] [-c] [-g] [-j] [-h] +./build.sh -p \${INSTALL_PREFIX} -t \${BUILD_TYPE} [-u] [-l] [-r] [-c] [-g] [-j] [-z] [-m] [-h] " exit 0 ;; @@ -106,6 +116,8 @@ CMAKE_CMD="cmake \ -DUSE_JFROG_CACHE=${USE_JFROG_CACHE} \ -DCUSTOMIZATION=${CUSTOMIZATION} \ -DFAISS_URL=${CUSTOMIZED_FAISS_URL} \ +-DMILVUS_CPU_VERSION=${CPU_VERSION} \ +-DBUILD_FAISS_WITH_MKL=${WITH_MKL} \ ../" echo ${CMAKE_CMD} ${CMAKE_CMD} @@ -147,4 +159,4 @@ else # compile and build make -j 8 install || exit 1 -fi \ No newline at end of file +fi diff --git a/core/cmake/DefineOptions.cmake b/core/cmake/DefineOptions.cmake index 167b6e9d66..b90f41fe8e 100644 --- a/core/cmake/DefineOptions.cmake +++ b/core/cmake/DefineOptions.cmake @@ -40,6 +40,11 @@ macro(define_option_string name description default) endif() endmacro() +#---------------------------------------------------------------------- +set_option_category("CPU version") + +define_option(MILVUS_CPU_VERSION "Build CPU version only" ON) + #---------------------------------------------------------------------- set_option_category("Thirdparty") diff --git a/core/conf/server_config.template b/core/conf/server_config.template index 8dfb30f534..07aaeb88e9 100644 --- a/core/conf/server_config.template +++ b/core/conf/server_config.template @@ -27,18 +27,21 @@ metric_config: port: 8080 # port prometheus uses to fetch metrics, must in range [1025, 65534] cache_config: + cpu_cache_capacity: 16 # GB, CPU memory used for cache, must be a positive integer cpu_cache_threshold: 0.85 # percentage of data that will be kept when cache cleanup is triggered, must be in range (0.0, 1.0] - gpu_cache_capacity: 4 # GB, GPU memory used for cache, must be a positive integer - gpu_cache_threshold: 0.85 # percentage of data that will be kept when cache cleanup is triggered, must be in range (0.0, 1.0] cache_insert_data: false # whether to load inserted data into cache, must be a boolean +# Uncomment the following config if you are using GPU version +# gpu_cache_capacity: 4 # GB, GPU memory used for cache, must be a positive integer +# gpu_cache_threshold: 0.85 # percentage of data that will be kept when cache cleanup is triggered, must be in range (0.0, 1.0] + engine_config: use_blas_threshold: 20 # if nq < use_blas_threshold, use SSE, faster with fluctuated response times # if nq >= use_blas_threshold, use OpenBlas, slower with stable response times gpu_search_threshold: 1000 # threshold beyond which the search computation is executed on GPUs only resource_config: - search_resources: # define the GPUs used for search computation, must be in format: gpux - - gpu0 - index_build_device: gpu0 # GPU used for building index, must be in format: gpux \ No newline at end of file + search_resources: # define the CPU / GPUs used for search computation, must be in format: cpu / gpux + - cpu + index_build_device: cpu # CPU / GPU used for building index, must be in format: cpu / gpux diff --git a/core/src/CMakeLists.txt b/core/src/CMakeLists.txt index ae3a458987..9e4065d646 100644 --- a/core/src/CMakeLists.txt +++ b/core/src/CMakeLists.txt @@ -20,11 +20,9 @@ include_directories(${MILVUS_SOURCE_DIR}) include_directories(${MILVUS_ENGINE_SRC}) -include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include) include_directories(${MILVUS_ENGINE_SRC}/grpc/gen-status) include_directories(${MILVUS_ENGINE_SRC}/grpc/gen-milvus) -#this statement must put here, since the INDEX_INCLUDE_DIRS is defined in code/CMakeList.txt add_subdirectory(index) set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE) @@ -109,35 +107,45 @@ set(boost_lib libboost_serialization.a ) -set(cuda_lib - ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so - cudart - cublas - ) - set(third_party_libs sqlite ${client_grpc_lib} yaml-cpp ${prometheus_lib} - ${cuda_lib} mysqlpp zlib ${boost_lib} ) -if (MILVUS_ENABLE_PROFILING STREQUAL "ON") +if (MILVUS_GPU_VERSION) + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(cuda_lib + ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so + cudart + cublas + ) set(third_party_libs ${third_party_libs} - gperftools - libunwind - ) + ${cuda_lib} + ) + aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/gpu wrapper_gpu_files) + set(engine_files ${engine_files} + ${wrapper_gpu_files} + ) +endif () + +if (MILVUS_ENABLE_PROFILING STREQUAL "ON") + set(third_party_libs ${third_party_libs} + gperftools + libunwind + ) endif () -link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") set(engine_libs pthread libgomp.a libgfortran.a + dl ) if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64") @@ -147,11 +155,11 @@ if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64") ) endif () -cuda_add_library(milvus_engine STATIC ${engine_files}) +add_library(milvus_engine STATIC ${engine_files}) target_link_libraries(milvus_engine knowhere - ${engine_libs} ${third_party_libs} + ${engine_libs} ) add_library(metrics STATIC ${metrics_files}) @@ -165,8 +173,6 @@ target_link_libraries(metrics ${metrics_lib}) set(server_libs milvus_engine - pthread - dl metrics ) diff --git a/core/src/config.h b/core/src/config.h new file mode 100644 index 0000000000..b031713328 --- /dev/null +++ b/core/src/config.h @@ -0,0 +1,3 @@ +#define MILVUS_VERSION "0.5.0" +#define BUILD_TYPE "Debug" +#define BUILD_TIME "2019-11-05 10:23.18" diff --git a/core/src/config.h.in b/core/src/config.h.in new file mode 100644 index 0000000000..9e8821881f --- /dev/null +++ b/core/src/config.h.in @@ -0,0 +1,3 @@ +#cmakedefine MILVUS_VERSION "@MILVUS_VERSION@" +#cmakedefine BUILD_TYPE "@BUILD_TYPE@" +#cmakedefine BUILD_TIME @BUILD_TIME@ \ No newline at end of file diff --git a/core/src/db/DB.h b/core/src/db/DB.h index 07fe30babd..a790fadb50 100644 --- a/core/src/db/DB.h +++ b/core/src/db/DB.h @@ -67,16 +67,15 @@ class DB { virtual Status Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - ResultIds& result_ids, ResultDistances& result_distances) = 0; + QueryResults& results) = 0; virtual Status Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) = 0; + const meta::DatesT& dates, QueryResults& results) = 0; virtual Status Query(const std::string& table_id, const std::vector& file_ids, uint64_t k, uint64_t nq, - uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids, - ResultDistances& result_distances) = 0; + uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) = 0; virtual Status Size(uint64_t& result) = 0; diff --git a/core/src/db/DBImpl.cpp b/core/src/db/DBImpl.cpp index fc31846bd3..6995de3d14 100644 --- a/core/src/db/DBImpl.cpp +++ b/core/src/db/DBImpl.cpp @@ -336,20 +336,20 @@ DBImpl::DropIndex(const std::string& table_id) { Status DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - ResultIds& result_ids, ResultDistances& result_distances) { + QueryResults& results) { if (shutting_down_.load(std::memory_order_acquire)) { return Status(DB_ERROR, "Milsvus server is shutdown!"); } meta::DatesT dates = {utils::GetDate()}; - Status result = Query(table_id, k, nq, nprobe, vectors, dates, result_ids, result_distances); + Status result = Query(table_id, k, nq, nprobe, vectors, dates, results); return result; } Status DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) { + const meta::DatesT& dates, QueryResults& results) { if (shutting_down_.load(std::memory_order_acquire)) { return Status(DB_ERROR, "Milsvus server is shutdown!"); } @@ -372,15 +372,14 @@ DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t npr } cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info before query - status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, result_ids, result_distances); + status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, results); cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info after query return status; } Status DBImpl::Query(const std::string& table_id, const std::vector& file_ids, uint64_t k, uint64_t nq, - uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids, - ResultDistances& result_distances) { + uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) { if (shutting_down_.load(std::memory_order_acquire)) { return Status(DB_ERROR, "Milsvus server is shutdown!"); } @@ -414,7 +413,7 @@ DBImpl::Query(const std::string& table_id, const std::vector& file_ } cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info before query - status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, result_ids, result_distances); + status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, results); cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info after query return status; } @@ -433,7 +432,7 @@ DBImpl::Size(uint64_t& result) { /////////////////////////////////////////////////////////////////////////////////////////////////////////////////// Status DBImpl::QueryAsync(const std::string& table_id, const meta::TableFilesSchema& files, uint64_t k, uint64_t nq, - uint64_t nprobe, const float* vectors, ResultIds& result_ids, ResultDistances& result_distances) { + uint64_t nprobe, const float* vectors, QueryResults& results) { server::CollectQueryMetrics metrics(nq); TimeRecorder rc(""); @@ -454,8 +453,7 @@ DBImpl::QueryAsync(const std::string& table_id, const meta::TableFilesSchema& fi } // step 3: construct results - result_ids = job->GetResultIds(); - result_distances = job->GetResultDistances(); + results = job->GetResult(); rc.ElapseFromBegin("Engine query totally cost"); return Status::OK(); diff --git a/core/src/db/DBImpl.h b/core/src/db/DBImpl.h index ad9c574bb1..e1e030cc32 100644 --- a/core/src/db/DBImpl.h +++ b/core/src/db/DBImpl.h @@ -91,16 +91,15 @@ class DBImpl : public DB { Status Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - ResultIds& result_ids, ResultDistances& result_distances) override; + QueryResults& results) override; Status Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors, - const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) override; + const meta::DatesT& dates, QueryResults& results) override; Status Query(const std::string& table_id, const std::vector& file_ids, uint64_t k, uint64_t nq, - uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids, - ResultDistances& result_distances) override; + uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) override; Status Size(uint64_t& result) override; @@ -108,7 +107,7 @@ class DBImpl : public DB { private: Status QueryAsync(const std::string& table_id, const meta::TableFilesSchema& files, uint64_t k, uint64_t nq, - uint64_t nprobe, const float* vectors, ResultIds& result_ids, ResultDistances& result_distances); + uint64_t nprobe, const float* vectors, QueryResults& results); void BackgroundTimerTask(); diff --git a/core/src/db/Types.h b/core/src/db/Types.h index cc2eab0383..94528a9a8a 100644 --- a/core/src/db/Types.h +++ b/core/src/db/Types.h @@ -19,7 +19,6 @@ #include "db/engine/ExecutionEngine.h" -#include #include #include #include @@ -27,13 +26,12 @@ namespace milvus { namespace engine { -using IDNumber = faiss::Index::idx_t; - +typedef int64_t IDNumber; typedef IDNumber* IDNumberPtr; typedef std::vector IDNumbers; -typedef std::vector ResultIds; -typedef std::vector ResultDistances; +typedef std::vector> QueryResult; +typedef std::vector QueryResults; struct TableIndex { int32_t engine_type_ = (int)EngineType::FAISS_IDMAP; diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index ee04191fef..19c699bda7 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -25,6 +25,7 @@ #include "utils/CommonUtil.h" #include "utils/Exception.h" #include "utils/Log.h" + #include "wrapper/ConfAdapter.h" #include "wrapper/ConfAdapterMgr.h" #include "wrapper/VecImpl.h" @@ -92,11 +93,19 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) { break; } case EngineType::FAISS_IVFFLAT: { +#ifdef MILVUS_CPU_VERSION + index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_CPU); +#else index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_MIX); +#endif break; } case EngineType::FAISS_IVFSQ8: { +#ifdef MILVUS_CPU_VERSION + index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_CPU); +#else index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_MIX); +#endif break; } case EngineType::NSG_MIX: { @@ -309,13 +318,30 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) { return Status::OK(); } #endif - try { - index_ = index_->CopyToGpu(device_id); - ENGINE_LOG_DEBUG << "CPU to GPU" << device_id; - } catch (std::exception& e) { - ENGINE_LOG_ERROR << e.what(); - return Status(DB_ERROR, e.what()); + + auto index = std::static_pointer_cast(cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_)); + bool already_in_cache = (index != nullptr); + if (already_in_cache) { + index_ = index; + } else { + if (index_ == nullptr) { + ENGINE_LOG_ERROR << "ExecutionEngineImpl: index is null, failed to copy to gpu"; + return Status(DB_ERROR, "index is null"); + } + + try { + index_ = index_->CopyToGpu(device_id); + ENGINE_LOG_DEBUG << "CPU to GPU" << device_id; + } catch (std::exception& e) { + ENGINE_LOG_ERROR << e.what(); + return Status(DB_ERROR, e.what()); + } } + + if (!already_in_cache) { + GpuCache(device_id); + } + return Status::OK(); } diff --git a/core/src/index/CMakeLists.txt b/core/src/index/CMakeLists.txt index f570752b68..86e5df8eca 100644 --- a/core/src/index/CMakeLists.txt +++ b/core/src/index/CMakeLists.txt @@ -19,12 +19,12 @@ cmake_minimum_required(VERSION 3.14) -message(STATUS "---------------core--------------") +message(STATUS "------------------------------KNOWHERE-----------------------------------") message(STATUS "Building using CMake version: ${CMAKE_VERSION}") -set(KNOWHERE_VERSION "0.1.0") +set(KNOWHERE_VERSION "0.5.0") string(REGEX MATCH "^[0-9]+\\.[0-9]+\\.[0-9]+" KNOWHERE_BASE_VERSION "${KNOWHERE_VERSION}") -project(knowhere VERSION "${KNOWHERE_BASE_VERSION}" LANGUAGES CUDA C CXX) +project(knowhere VERSION "${KNOWHERE_BASE_VERSION}" LANGUAGES C CXX) set(CMAKE_CXX_STANDARD 14) set(KNOWHERE_VERSION_MAJOR "${knowhere_VERSION_MAJOR}") @@ -45,17 +45,6 @@ if(NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE Release) endif(NOT CMAKE_BUILD_TYPE) -if(CMAKE_BUILD_TYPE STREQUAL "Release") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp") - set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3") -else() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp") - set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g") -endif() -MESSAGE(STATUS "CMAKE_CXX_FLAGS" ${CMAKE_CXX_FLAGS}) - -find_package(CUDA) - if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)") message(STATUS "building milvus_engine on x86 architecture") set(KNOWHERE_BUILD_ARCH x86_64) @@ -77,15 +66,39 @@ message(STATUS "Build type = ${BUILD_TYPE}") set(INDEX_SOURCE_DIR ${PROJECT_SOURCE_DIR}) set(INDEX_BINARY_DIR ${PROJECT_BINARY_DIR}) -message(STATUS "Core source dir: ${PROJECT_SOURCE_DIR}") -message(STATUS "Core binary dir: ${PROJECT_BINARY_DIR}") set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${INDEX_SOURCE_DIR}/cmake") include(ExternalProject) include(DefineOptionsCore) include(BuildUtilsCore) + +set(KNOWHERE_GPU_VERSION false) +if (MILVUS_CPU_VERSION OR KNOWHERE_CPU_VERSION) + message(STATUS "Building Knowhere CPU version") + add_compile_definitions("MILVUS_CPU_VERSION") +else () + message(STATUS "Building Knowhere GPU version") + add_compile_definitions("MILVUS_GPU_VERSION") + set(KNOWHERE_GPU_VERSION true) + enable_language(CUDA) + find_package(CUDA 10 REQUIRED) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda") +endif () + include(ThirdPartyPackagesCore) +if (CMAKE_BUILD_TYPE STREQUAL "Release") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp") + if (KNOWHERE_GPU_VERSION) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3") + endif () +else () + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp") + if (KNOWHERE_GPU_VERSION) + set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g") + endif () +endif () + add_subdirectory(knowhere) if (BUILD_COVERAGE STREQUAL "ON") @@ -94,7 +107,7 @@ endif() set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE) -if(BUILD_UNIT_TEST STREQUAL "ON") +if(KNOWHERE_BUILD_TESTS) add_subdirectory(unittest) endif() diff --git a/core/src/index/cmake/DefineOptionsCore.cmake b/core/src/index/cmake/DefineOptionsCore.cmake index cf5e8ea328..ba29fea207 100644 --- a/core/src/index/cmake/DefineOptionsCore.cmake +++ b/core/src/index/cmake/DefineOptionsCore.cmake @@ -40,6 +40,15 @@ macro(define_option_string name description default) endif() endmacro() +#---------------------------------------------------------------------- +set_option_category("CPU version") + +if(MILVUS_CPU_VERSION) + define_option(KNOWHERE_CPU_VERSION "Build CPU version only" ON) +else() + define_option(KNOWHERE_CPU_VERSION "Build CPU version only" OFF) +endif() + #---------------------------------------------------------------------- set_option_category("Thirdparty") @@ -70,7 +79,7 @@ define_option(KNOWHERE_WITH_FAISS "Build with FAISS library" ON) define_option(KNOWHERE_WITH_FAISS_GPU_VERSION "Build with FAISS GPU version" ON) -define_option(KNOWHERE_WITH_OPENBLAS "Build with OpenBLAS library" ON) +define_option(BUILD_FAISS_WITH_MKL "Build FAISS with MKL" OFF) #---------------------------------------------------------------------- if(MSVC) diff --git a/core/src/index/cmake/ThirdPartyPackagesCore.cmake b/core/src/index/cmake/ThirdPartyPackagesCore.cmake index 66ad5e9bbe..54d6ab568b 100644 --- a/core/src/index/cmake/ThirdPartyPackagesCore.cmake +++ b/core/src/index/cmake/ThirdPartyPackagesCore.cmake @@ -26,24 +26,24 @@ set(KNOWHERE_THIRDPARTY_DEPENDENCIES message(STATUS "Using ${KNOWHERE_DEPENDENCY_SOURCE} approach to find dependencies") # For each dependency, set dependency source to global default, if unset -foreach(DEPENDENCY ${KNOWHERE_THIRDPARTY_DEPENDENCIES}) - if("${${DEPENDENCY}_SOURCE}" STREQUAL "") +foreach (DEPENDENCY ${KNOWHERE_THIRDPARTY_DEPENDENCIES}) + if ("${${DEPENDENCY}_SOURCE}" STREQUAL "") set(${DEPENDENCY}_SOURCE ${KNOWHERE_DEPENDENCY_SOURCE}) - endif() -endforeach() + endif () +endforeach () macro(build_dependency DEPENDENCY_NAME) - if("${DEPENDENCY_NAME}" STREQUAL "ARROW") + if ("${DEPENDENCY_NAME}" STREQUAL "ARROW") build_arrow() - elseif("${DEPENDENCY_NAME}" STREQUAL "LAPACK") + elseif ("${DEPENDENCY_NAME}" STREQUAL "LAPACK") build_lapack() elseif ("${DEPENDENCY_NAME}" STREQUAL "GTest") build_gtest() elseif ("${DEPENDENCY_NAME}" STREQUAL "OpenBLAS") build_openblas() - elseif("${DEPENDENCY_NAME}" STREQUAL "FAISS") + elseif ("${DEPENDENCY_NAME}" STREQUAL "FAISS") build_faiss() - else() + else () message(FATAL_ERROR "Unknown thirdparty dependency to build: ${DEPENDENCY_NAME}") endif () endmacro() @@ -51,7 +51,7 @@ endmacro() macro(resolve_dependency DEPENDENCY_NAME) if (${DEPENDENCY_NAME}_SOURCE STREQUAL "AUTO") #message(STATUS "Finding ${DEPENDENCY_NAME} package") - #message(STATUS "${DEPENDENCY_NAME} package not found") + #message(STATUS "${DEPENDENCY_NAME} package not found") build_dependency(${DEPENDENCY_NAME}) elseif (${DEPENDENCY_NAME}_SOURCE STREQUAL "BUNDLED") build_dependency(${DEPENDENCY_NAME}) @@ -64,28 +64,28 @@ endmacro() # Identify OS if (UNIX) if (APPLE) - set (CMAKE_OS_NAME "osx" CACHE STRING "Operating system name" FORCE) + set(CMAKE_OS_NAME "osx" CACHE STRING "Operating system name" FORCE) else (APPLE) ## Check for Debian GNU/Linux ________________ - find_file (DEBIAN_FOUND debian_version debconf.conf + find_file(DEBIAN_FOUND debian_version debconf.conf PATHS /etc ) if (DEBIAN_FOUND) - set (CMAKE_OS_NAME "debian" CACHE STRING "Operating system name" FORCE) + set(CMAKE_OS_NAME "debian" CACHE STRING "Operating system name" FORCE) endif (DEBIAN_FOUND) ## Check for Fedora _________________________ - find_file (FEDORA_FOUND fedora-release + find_file(FEDORA_FOUND fedora-release PATHS /etc ) if (FEDORA_FOUND) - set (CMAKE_OS_NAME "fedora" CACHE STRING "Operating system name" FORCE) + set(CMAKE_OS_NAME "fedora" CACHE STRING "Operating system name" FORCE) endif (FEDORA_FOUND) ## Check for RedHat _________________________ - find_file (REDHAT_FOUND redhat-release inittab.RH + find_file(REDHAT_FOUND redhat-release inittab.RH PATHS /etc ) if (REDHAT_FOUND) - set (CMAKE_OS_NAME "redhat" CACHE STRING "Operating system name" FORCE) + set(CMAKE_OS_NAME "redhat" CACHE STRING "Operating system name" FORCE) endif (REDHAT_FOUND) ## Extra check for Ubuntu ____________________ if (DEBIAN_FOUND) @@ -94,18 +94,18 @@ if (UNIX) ## a first superficial inspection a system will ## be considered as Debian, which signifies an ## extra check is required. - find_file (UBUNTU_EXTRA legal issue + find_file(UBUNTU_EXTRA legal issue PATHS /etc ) if (UBUNTU_EXTRA) ## Scan contents of file - file (STRINGS ${UBUNTU_EXTRA} UBUNTU_FOUND + file(STRINGS ${UBUNTU_EXTRA} UBUNTU_FOUND REGEX Ubuntu ) ## Check result of string search if (UBUNTU_FOUND) - set (CMAKE_OS_NAME "ubuntu" CACHE STRING "Operating system name" FORCE) - set (DEBIAN_FOUND FALSE) + set(CMAKE_OS_NAME "ubuntu" CACHE STRING "Operating system name" FORCE) + set(DEBIAN_FOUND FALSE) endif (UBUNTU_FOUND) endif (UBUNTU_EXTRA) endif (DEBIAN_FOUND) @@ -119,17 +119,17 @@ set(THIRDPARTY_DIR "${INDEX_SOURCE_DIR}/thirdparty") # ---------------------------------------------------------------------- # JFrog -if(NOT DEFINED USE_JFROG_CACHE) +if (NOT DEFINED USE_JFROG_CACHE) set(USE_JFROG_CACHE "OFF") -endif() -if(USE_JFROG_CACHE STREQUAL "ON") +endif () +if (USE_JFROG_CACHE STREQUAL "ON") set(JFROG_ARTFACTORY_CACHE_URL "${JFROG_ARTFACTORY_URL}/milvus/thirdparty/cache/${CMAKE_OS_NAME}/${KNOWHERE_BUILD_ARCH}/${BUILD_TYPE}") set(THIRDPARTY_PACKAGE_CACHE "${THIRDPARTY_DIR}/cache") - if(NOT EXISTS ${THIRDPARTY_PACKAGE_CACHE}) + if (NOT EXISTS ${THIRDPARTY_PACKAGE_CACHE}) message(STATUS "Will create cached directory: ${THIRDPARTY_PACKAGE_CACHE}") file(MAKE_DIRECTORY ${THIRDPARTY_PACKAGE_CACHE}) - endif() -endif() + endif () +endif () macro(resolve_dependency DEPENDENCY_NAME) if (${DEPENDENCY_NAME}_SOURCE STREQUAL "AUTO") @@ -150,11 +150,11 @@ string(TOUPPER ${CMAKE_BUILD_TYPE} UPPERCASE_BUILD_TYPE) set(EP_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}}") set(EP_C_FLAGS "${CMAKE_C_FLAGS} ${CMAKE_C_FLAGS_${UPPERCASE_BUILD_TYPE}}") -if(NOT MSVC) +if (NOT MSVC) # Set -fPIC on all external projects set(EP_CXX_FLAGS "${EP_CXX_FLAGS} -fPIC") set(EP_C_FLAGS "${EP_C_FLAGS} -fPIC") -endif() +endif () # CC/CXX environment variables are captured on the first invocation of the # builder (e.g make or ninja) instead of when CMake is invoked into to build @@ -164,13 +164,13 @@ endif() set(EP_COMMON_TOOLCHAIN -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER} -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}) -if(CMAKE_AR) +if (CMAKE_AR) set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_AR=${CMAKE_AR}) -endif() +endif () -if(CMAKE_RANLIB) +if (CMAKE_RANLIB) set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_RANLIB=${CMAKE_RANLIB}) -endif() +endif () # External projects are still able to override the following declarations. # cmake command line will favor the last defined variable when a duplicate is @@ -184,18 +184,18 @@ set(EP_COMMON_CMAKE_ARGS -DCMAKE_CXX_FLAGS=${EP_CXX_FLAGS} -DCMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}=${EP_CXX_FLAGS}) -if(NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD) +if (NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD) set(EP_LOG_OPTIONS LOG_CONFIGURE 1 LOG_BUILD 1 LOG_INSTALL 1 LOG_DOWNLOAD 1) -else() +else () set(EP_LOG_OPTIONS) -endif() +endif () # Ensure that a default make is set -if("${MAKE}" STREQUAL "") - if(NOT MSVC) +if ("${MAKE}" STREQUAL "") + if (NOT MSVC) find_program(MAKE make) - endif() -endif() + endif () +endif () set(MAKE_BUILD_ARGS "-j8") @@ -212,32 +212,32 @@ find_package(Threads REQUIRED) # Read toolchain versions from cpp/thirdparty/versions.txt file(STRINGS "${THIRDPARTY_DIR}/versions.txt" TOOLCHAIN_VERSIONS_TXT) -foreach(_VERSION_ENTRY ${TOOLCHAIN_VERSIONS_TXT}) +foreach (_VERSION_ENTRY ${TOOLCHAIN_VERSIONS_TXT}) # Exclude comments - if(NOT _VERSION_ENTRY MATCHES "^[^#][A-Za-z0-9-_]+_VERSION=") + if (NOT _VERSION_ENTRY MATCHES "^[^#][A-Za-z0-9-_]+_VERSION=") continue() - endif() + endif () string(REGEX MATCH "^[^=]*" _LIB_NAME ${_VERSION_ENTRY}) string(REPLACE "${_LIB_NAME}=" "" _LIB_VERSION ${_VERSION_ENTRY}) # Skip blank or malformed lines - if(${_LIB_VERSION} STREQUAL "") + if (${_LIB_VERSION} STREQUAL "") continue() - endif() + endif () # For debugging #message(STATUS "${_LIB_NAME}: ${_LIB_VERSION}") set(${_LIB_NAME} "${_LIB_VERSION}") -endforeach() +endforeach () -if(CUSTOMIZATION) +if (CUSTOMIZATION) execute_process(COMMAND wget -q --method HEAD ${FAISS_URL} RESULT_VARIABLE return_code) message(STATUS "Check the remote cache file ${FAISS_URL}. return code = ${return_code}") if (NOT return_code EQUAL 0) MESSAGE(FATAL_ERROR "Can't access to ${FAISS_URL}") - else() + else () set(FAISS_SOURCE_URL ${FAISS_URL}) # set(FAISS_MD5 "a589663865a8558205533c8ac414278c") # set(FAISS_MD5 "57da9c4f599cc8fa4260488b1c96e1cc") # commit-id 6dbdf75987c34a2c853bd172ea0d384feea8358c branch-0.2.0 @@ -246,36 +246,35 @@ if(CUSTOMIZATION) # set(FAISS_MD5 "c89ea8e655f5cdf58f42486f13614714") # commit-id 9c28a1cbb88f41fa03b03d7204106201ad33276b branch-0.2.1 # set(FAISS_MD5 "87fdd86351ffcaf3f80dc26ade63c44b") # commit-id 841a156e67e8e22cd8088e1b58c00afbf2efc30b branch-0.2.1 # set(FAISS_MD5 "f3b2ce3364c3fa7febd3aa7fdd0fe380") # commit-id 694e03458e6b69ce8a62502f71f69a614af5af8f branch-0.3.0 - # set(FAISS_MD5 "bb30722c22390ce5f6759ccb216c1b2a") # commit-id d324db297475286afe107847c7fb7a0f9dc7e90e branch-0.3.0 - set(FAISS_MD5 "2293cdb209c3718e3b19f3edae8b32b3") # commit-id a13c1205dc52977a9ad3b33a14efa958604a8bff branch-0.3.0 - endif() -else() + set(FAISS_MD5 "bb30722c22390ce5f6759ccb216c1b2a") # commit-id d324db297475286afe107847c7fb7a0f9dc7e90e branch-0.3.0 + endif () +else () set(FAISS_SOURCE_URL "https://github.com/milvus-io/faiss/archive/1.6.0.tar.gz") set(FAISS_MD5 "eb96d84f98b078a9eec04a796f5c792e") -endif() +endif () message(STATUS "FAISS URL = ${FAISS_SOURCE_URL}") -if(DEFINED ENV{KNOWHERE_ARROW_URL}) +if (DEFINED ENV{KNOWHERE_ARROW_URL}) set(ARROW_SOURCE_URL "$ENV{KNOWHERE_ARROW_URL}") -else() +else () set(ARROW_SOURCE_URL "https://github.com/apache/arrow.git" ) -endif() +endif () if (DEFINED ENV{KNOWHERE_GTEST_URL}) set(GTEST_SOURCE_URL "$ENV{KNOWHERE_GTEST_URL}") else () set(GTEST_SOURCE_URL "https://github.com/google/googletest/archive/release-${GTEST_VERSION}.tar.gz") -endif() +endif () set(GTEST_MD5 "2e6fbeb6a91310a16efe181886c59596") -if(DEFINED ENV{KNOWHERE_LAPACK_URL}) +if (DEFINED ENV{KNOWHERE_LAPACK_URL}) set(LAPACK_SOURCE_URL "$ENV{KNOWHERE_LAPACK_URL}") -else() +else () set(LAPACK_SOURCE_URL "https://github.com/Reference-LAPACK/lapack/archive/${LAPACK_VERSION}.tar.gz") -endif() +endif () set(LAPACK_MD5 "96591affdbf58c450d45c1daa540dbd2") if (DEFINED ENV{KNOWHERE_OPENBLAS_URL}) @@ -283,7 +282,7 @@ if (DEFINED ENV{KNOWHERE_OPENBLAS_URL}) else () set(OPENBLAS_SOURCE_URL "https://github.com/xianyi/OpenBLAS/archive/${OPENBLAS_VERSION}.tar.gz") -endif() +endif () set(OPENBLAS_MD5 "8a110a25b819a4b94e8a9580702b6495") # ---------------------------------------------------------------------- @@ -293,10 +292,10 @@ set(ARROW_PREFIX "${INDEX_BINARY_DIR}/arrow_ep-prefix/src/arrow_ep/cpp") macro(build_arrow) message(STATUS "Building Apache ARROW-${ARROW_VERSION} from source") set(ARROW_STATIC_LIB_NAME arrow) - set(ARROW_STATIC_LIB + set(ARROW_STATIC_LIB "${ARROW_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}${ARROW_STATIC_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX}" ) - set(ARROW_INCLUDE_DIR "${ARROW_PREFIX}/include") + set(ARROW_INCLUDE_DIR "${ARROW_PREFIX}/include") set(ARROW_CMAKE_ARGS ${EP_COMMON_CMAKE_ARGS} @@ -326,10 +325,10 @@ macro(build_arrow) -DBOOST_SOURCE=AUTO #try to find BOOST in the system default locations and build from source if not found ) - - if(USE_JFROG_CACHE STREQUAL "ON") + + if (USE_JFROG_CACHE STREQUAL "ON") execute_process(COMMAND sh -c "git ls-remote --heads --tags ${ARROW_SOURCE_URL} ${ARROW_VERSION} | cut -f 1" OUTPUT_VARIABLE ARROW_LAST_COMMIT_ID) - if(${ARROW_LAST_COMMIT_ID} MATCHES "^[^#][a-z0-9]+") + if (${ARROW_LAST_COMMIT_ID} MATCHES "^[^#][a-z0-9]+") string(MD5 ARROW_COMBINE_MD5 "${ARROW_LAST_COMMIT_ID}") set(ARROW_CACHE_PACKAGE_NAME "arrow_${ARROW_COMBINE_MD5}.tar.gz") set(ARROW_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${ARROW_CACHE_PACKAGE_NAME}") @@ -359,18 +358,18 @@ macro(build_arrow) ) ExternalProject_Create_Cache(arrow_ep ${ARROW_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/arrow_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${ARROW_CACHE_URL}) - else() + else () file(DOWNLOAD ${ARROW_CACHE_URL} ${ARROW_CACHE_PACKAGE_PATH} STATUS status) list(GET status 0 status_code) message(STATUS "DOWNLOADING FROM ${ARROW_CACHE_URL} TO ${ARROW_CACHE_PACKAGE_PATH}. STATUS = ${status_code}") if (status_code EQUAL 0) ExternalProject_Use_Cache(arrow_ep ${ARROW_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR}) - endif() - endif() - else() + endif () + endif () + else () message(FATAL_ERROR "The last commit ID of \"${ARROW_SOURCE_URL}\" repository don't match!") - endif() - else() + endif () + else () externalproject_add(arrow_ep GIT_REPOSITORY ${ARROW_SOURCE_URL} @@ -390,14 +389,14 @@ macro(build_arrow) BUILD_BYPRODUCTS "${ARROW_STATIC_LIB}" ) - endif() + endif () file(MAKE_DIRECTORY "${ARROW_PREFIX}/include") add_library(arrow STATIC IMPORTED) set_target_properties(arrow PROPERTIES IMPORTED_LOCATION "${ARROW_STATIC_LIB}" INTERFACE_INCLUDE_DIRECTORIES "${ARROW_INCLUDE_DIR}") - add_dependencies(arrow arrow_ep) + add_dependencies(arrow arrow_ep) set(JEMALLOC_PREFIX "${INDEX_BINARY_DIR}/arrow_ep-prefix/src/arrow_ep-build/jemalloc_ep-prefix/src/jemalloc_ep") @@ -408,13 +407,13 @@ macro(build_arrow) endmacro() -if(KNOWHERE_WITH_ARROW AND NOT TARGET arrow_ep) +if (KNOWHERE_WITH_ARROW AND NOT TARGET arrow_ep) resolve_dependency(ARROW) link_directories(SYSTEM ${ARROW_PREFIX}/lib/) include_directories(SYSTEM ${ARROW_INCLUDE_DIR}) -endif() +endif () # ---------------------------------------------------------------------- # OpenBLAS @@ -428,7 +427,7 @@ macro(build_openblas) set(OPENBLAS_REAL_STATIC_LIB "${OPENBLAS_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}openblas_haswellp-r0.3.6${CMAKE_STATIC_LIBRARY_SUFFIX}") - if(USE_JFROG_CACHE STREQUAL "ON") + if (USE_JFROG_CACHE STREQUAL "ON") set(OPENBLAS_CACHE_PACKAGE_NAME "openblas_${OPENBLAS_MD5}.tar.gz") set(OPENBLAS_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${OPENBLAS_CACHE_PACKAGE_NAME}") set(OPENBLAS_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${OPENBLAS_CACHE_PACKAGE_NAME}") @@ -455,15 +454,15 @@ macro(build_openblas) ${OPENBLAS_STATIC_LIB}) ExternalProject_Create_Cache(openblas_ep ${OPENBLAS_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/openblas_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${OPENBLAS_CACHE_URL}) - else() + else () file(DOWNLOAD ${OPENBLAS_CACHE_URL} ${OPENBLAS_CACHE_PACKAGE_PATH} STATUS status) list(GET status 0 status_code) message(STATUS "DOWNLOADING FROM ${OPENBLAS_CACHE_URL} TO ${OPENBLAS_CACHE_PACKAGE_PATH}. STATUS = ${status_code}") if (status_code EQUAL 0) ExternalProject_Use_Cache(openblas_ep ${OPENBLAS_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR}) - endif() - endif() - else() + endif () + endif () + else () externalproject_add(openblas_ep URL ${OPENBLAS_SOURCE_URL} @@ -481,7 +480,7 @@ macro(build_openblas) install BUILD_BYPRODUCTS ${OPENBLAS_STATIC_LIB}) - endif() + endif () file(MAKE_DIRECTORY "${OPENBLAS_INCLUDE_DIR}") add_library(openblas STATIC IMPORTED) @@ -510,7 +509,7 @@ macro(build_lapack) "-DCMAKE_INSTALL_PREFIX=${LAPACK_PREFIX}" -DCMAKE_INSTALL_LIBDIR=lib) - if(USE_JFROG_CACHE STREQUAL "ON") + if (USE_JFROG_CACHE STREQUAL "ON") set(LAPACK_CACHE_PACKAGE_NAME "lapack_${LAPACK_MD5}.tar.gz") set(LAPACK_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${LAPACK_CACHE_PACKAGE_NAME}") set(LAPACK_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${LAPACK_CACHE_PACKAGE_NAME}") @@ -531,15 +530,15 @@ macro(build_lapack) ${LAPACK_STATIC_LIB}) ExternalProject_Create_Cache(lapack_ep ${LAPACK_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/lapack_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${LAPACK_CACHE_URL}) - else() + else () file(DOWNLOAD ${LAPACK_CACHE_URL} ${LAPACK_CACHE_PACKAGE_PATH} STATUS status) list(GET status 0 status_code) message(STATUS "DOWNLOADING FROM ${LAPACK_CACHE_URL} TO ${LAPACK_CACHE_PACKAGE_PATH}. STATUS = ${status_code}") if (status_code EQUAL 0) ExternalProject_Use_Cache(lapack_ep ${LAPACK_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR}) - endif() - endif() - else() + endif () + endif () + else () externalproject_add(lapack_ep URL ${LAPACK_SOURCE_URL} @@ -551,7 +550,7 @@ macro(build_lapack) ${MAKE_BUILD_ARGS} BUILD_BYPRODUCTS ${LAPACK_STATIC_LIB}) - endif() + endif () file(MAKE_DIRECTORY "${LAPACK_INCLUDE_DIR}") add_library(lapack STATIC IMPORTED) @@ -571,13 +570,13 @@ macro(build_gtest) set(GTEST_VENDORED TRUE) set(GTEST_CMAKE_CXX_FLAGS "${EP_CXX_FLAGS}") - if(APPLE) + if (APPLE) set(GTEST_CMAKE_CXX_FLAGS ${GTEST_CMAKE_CXX_FLAGS} -DGTEST_USE_OWN_TR1_TUPLE=1 -Wno-unused-value -Wno-ignored-attributes) - endif() + endif () set(GTEST_PREFIX "${INDEX_BINARY_DIR}/googletest_ep-prefix/src/googletest_ep") set(GTEST_INCLUDE_DIR "${GTEST_PREFIX}/include") @@ -596,10 +595,10 @@ macro(build_gtest) set(GMOCK_INCLUDE_DIR "${GTEST_PREFIX}/include") set(GMOCK_STATIC_LIB "${GTEST_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}gmock${CMAKE_STATIC_LIBRARY_SUFFIX}" - ) + ) - if(USE_JFROG_CACHE STREQUAL "ON") + if (USE_JFROG_CACHE STREQUAL "ON") set(GTEST_CACHE_PACKAGE_NAME "googletest_${GTEST_MD5}.tar.gz") set(GTEST_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${GTEST_CACHE_PACKAGE_NAME}") set(GTEST_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${GTEST_CACHE_PACKAGE_NAME}") @@ -622,15 +621,15 @@ macro(build_gtest) ${EP_LOG_OPTIONS}) ExternalProject_Create_Cache(googletest_ep ${GTEST_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/googletest_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${GTEST_CACHE_URL}) - else() + else () file(DOWNLOAD ${GTEST_CACHE_URL} ${GTEST_CACHE_PACKAGE_PATH} STATUS status) list(GET status 0 status_code) message(STATUS "DOWNLOADING FROM ${GTEST_CACHE_URL} TO ${GTEST_CACHE_PACKAGE_PATH}. STATUS = ${status_code}") if (status_code EQUAL 0) ExternalProject_Use_Cache(googletest_ep ${GTEST_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR}) - endif() - endif() - else() + endif () + endif () + else () ExternalProject_Add(googletest_ep URL ${GTEST_SOURCE_URL} @@ -644,20 +643,20 @@ macro(build_gtest) CMAKE_ARGS ${GTEST_CMAKE_ARGS} ${EP_LOG_OPTIONS}) - endif() + endif () # The include directory must exist before it is referenced by a target. file(MAKE_DIRECTORY "${GTEST_INCLUDE_DIR}") add_library(gtest STATIC IMPORTED) set_target_properties(gtest - PROPERTIES IMPORTED_LOCATION "${GTEST_STATIC_LIB}" - INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}") + PROPERTIES IMPORTED_LOCATION "${GTEST_STATIC_LIB}" + INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}") add_library(gtest_main STATIC IMPORTED) set_target_properties(gtest_main - PROPERTIES IMPORTED_LOCATION "${GTEST_MAIN_STATIC_LIB}" - INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}") + PROPERTIES IMPORTED_LOCATION "${GTEST_MAIN_STATIC_LIB}" + INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}") add_library(gmock STATIC IMPORTED) set_target_properties(gmock @@ -673,44 +672,88 @@ endmacro() if (KNOWHERE_BUILD_TESTS AND NOT TARGET googletest_ep) resolve_dependency(GTest) - if(NOT GTEST_VENDORED) - endif() + if (NOT GTEST_VENDORED) + endif () # TODO: Don't use global includes but rather target_include_directories get_target_property(GTEST_INCLUDE_DIR gtest INTERFACE_INCLUDE_DIRECTORIES) link_directories(SYSTEM "${GTEST_PREFIX}/lib") include_directories(SYSTEM ${GTEST_INCLUDE_DIR}) -endif() +endif () # ---------------------------------------------------------------------- # FAISS macro(build_faiss) message(STATUS "Building FAISS-${FAISS_VERSION} from source") + + if (NOT DEFINED BUILD_FAISS_WITH_MKL) + set(BUILD_FAISS_WITH_MKL OFF) + endif () + + if (EXISTS "/proc/cpuinfo") + FILE(READ /proc/cpuinfo PROC_CPUINFO) + + SET(VENDOR_ID_RX "vendor_id[ \t]*:[ \t]*([a-zA-Z]+)\n") + STRING(REGEX MATCH "${VENDOR_ID_RX}" VENDOR_ID "${PROC_CPUINFO}") + STRING(REGEX REPLACE "${VENDOR_ID_RX}" "\\1" VENDOR_ID "${VENDOR_ID}") + + if (NOT ${VENDOR_ID} STREQUAL "GenuineIntel") + set(BUILD_FAISS_WITH_MKL OFF) + endif () + endif () + set(FAISS_PREFIX "${INDEX_BINARY_DIR}/faiss_ep-prefix/src/faiss_ep") set(FAISS_INCLUDE_DIR "${FAISS_PREFIX}/include") set(FAISS_STATIC_LIB "${FAISS_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}faiss${CMAKE_STATIC_LIBRARY_SUFFIX}") - set(FAISS_CONFIGURE_ARGS "--prefix=${FAISS_PREFIX}" "CFLAGS=${EP_C_FLAGS}" "CXXFLAGS=${EP_CXX_FLAGS}" - "LDFLAGS=-L${OPENBLAS_PREFIX}/lib -L${LAPACK_PREFIX}/lib -lopenblas -llapack" --without-python) - - if(${KNOWHERE_WITH_FAISS_GPU_VERSION} STREQUAL "ON") + set(FAISS_CFLAGS ${EP_C_FLAGS}) + set(FAISS_CXXFLAGS ${EP_CXX_FLAGS}) + + if (BUILD_FAISS_WITH_MKL) + + find_path(MKL_LIB_PATH + NAMES "libmkl_intel_ilp64.a" "libmkl_gnu_thread.a" "libmkl_core.a" + PATH_SUFFIXES "intel/compilers_and_libraries_${MKL_VERSION}/linux/mkl/lib/intel64/") + if (${MKL_LIB_PATH} STREQUAL "MKL_LIB_PATH-NOTFOUND") + message(FATAL_ERROR "Could not find MKL libraries") + endif () + message(STATUS "Build Faiss with MKL. MKL lib path = ${MKL_LIB_PATH}") + + set(MKL_LIBS + ${MKL_LIB_PATH}/libmkl_intel_ilp64.a + ${MKL_LIB_PATH}/libmkl_gnu_thread.a + ${MKL_LIB_PATH}/libmkl_core.a + ) + + set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS} + "CPPFLAGS=-DFINTEGER=long -DMKL_ILP64 -m64 -I${MKL_LIB_PATH}/../../include" + "LDFLAGS=-L${MKL_LIB_PATH}" + ) + + else () + message(STATUS "Build Faiss with OpenBlas/LAPACK") + set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS} + "LDFLAGS=-L${OPENBLAS_PREFIX}/lib -L${LAPACK_PREFIX}/lib") + endif () + + if (KNOWHERE_GPU_VERSION) set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS} "--with-cuda=${CUDA_TOOLKIT_ROOT_DIR}" "--with-cuda-arch=-gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75" ) - else() + else () set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS} --without-cuda) - endif() + endif () - if(USE_JFROG_CACHE STREQUAL "ON") + if (USE_JFROG_CACHE STREQUAL "ON") string(MD5 FAISS_COMBINE_MD5 "${FAISS_MD5}${LAPACK_MD5}${OPENBLAS_MD5}") set(FAISS_CACHE_PACKAGE_NAME "faiss_${FAISS_COMBINE_MD5}.tar.gz") set(FAISS_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${FAISS_CACHE_PACKAGE_NAME}") @@ -735,18 +778,20 @@ macro(build_faiss) BUILD_BYPRODUCTS ${FAISS_STATIC_LIB}) - ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep) + if (NOT BUILD_FAISS_WITH_MKL) + ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep) + endif () ExternalProject_Create_Cache(faiss_ep ${FAISS_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/faiss_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${FAISS_CACHE_URL}) - else() + else () file(DOWNLOAD ${FAISS_CACHE_URL} ${FAISS_CACHE_PACKAGE_PATH} STATUS status) list(GET status 0 status_code) message(STATUS "DOWNLOADING FROM ${FAISS_CACHE_URL} TO ${FAISS_CACHE_PACKAGE_PATH}. STATUS = ${status_code}") if (status_code EQUAL 0) ExternalProject_Use_Cache(faiss_ep ${FAISS_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR}) - endif() - endif() - else() + endif () + endif () + else () externalproject_add(faiss_ep URL ${FAISS_SOURCE_URL} @@ -763,35 +808,54 @@ macro(build_faiss) BUILD_BYPRODUCTS ${FAISS_STATIC_LIB}) - ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep) - endif() + if (NOT BUILD_FAISS_WITH_MKL) + ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep) + endif () + + endif () file(MAKE_DIRECTORY "${FAISS_INCLUDE_DIR}") add_library(faiss STATIC IMPORTED) + set_target_properties( faiss - PROPERTIES IMPORTED_LOCATION "${FAISS_STATIC_LIB}" + PROPERTIES + IMPORTED_LOCATION "${FAISS_STATIC_LIB}" INTERFACE_INCLUDE_DIRECTORIES "${FAISS_INCLUDE_DIR}" - INTERFACE_LINK_LIBRARIES "openblas;lapack" ) + ) + if (BUILD_FAISS_WITH_MKL) + set_target_properties( + faiss + PROPERTIES + INTERFACE_LINK_LIBRARIES "${MKL_LIBS}") + else () + set_target_properties( + faiss + PROPERTIES + INTERFACE_LINK_LIBRARIES "openblas;lapack") + endif () + add_dependencies(faiss faiss_ep) endmacro() -if(KNOWHERE_WITH_FAISS AND NOT TARGET faiss_ep) +if (KNOWHERE_WITH_FAISS AND NOT TARGET faiss_ep) - resolve_dependency(OpenBLAS) - get_target_property(OPENBLAS_INCLUDE_DIR openblas INTERFACE_INCLUDE_DIRECTORIES) - include_directories(SYSTEM "${OPENBLAS_INCLUDE_DIR}") - link_directories(SYSTEM ${OPENBLAS_PREFIX}/lib) + if (NOT BUILD_FAISS_WITH_MKL) + resolve_dependency(OpenBLAS) + get_target_property(OPENBLAS_INCLUDE_DIR openblas INTERFACE_INCLUDE_DIRECTORIES) + include_directories(SYSTEM "${OPENBLAS_INCLUDE_DIR}") + link_directories(SYSTEM ${OPENBLAS_PREFIX}/lib) - resolve_dependency(LAPACK) - get_target_property(LAPACK_INCLUDE_DIR lapack INTERFACE_INCLUDE_DIRECTORIES) - include_directories(SYSTEM "${LAPACK_INCLUDE_DIR}") - link_directories(SYSTEM "${LAPACK_PREFIX}/lib") + resolve_dependency(LAPACK) + get_target_property(LAPACK_INCLUDE_DIR lapack INTERFACE_INCLUDE_DIRECTORIES) + include_directories(SYSTEM "${LAPACK_INCLUDE_DIR}") + link_directories(SYSTEM "${LAPACK_PREFIX}/lib") + endif () resolve_dependency(FAISS) get_target_property(FAISS_INCLUDE_DIR faiss INTERFACE_INCLUDE_DIRECTORIES) include_directories(SYSTEM "${FAISS_INCLUDE_DIR}") link_directories(SYSTEM ${FAISS_PREFIX}/lib/) -endif() +endif () diff --git a/core/src/index/knowhere/CMakeLists.txt b/core/src/index/knowhere/CMakeLists.txt index bece9058a9..11c79e5466 100644 --- a/core/src/index/knowhere/CMakeLists.txt +++ b/core/src/index/knowhere/CMakeLists.txt @@ -1,6 +1,3 @@ -include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) -link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64) - include_directories(${INDEX_SOURCE_DIR}/knowhere) include_directories(${INDEX_SOURCE_DIR}/thirdparty) include_directories(${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService) @@ -19,9 +16,9 @@ file(GLOB SRC_FILES ${SPTAG_SOURCE_DIR}/AnnService/src/Core/KDT/*.cpp ${SPTAG_SOURCE_DIR}/AnnService/src/Helper/*.cpp) -if(NOT TARGET SPTAGLibStatic) +if (NOT TARGET SPTAGLibStatic) add_library(SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES}) -endif() +endif () set(external_srcs knowhere/adapter/SptagAdapter.cpp @@ -36,19 +33,13 @@ set(index_srcs knowhere/index/vector_index/IndexKDT.cpp knowhere/index/vector_index/IndexIDMAP.cpp knowhere/index/vector_index/IndexIVF.cpp - knowhere/index/vector_index/IndexGPUIVF.cpp knowhere/index/vector_index/helpers/KDTParameterMgr.cpp knowhere/index/vector_index/IndexNSG.cpp knowhere/index/vector_index/nsg/NSG.cpp knowhere/index/vector_index/nsg/NSGIO.cpp knowhere/index/vector_index/nsg/NSGHelper.cpp - knowhere/index/vector_index/helpers/Cloner.cpp - knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp knowhere/index/vector_index/IndexIVFSQ.cpp - knowhere/index/vector_index/IndexGPUIVFSQ.cpp - knowhere/index/vector_index/IndexIVFSQHybrid.cpp knowhere/index/vector_index/IndexIVFPQ.cpp - knowhere/index/vector_index/IndexGPUIVFPQ.cpp knowhere/index/vector_index/FaissBaseIndex.cpp knowhere/index/vector_index/helpers/FaissIO.cpp knowhere/index/vector_index/helpers/IndexParameter.cpp @@ -57,24 +48,56 @@ set(index_srcs set(depend_libs SPTAGLibStatic faiss - openblas - lapack arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a - cudart - cublas gomp gfortran pthread ) +if (BUILD_FAISS_WITH_MKL) + set(depend_libs ${depend_libs} + "-Wl,--start-group \ + ${MKL_LIB_PATH}/libmkl_intel_ilp64.a \ + ${MKL_LIB_PATH}/libmkl_gnu_thread.a \ + ${MKL_LIB_PATH}/libmkl_core.a \ + -Wl,--end-group -lgomp -lpthread -lm -ldl" + ) +else () + set(depend_libs ${depend_libs} + lapack + openblas) +endif () -if(NOT TARGET knowhere) +if (KNOWHERE_GPU_VERSION) + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(cuda_lib + cudart + cublas + ) + set(depend_libs ${depend_libs} + ${cuda_lib} + ) + + set(index_srcs ${index_srcs} + knowhere/index/vector_index/IndexGPUIVF.cpp + knowhere/index/vector_index/helpers/Cloner.cpp + knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp + knowhere/index/vector_index/IndexGPUIVFSQ.cpp + knowhere/index/vector_index/IndexIVFSQHybrid.cpp + knowhere/index/vector_index/IndexGPUIVFPQ.cpp + knowhere/index/vector_index/IndexGPUIDMAP.cpp + ) + +endif () + +if (NOT TARGET knowhere) add_library( knowhere STATIC ${external_srcs} ${index_srcs} ) -endif() +endif () target_link_libraries( knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp new file mode 100644 index 0000000000..5c7edd73fe --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp @@ -0,0 +1,113 @@ +#include "IndexGPUIDMAP.h" + +#include +#include +#include +#include + +#ifdef MILVUS_GPU_VERSION + +#include + +#endif + +#include "knowhere/adapter/VectorAdapter.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/helpers/FaissIO.h" + +namespace knowhere { + + VectorIndexPtr + GPUIDMAP::CopyGpuToCpu(const Config &config) { + std::lock_guard lk(mutex_); + + faiss::Index *device_index = index_.get(); + faiss::Index *host_index = faiss::gpu::index_gpu_to_cpu(device_index); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); + } + + VectorIndexPtr + GPUIDMAP::Clone() { + auto cpu_idx = CopyGpuToCpu(Config()); + + if (auto idmap = std::dynamic_pointer_cast(cpu_idx)) { + return idmap->CopyCpuToGpu(gpu_id_, Config()); + } else { + KNOWHERE_THROW_MSG("IndexType not Support GpuClone"); + } + } + + BinarySet + GPUIDMAP::SerializeImpl() { + try { + MemoryIOWriter writer; + { + faiss::Index *index = index_.get(); + faiss::Index *host_index = faiss::gpu::index_gpu_to_cpu(index); + + faiss::write_index(host_index, &writer); + delete host_index; + } + auto data = std::make_shared(); + data.reset(writer.data_); + + BinarySet res_set; + res_set.Append("IVF", data, writer.rp); + + return res_set; + } catch (std::exception &e) { + KNOWHERE_THROW_MSG(e.what()); + } + } + + void + GPUIDMAP::LoadImpl(const BinarySet &index_binary) { + auto binary = index_binary.GetByName("IVF"); + MemoryIOReader reader; + { + reader.total = binary->size; + reader.data_ = binary->data.get(); + + faiss::Index *index = faiss::read_index(&reader); + + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { + ResScope rs(res, gpu_id_, false); + auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index); + index_.reset(device_index); + res_ = res; + } else { + KNOWHERE_THROW_MSG("Load error, can't get gpu resource"); + } + + delete index; + } + } + + VectorIndexPtr + GPUIDMAP::CopyGpuToGpu(const int64_t &device_id, const Config &config) { + auto cpu_index = CopyGpuToCpu(config); + return std::static_pointer_cast(cpu_index)->CopyCpuToGpu(device_id, config); + } + + float * + GPUIDMAP::GetRawVectors() { + KNOWHERE_THROW_MSG("Not support"); + } + + int64_t * + GPUIDMAP::GetRawIds() { + KNOWHERE_THROW_MSG("Not support"); + } + + void + GPUIDMAP::search_impl(int64_t n, const float *data, int64_t k, float *distances, int64_t *labels, + const Config &cfg) { + ResScope rs(res_, gpu_id_); + index_->search(n, (float *) data, k, distances, labels); + } + +} // knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h new file mode 100644 index 0000000000..3272067cf0 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h @@ -0,0 +1,47 @@ +#pragma once + +#include "IndexGPUIVF.h" +#include "IndexIVF.h" +#include "IndexIDMAP.h" + +#include +#include + +namespace knowhere { + + class GPUIDMAP : public IDMAP, public GPUIndex { + public: + explicit GPUIDMAP(std::shared_ptr index, const int64_t &device_id, ResPtr &res) + : IDMAP(std::move(index)), GPUIndex(device_id, res) { + } + + VectorIndexPtr + CopyGpuToCpu(const Config &config) override; + + float * + GetRawVectors() override; + + int64_t * + GetRawIds() override; + + VectorIndexPtr + Clone() override; + + VectorIndexPtr + CopyGpuToGpu(const int64_t &device_id, const Config &config) override; + + protected: + void + search_impl(int64_t n, const float *data, int64_t k, float *distances, int64_t *labels, + const Config &cfg) override; + + BinarySet + SerializeImpl() override; + + void + LoadImpl(const BinarySet &index_binary) override; + }; + + using GPUIDMAPPtr = std::shared_ptr; + +} // knowhere \ No newline at end of file diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp index f926951736..f02752abf2 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp @@ -17,10 +17,18 @@ #include #include -#include + #include +#include +#include #include +#ifdef MILVUS_GPU_VERSION + +#include + +#endif + #include #include "knowhere/adapter/VectorAdapter.h" @@ -28,244 +36,167 @@ #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +#ifdef MILVUS_GPU_VERSION + +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#include "knowhere/index/vector_index/IndexGPUIDMAP.h" + +#endif + namespace knowhere { -BinarySet -IDMAP::Serialize() { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - - std::lock_guard lk(mutex_); - return SerializeImpl(); -} - -void -IDMAP::Load(const BinarySet& index_binary) { - std::lock_guard lk(mutex_); - LoadImpl(index_binary); -} - -DatasetPtr -IDMAP::Search(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - - config->CheckValid(); - // auto metric_type = config["metric_type"].as_string() == "L2" ? - // faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; - // index_->metric_type = metric_type; - - GETTENSOR(dataset) - - auto elems = rows * config->k; - auto res_ids = (int64_t*)malloc(sizeof(int64_t) * elems); - auto res_dis = (float*)malloc(sizeof(float) * elems); - - search_impl(rows, (float*)p_data, config->k, res_dis, res_ids, Config()); - - auto id_buf = MakeMutableBufferSmart((uint8_t*)res_ids, sizeof(int64_t) * elems); - auto dist_buf = MakeMutableBufferSmart((uint8_t*)res_dis, sizeof(float) * elems); - - std::vector id_bufs{nullptr, id_buf}; - std::vector dist_bufs{nullptr, dist_buf}; - - auto int64_type = std::make_shared(); - auto float_type = std::make_shared(); - - auto id_array_data = arrow::ArrayData::Make(int64_type, elems, id_bufs); - auto dist_array_data = arrow::ArrayData::Make(float_type, elems, dist_bufs); - - auto ids = std::make_shared>(id_array_data); - auto dists = std::make_shared>(dist_array_data); - std::vector array{ids, dists}; - - return std::make_shared(array, nullptr); -} - -void -IDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) { - index_->search(n, (float*)data, k, distances, labels); -} - -void -IDMAP::Add(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - - std::lock_guard lk(mutex_); - GETTENSOR(dataset) - - // TODO: magic here. - auto array = dataset->array()[0]; - auto p_ids = array->data()->GetValues(1, 0); - - index_->add_with_ids(rows, (float*)p_data, p_ids); -} - -int64_t -IDMAP::Count() { - return index_->ntotal; -} - -int64_t -IDMAP::Dimension() { - return index_->d; -} - -// TODO(linxj): return const pointer -float* -IDMAP::GetRawVectors() { - try { - auto file_index = dynamic_cast(index_.get()); - auto flat_index = dynamic_cast(file_index->index); - return flat_index->xb.data(); - } catch (std::exception& e) { - KNOWHERE_THROW_MSG(e.what()); - } -} - -// TODO(linxj): return const pointer -int64_t* -IDMAP::GetRawIds() { - try { - auto file_index = dynamic_cast(index_.get()); - return file_index->id_map.data(); - } catch (std::exception& e) { - KNOWHERE_THROW_MSG(e.what()); - } -} - -const char* type = "IDMap,Flat"; - -void -IDMAP::Train(const Config& config) { - config->CheckValid(); - - auto index = faiss::index_factory(config->d, type, GetMetricType(config->metric_type)); - index_.reset(index); -} - -VectorIndexPtr -IDMAP::Clone() { - std::lock_guard lk(mutex_); - - auto clone_index = faiss::clone_index(index_.get()); - std::shared_ptr new_index; - new_index.reset(clone_index); - return std::make_shared(new_index); -} - -VectorIndexPtr -IDMAP::CopyCpuToGpu(const int64_t& device_id, const Config& config) { - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { - ResScope rs(res, device_id, false); - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); - - std::shared_ptr device_index; - device_index.reset(gpu_index); - return std::make_shared(device_index, device_id, res); - } else { - KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); - } -} - -void -IDMAP::Seal() { - // do nothing -} - -VectorIndexPtr -GPUIDMAP::CopyGpuToCpu(const Config& config) { - std::lock_guard lk(mutex_); - - faiss::Index* device_index = index_.get(); - faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); - - std::shared_ptr new_index; - new_index.reset(host_index); - return std::make_shared(new_index); -} - -VectorIndexPtr -GPUIDMAP::Clone() { - auto cpu_idx = CopyGpuToCpu(Config()); - - if (auto idmap = std::dynamic_pointer_cast(cpu_idx)) { - return idmap->CopyCpuToGpu(gpu_id_, Config()); - } else { - KNOWHERE_THROW_MSG("IndexType not Support GpuClone"); - } -} - -BinarySet -GPUIDMAP::SerializeImpl() { - try { - MemoryIOWriter writer; - { - faiss::Index* index = index_.get(); - faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(index); - - faiss::write_index(host_index, &writer); - delete host_index; + BinarySet + IDMAP::Serialize() { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); } - auto data = std::make_shared(); - data.reset(writer.data_); - BinarySet res_set; - res_set.Append("IVF", data, writer.rp); - - return res_set; - } catch (std::exception& e) { - KNOWHERE_THROW_MSG(e.what()); + std::lock_guard lk(mutex_); + return SerializeImpl(); } -} -void -GPUIDMAP::LoadImpl(const BinarySet& index_binary) { - auto binary = index_binary.GetByName("IVF"); - MemoryIOReader reader; - { - reader.total = binary->size; - reader.data_ = binary->data.get(); + void + IDMAP::Load(const BinarySet &index_binary) { + std::lock_guard lk(mutex_); + LoadImpl(index_binary); + } - faiss::Index* index = faiss::read_index(&reader); + DatasetPtr + IDMAP::Search(const DatasetPtr &dataset, const Config &config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { - ResScope rs(res, gpu_id_, false); - auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index); - index_.reset(device_index); - res_ = res; + config->CheckValid(); + // auto metric_type = config["metric_type"].as_string() == "L2" ? + // faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; + // index_->metric_type = metric_type; + + GETTENSOR(dataset) + + auto elems = rows * config->k; + auto res_ids = (int64_t *) malloc(sizeof(int64_t) * elems); + auto res_dis = (float *) malloc(sizeof(float) * elems); + + search_impl(rows, (float *) p_data, config->k, res_dis, res_ids, Config()); + + auto id_buf = MakeMutableBufferSmart((uint8_t *) res_ids, sizeof(int64_t) * elems); + auto dist_buf = MakeMutableBufferSmart((uint8_t *) res_dis, sizeof(float) * elems); + + std::vector id_bufs{nullptr, id_buf}; + std::vector dist_bufs{nullptr, dist_buf}; + + auto int64_type = std::make_shared(); + auto float_type = std::make_shared(); + + auto id_array_data = arrow::ArrayData::Make(int64_type, elems, id_bufs); + auto dist_array_data = arrow::ArrayData::Make(float_type, elems, dist_bufs); + + auto ids = std::make_shared>(id_array_data); + auto dists = std::make_shared>(dist_array_data); + std::vector array{ids, dists}; + + return std::make_shared(array, nullptr); + } + + void + IDMAP::search_impl(int64_t n, const float *data, int64_t k, float *distances, int64_t *labels, const Config &cfg) { + index_->search(n, (float *) data, k, distances, labels); + } + + void + IDMAP::Add(const DatasetPtr &dataset, const Config &config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } + + std::lock_guard lk(mutex_); + GETTENSOR(dataset) + + // TODO: magic here. + auto array = dataset->array()[0]; + auto p_ids = array->data()->GetValues(1, 0); + + index_->add_with_ids(rows, (float *) p_data, p_ids); + } + + int64_t + IDMAP::Count() { + return index_->ntotal; + } + + int64_t + IDMAP::Dimension() { + return index_->d; + } + +// TODO(linxj): return const pointer + float * + IDMAP::GetRawVectors() { + try { + auto file_index = dynamic_cast(index_.get()); + auto flat_index = dynamic_cast(file_index->index); + return flat_index->xb.data(); + } catch (std::exception &e) { + KNOWHERE_THROW_MSG(e.what()); + } + } + +// TODO(linxj): return const pointer + int64_t * + IDMAP::GetRawIds() { + try { + auto file_index = dynamic_cast(index_.get()); + return file_index->id_map.data(); + } catch (std::exception &e) { + KNOWHERE_THROW_MSG(e.what()); + } + } + + const char *type = "IDMap,Flat"; + + void + IDMAP::Train(const Config &config) { + config->CheckValid(); + + auto index = faiss::index_factory(config->d, type, GetMetricType(config->metric_type)); + index_.reset(index); + } + + VectorIndexPtr + IDMAP::Clone() { + std::lock_guard lk(mutex_); + + auto clone_index = faiss::clone_index(index_.get()); + std::shared_ptr new_index; + new_index.reset(clone_index); + return std::make_shared(new_index); + } + + VectorIndexPtr + IDMAP::CopyCpuToGpu(const int64_t &device_id, const Config &config) { + +#ifdef MILVUS_GPU_VERSION + + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); + + std::shared_ptr device_index; + device_index.reset(gpu_index); + return std::make_shared(device_index, device_id, res); } else { - KNOWHERE_THROW_MSG("Load error, can't get gpu resource"); + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); } +#else + KNOWHERE_THROW_MSG("Calling IDMAP::CopyCpuToGpu when we are using CPU version"); +#endif - delete index; } -} -VectorIndexPtr -GPUIDMAP::CopyGpuToGpu(const int64_t& device_id, const Config& config) { - auto cpu_index = CopyGpuToCpu(config); - return std::static_pointer_cast(cpu_index)->CopyCpuToGpu(device_id, config); -} - -float* -GPUIDMAP::GetRawVectors() { - KNOWHERE_THROW_MSG("Not support"); -} - -int64_t* -GPUIDMAP::GetRawIds() { - KNOWHERE_THROW_MSG("Not support"); -} - -void -GPUIDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) { - ResScope rs(res_, gpu_id_); - index_->search(n, (float*)data, k, distances, labels); -} + void + IDMAP::Seal() { + // do nothing + } } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h index ec1cbb9e77..0f66e8fac0 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h @@ -17,7 +17,6 @@ #pragma once -#include "IndexGPUIVF.h" #include "IndexIVF.h" #include @@ -67,32 +66,4 @@ class IDMAP : public VectorIndex, public FaissBaseIndex { using IDMAPPtr = std::shared_ptr; -class GPUIDMAP : public IDMAP, public GPUIndex { - public: - explicit GPUIDMAP(std::shared_ptr index, const int64_t& device_id, ResPtr& res) - : IDMAP(std::move(index)), GPUIndex(device_id, res) { - } - - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; - float* - GetRawVectors() override; - int64_t* - GetRawIds() override; - VectorIndexPtr - Clone() override; - VectorIndexPtr - CopyGpuToGpu(const int64_t& device_id, const Config& config) override; - - protected: - void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; - BinarySet - SerializeImpl() override; - void - LoadImpl(const BinarySet& index_binary) override; -}; - -using GPUIDMAPPtr = std::shared_ptr; - } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp index 6da5db38ec..1b1d1a86f0 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp @@ -15,11 +15,19 @@ // specific language governing permissions and limitations // under the License. +#include #include #include #include #include +#include +#include +#include +#include +#ifdef MILVUS_GPU_VERSION +#include #include +#endif #include #include @@ -29,7 +37,9 @@ #include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Log.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/IndexGPUIVF.h" +#endif #include "knowhere/index/vector_index/IndexIVF.h" namespace knowhere { @@ -221,16 +231,18 @@ IVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, int6 faiss::ivflib::search_with_parameters(index_.get(), n, (float*)data, k, distances, labels, params.get()); stdclock::time_point after = stdclock::now(); double search_cost = (std::chrono::duration(after - before)).count(); - KNOWHERE_LOG_DEBUG << "K=" << k << " NQ=" << n << " NL=" << faiss::indexIVF_stats.nlist - << " ND=" << faiss::indexIVF_stats.ndis << " NH=" << faiss::indexIVF_stats.nheap_updates - << " Q=" << faiss::indexIVF_stats.quantization_time - << " S=" << faiss::indexIVF_stats.search_time; + KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost + << ", quantization cost: " << faiss::indexIVF_stats.quantization_time + << ", data search cost: " << faiss::indexIVF_stats.search_time; faiss::indexIVF_stats.quantization_time = 0; faiss::indexIVF_stats.search_time = 0; } VectorIndexPtr IVF::CopyCpuToGpu(const int64_t& device_id, const Config& config) { + +#ifdef MILVUS_GPU_VERSION + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); @@ -241,6 +253,10 @@ IVF::CopyCpuToGpu(const int64_t& device_id, const Config& config) { } else { KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); } + +#else + KNOWHERE_THROW_MSG("Calling IVF::CopyCpuToGpu when we are using CPU version"); +#endif } VectorIndexPtr diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp index 6e9a1d94da..4abeaf4385 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. +#ifdef MILVUS_GPU_VERSION +#include #include +#endif #include + #include #include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" + +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif +#include "knowhere/index/vector_index/IndexIVFSQ.h" namespace knowhere { @@ -54,6 +62,9 @@ IVFSQ::Clone_impl(const std::shared_ptr& index) { VectorIndexPtr IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { + +#ifdef MILVUS_GPU_VERSION + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); @@ -65,6 +76,10 @@ IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { } else { KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); } + +#else + KNOWHERE_THROW_MSG("Calling IVFSQ::CopyCpuToGpu when we are using CPU version"); +#endif } } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp index f5519b8240..8f6d93d7ff 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp @@ -19,8 +19,10 @@ #include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" +#endif + #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/nsg/NSG.h" #include "knowhere/index/vector_index/nsg/NSGIO.h" @@ -117,7 +119,11 @@ NSG::Train(const DatasetPtr& dataset, const Config& config) { } // TODO(linxj): dev IndexFactory, support more IndexType +#ifdef MILVUS_GPU_VERSION auto preprocess_index = std::make_shared(build_cfg->gpu_id); +#else + auto preprocess_index = std::make_shared(); +#endif auto model = preprocess_index->Train(dataset, config); preprocess_index->set_index_model(model); preprocess_index->AddWithoutIds(dataset, config); diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt index 2e84908cd7..145278a636 100644 --- a/core/src/index/unittest/CMakeLists.txt +++ b/core/src/index/unittest/CMakeLists.txt @@ -2,26 +2,32 @@ include_directories(${INDEX_SOURCE_DIR}/thirdparty) include_directories(${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService) include_directories(${INDEX_SOURCE_DIR}/knowhere) include_directories(${INDEX_SOURCE_DIR}) -include_directories(/usr/local/cuda/include) -link_directories(/usr/local/cuda/lib64) - -message(STATUS "arrow prefix: ${ARROW_PREFIX}") -message(STATUS "libjemalloc_pic path: ${ARROW_PREFIX}/lib/libjemalloc_pic.a") set(depend_libs gtest gmock gtest_main gmock_main - faiss openblas lapack + faiss arrow "${ARROW_PREFIX}/lib/libjemalloc_pic.a" ) +if (BUILD_FAISS_WITH_MKL) + set(depend_libs ${depend_libs} + "-Wl,--start-group \ + ${MKL_LIB_PATH}/libmkl_intel_ilp64.a \ + ${MKL_LIB_PATH}/libmkl_gnu_thread.a \ + ${MKL_LIB_PATH}/libmkl_core.a \ + -Wl,--end-group -lgomp -lpthread -lm -ldl" + ) +else () + set(depend_libs ${depend_libs} + lapack + openblas) +endif () set(basic_libs - cudart cublas gomp gfortran pthread ) set(util_srcs ${MILVUS_ENGINE_SRC}/external/easyloggingpp/easylogging++.cc - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/adapter/Structure.cpp @@ -31,32 +37,49 @@ set(util_srcs ${INDEX_SOURCE_DIR}/unittest/utils.cpp ) +if (KNOWHERE_GPU_VERSION) + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(cuda_lib + cudart + cublas + ) + set(basic_libs ${basic_libs} + ${cuda_lib} + ) + set(util_srcs ${util_srcs} + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp + ) +endif () + # set(ivf_srcs - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVF.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp ) -if(NOT TARGET test_ivf) +if (KNOWHERE_GPU_VERSION) + set(ivf_srcs ${ivf_srcs} + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp + ) +endif () +if (NOT TARGET test_ivf) add_executable(test_ivf test_ivf.cpp ${ivf_srcs} ${util_srcs}) -endif() +endif () target_link_libraries(test_ivf ${depend_libs} ${unittest_libs} ${basic_libs}) # -set(idmap_srcs - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp - ) -if(NOT TARGET test_idmap) - add_executable(test_idmap test_idmap.cpp ${idmap_srcs} ${ivf_srcs} ${util_srcs}) -endif() +if (NOT TARGET test_idmap) + add_executable(test_idmap test_idmap.cpp ${ivf_srcs} ${util_srcs}) +endif () target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs}) # @@ -66,25 +89,28 @@ set(kdt_srcs ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexKDT.cpp ) -if(NOT TARGET test_kdt) +if (NOT TARGET test_kdt) add_executable(test_kdt test_kdt.cpp ${kdt_srcs} ${util_srcs}) -endif() +endif () target_link_libraries(test_kdt SPTAGLibStatic ${depend_libs} ${unittest_libs} ${basic_libs}) -add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs}) -target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs}) +if (KNOWHERE_GPU_VERSION) + add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs}) + target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs}) -add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs}) -target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs}) + add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs}) + target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs}) +endif () install(TARGETS test_ivf DESTINATION unittest) install(TARGETS test_idmap DESTINATION unittest) install(TARGETS test_kdt DESTINATION unittest) -install(TARGETS test_gpuresource DESTINATION unittest) -install(TARGETS test_customized_index DESTINATION unittest) - +if (KNOWHERE_GPU_VERSION) + install(TARGETS test_gpuresource DESTINATION unittest) + install(TARGETS test_customized_index DESTINATION unittest) +endif () #add_subdirectory(faiss_ori) #add_subdirectory(faiss_benchmark) add_subdirectory(test_nsg) diff --git a/core/src/index/unittest/Helper.h b/core/src/index/unittest/Helper.h index 8d4bb0f4ae..074c9548aa 100644 --- a/core/src/index/unittest/Helper.h +++ b/core/src/index/unittest/Helper.h @@ -18,13 +18,16 @@ #include #include -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" + +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#endif int DEVICEID = 0; constexpr int64_t DIM = 128; @@ -36,22 +39,25 @@ constexpr int64_t TEMPMEM = 1024 * 1024 * 300; constexpr int64_t RESNUM = 2; knowhere::IVFIndexPtr -IndexFactory(const std::string& type) { +IndexFactory(const std::string &type) { if (type == "IVF") { return std::make_shared(); } else if (type == "IVFPQ") { return std::make_shared(); - } else if (type == "GPUIVF") { + } else if (type == "IVFSQ") { + return std::make_shared(); + } +#ifdef MILVUS_GPU_VERSION + else if (type == "GPUIVF") { return std::make_shared(DEVICEID); } else if (type == "GPUIVFPQ") { return std::make_shared(DEVICEID); - } else if (type == "IVFSQ") { - return std::make_shared(); } else if (type == "GPUIVFSQ") { return std::make_shared(DEVICEID); } else if (type == "IVFSQHybrid") { return std::make_shared(DEVICEID); } +#endif } enum class ParameterType { @@ -61,15 +67,15 @@ enum class ParameterType { }; class ParamGenerator { - public: - static ParamGenerator& +public: + static ParamGenerator & GetInstance() { static ParamGenerator instance; return instance; } knowhere::Config - Gen(const ParameterType& type) { + Gen(const ParameterType &type) { if (type == ParameterType::ivf) { auto tempconf = std::make_shared(); tempconf->d = DIM; @@ -107,14 +113,18 @@ class ParamGenerator { #include class TestGpuIndexBase : public ::testing::Test { - protected: +protected: void SetUp() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); +#endif } void TearDown() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); +#endif } }; diff --git a/core/src/index/unittest/faiss_benchmark/CMakeLists.txt b/core/src/index/unittest/faiss_benchmark/CMakeLists.txt index 556364b68a..72eb7e7a7e 100644 --- a/core/src/index/unittest/faiss_benchmark/CMakeLists.txt +++ b/core/src/index/unittest/faiss_benchmark/CMakeLists.txt @@ -1,24 +1,50 @@ -include_directories(${INDEX_SOURCE_DIR}/thirdparty) -include_directories(${INDEX_SOURCE_DIR}/include) -include_directories(/usr/local/cuda/include) -include_directories(/usr/local/hdf5/include) +if (KNOWHERE_GPU_VERSION) -link_directories(/usr/local/cuda/lib64) -link_directories(/usr/local/hdf5/lib) + include_directories(${INDEX_SOURCE_DIR}/thirdparty) + include_directories(${INDEX_SOURCE_DIR}/include) + include_directories(/usr/local/cuda/include) + include_directories(/usr/local/hdf5/include) -set(unittest_libs - gtest gmock gtest_main gmock_main) + link_directories(/usr/local/cuda/lib64) + link_directories(/usr/local/hdf5/lib) -set(depend_libs - faiss openblas lapack hdf5 - arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a - ) + set(unittest_libs + gtest gmock gtest_main gmock_main) -set(basic_libs - cudart cublas - gomp gfortran pthread - ) + set(depend_libs + faiss hdf5 + arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a + ) + if (BUILD_FAISS_WITH_MKL) + set(depend_libs ${depend_libs} + "-Wl,--start-group \ + ${MKL_LIB_PATH}/libmkl_intel_ilp64.a \ + ${MKL_LIB_PATH}/libmkl_gnu_thread.a \ + ${MKL_LIB_PATH}/libmkl_core.a \ + -Wl,--end-group -lgomp -lpthread -lm -ldl" + ) + else () + set(depend_libs ${depend_libs} + lapack + openblas) + endif () -add_executable(test_faiss_benchmark faiss_benchmark_test.cpp) -target_link_libraries(test_faiss_benchmark ${depend_libs} ${unittest_libs} ${basic_libs}) -install(TARGETS test_faiss_benchmark DESTINATION unittest) + set(basic_libs + gomp gfortran pthread + ) + + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(cuda_lib + cudart + cublas + ) + set(basic_libs ${basic_libs} + ${cuda_lib} + ) + + add_executable(test_faiss_benchmark faiss_benchmark_test.cpp) + target_link_libraries(test_faiss_benchmark ${depend_libs} ${unittest_libs} ${basic_libs}) + install(TARGETS test_faiss_benchmark DESTINATION unittest) + +endif () diff --git a/core/src/index/unittest/faiss_ori/CMakeLists.txt b/core/src/index/unittest/faiss_ori/CMakeLists.txt index d01463aaab..8216764ab7 100644 --- a/core/src/index/unittest/faiss_ori/CMakeLists.txt +++ b/core/src/index/unittest/faiss_ori/CMakeLists.txt @@ -1,26 +1,49 @@ -include_directories(${INDEX_SOURCE_DIR}/thirdparty) -include_directories(${INDEX_SOURCE_DIR}/include) -include_directories(/usr/local/cuda/include) -link_directories(/usr/local/cuda/lib64) +if (KNOWHERE_GPU_VERSION) -set(unittest_libs - gtest gmock gtest_main gmock_main) + include_directories(${INDEX_SOURCE_DIR}/thirdparty) + include_directories(${INDEX_SOURCE_DIR}/include) -set(depend_libs - faiss openblas lapack - arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a - ) + set(unittest_libs + gtest gmock gtest_main gmock_main) -set(basic_libs - cudart cublas - gomp gfortran pthread - ) + set(depend_libs + faiss + arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a + ) + if (BUILD_FAISS_WITH_MKL) + set(depend_libs ${depend_libs} + "-Wl,--start-group \ + ${MKL_LIB_PATH}/libmkl_intel_ilp64.a \ + ${MKL_LIB_PATH}/libmkl_gnu_thread.a \ + ${MKL_LIB_PATH}/libmkl_core.a \ + -Wl,--end-group -lgomp -lpthread -lm -ldl" + ) + else () + set(depend_libs ${depend_libs} + lapack + openblas) + endif () + set(basic_libs + gomp gfortran pthread + ) -# -if(NOT TARGET test_gpu) - add_executable(test_gpu gpuresource_test.cpp) -endif() -target_link_libraries(test_gpu ${depend_libs} ${unittest_libs} ${basic_libs}) + include_directories(${CUDA_INCLUDE_DIRS}) + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(cuda_lib + cudart + cublas + ) + set(basic_libs ${basic_libs} + ${cuda_lib} + ) -install(TARGETS test_gpu DESTINATION unittest) \ No newline at end of file + # + if (NOT TARGET test_gpu) + add_executable(test_gpu gpuresource_test.cpp) + endif () + target_link_libraries(test_gpu ${depend_libs} ${unittest_libs} ${basic_libs}) + + install(TARGETS test_gpu DESTINATION unittest) + +endif () \ No newline at end of file diff --git a/core/src/index/unittest/test_idmap.cpp b/core/src/index/unittest/test_idmap.cpp index d1ff3ee046..b6e51361b9 100644 --- a/core/src/index/unittest/test_idmap.cpp +++ b/core/src/index/unittest/test_idmap.cpp @@ -21,8 +21,10 @@ #include "knowhere/adapter/Structure.h" #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/Cloner.h" - +#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#endif #include "Helper.h" #include "unittest/utils.h" @@ -116,6 +118,7 @@ TEST_F(IDMAPTest, idmap_serialize) { } } +#ifdef MILVUS_GPU_VERSION TEST_F(IDMAPTest, copy_test) { ASSERT_TRUE(!xb.empty()); @@ -175,3 +178,4 @@ TEST_F(IDMAPTest, copy_test) { AssertAnns(device_result, nq, k); } } +#endif diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp index 8b17e08272..953884bc2f 100644 --- a/core/src/index/unittest/test_ivf.cpp +++ b/core/src/index/unittest/test_ivf.cpp @@ -20,13 +20,24 @@ #include #include +#ifdef MILVUS_GPU_VERSION #include +#endif #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" + #include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" + +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/helpers/Cloner.h" +#endif #include "unittest/Helper.h" #include "unittest/utils.h" @@ -36,11 +47,12 @@ using ::testing::TestWithParam; using ::testing::Values; class IVFTest : public DataGen, public TestWithParam<::std::tuple> { - protected: +protected: void SetUp() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); - +#endif ParameterType parameter_type; std::tie(index_type, parameter_type) = GetParam(); // Init_with_default(); @@ -54,7 +66,9 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple(bin->data.get()), bin->size); + writer(static_cast(bin->data.get()), bin->size); FileIOReader reader(filename); reader(ret, bin->size); @@ -148,6 +165,7 @@ TEST_P(IVFTest, ivf_serialize) { } } +#ifdef MILVUS_GPU_VERSION TEST_P(IVFTest, clone_test) { assert(!xb.empty()); @@ -198,18 +216,18 @@ TEST_P(IVFTest, clone_test) { auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ - auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - auto clone_result = clone_index->Search(query_dataset, conf); - AssertEqual(result, clone_result); - std::cout << "clone G <=> C [" << index_type << "] success" << std::endl; - }); + auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + auto clone_result = clone_index->Search(query_dataset, conf); + AssertEqual(result, clone_result); + std::cout << "clone G <=> C [" << index_type << "] success" << std::endl; + }); } else { EXPECT_THROW( - { - std::cout << "clone G <=> C [" << index_type << "] failed" << std::endl; - auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - }, - knowhere::KnowhereException); + { + std::cout << "clone G <=> C [" << index_type << "] failed" << std::endl; + auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + }, + knowhere::KnowhereException); } } @@ -223,22 +241,24 @@ TEST_P(IVFTest, clone_test) { auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); - auto clone_result = clone_index->Search(query_dataset, conf); - AssertEqual(result, clone_result); - std::cout << "clone C <=> G [" << index_type << "] success" << std::endl; - }); + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); + auto clone_result = clone_index->Search(query_dataset, conf); + AssertEqual(result, clone_result); + std::cout << "clone C <=> G [" << index_type << "] success" << std::endl; + }); } else { EXPECT_THROW( - { - std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl; - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); - }, - knowhere::KnowhereException); + { + std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl; + auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); + }, + knowhere::KnowhereException); } } } +#endif +#ifdef MILVUS_GPU_VERSION #ifdef CUSTOMIZATION TEST_P(IVFTest, gpu_seal_test) { std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"}; @@ -271,5 +291,5 @@ TEST_P(IVFTest, gpu_seal_test) { auto with_seal = tc.RecordSection("With seal"); ASSERT_GE(without_seal, with_seal); } - +#endif #endif diff --git a/core/src/index/unittest/test_nsg/CMakeLists.txt b/core/src/index/unittest/test_nsg/CMakeLists.txt index 3d22051d82..01227ea745 100644 --- a/core/src/index/unittest/test_nsg/CMakeLists.txt +++ b/core/src/index/unittest/test_nsg/CMakeLists.txt @@ -4,14 +4,13 @@ add_definitions(-std=c++11 -O3 -lboost -march=native -Wall -DINFO) -find_package(OpenMP) +find_package(OpenMP REQUIRED) if (OPENMP_FOUND) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") else () message(FATAL_ERROR "no OpenMP supprot") endif () -message(${OpenMP_CXX_FLAGS}) include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg) aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg nsg_src) diff --git a/core/src/index/unittest/test_nsg/test_nsg.cpp b/core/src/index/unittest/test_nsg/test_nsg.cpp index 11b9becce4..47c014e691 100644 --- a/core/src/index/unittest/test_nsg/test_nsg.cpp +++ b/core/src/index/unittest/test_nsg/test_nsg.cpp @@ -21,7 +21,9 @@ #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/FaissBaseIndex.h" #include "knowhere/index/vector_index/IndexNSG.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif #include "knowhere/index/vector_index/nsg/NSGIO.h" #include "unittest/utils.h" @@ -37,7 +39,9 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test { void SetUp() override { // Init_with_default(); +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2); +#endif Generate(256, 1000000 / 100, 1); index_ = std::make_shared(); @@ -60,7 +64,9 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test { void TearDown() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); +#endif } protected: diff --git a/core/src/main.cpp b/core/src/main.cpp index 85f91469bb..cd00024afe 100644 --- a/core/src/main.cpp +++ b/core/src/main.cpp @@ -25,14 +25,14 @@ #include "external/easyloggingpp/easylogging++.h" #include "metrics/Metrics.h" #include "server/Server.h" -#include "src/version.h" +#include "src/config.h" #include "utils/CommonUtil.h" #include "utils/SignalUtil.h" INITIALIZE_EASYLOGGINGPP void -print_help(const std::string& app_name) { +print_help(const std::string &app_name) { std::cout << std::endl << "Usage: " << app_name << " [OPTIONS]" << std::endl << std::endl; std::cout << " Options:" << std::endl; std::cout << " -h --help Print this help" << std::endl; @@ -52,19 +52,24 @@ print_banner() { std::cout << std::endl; std::cout << "Welcome to Milvus!" << std::endl; std::cout << "Milvus " << BUILD_TYPE << " version: v" << MILVUS_VERSION << ", built at " << BUILD_TIME << std::endl; +#ifdef MILVUS_CPU_VERSION + std::cout << "You are using Milvus CPU version" << std::endl; +#else + std::cout << "You are using Milvus GPU version" << std::endl; +#endif std::cout << std::endl; } int -main(int argc, char* argv[]) { +main(int argc, char *argv[]) { print_banner(); - static struct option long_options[] = {{"conf_file", required_argument, nullptr, 'c'}, + static struct option long_options[] = {{"conf_file", required_argument, nullptr, 'c'}, {"log_conf_file", required_argument, nullptr, 'l'}, - {"help", no_argument, nullptr, 'h'}, - {"daemon", no_argument, nullptr, 'd'}, - {"pid_file", required_argument, nullptr, 'p'}, - {nullptr, 0, nullptr, 0}}; + {"help", no_argument, nullptr, 'h'}, + {"daemon", no_argument, nullptr, 'd'}, + {"pid_file", required_argument, nullptr, 'p'}, + {nullptr, 0, nullptr, 0}}; int option_index = 0; int64_t start_daemonized = 0; @@ -73,7 +78,7 @@ main(int argc, char* argv[]) { std::string pid_filename; std::string app_name = argv[0]; - milvus::server::Server& server = milvus::server::Server::GetInstance(); + milvus::server::Server &server = milvus::server::Server::GetInstance(); milvus::Status s; if (argc < 2) { @@ -85,21 +90,21 @@ main(int argc, char* argv[]) { while ((value = getopt_long(argc, argv, "c:l:p:dh", long_options, &option_index)) != -1) { switch (value) { case 'c': { - char* config_filename_ptr = strdup(optarg); + char *config_filename_ptr = strdup(optarg); config_filename = config_filename_ptr; free(config_filename_ptr); std::cout << "Loading configuration from: " << config_filename << std::endl; break; } case 'l': { - char* log_filename_ptr = strdup(optarg); + char *log_filename_ptr = strdup(optarg); log_config_file = log_filename_ptr; free(log_filename_ptr); - std::cout << "Initial log config from: " << log_config_file << std::endl; + std::cout << "Initializing log config from: " << log_config_file << std::endl; break; } case 'p': { - char* pid_filename_ptr = strdup(optarg); + char *pid_filename_ptr = strdup(optarg); pid_filename = pid_filename_ptr; free(pid_filename_ptr); std::cout << pid_filename << std::endl; @@ -142,7 +147,7 @@ main(int argc, char* argv[]) { return EXIT_SUCCESS; -FAIL: + FAIL: std::cout << "Milvus server exit..." << std::endl; return EXIT_FAILURE; } diff --git a/core/src/metrics/SystemInfo.cpp b/core/src/metrics/SystemInfo.cpp index 154f7b0797..b0d1577d36 100644 --- a/core/src/metrics/SystemInfo.cpp +++ b/core/src/metrics/SystemInfo.cpp @@ -19,7 +19,6 @@ #include "utils/Log.h" #include -#include #include #include #include @@ -29,6 +28,10 @@ #include #include +#ifdef MILVUS_GPU_VERSION +#include +#endif + namespace milvus { namespace server { @@ -60,6 +63,7 @@ SystemInfo::Init() { total_ram_ = GetPhysicalMemory(); fclose(file); +#ifdef MILVUS_GPU_VERSION // initialize GPU information nvmlReturn_t nvmlresult; nvmlresult = nvmlInit(); @@ -72,6 +76,7 @@ SystemInfo::Init() { SERVER_LOG_ERROR << "Unable to get devidce number"; return; } +#endif // initialize network traffic information std::pair in_and_out_octets = Octets(); @@ -209,10 +214,14 @@ SystemInfo::CPUPercent() { std::vector SystemInfo::GPUMemoryTotal() { + // get GPU usage percent if (!initialized_) Init(); std::vector result; + +#ifdef MILVUS_GPU_VERSION + nvmlMemory_t nvmlMemory; for (int i = 0; i < num_device_; ++i) { nvmlDevice_t device; @@ -220,6 +229,8 @@ SystemInfo::GPUMemoryTotal() { nvmlDeviceGetMemoryInfo(device, &nvmlMemory); result.push_back(nvmlMemory.total); } +#endif + return result; } @@ -228,6 +239,9 @@ SystemInfo::GPUTemperature() { if (!initialized_) Init(); std::vector result; + +#ifdef MILVUS_GPU_VERSION + for (int i = 0; i < num_device_; i++) { nvmlDevice_t device; nvmlDeviceGetHandleByIndex(i, &device); @@ -235,6 +249,9 @@ SystemInfo::GPUTemperature() { nvmlDeviceGetTemperature(device, NVML_TEMPERATURE_GPU, &temp); result.push_back(temp); } + +#endif + return result; } @@ -283,6 +300,9 @@ SystemInfo::GPUMemoryUsed() { Init(); std::vector result; + +#ifdef MILVUS_GPU_VERSION + nvmlMemory_t nvmlMemory; for (int i = 0; i < num_device_; ++i) { nvmlDevice_t device; @@ -290,6 +310,9 @@ SystemInfo::GPUMemoryUsed() { nvmlDeviceGetMemoryInfo(device, &nvmlMemory); result.push_back(nvmlMemory.used); } + +#endif + return result; } diff --git a/core/src/scheduler/JobMgr.cpp b/core/src/scheduler/JobMgr.cpp index 4404d95763..794f6a0f37 100644 --- a/core/src/scheduler/JobMgr.cpp +++ b/core/src/scheduler/JobMgr.cpp @@ -104,25 +104,20 @@ JobMgr::build_task(const JobPtr& job) { void JobMgr::calculate_path(const TaskPtr& task) { - if (task->type_ == TaskType::SearchTask) { - if (task->label()->Type() != TaskLabelType::SPECIFIED_RESOURCE) { - return; - } - - std::vector path; - auto spec_label = std::static_pointer_cast(task->label()); - auto src = res_mgr_->GetDiskResources()[0]; - auto dest = spec_label->resource(); - ShortestPath(src.lock(), dest.lock(), res_mgr_, path); - task->path() = Path(path, path.size() - 1); - } else if (task->type_ == TaskType::BuildIndexTask) { - auto spec_label = std::static_pointer_cast(task->label()); - auto src = res_mgr_->GetDiskResources()[0]; - auto dest = spec_label->resource(); - std::vector path; - ShortestPath(src.lock(), dest.lock(), res_mgr_, path); - task->path() = Path(path, path.size() - 1); + if (task->type_ != TaskType::SearchTask) { + return; } + + if (task->label()->Type() != TaskLabelType::SPECIFIED_RESOURCE) { + return; + } + + std::vector path; + auto spec_label = std::static_pointer_cast(task->label()); + auto src = res_mgr_->GetDiskResources()[0]; + auto dest = spec_label->resource(); + ShortestPath(src.lock(), dest.lock(), res_mgr_, path); + task->path() = Path(path, path.size() - 1); } } // namespace scheduler diff --git a/core/src/scheduler/SchedInst.cpp b/core/src/scheduler/SchedInst.cpp index f0c00c2d2a..61e0c09759 100644 --- a/core/src/scheduler/SchedInst.cpp +++ b/core/src/scheduler/SchedInst.cpp @@ -18,7 +18,6 @@ #include "scheduler/SchedInst.h" #include "ResourceFactory.h" #include "Utils.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" #include "server/Config.h" #include @@ -55,8 +54,8 @@ load_simple_config() { // get resources auto gpu_ids = get_gpu_pool(); - int32_t index_build_device_id; - config.GetResourceConfigIndexBuildDevice(index_build_device_id); + int32_t build_gpu_id; + config.GetResourceConfigIndexBuildDevice(build_gpu_id); // create and connect ResMgrInst::GetInstance()->Add(ResourceFactory::Create("disk", "DISK", 0, true, false)); @@ -70,15 +69,15 @@ load_simple_config() { for (auto& gpu_id : gpu_ids) { ResMgrInst::GetInstance()->Add(ResourceFactory::Create(std::to_string(gpu_id), "GPU", gpu_id, true, true)); ResMgrInst::GetInstance()->Connect("cpu", std::to_string(gpu_id), pcie); - if (index_build_device_id == gpu_id) { + if (build_gpu_id == gpu_id) { find_build_gpu_id = true; } } - if (not find_build_gpu_id && index_build_device_id != server::CPU_DEVICE_ID) { + if (not find_build_gpu_id) { ResMgrInst::GetInstance()->Add( - ResourceFactory::Create(std::to_string(index_build_device_id), "GPU", index_build_device_id, true, true)); - ResMgrInst::GetInstance()->Connect("cpu", std::to_string(index_build_device_id), pcie); + ResourceFactory::Create(std::to_string(build_gpu_id), "GPU", build_gpu_id, true, true)); + ResMgrInst::GetInstance()->Connect("cpu", std::to_string(build_gpu_id), pcie); } } diff --git a/core/src/scheduler/SchedInst.h b/core/src/scheduler/SchedInst.h index d51611af26..a3048069f9 100644 --- a/core/src/scheduler/SchedInst.h +++ b/core/src/scheduler/SchedInst.h @@ -106,6 +106,7 @@ class OptimizerInst { has_cpu = true; } } + std::vector pass_list; pass_list.push_back(std::make_shared()); pass_list.push_back(std::make_shared()); diff --git a/core/src/scheduler/TaskCreator.cpp b/core/src/scheduler/TaskCreator.cpp index 9f3bc2ae9a..40cfa9aac6 100644 --- a/core/src/scheduler/TaskCreator.cpp +++ b/core/src/scheduler/TaskCreator.cpp @@ -70,15 +70,8 @@ TaskCreator::Create(const DeleteJobPtr& job) { std::vector TaskCreator::Create(const BuildIndexJobPtr& job) { std::vector tasks; - server::Config& config = server::Config::GetInstance(); - int32_t build_index_id; - Status stat = config.GetResourceConfigIndexBuildDevice(build_index_id); - ResourcePtr res_ptr; - if (build_index_id == server::CPU_DEVICE_ID) { - res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); - } else { - res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, build_index_id); - } + // TODO(yukun): remove "disk" hardcode here + ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("disk"); for (auto& to_index_file : job->to_index_files()) { auto label = std::make_shared(std::weak_ptr(res_ptr)); diff --git a/core/src/scheduler/Utils.cpp b/core/src/scheduler/Utils.cpp index 998e545ba5..2fd573e47a 100644 --- a/core/src/scheduler/Utils.cpp +++ b/core/src/scheduler/Utils.cpp @@ -19,7 +19,9 @@ #include "server/Config.h" #include "utils/Log.h" +#ifdef MILVUS_GPU_VERSION #include +#endif #include #include #include @@ -38,7 +40,9 @@ get_current_timestamp() { uint64_t get_num_gpu() { int n_devices = 0; +#ifdef MILVUS_GPU_VERSION cudaGetDeviceCount(&n_devices); +#endif return n_devices; } diff --git a/core/src/scheduler/action/PushTaskToNeighbour.cpp b/core/src/scheduler/action/PushTaskToNeighbour.cpp index 9aed678937..b8a4a1164b 100644 --- a/core/src/scheduler/action/PushTaskToNeighbour.cpp +++ b/core/src/scheduler/action/PushTaskToNeighbour.cpp @@ -138,41 +138,73 @@ Action::SpecifiedResourceLabelTaskScheduler(const ResourceMgrPtr& res_mgr, Resou std::shared_ptr event) { auto task_item = event->task_table_item_; auto task = event->task_table_item_->task; - // if (resource->type() == ResourceType::DISK) { - // // step 1: calculate shortest path per resource, from disk to compute resource - // auto compute_resources = res_mgr->GetComputeResources(); - // std::vector> paths; - // std::vector transport_costs; - // for (auto& res : compute_resources) { - // std::vector path; - // uint64_t transport_cost = ShortestPath(resource, res, res_mgr, path); - // transport_costs.push_back(transport_cost); - // paths.emplace_back(path); - // } - // if (task->job_.lock()->type() == JobType::BUILD) { - // // step2: Read device id in config - // // get build index gpu resource - // server::Config& config = server::Config::GetInstance(); - // int32_t build_index_gpu; - // Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu); - // - // bool find_gpu_res = false; - // if (res_mgr->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) { - // for (uint64_t i = 0; i < compute_resources.size(); ++i) { - // if (compute_resources[i]->name() == - // res_mgr->GetResource(ResourceType::GPU, build_index_gpu)->name()) { - // find_gpu_res = true; - // Path task_path(paths[i], paths[i].size() - 1); - // task->path() = task_path; - // break; - // } - // } - // } - // if (not find_gpu_res) { - // task->path() = Path(paths[0], paths[0].size() - 1); - // } - // } - // } + if (resource->type() == ResourceType::DISK) { + // step 1: calculate shortest path per resource, from disk to compute resource + auto compute_resources = res_mgr->GetComputeResources(); + std::vector> paths; + std::vector transport_costs; + for (auto& res : compute_resources) { + std::vector path; + uint64_t transport_cost = ShortestPath(resource, res, res_mgr, path); + transport_costs.push_back(transport_cost); + paths.emplace_back(path); + } + // if (task->job_.lock()->type() == JobType::SEARCH) { + // auto label = task->label(); + // auto spec_label = std::static_pointer_cast(label); + // if (spec_label->resource().lock()->type() == ResourceType::CPU) { + // std::vector spec_path; + // spec_path.push_back(spec_label->resource().lock()->name()); + // spec_path.push_back(resource->name()); + // task->path() = Path(spec_path, spec_path.size() - 1); + // } else { + // // step 2: select min cost, cost(resource) = avg_cost * task_to_do + transport_cost + // uint64_t min_cost = std::numeric_limits::max(); + // uint64_t min_cost_idx = 0; + // for (uint64_t i = 0; i < compute_resources.size(); ++i) { + // if (compute_resources[i]->TotalTasks() == 0) { + // min_cost_idx = i; + // break; + // } + // uint64_t cost = compute_resources[i]->TaskAvgCost() * + // compute_resources[i]->NumOfTaskToExec() + + // transport_costs[i]; + // if (min_cost > cost) { + // min_cost = cost; + // min_cost_idx = i; + // } + // } + // + // // step 3: set path in task + // Path task_path(paths[min_cost_idx], paths[min_cost_idx].size() - 1); + // task->path() = task_path; + // } + // + // } else + if (task->job_.lock()->type() == JobType::BUILD) { + // step2: Read device id in config + // get build index gpu resource + server::Config& config = server::Config::GetInstance(); + int32_t build_index_gpu; + Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu); + + bool find_gpu_res = false; + if (res_mgr->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) { + for (uint64_t i = 0; i < compute_resources.size(); ++i) { + if (compute_resources[i]->name() == + res_mgr->GetResource(ResourceType::GPU, build_index_gpu)->name()) { + find_gpu_res = true; + Path task_path(paths[i], paths[i].size() - 1); + task->path() = task_path; + break; + } + } + } + if (not find_gpu_res) { + task->path() = Path(paths[0], paths[0].size() - 1); + } + } + } if (resource->name() == task->path().Last()) { resource->WakeupExecutor(); diff --git a/core/src/scheduler/job/BuildIndexJob.cpp b/core/src/scheduler/job/BuildIndexJob.cpp index 3247383db3..4c4c3b5054 100644 --- a/core/src/scheduler/job/BuildIndexJob.cpp +++ b/core/src/scheduler/job/BuildIndexJob.cpp @@ -50,10 +50,7 @@ void BuildIndexJob::BuildIndexDone(size_t to_index_id) { std::unique_lock lock(mutex_); to_index_files_.erase(to_index_id); - if (to_index_files_.empty()) { - cv_.notify_all(); - } - + cv_.notify_all(); SERVER_LOG_DEBUG << "BuildIndexJob " << id() << " finish index file: " << to_index_id; } diff --git a/core/src/scheduler/job/SearchJob.cpp b/core/src/scheduler/job/SearchJob.cpp index ec93c69f55..47c825c122 100644 --- a/core/src/scheduler/job/SearchJob.cpp +++ b/core/src/scheduler/job/SearchJob.cpp @@ -49,21 +49,13 @@ void SearchJob::SearchDone(size_t index_id) { std::unique_lock lock(mutex_); index_files_.erase(index_id); - if (index_files_.empty()) { - cv_.notify_all(); - } - + cv_.notify_all(); SERVER_LOG_DEBUG << "SearchJob " << id() << " finish index file: " << index_id; } -ResultIds& -SearchJob::GetResultIds() { - return result_ids_; -} - -ResultDistances& -SearchJob::GetResultDistances() { - return result_distances_; +ResultSet& +SearchJob::GetResult() { + return result_; } Status& diff --git a/core/src/scheduler/job/SearchJob.h b/core/src/scheduler/job/SearchJob.h index ff5ab34131..1e586090b9 100644 --- a/core/src/scheduler/job/SearchJob.h +++ b/core/src/scheduler/job/SearchJob.h @@ -29,7 +29,6 @@ #include #include "Job.h" -#include "db/Types.h" #include "db/meta/MetaTypes.h" namespace milvus { @@ -38,9 +37,9 @@ namespace scheduler { using engine::meta::TableFileSchemaPtr; using Id2IndexMap = std::unordered_map; - -using ResultIds = engine::ResultIds; -using ResultDistances = engine::ResultDistances; +using IdDistPair = std::pair; +using Id2DistVec = std::vector; +using ResultSet = std::vector; class SearchJob : public Job { public: @@ -56,11 +55,8 @@ class SearchJob : public Job { void SearchDone(size_t index_id); - ResultIds& - GetResultIds(); - - ResultDistances& - GetResultDistances(); + ResultSet& + GetResult(); Status& GetStatus(); @@ -94,11 +90,6 @@ class SearchJob : public Job { return index_files_; } - std::mutex& - mutex() { - return mutex_; - } - private: uint64_t topk_ = 0; uint64_t nq_ = 0; @@ -108,8 +99,7 @@ class SearchJob : public Job { Id2IndexMap index_files_; // TODO: column-base better ? - ResultIds result_ids_; - ResultDistances result_distances_; + ResultSet result_; Status status_; std::mutex mutex_; diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.cpp b/core/src/scheduler/optimizer/OnlyGPUPass.cpp index e5d3c71fd3..3fcda0e8a3 100644 --- a/core/src/scheduler/optimizer/OnlyGPUPass.cpp +++ b/core/src/scheduler/optimizer/OnlyGPUPass.cpp @@ -46,7 +46,7 @@ OnlyGPUPass::Run(const TaskPtr& task) { auto label = std::make_shared(std::weak_ptr(res_ptr)); task->label() = label; - specified_gpu_id_ = (specified_gpu_id_ + 1) % gpu_id.size(); + specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size(); return true; } diff --git a/core/src/scheduler/task/SearchTask.cpp b/core/src/scheduler/task/SearchTask.cpp index 08bc6525aa..1bf1caff76 100644 --- a/core/src/scheduler/task/SearchTask.cpp +++ b/core/src/scheduler/task/SearchTask.cpp @@ -219,11 +219,8 @@ XSearchTask::Execute() { // step 3: pick up topk result auto spec_k = index_engine_->Count() < topk ? index_engine_->Count() : topk; - { - std::unique_lock lock(search_job->mutex()); - XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, metric_l2, - search_job->GetResultIds(), search_job->GetResultDistances()); - } + XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, metric_l2, + search_job->GetResult()); span = rc.RecordSection(hdr + ", reduce topk"); // search_job->AccumReduceCost(span); @@ -243,69 +240,71 @@ XSearchTask::Execute() { } void -XSearchTask::MergeTopkToResultSet(const scheduler::ResultIds& src_ids, const scheduler::ResultDistances& src_distances, - size_t src_k, size_t nq, size_t topk, bool ascending, scheduler::ResultIds& tar_ids, - scheduler::ResultDistances& tar_distances) { - if (src_ids.empty()) { - return; +XSearchTask::MergeTopkToResultSet(const std::vector& input_ids, const std::vector& input_distance, + uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending, + scheduler::ResultSet& result) { + if (result.empty()) { + result.resize(nq); } - size_t tar_k = tar_ids.size() / nq; - size_t buf_k = std::min(topk, src_k + tar_k); - - scheduler::ResultIds buf_ids(nq * buf_k, -1); - scheduler::ResultDistances buf_distances(nq * buf_k, 0.0); - for (uint64_t i = 0; i < nq; i++) { - size_t buf_k_j = 0, src_k_j = 0, tar_k_j = 0; - size_t buf_idx, src_idx, tar_idx; + scheduler::Id2DistVec result_buf; + auto& result_i = result[i]; - size_t buf_k_multi_i = buf_k * i; - size_t src_k_multi_i = topk * i; - size_t tar_k_multi_i = tar_k * i; - - while (buf_k_j < buf_k && src_k_j < src_k && tar_k_j < tar_k) { - src_idx = src_k_multi_i + src_k_j; - tar_idx = tar_k_multi_i + tar_k_j; - buf_idx = buf_k_multi_i + buf_k_j; - - if ((ascending && src_distances[src_idx] < tar_distances[tar_idx]) || - (!ascending && src_distances[src_idx] > tar_distances[tar_idx])) { - buf_ids[buf_idx] = src_ids[src_idx]; - buf_distances[buf_idx] = src_distances[src_idx]; - src_k_j++; - } else { - buf_ids[buf_idx] = tar_ids[tar_idx]; - buf_distances[buf_idx] = tar_distances[tar_idx]; - tar_k_j++; + if (result[i].empty()) { + result_buf.resize(input_k, scheduler::IdDistPair(-1, 0.0)); + uint64_t input_k_multi_i = topk * i; + for (auto k = 0; k < input_k; ++k) { + uint64_t idx = input_k_multi_i + k; + auto& result_buf_item = result_buf[k]; + result_buf_item.first = input_ids[idx]; + result_buf_item.second = input_distance[idx]; } - buf_k_j++; - } - - if (buf_k_j < buf_k) { - if (src_k_j < src_k) { - while (buf_k_j < buf_k && src_k_j < src_k) { - buf_idx = buf_k_multi_i + buf_k_j; - src_idx = src_k_multi_i + src_k_j; - buf_ids[buf_idx] = src_ids[src_idx]; - buf_distances[buf_idx] = src_distances[src_idx]; - src_k_j++; - buf_k_j++; + } else { + size_t tar_size = result_i.size(); + uint64_t output_k = std::min(topk, input_k + tar_size); + result_buf.resize(output_k, scheduler::IdDistPair(-1, 0.0)); + size_t buf_k = 0, src_k = 0, tar_k = 0; + uint64_t src_idx; + uint64_t input_k_multi_i = topk * i; + while (buf_k < output_k && src_k < input_k && tar_k < tar_size) { + src_idx = input_k_multi_i + src_k; + auto& result_buf_item = result_buf[buf_k]; + auto& result_item = result_i[tar_k]; + if ((ascending && input_distance[src_idx] < result_item.second) || + (!ascending && input_distance[src_idx] > result_item.second)) { + result_buf_item.first = input_ids[src_idx]; + result_buf_item.second = input_distance[src_idx]; + src_k++; + } else { + result_buf_item = result_item; + tar_k++; } - } else { - while (buf_k_j < buf_k && tar_k_j < tar_k) { - buf_idx = buf_k_multi_i + buf_k_j; - tar_idx = tar_k_multi_i + tar_k_j; - buf_ids[buf_idx] = tar_ids[tar_idx]; - buf_distances[buf_idx] = tar_distances[tar_idx]; - tar_k_j++; - buf_k_j++; + buf_k++; + } + + if (buf_k < output_k) { + if (src_k < input_k) { + while (buf_k < output_k && src_k < input_k) { + src_idx = input_k_multi_i + src_k; + auto& result_buf_item = result_buf[buf_k]; + result_buf_item.first = input_ids[src_idx]; + result_buf_item.second = input_distance[src_idx]; + src_k++; + buf_k++; + } + } else { + while (buf_k < output_k && tar_k < tar_size) { + result_buf[buf_k] = result_i[tar_k]; + tar_k++; + buf_k++; + } } } } + + result_i.swap(result_buf); } - tar_ids.swap(buf_ids); - tar_distances.swap(buf_distances); } // void diff --git a/core/src/scheduler/task/SearchTask.h b/core/src/scheduler/task/SearchTask.h index bd51137341..bbc8b5bd8f 100644 --- a/core/src/scheduler/task/SearchTask.h +++ b/core/src/scheduler/task/SearchTask.h @@ -39,9 +39,8 @@ class XSearchTask : public Task { public: static void - MergeTopkToResultSet(const scheduler::ResultIds& src_ids, const scheduler::ResultDistances& src_distances, - size_t src_k, size_t nq, size_t topk, bool ascending, scheduler::ResultIds& tar_ids, - scheduler::ResultDistances& tar_distances); + MergeTopkToResultSet(const std::vector& input_ids, const std::vector& input_distance, + uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending, scheduler::ResultSet& result); // static void // MergeTopkArray(std::vector& tar_ids, std::vector& tar_distance, uint64_t& tar_input_k, diff --git a/core/src/sdk/grpc/ClientProxy.cpp b/core/src/sdk/grpc/ClientProxy.cpp index c726cfc532..91a11adf8c 100644 --- a/core/src/sdk/grpc/ClientProxy.cpp +++ b/core/src/sdk/grpc/ClientProxy.cpp @@ -17,7 +17,7 @@ #include "sdk/grpc/ClientProxy.h" #include "grpc/gen-milvus/milvus.grpc.pb.h" -#include "src/version.h" +#include "src/config.h" #include #include diff --git a/core/src/server/Config.cpp b/core/src/server/Config.cpp index cc88dccffa..81a0fd7042 100644 --- a/core/src/server/Config.cpp +++ b/core/src/server/Config.cpp @@ -162,6 +162,7 @@ Config::ValidateConfig() { return s; } +#ifdef MILVUS_GPU_VERSION int64_t cache_gpu_cache_capacity; s = GetCacheConfigGpuCacheCapacity(cache_gpu_cache_capacity); if (!s.ok()) { @@ -173,6 +174,7 @@ Config::ValidateConfig() { if (!s.ok()) { return s; } +#endif bool cache_insert_data; s = GetCacheConfigCacheInsertData(cache_insert_data); @@ -401,7 +403,8 @@ Status Config::CheckServerConfigDeployMode(const std::string& value) { if (value != "single" && value != "cluster_readonly" && value != "cluster_writable") { return Status(SERVER_INVALID_ARGUMENT, - "server_config.deploy_mode is not one of single, cluster_readonly, and cluster_writable."); + "server_config.deploy_mode is not one of " + "single, cluster_readonly, and cluster_writable."); } return Status::OK(); } @@ -589,18 +592,15 @@ Config::CheckCacheConfigGpuCacheCapacity(const std::string& value) { return Status(SERVER_INVALID_ARGUMENT, msg); } else { uint64_t gpu_cache_capacity = std::stoi(value) * GB; - int device_id; - Status s = GetResourceConfigIndexBuildDevice(device_id); + int gpu_index; + Status s = GetResourceConfigIndexBuildDevice(gpu_index); if (!s.ok()) { return s; } - if (device_id == server::CPU_DEVICE_ID) - return Status::OK(); - size_t gpu_memory; - if (!ValidationUtil::GetGpuMemory(device_id, gpu_memory).ok()) { - std::string msg = "Fail to get GPU memory for GPU device: " + std::to_string(device_id); + if (!ValidationUtil::GetGpuMemory(gpu_index, gpu_memory).ok()) { + std::string msg = "Fail to get GPU memory for GPU device: " + std::to_string(gpu_index); return Status(SERVER_UNEXPECTED_ERROR, msg); } else if (gpu_cache_capacity >= gpu_memory) { std::string msg = "Invalid gpu cache capacity: " + value + @@ -689,21 +689,31 @@ Config::CheckResourceConfigMode(const std::string& value) { } Status -CheckGpuDevice(const std::string& value) { - const std::regex pat("gpu(\\d+)"); - std::cmatch m; - if (!std::regex_match(value.c_str(), m, pat)) { - std::string msg = "Invalid gpu device: " + value + - ". Possible reason: resource_config.search_resources does not match your hardware."; +CheckResource(const std::string& value) { + std::string s = value; + std::transform(s.begin(), s.end(), s.begin(), ::tolower); +#ifdef MILVUS_CPU_VERSION + if (s != "cpu") { + return Status(SERVER_INVALID_ARGUMENT, "Invalid CPU resource: " + s); + } +#else + const std::regex pat("cpu|gpu(\\d+)"); + std::smatch m; + if (!std::regex_match(s, m, pat)) { + std::string msg = "Invalid search resource: " + value + + ". Possible reason: resource_config.search_resources is not in the format of cpux or gpux"; return Status(SERVER_INVALID_ARGUMENT, msg); } - int32_t gpu_index = std::stoi(value.substr(3)); - if (!ValidationUtil::ValidateGpuIndex(gpu_index).ok()) { - std::string msg = "Invalid gpu device: " + value + + if (s.compare(0, 3, "gpu") == 0) { + int32_t gpu_index = std::stoi(s.substr(3)); + if (!ValidationUtil::ValidateGpuIndex(gpu_index).ok()) { + std::string msg = "Invalid search resource: " + value + ". Possible reason: resource_config.search_resources does not match your hardware."; - return Status(SERVER_INVALID_ARGUMENT, msg); + return Status(SERVER_INVALID_ARGUMENT, msg); + } } +#endif return Status::OK(); } @@ -716,14 +726,10 @@ Config::CheckResourceConfigSearchResources(const std::vector& value return Status(SERVER_INVALID_ARGUMENT, msg); } - for (auto& device : value) { - if (device == "cpu") { - continue; - } - if (!CheckGpuDevice(device).ok()) { - std::string msg = "Invalid search resource: " + device + - ". Possible reason: resource_config.search_resources does not match your hardware."; - return Status(SERVER_INVALID_ARGUMENT, msg); + for (auto& resource : value) { + auto status = CheckResource(resource); + if (!status.ok()) { + return Status(SERVER_INVALID_ARGUMENT, status.message()); } } return Status::OK(); @@ -731,13 +737,9 @@ Config::CheckResourceConfigSearchResources(const std::vector& value Status Config::CheckResourceConfigIndexBuildDevice(const std::string& value) { - if (value == "cpu") { - return Status::OK(); - } - if (!CheckGpuDevice(value).ok()) { - std::string msg = "Invalid index build device: " + value + - ". Possible reason: resource_config.index_build_device does not match your hardware."; - return Status(SERVER_INVALID_ARGUMENT, msg); + auto status = CheckResource(value); + if (!status.ok()) { + return Status(SERVER_INVALID_ARGUMENT, status.message()); } return Status::OK(); } @@ -1015,11 +1017,12 @@ Config::GetResourceConfigIndexBuildDevice(int32_t& value) { return s; } - if (str == "cpu") { - value = CPU_DEVICE_ID; - } else { + if (str != "cpu") { value = std::stoi(str.substr(3)); } + else { + value = -1; + } return Status::OK(); } diff --git a/core/src/server/Config.h b/core/src/server/Config.h index 45591fb5ec..e8a396d568 100644 --- a/core/src/server/Config.h +++ b/core/src/server/Config.h @@ -93,9 +93,7 @@ static const char* CONFIG_RESOURCE_MODE = "mode"; static const char* CONFIG_RESOURCE_MODE_DEFAULT = "simple"; static const char* CONFIG_RESOURCE_SEARCH_RESOURCES = "search_resources"; static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE = "index_build_device"; -static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE_DEFAULT = "gpu0"; - -const int32_t CPU_DEVICE_ID = -1; +static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE_DEFAULT = "cpu"; class Config { public: diff --git a/core/src/server/Server.cpp b/core/src/server/Server.cpp index 5676504722..eb2b077829 100644 --- a/core/src/server/Server.cpp +++ b/core/src/server/Server.cpp @@ -25,7 +25,7 @@ #include "server/DBWrapper.h" #include "server/Server.h" #include "server/grpc_impl/GrpcServer.h" -#include "src/version.h" +#include "src/config.h" #include "utils/Log.h" #include "utils/LogUtil.h" #include "utils/SignalUtil.h" diff --git a/core/src/server/grpc_impl/GrpcRequestTask.cpp b/core/src/server/grpc_impl/GrpcRequestTask.cpp index 77f262bda6..0816d45750 100644 --- a/core/src/server/grpc_impl/GrpcRequestTask.cpp +++ b/core/src/server/grpc_impl/GrpcRequestTask.cpp @@ -28,7 +28,7 @@ #include "scheduler/SchedInst.h" #include "server/DBWrapper.h" #include "server/Server.h" -#include "src/version.h" +#include "src/config.h" #include "utils/CommonUtil.h" #include "utils/Log.h" #include "utils/TimeRecorder.h" @@ -637,8 +637,7 @@ SearchTask::OnExecute() { rc.RecordSection("prepare vector data"); // step 6: search vectors - engine::ResultIds result_ids; - engine::ResultDistances result_distances; + engine::QueryResults results; auto record_count = (uint64_t)search_param_->query_record_array().size(); #ifdef MILVUS_ENABLE_PROFILING @@ -648,11 +647,11 @@ SearchTask::OnExecute() { #endif if (file_id_array_.empty()) { - status = DBWrapper::DB()->Query(table_name_, (size_t)top_k, record_count, nprobe, vec_f.data(), dates, - result_ids, result_distances); + status = + DBWrapper::DB()->Query(table_name_, (size_t)top_k, record_count, nprobe, vec_f.data(), dates, results); } else { status = DBWrapper::DB()->Query(table_name_, file_id_array_, (size_t)top_k, record_count, nprobe, - vec_f.data(), dates, result_ids, result_distances); + vec_f.data(), dates, results); } #ifdef MILVUS_ENABLE_PROFILING @@ -664,20 +663,23 @@ SearchTask::OnExecute() { return status; } - if (result_ids.empty()) { + if (results.empty()) { return Status::OK(); // empty table } - size_t result_k = result_ids.size() / record_count; + if (results.size() != record_count) { + std::string msg = "Search " + std::to_string(record_count) + " vectors but only return " + + std::to_string(results.size()) + " results"; + return Status(SERVER_ILLEGAL_SEARCH_RESULT, msg); + } // step 7: construct result array - for (size_t i = 0; i < record_count; i++) { + for (auto& result : results) { ::milvus::grpc::TopKQueryResult* topk_query_result = topk_result_list->add_topk_query_result(); - for (size_t j = 0; j < result_k; j++) { + for (auto& pair : result) { ::milvus::grpc::QueryResult* grpc_result = topk_query_result->add_query_result_arrays(); - size_t idx = i * result_k + j; - grpc_result->set_id(result_ids[idx]); - grpc_result->set_distance(result_distances[idx]); + grpc_result->set_id(pair.first); + grpc_result->set_distance(pair.second); } } diff --git a/core/src/utils/ValidationUtil.cpp b/core/src/utils/ValidationUtil.cpp index 68088d6c93..347ba44e85 100644 --- a/core/src/utils/ValidationUtil.cpp +++ b/core/src/utils/ValidationUtil.cpp @@ -20,7 +20,9 @@ #include "db/engine/ExecutionEngine.h" #include +#ifdef MILVUS_GPU_VERSION #include +#endif #include #include #include @@ -168,6 +170,8 @@ ValidationUtil::ValidateSearchNprobe(int64_t nprobe, const engine::meta::TableSc Status ValidationUtil::ValidateGpuIndex(uint32_t gpu_index) { + +#ifdef MILVUS_GPU_VERSION int num_devices = 0; auto cuda_err = cudaGetDeviceCount(&num_devices); if (cuda_err != cudaSuccess) { @@ -181,12 +185,16 @@ ValidationUtil::ValidateGpuIndex(uint32_t gpu_index) { SERVER_LOG_ERROR << msg; return Status(SERVER_INVALID_ARGUMENT, msg); } +#endif return Status::OK(); } Status ValidationUtil::GetGpuMemory(uint32_t gpu_index, size_t& memory) { + +#ifdef MILVUS_GPU_VERSION + cudaDeviceProp deviceProp; auto cuda_err = cudaGetDeviceProperties(&deviceProp, gpu_index); if (cuda_err) { @@ -196,6 +204,8 @@ ValidationUtil::GetGpuMemory(uint32_t gpu_index, size_t& memory) { } memory = deviceProp.totalGlobalMem; +#endif + return Status::OK(); } diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp index 4ac8e22f52..2dcf6bab7e 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/wrapper/ConfAdapter.cpp @@ -109,7 +109,7 @@ IVFSQConfAdapter::Match(const TempMetaConf& metaconf) { conf->nlist = MatchNlist(metaconf.size, metaconf.nlist); conf->d = metaconf.dim; conf->metric_type = metaconf.metric_type; - conf->gpu_id = metaconf.gpu_id; + conf->gpu_id = conf->gpu_id; conf->nbits = 8; MatchBase(conf); return conf; diff --git a/core/src/wrapper/KnowhereResource.cpp b/core/src/wrapper/KnowhereResource.cpp index 650ae727c1..de3b909b08 100644 --- a/core/src/wrapper/KnowhereResource.cpp +++ b/core/src/wrapper/KnowhereResource.cpp @@ -16,7 +16,9 @@ // under the License. #include "wrapper/KnowhereResource.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif #include "server/Config.h" #include @@ -32,6 +34,9 @@ constexpr int64_t M_BYTE = 1024 * 1024; Status KnowhereResource::Initialize() { + +#ifdef MILVUS_GPU_VERSION + struct GpuResourceSetting { int64_t pinned_memory = 300 * M_BYTE; int64_t temp_memory = 300 * M_BYTE; @@ -73,12 +78,16 @@ KnowhereResource::Initialize() { iter->second.temp_memory, iter->second.resource_num); } +#endif + return Status::OK(); } Status KnowhereResource::Finalize() { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource. +#endif return Status::OK(); } diff --git a/core/src/wrapper/VecImpl.cpp b/core/src/wrapper/VecImpl.cpp index c97900f839..8bf979e668 100644 --- a/core/src/wrapper/VecImpl.cpp +++ b/core/src/wrapper/VecImpl.cpp @@ -18,11 +18,16 @@ #include "wrapper/VecImpl.h" #include "DataTransfer.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/index/vector_index/IndexIDMAP.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" -#include "knowhere/index/vector_index/helpers/Cloner.h" #include "utils/Log.h" +#include "wrapper/WrapperException.h" + +#ifdef MILVUS_GPU_VERSION + +#include +#include + +#endif /* * no parameter check in this layer. @@ -30,326 +35,216 @@ */ namespace milvus { -namespace engine { + namespace engine { -Status -VecIndexImpl::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) { - try { - dim = cfg->d; - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); + Status + VecIndexImpl::BuildAll(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg, + const int64_t &nt, + const float *xt) { + try { + dim = cfg->d; + auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - auto preprocessor = index_->BuildPreprocessor(dataset, cfg); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(dataset, cfg); - index_->set_index_model(model); - index_->Add(dataset, cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} + auto preprocessor = index_->BuildPreprocessor(dataset, cfg); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(dataset, cfg); + index_->set_index_model(model); + index_->Add(dataset, cfg); + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); + } -Status -VecIndexImpl::Add(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg) { - try { - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); + Status + VecIndexImpl::Add(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg) { + try { + auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - index_->Add(dataset, cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} + index_->Add(dataset, cfg); + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); + } -Status -VecIndexImpl::Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, const Config& cfg) { - try { - auto k = cfg->k; - auto dataset = GenDataset(nq, dim, xq); + Status + VecIndexImpl::Search(const int64_t &nq, const float *xq, float *dist, int64_t *ids, const Config &cfg) { + try { + auto k = cfg->k; + auto dataset = GenDataset(nq, dim, xq); - Config search_cfg = cfg; + Config search_cfg = cfg; - auto res = index_->Search(dataset, search_cfg); - auto ids_array = res->array()[0]; - auto dis_array = res->array()[1]; + auto res = index_->Search(dataset, search_cfg); + auto ids_array = res->array()[0]; + auto dis_array = res->array()[1]; - //{ - // auto& ids = ids_array; - // auto& dists = dis_array; - // std::stringstream ss_id; - // std::stringstream ss_dist; - // for (auto i = 0; i < 10; i++) { - // for (auto j = 0; j < k; ++j) { - // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - // } - // ss_id << std::endl; - // ss_dist << std::endl; - // } - // std::cout << "id\n" << ss_id.str() << std::endl; - // std::cout << "dist\n" << ss_dist.str() << std::endl; - //} + //{ + // auto& ids = ids_array; + // auto& dists = dis_array; + // std::stringstream ss_id; + // std::stringstream ss_dist; + // for (auto i = 0; i < 10; i++) { + // for (auto j = 0; j < k; ++j) { + // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; + // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; + // } + // ss_id << std::endl; + // ss_dist << std::endl; + // } + // std::cout << "id\n" << ss_id.str() << std::endl; + // std::cout << "dist\n" << ss_dist.str() << std::endl; + //} - auto p_ids = ids_array->data()->GetValues(1, 0); - auto p_dist = dis_array->data()->GetValues(1, 0); + auto p_ids = ids_array->data()->GetValues(1, 0); + auto p_dist = dis_array->data()->GetValues(1, 0); - // TODO(linxj): avoid copy here. - memcpy(ids, p_ids, sizeof(int64_t) * nq * k); - memcpy(dist, p_dist, sizeof(float) * nq * k); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} + // TODO(linxj): avoid copy here. + memcpy(ids, p_ids, sizeof(int64_t) * nq * k); + memcpy(dist, p_dist, sizeof(float) * nq * k); + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); + } -knowhere::BinarySet -VecIndexImpl::Serialize() { - type = ConvertToCpuIndexType(type); - return index_->Serialize(); -} - -Status -VecIndexImpl::Load(const knowhere::BinarySet& index_binary) { - index_->Load(index_binary); - dim = Dimension(); - return Status::OK(); -} - -int64_t -VecIndexImpl::Dimension() { - return index_->Dimension(); -} - -int64_t -VecIndexImpl::Count() { - return index_->Count(); -} - -IndexType -VecIndexImpl::GetType() { - return type; -} - -VecIndexPtr -VecIndexImpl::CopyToGpu(const int64_t& device_id, const Config& cfg) { - // TODO(linxj): exception handle - auto gpu_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, cfg); - auto new_index = std::make_shared(gpu_index, ConvertToGpuIndexType(type)); - new_index->dim = dim; - return new_index; -} - -VecIndexPtr -VecIndexImpl::CopyToCpu(const Config& cfg) { - // TODO(linxj): exception handle - auto cpu_index = knowhere::cloner::CopyGpuToCpu(index_, cfg); - auto new_index = std::make_shared(cpu_index, ConvertToCpuIndexType(type)); - new_index->dim = dim; - return new_index; -} - -VecIndexPtr -VecIndexImpl::Clone() { - // TODO(linxj): exception handle - auto clone_index = std::make_shared(index_->Clone(), type); - clone_index->dim = dim; - return clone_index; -} - -int64_t -VecIndexImpl::GetDeviceId() { - if (auto device_idx = std::dynamic_pointer_cast(index_)) { - return device_idx->GetGpuDevice(); - } - // else - return -1; // -1 == cpu -} - -float* -BFIndex::GetRawVectors() { - auto raw_index = std::dynamic_pointer_cast(index_); - if (raw_index) { - return raw_index->GetRawVectors(); - } - return nullptr; -} - -int64_t* -BFIndex::GetRawIds() { - return std::static_pointer_cast(index_)->GetRawIds(); -} - -ErrorCode -BFIndex::Build(const Config& cfg) { - try { - dim = cfg->d; - std::static_pointer_cast(index_)->Train(cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_UNEXPECTED_ERROR; - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_ERROR; - } - return KNOWHERE_SUCCESS; -} - -Status -BFIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) { - try { - dim = cfg->d; - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - - std::static_pointer_cast(index_)->Train(cfg); - index_->Add(dataset, cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} - -// TODO(linxj): add lock here. -Status -IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) { - try { - dim = cfg->d; - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - - auto preprocessor = index_->BuildPreprocessor(dataset, cfg); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(dataset, cfg); - index_->set_index_model(model); - index_->Add(dataset, cfg); - - if (auto device_index = std::dynamic_pointer_cast(index_)) { - auto host_index = device_index->CopyGpuToCpu(Config()); - index_ = host_index; + knowhere::BinarySet + VecIndexImpl::Serialize() { type = ConvertToCpuIndexType(type); - } else { - WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed"; - return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed"); + return index_->Serialize(); } - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} -Status -IVFMixIndex::Load(const knowhere::BinarySet& index_binary) { - index_->Load(index_binary); - dim = Dimension(); - return Status::OK(); -} - -knowhere::QuantizerPtr -IVFHybridIndex::LoadQuantizer(const Config& conf) { - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - return new_idx->LoadQuantizer(conf); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - } -} - -Status -IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) { - try { - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - new_idx->SetQuantizer(q); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - return Status(KNOWHERE_ERROR, "not support"); + Status + VecIndexImpl::Load(const knowhere::BinarySet &index_binary) { + index_->Load(index_binary); + dim = Dimension(); + return Status::OK(); } - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} -Status -IVFHybridIndex::UnsetQuantizer() { - try { - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - new_idx->UnsetQuantizer(); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - return Status(KNOWHERE_ERROR, "not support"); + int64_t + VecIndexImpl::Dimension() { + return index_->Dimension(); } - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return Status(KNOWHERE_ERROR, e.what()); - } - return Status::OK(); -} -VecIndexPtr -IVFHybridIndex::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { - try { - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - return std::make_shared(new_idx->LoadData(q, conf), type); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + int64_t + VecIndexImpl::Count() { + return index_->Count(); } - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - } - return nullptr; -} -std::pair -IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) { - try { - // TODO(linxj): Hardcode here - if (auto hybrid_idx = std::dynamic_pointer_cast(index_)) { - auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg); - auto new_idx = std::make_shared(pair.first, type); - return std::make_pair(new_idx, pair.second); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + IndexType + VecIndexImpl::GetType() { + return type; } - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - } - return std::make_pair(nullptr, nullptr); -} -} // namespace engine + VecIndexPtr + VecIndexImpl::CopyToGpu(const int64_t &device_id, const Config &cfg) { + // TODO(linxj): exception handle +#ifdef MILVUS_GPU_VERSION + auto gpu_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, cfg); + auto new_index = std::make_shared(gpu_index, ConvertToGpuIndexType(type)); + new_index->dim = dim; + return new_index; +#else + WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToGpu when we are using CPU version"; + throw WrapperException("Calling VecIndexImpl::CopyToGpu when we are using CPU version"); +#endif + } + + VecIndexPtr + VecIndexImpl::CopyToCpu(const Config &cfg) { + // TODO(linxj): exception handle +#ifdef MILVUS_GPU_VERSION + auto cpu_index = knowhere::cloner::CopyGpuToCpu(index_, cfg); + auto new_index = std::make_shared(cpu_index, ConvertToCpuIndexType(type)); + new_index->dim = dim; + return new_index; +#else + WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToCpu when we are using CPU version"; + throw WrapperException("Calling VecIndexImpl::CopyToCpu when we are using CPU version"); +#endif + + } + + VecIndexPtr + VecIndexImpl::Clone() { + // TODO(linxj): exception handle + auto clone_index = std::make_shared(index_->Clone(), type); + clone_index->dim = dim; + return clone_index; + } + + int64_t + VecIndexImpl::GetDeviceId() { +#ifdef MILVUS_GPU_VERSION + if (auto device_idx = std::dynamic_pointer_cast(index_)) { + return device_idx->GetGpuDevice(); + } +#else + // else + return -1; // -1 == cpu +#endif + } + + float * + BFIndex::GetRawVectors() { + auto raw_index = std::dynamic_pointer_cast(index_); + if (raw_index) { + return raw_index->GetRawVectors(); + } + return nullptr; + } + + int64_t * + BFIndex::GetRawIds() { + return std::static_pointer_cast(index_)->GetRawIds(); + } + + ErrorCode + BFIndex::Build(const Config &cfg) { + try { + dim = cfg->d; + std::static_pointer_cast(index_)->Train(cfg); + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return KNOWHERE_UNEXPECTED_ERROR; + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return KNOWHERE_ERROR; + } + return KNOWHERE_SUCCESS; + } + + Status + BFIndex::BuildAll(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg, const int64_t &nt, + const float *xt) { + try { + dim = cfg->d; + auto dataset = GenDatasetWithIds(nb, dim, xb, ids); + + std::static_pointer_cast(index_)->Train(cfg); + index_->Add(dataset, cfg); + } catch (knowhere::KnowhereException &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception &e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); + } + + } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/VecImpl.h b/core/src/wrapper/VecImpl.h index 22d734cf92..2b6f07827e 100644 --- a/core/src/wrapper/VecImpl.h +++ b/core/src/wrapper/VecImpl.h @@ -77,41 +77,6 @@ class VecIndexImpl : public VecIndex { std::shared_ptr index_ = nullptr; }; -class IVFMixIndex : public VecIndexImpl { - public: - explicit IVFMixIndex(std::shared_ptr index, const IndexType& type) - : VecIndexImpl(std::move(index), type) { - } - - Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) override; - - Status - Load(const knowhere::BinarySet& index_binary) override; -}; - -class IVFHybridIndex : public IVFMixIndex { - public: - explicit IVFHybridIndex(std::shared_ptr index, const IndexType& type) - : IVFMixIndex(std::move(index), type) { - } - - knowhere::QuantizerPtr - LoadQuantizer(const Config& conf) override; - - Status - SetQuantizer(const knowhere::QuantizerPtr& q) override; - - Status - UnsetQuantizer() override; - std::pair - CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) override; - - VecIndexPtr - LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override; -}; - class BFIndex : public VecIndexImpl { public: explicit BFIndex(std::shared_ptr index) diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp index abf97e69e5..d9e252d7ed 100644 --- a/core/src/wrapper/VecIndex.cpp +++ b/core/src/wrapper/VecIndex.cpp @@ -18,19 +18,23 @@ #include "wrapper/VecIndex.h" #include "VecImpl.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/IndexKDT.h" #include "knowhere/index/vector_index/IndexNSG.h" #include "utils/Log.h" +#ifdef MILVUS_GPU_VERSION #include +#include "wrapper/gpu/GPUVecImpl.h" +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#endif namespace milvus { namespace engine { @@ -119,43 +123,46 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { index = std::make_shared(); break; } - case IndexType::FAISS_IVFFLAT_GPU: { - index = std::make_shared(gpu_device); - break; - } - case IndexType::FAISS_IVFFLAT_MIX: { - index = std::make_shared(gpu_device); - return std::make_shared(index, IndexType::FAISS_IVFFLAT_MIX); - } case IndexType::FAISS_IVFPQ_CPU: { index = std::make_shared(); break; } - case IndexType::FAISS_IVFPQ_GPU: { - index = std::make_shared(gpu_device); - break; - } case IndexType::SPTAG_KDT_RNT_CPU: { index = std::make_shared(); break; } - case IndexType::FAISS_IVFSQ8_MIX: { - index = std::make_shared(gpu_device); - return std::make_shared(index, IndexType::FAISS_IVFSQ8_MIX); - } case IndexType::FAISS_IVFSQ8_CPU: { index = std::make_shared(); break; } + +#ifdef MILVUS_GPU_VERSION + case IndexType::FAISS_IVFFLAT_GPU: { + index = std::make_shared(gpu_device); + break; + } + case IndexType::FAISS_IVFPQ_GPU: { + index = std::make_shared(gpu_device); + break; + } + case IndexType::FAISS_IVFSQ8_MIX: { + index = std::make_shared(gpu_device); + return std::make_shared(index, IndexType::FAISS_IVFSQ8_MIX); + } case IndexType::FAISS_IVFSQ8_GPU: { index = std::make_shared(gpu_device); break; } + case IndexType::FAISS_IVFFLAT_MIX: { + index = std::make_shared(gpu_device); + return std::make_shared(index, IndexType::FAISS_IVFFLAT_MIX); + } #ifdef CUSTOMIZATION case IndexType::FAISS_IVFSQ8_HYBRID: { index = std::make_shared(gpu_device); return std::make_shared(index, IndexType::FAISS_IVFSQ8_HYBRID); } +#endif #endif case IndexType::NSG_MIX: { // TODO(linxj): bug. index = std::make_shared(gpu_device); diff --git a/core/src/version.h.macro b/core/src/wrapper/WrapperException.cpp similarity index 74% rename from core/src/version.h.macro rename to core/src/wrapper/WrapperException.cpp index 454d8a990a..611a3571da 100644 --- a/core/src/version.h.macro +++ b/core/src/wrapper/WrapperException.cpp @@ -15,8 +15,16 @@ // specific language governing permissions and limitations // under the License. -#pragma once +#include "wrapper/WrapperException.h" -#define MILVUS_VERSION "@MILVUS_VERSION@" -#define BUILD_TYPE "@BUILD_TYPE@" -#define BUILD_TIME @BUILD_TIME@ \ No newline at end of file +namespace milvus { +namespace engine { + +WrapperException::WrapperException(const std::string &msg) : msg(msg) {} + +const char *WrapperException::what() const noexcept { + return msg.c_str(); +} + +} // namespace engine +} // namespace milvus diff --git a/core/src/wrapper/WrapperException.h b/core/src/wrapper/WrapperException.h new file mode 100644 index 0000000000..c2e2748a43 --- /dev/null +++ b/core/src/wrapper/WrapperException.h @@ -0,0 +1,36 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +namespace milvus { +namespace engine { + +class WrapperException : public std::exception { + public: + explicit WrapperException(const std::string& msg); + + const char* what() const noexcept override; + + const std::string msg; +}; + +} // namespace engine +} // namespace milvus diff --git a/core/src/wrapper/gpu/GPUVecImpl.cpp b/core/src/wrapper/gpu/GPUVecImpl.cpp new file mode 100644 index 0000000000..4ff5e665cc --- /dev/null +++ b/core/src/wrapper/gpu/GPUVecImpl.cpp @@ -0,0 +1,164 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "wrapper/VecImpl.h" +#include "GPUVecImpl.h" +#include "src/wrapper/DataTransfer.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" +#include "utils/Log.h" + +/* + * no parameter check in this layer. + * only responible for index combination + */ + +namespace milvus { +namespace engine { + + +// TODO(linxj): add lock here. +Status +IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, + const float* xt) { + try { + dim = cfg->d; + auto dataset = GenDatasetWithIds(nb, dim, xb, ids); + + auto preprocessor = index_->BuildPreprocessor(dataset, cfg); + index_->set_preprocessor(preprocessor); + auto model = index_->Train(dataset, cfg); + index_->set_index_model(model); + index_->Add(dataset, cfg); + + if (auto device_index = std::dynamic_pointer_cast(index_)) { + auto host_index = device_index->CopyGpuToCpu(Config()); + index_ = host_index; + type = ConvertToCpuIndexType(type); + } else { + WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed"; + return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed"); + } + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); +} + +Status +IVFMixIndex::Load(const knowhere::BinarySet& index_binary) { + index_->Load(index_binary); + dim = Dimension(); + return Status::OK(); +} + +knowhere::QuantizerPtr +IVFHybridIndex::LoadQuantizer(const Config& conf) { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + return new_idx->LoadQuantizer(conf); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + } +} + +Status +IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + new_idx->SetQuantizer(q); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + return Status(KNOWHERE_ERROR, "not support"); + } + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); +} + +Status +IVFHybridIndex::UnsetQuantizer() { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + new_idx->UnsetQuantizer(); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + return Status(KNOWHERE_ERROR, "not support"); + } + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + return Status(KNOWHERE_ERROR, e.what()); + } + return Status::OK(); +} + +VecIndexPtr +IVFHybridIndex::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { + try { + // TODO(linxj): Hardcode here + if (auto new_idx = std::dynamic_pointer_cast(index_)) { + return std::make_shared(new_idx->LoadData(q, conf), type); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + } + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + } + return nullptr; +} + +std::pair +IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) { + try { + // TODO(linxj): Hardcode here + if (auto hybrid_idx = std::dynamic_pointer_cast(index_)) { + auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg); + auto new_idx = std::make_shared(pair.first, type); + return std::make_pair(new_idx, pair.second); + } else { + WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); + } + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + } + return std::make_pair(nullptr, nullptr); +} + +} // namespace engine +} // namespace milvus diff --git a/core/src/wrapper/gpu/GPUVecImpl.h b/core/src/wrapper/gpu/GPUVecImpl.h new file mode 100644 index 0000000000..39df216185 --- /dev/null +++ b/core/src/wrapper/gpu/GPUVecImpl.h @@ -0,0 +1,66 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +#include "wrapper/VecIndex.h" +#include "knowhere/index/vector_index/VectorIndex.h" + +namespace milvus { +namespace engine { + +class IVFMixIndex : public VecIndexImpl { +public: + explicit IVFMixIndex(std::shared_ptr index, const IndexType &type) + : VecIndexImpl(std::move(index), type) { + } + + Status + BuildAll(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg, const int64_t &nt, + const float *xt) override; + + Status + Load(const knowhere::BinarySet &index_binary) override; +}; + +class IVFHybridIndex : public IVFMixIndex { +public: + explicit IVFHybridIndex(std::shared_ptr index, const IndexType &type) + : IVFMixIndex(std::move(index), type) { + } + + knowhere::QuantizerPtr + LoadQuantizer(const Config &conf) override; + + Status + SetQuantizer(const knowhere::QuantizerPtr &q) override; + + Status + UnsetQuantizer() override; + + std::pair + CopyToGpuWithQuantizer(const int64_t &device_id, const Config &cfg) override; + + VecIndexPtr + LoadData(const knowhere::QuantizerPtr &q, const Config &conf) override; +}; + +} // namespace engine +} // namespace milvus diff --git a/core/thirdparty/versions.txt b/core/thirdparty/versions.txt index 4faaf119e4..68023d4072 100644 --- a/core/thirdparty/versions.txt +++ b/core/thirdparty/versions.txt @@ -9,5 +9,6 @@ LIBUNWIND_VERSION=1.3.1 GPERFTOOLS_VERSION=2.7 GRPC_VERSION=master ZLIB_VERSION=v1.2.11 +MKL_VERSION=2019.5.281 # vim: set filetype=sh: diff --git a/core/ubuntu_build_deps.sh b/core/ubuntu_build_deps.sh index e454a147ac..911046aa1f 100755 --- a/core/ubuntu_build_deps.sh +++ b/core/ubuntu_build_deps.sh @@ -1,5 +1,11 @@ #!/bin/bash +wget -P /tmp https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB +sudo apt-key add /tmp/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB + +sudo sh -c 'echo deb https://apt.repos.intel.com/mkl all main > /etc/apt/sources.list.d/intel-mkl.list' +sudo apt-get -y update && sudo apt-get -y install intel-mkl-gnu-2019.5-281 intel-mkl-core-2019.5-281 + sudo apt-get install -y gfortran libmysqlclient-dev mysql-client libcurl4-openssl-dev libboost-system-dev \ libboost-filesystem-dev libboost-serialization-dev libboost-regex-dev diff --git a/core/unittest/CMakeLists.txt b/core/unittest/CMakeLists.txt index 62b5bdf256..1ec300a7bc 100644 --- a/core/unittest/CMakeLists.txt +++ b/core/unittest/CMakeLists.txt @@ -16,7 +16,6 @@ # specific language governing permissions and limitations # under the License. #------------------------------------------------------------------------------- -include_directories("${CUDA_TOOLKIT_ROOT_DIR}/include") foreach(dir ${INDEX_INCLUDE_DIRS}) include_directories(${dir}) @@ -26,8 +25,6 @@ include_directories(${MILVUS_SOURCE_DIR}) include_directories(${MILVUS_ENGINE_SRC}) include_directories(${CMAKE_CURRENT_SOURCE_DIR}) -link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") - aux_source_directory(${MILVUS_ENGINE_SRC}/cache cache_files) aux_source_directory(${MILVUS_ENGINE_SRC}/config config_files) aux_source_directory(${MILVUS_ENGINE_SRC}/metrics metrics_files) @@ -116,11 +113,22 @@ set(unittest_libs prometheus-cpp-core dl z - ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so - cudart - cublas ) +if (MILVUS_GPU_VERSION) + include_directories("${CUDA_INCLUDE_DIRS}") + link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") + set(unittest_libs ${unittest_libs} + ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so + cudart + cublas + ) + aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/gpu wrapper_gpu_files) + set(common_files ${common_files} + ${wrapper_gpu_files} + ) +endif() + add_subdirectory(db) add_subdirectory(wrapper) add_subdirectory(metrics) diff --git a/core/unittest/db/CMakeLists.txt b/core/unittest/db/CMakeLists.txt index 3954dd8656..83f242a31a 100644 --- a/core/unittest/db/CMakeLists.txt +++ b/core/unittest/db/CMakeLists.txt @@ -20,7 +20,7 @@ aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} test_files) -cuda_add_executable(test_db +add_executable(test_db ${common_files} ${test_files} ) diff --git a/core/unittest/db/test_db.cpp b/core/unittest/db/test_db.cpp index f9e8da9c0f..b869d17388 100644 --- a/core/unittest/db/test_db.cpp +++ b/core/unittest/db/test_db.cpp @@ -175,8 +175,7 @@ TEST_F(DBTest, DB_TEST) { BuildVectors(qb, qxb); std::thread search([&]() { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; + milvus::engine::QueryResults results; int k = 10; std::this_thread::sleep_for(std::chrono::seconds(2)); @@ -191,17 +190,17 @@ TEST_F(DBTest, DB_TEST) { prev_count = count; START_TIMER; - stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), result_ids, result_distances); + stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), results); ss << "Search " << j << " With Size " << count / milvus::engine::M << " M"; STOP_TIMER(ss.str()); ASSERT_TRUE(stat.ok()); - for (auto i = 0; i < qb; ++i) { - ASSERT_EQ(result_ids[i*k], target_ids[i]); + for (auto k = 0; k < qb; ++k) { + ASSERT_EQ(results[k][0].first, target_ids[k]); ss.str(""); - ss << "Result [" << i << "]:"; - for (auto t = 0; t < k; t++) { - ss << result_ids[i * k + t] << " "; + ss << "Result [" << k << "]:"; + for (auto result : results[k]) { + ss << result.first << " "; } /* LOG(DEBUG) << ss.str(); */ } @@ -285,18 +284,16 @@ TEST_F(DBTest, SEARCH_TEST) { db_->CreateIndex(TABLE_NAME, index); // wait until build index finish { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results); ASSERT_TRUE(stat.ok()); } {//search by specify index file milvus::engine::meta::DatesT dates; std::vector file_ids = {"1", "2", "3", "4", "5", "6"}; - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results); ASSERT_TRUE(stat.ok()); } @@ -306,25 +303,22 @@ TEST_F(DBTest, SEARCH_TEST) { db_->CreateIndex(TABLE_NAME, index); // wait until build index finish { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results); ASSERT_TRUE(stat.ok()); } { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), result_ids, result_distances); + milvus::engine::QueryResults large_nq_results; + stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), large_nq_results); ASSERT_TRUE(stat.ok()); } {//search by specify index file milvus::engine::meta::DatesT dates; std::vector file_ids = {"1", "2", "3", "4", "5", "6"}; - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results); ASSERT_TRUE(stat.ok()); } @@ -397,12 +391,11 @@ TEST_F(DBTest, SHUTDOWN_TEST) { ASSERT_FALSE(stat.ok()); milvus::engine::meta::DatesT dates; - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(table_info.table_id_, 1, 1, 1, nullptr, dates, result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(table_info.table_id_, 1, 1, 1, nullptr, dates, results); ASSERT_FALSE(stat.ok()); std::vector file_ids; - stat = db_->Query(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, result_ids, result_distances); + stat = db_->Query(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, results); ASSERT_FALSE(stat.ok()); stat = db_->DeleteTable(table_info.table_id_, dates); diff --git a/core/unittest/db/test_db_mysql.cpp b/core/unittest/db/test_db_mysql.cpp index 30a616e662..ae1da8012a 100644 --- a/core/unittest/db/test_db_mysql.cpp +++ b/core/unittest/db/test_db_mysql.cpp @@ -81,8 +81,7 @@ TEST_F(MySqlDBTest, DB_TEST) { ASSERT_EQ(target_ids.size(), qb); std::thread search([&]() { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; + milvus::engine::QueryResults results; int k = 10; std::this_thread::sleep_for(std::chrono::seconds(5)); @@ -97,25 +96,25 @@ TEST_F(MySqlDBTest, DB_TEST) { prev_count = count; START_TIMER; - stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), result_ids, result_distances); + stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), results); ss << "Search " << j << " With Size " << count / milvus::engine::M << " M"; STOP_TIMER(ss.str()); ASSERT_TRUE(stat.ok()); - for (auto i = 0; i < qb; ++i) { + for (auto k = 0; k < qb; ++k) { // std::cout << results[k][0].first << " " << target_ids[k] << std::endl; // ASSERT_EQ(results[k][0].first, target_ids[k]); bool exists = false; - for (auto t = 0; t < k; t++) { - if (result_ids[i * k + t] == target_ids[i]) { + for (auto &result : results[k]) { + if (result.first == target_ids[k]) { exists = true; } } ASSERT_TRUE(exists); ss.str(""); - ss << "Result [" << i << "]:"; - for (auto t = 0; t < k; t++) { - ss << result_ids[i * k + t] << " "; + ss << "Result [" << k << "]:"; + for (auto result : results[k]) { + ss << result.first << " "; } /* LOG(DEBUG) << ss.str(); */ } @@ -189,9 +188,8 @@ TEST_F(MySqlDBTest, SEARCH_TEST) { sleep(2); // wait until build index finish - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances); + milvus::engine::QueryResults results; + stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results); ASSERT_TRUE(stat.ok()); } diff --git a/core/unittest/db/test_mem.cpp b/core/unittest/db/test_mem.cpp index 939e61246c..e05811ff9e 100644 --- a/core/unittest/db/test_mem.cpp +++ b/core/unittest/db/test_mem.cpp @@ -259,11 +259,10 @@ TEST_F(MemManagerTest2, SERIAL_INSERT_SEARCH_TEST) { int topk = 10, nprobe = 10; for (auto& pair : search_vectors) { auto& search = pair.second; - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; - stat = db_->Query(GetTableName(), topk, 1, nprobe, search.data(), result_ids, result_distances); - ASSERT_EQ(result_ids[0], pair.first); - ASSERT_LT(result_distances[0], 1e-4); + milvus::engine::QueryResults results; + stat = db_->Query(GetTableName(), topk, 1, nprobe, search.data(), results); + ASSERT_EQ(results[0][0].first, pair.first); + ASSERT_LT(results[0][0].second, 1e-4); } } @@ -315,8 +314,7 @@ TEST_F(MemManagerTest2, CONCURRENT_INSERT_SEARCH_TEST) { BuildVectors(qb, qxb); std::thread search([&]() { - milvus::engine::ResultIds result_ids; - milvus::engine::ResultDistances result_distances; + milvus::engine::QueryResults results; int k = 10; std::this_thread::sleep_for(std::chrono::seconds(2)); @@ -331,17 +329,17 @@ TEST_F(MemManagerTest2, CONCURRENT_INSERT_SEARCH_TEST) { prev_count = count; START_TIMER; - stat = db_->Query(GetTableName(), k, qb, 10, qxb.data(), result_ids, result_distances); + stat = db_->Query(GetTableName(), k, qb, 10, qxb.data(), results); ss << "Search " << j << " With Size " << count / milvus::engine::M << " M"; STOP_TIMER(ss.str()); ASSERT_TRUE(stat.ok()); - for (auto i = 0; i < qb; ++i) { - ASSERT_EQ(result_ids[i * k], target_ids[i]); + for (auto k = 0; k < qb; ++k) { + ASSERT_EQ(results[k][0].first, target_ids[k]); ss.str(""); - ss << "Result [" << i << "]:"; - for (auto t = 0; t < k; t++) { - ss << result_ids[i * k + t] << " "; + ss << "Result [" << k << "]:"; + for (auto result : results[k]) { + ss << result.first << " "; } /* LOG(DEBUG) << ss.str(); */ } diff --git a/core/unittest/db/test_search.cpp b/core/unittest/db/test_search.cpp index 1d1d9a677a..b8cf08b3e2 100644 --- a/core/unittest/db/test_search.cpp +++ b/core/unittest/db/test_search.cpp @@ -19,7 +19,6 @@ #include #include -#include "scheduler/job/SearchJob.h" #include "scheduler/task/SearchTask.h" #include "utils/TimeRecorder.h" #include "utils/ThreadPool.h" @@ -29,80 +28,74 @@ namespace { namespace ms = milvus::scheduler; void -BuildResult(ms::ResultIds& output_ids, - ms::ResultDistances & output_distances, - size_t input_k, - size_t topk, - size_t nq, +BuildResult(std::vector& output_ids, + std::vector& output_distance, + uint64_t input_k, + uint64_t topk, + uint64_t nq, bool ascending) { output_ids.clear(); output_ids.resize(nq * topk); - output_distances.clear(); - output_distances.resize(nq * topk); + output_distance.clear(); + output_distance.resize(nq * topk); - for (size_t i = 0; i < nq; i++) { + for (uint64_t i = 0; i < nq; i++) { //insert valid items - for (size_t j = 0; j < input_k; j++) { + for (uint64_t j = 0; j < input_k; j++) { output_ids[i * topk + j] = (int64_t)(drand48() * 100000); - output_distances[i * topk + j] = ascending ? (j + drand48()) : ((input_k - j) + drand48()); + output_distance[i * topk + j] = ascending ? (j + drand48()) : ((input_k - j) + drand48()); } //insert invalid items - for (size_t j = input_k; j < topk; j++) { + for (uint64_t j = input_k; j < topk; j++) { output_ids[i * topk + j] = -1; - output_distances[i * topk + j] = -1.0; + output_distance[i * topk + j] = -1.0; } } } void -CopyResult(ms::ResultIds& output_ids, - ms::ResultDistances& output_distances, - size_t output_topk, - ms::ResultIds& input_ids, - ms::ResultDistances& input_distances, - size_t input_topk, - size_t nq) { +CopyResult(std::vector& output_ids, + std::vector& output_distance, + uint64_t output_topk, + std::vector& input_ids, + std::vector& input_distance, + uint64_t input_topk, + uint64_t nq) { ASSERT_TRUE(input_ids.size() >= nq * input_topk); - ASSERT_TRUE(input_distances.size() >= nq * input_topk); + ASSERT_TRUE(input_distance.size() >= nq * input_topk); ASSERT_TRUE(output_topk <= input_topk); output_ids.clear(); output_ids.resize(nq * output_topk); - output_distances.clear(); - output_distances.resize(nq * output_topk); + output_distance.clear(); + output_distance.resize(nq * output_topk); - for (size_t i = 0; i < nq; i++) { - for (size_t j = 0; j < output_topk; j++) { + for (uint64_t i = 0; i < nq; i++) { + for (uint64_t j = 0; j < output_topk; j++) { output_ids[i * output_topk + j] = input_ids[i * input_topk + j]; - output_distances[i * output_topk + j] = input_distances[i * input_topk + j]; + output_distance[i * output_topk + j] = input_distance[i * input_topk + j]; } } } void -CheckTopkResult(const ms::ResultIds& input_ids_1, - const ms::ResultDistances& input_distances_1, - size_t input_k_1, - const ms::ResultIds& input_ids_2, - const ms::ResultDistances& input_distances_2, - size_t input_k_2, - size_t topk, - size_t nq, +CheckTopkResult(const std::vector& input_ids_1, + const std::vector& input_distance_1, + const std::vector& input_ids_2, + const std::vector& input_distance_2, + uint64_t topk, + uint64_t nq, bool ascending, - const ms::ResultIds& result_ids, - const ms::ResultDistances& result_distances) { - ASSERT_EQ(result_ids.size(), result_distances.size()); - ASSERT_EQ(input_ids_1.size(), input_distances_1.size()); - ASSERT_EQ(input_ids_2.size(), input_distances_2.size()); + const milvus::scheduler::ResultSet& result) { + ASSERT_EQ(result.size(), nq); + ASSERT_EQ(input_ids_1.size(), input_distance_1.size()); + ASSERT_EQ(input_ids_2.size(), input_distance_2.size()); - size_t result_k = result_distances.size() / nq; - ASSERT_EQ(result_k, std::min(topk, input_k_1 + input_k_2)); - - for (size_t i = 0; i < nq; i++) { + for (int64_t i = 0; i < nq; i++) { std::vector - src_vec(input_distances_1.begin() + i * topk, input_distances_1.begin() + (i + 1) * topk); + src_vec(input_distance_1.begin() + i * topk, input_distance_1.begin() + (i + 1) * topk); src_vec.insert(src_vec.end(), - input_distances_2.begin() + i * topk, - input_distances_2.begin() + (i + 1) * topk); + input_distance_2.begin() + i * topk, + input_distance_2.begin() + (i + 1) * topk); if (ascending) { std::sort(src_vec.begin(), src_vec.end()); } else { @@ -118,16 +111,15 @@ CheckTopkResult(const ms::ResultIds& input_ids_1, ++iter; } - size_t n = std::min(topk, result_ids.size() / nq); - for (size_t j = 0; j < n; j++) { - size_t idx = i * n + j; - if (result_ids[idx] < 0) { + uint64_t n = std::min(topk, result[i].size()); + for (uint64_t j = 0; j < n; j++) { + if (result[i][j].first < 0) { continue; } - if (src_vec[j] != result_distances[idx]) { - std::cout << src_vec[j] << " " << result_distances[idx] << std::endl; + if (src_vec[j] != result[i][j].second) { + std::cout << src_vec[j] << " " << result[i][j].second << std::endl; } - ASSERT_TRUE(src_vec[j] == result_distances[idx]); + ASSERT_TRUE(src_vec[j] == result[i][j].second); } } } @@ -135,21 +127,20 @@ CheckTopkResult(const ms::ResultIds& input_ids_1, } // namespace void -MergeTopkToResultSetTest(size_t topk_1, size_t topk_2, size_t nq, size_t topk, bool ascending) { - ms::ResultIds ids1, ids2; - ms::ResultDistances dist1, dist2; - ms::ResultIds result_ids; - ms::ResultDistances result_distances; +MergeTopkToResultSetTest(uint64_t topk_1, uint64_t topk_2, uint64_t nq, uint64_t topk, bool ascending) { + std::vector ids1, ids2; + std::vector dist1, dist2; + ms::ResultSet result; BuildResult(ids1, dist1, topk_1, topk, nq, ascending); BuildResult(ids2, dist2, topk_2, topk, nq, ascending); - ms::XSearchTask::MergeTopkToResultSet(ids1, dist1, topk_1, nq, topk, ascending, result_ids, result_distances); - ms::XSearchTask::MergeTopkToResultSet(ids2, dist2, topk_2, nq, topk, ascending, result_ids, result_distances); - CheckTopkResult(ids1, dist1, topk_1, ids2, dist2, topk_2, topk, nq, ascending, result_ids, result_distances); + ms::XSearchTask::MergeTopkToResultSet(ids1, dist1, topk_1, nq, topk, ascending, result); + ms::XSearchTask::MergeTopkToResultSet(ids2, dist2, topk_2, nq, topk, ascending, result); + CheckTopkResult(ids1, dist1, ids2, dist2, topk, nq, ascending, result); } TEST(DBSearchTest, MERGE_RESULT_SET_TEST) { - size_t NQ = 15; - size_t TOP_K = 64; + uint64_t NQ = 15; + uint64_t TOP_K = 64; /* test1, id1/dist1 valid, id2/dist2 empty */ MergeTopkToResultSetTest(TOP_K, 0, NQ, TOP_K, true); @@ -168,21 +159,21 @@ TEST(DBSearchTest, MERGE_RESULT_SET_TEST) { MergeTopkToResultSetTest(TOP_K / 2, TOP_K / 3, NQ, TOP_K, false); } -//void MergeTopkArrayTest(size_t topk_1, size_t topk_2, size_t nq, size_t topk, bool ascending) { +//void MergeTopkArrayTest(uint64_t topk_1, uint64_t topk_2, uint64_t nq, uint64_t topk, bool ascending) { // std::vector ids1, ids2; // std::vector dist1, dist2; // ms::ResultSet result; // BuildResult(ids1, dist1, topk_1, topk, nq, ascending); // BuildResult(ids2, dist2, topk_2, topk, nq, ascending); -// size_t result_topk = std::min(topk, topk_1 + topk_2); +// uint64_t result_topk = std::min(topk, topk_1 + topk_2); // ms::XSearchTask::MergeTopkArray(ids1, dist1, topk_1, ids2, dist2, topk_2, nq, topk, ascending); // if (ids1.size() != result_topk * nq) { // std::cout << ids1.size() << " " << result_topk * nq << std::endl; // } // ASSERT_TRUE(ids1.size() == result_topk * nq); // ASSERT_TRUE(dist1.size() == result_topk * nq); -// for (size_t i = 0; i < nq; i++) { -// for (size_t k = 1; k < result_topk; k++) { +// for (uint64_t i = 0; i < nq; i++) { +// for (uint64_t k = 1; k < result_topk; k++) { // float f0 = dist1[i * topk + k - 1]; // float f1 = dist1[i * topk + k]; // if (ascending) { @@ -201,8 +192,8 @@ TEST(DBSearchTest, MERGE_RESULT_SET_TEST) { //} //TEST(DBSearchTest, MERGE_ARRAY_TEST) { -// size_t NQ = 15; -// size_t TOP_K = 64; +// uint64_t NQ = 15; +// uint64_t TOP_K = 64; // // /* test1, id1/dist1 valid, id2/dist2 empty */ // MergeTopkArrayTest(TOP_K, 0, NQ, TOP_K, true); @@ -231,23 +222,23 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) { int32_t index_file_num = 478; /* sift1B dataset, index files num */ bool ascending = true; - std::vector thread_vec = {4}; - std::vector nq_vec = {1000}; - std::vector topk_vec = {64}; - size_t NQ = nq_vec[nq_vec.size() - 1]; - size_t TOPK = topk_vec[topk_vec.size() - 1]; + std::vector thread_vec = {4, 8}; + std::vector nq_vec = {1, 10, 100}; + std::vector topk_vec = {1, 4, 16, 64}; + int32_t NQ = nq_vec[nq_vec.size() - 1]; + int32_t TOPK = topk_vec[topk_vec.size() - 1]; - std::vector id_vec; - std::vector dist_vec; - ms::ResultIds input_ids; - ms::ResultDistances input_distances; + std::vector> id_vec; + std::vector> dist_vec; + std::vector input_ids; + std::vector input_distance; int32_t i, k, step; /* generate testing data */ for (i = 0; i < index_file_num; i++) { - BuildResult(input_ids, input_distances, TOPK, TOPK, NQ, ascending); + BuildResult(input_ids, input_distance, TOPK, TOPK, NQ, ascending); id_vec.push_back(input_ids); - dist_vec.push_back(input_distances); + dist_vec.push_back(input_distance); } for (int32_t max_thread_num : thread_vec) { @@ -256,11 +247,10 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) { for (int32_t nq : nq_vec) { for (int32_t top_k : topk_vec) { - ms::ResultIds final_result_ids, final_result_ids_2, final_result_ids_3; - ms::ResultDistances final_result_distances, final_result_distances_2, final_result_distances_3; + ms::ResultSet final_result, final_result_2, final_result_3; - std::vector id_vec_1(index_file_num); - std::vector dist_vec_1(index_file_num); + std::vector> id_vec_1(index_file_num); + std::vector> dist_vec_1(index_file_num); for (i = 0; i < index_file_num; i++) { CopyResult(id_vec_1[i], dist_vec_1[i], top_k, id_vec[i], dist_vec[i], TOPK, nq); } @@ -278,10 +268,8 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) { nq, top_k, ascending, - final_result_ids, - final_result_distances); - ASSERT_EQ(final_result_ids.size(), nq * top_k); - ASSERT_EQ(final_result_distances.size(), nq * top_k); + final_result); + ASSERT_EQ(final_result.size(), nq); } rc1.RecordSection("reduce done"); @@ -290,7 +278,7 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) { // /* method-2 */ // std::vector> id_vec_2(index_file_num); // std::vector> dist_vec_2(index_file_num); -// std::vector k_vec_2(index_file_num); +// std::vector k_vec_2(index_file_num); // for (i = 0; i < index_file_num; i++) { // CopyResult(id_vec_2[i], dist_vec_2[i], top_k, id_vec[i], dist_vec[i], TOPK, nq); // k_vec_2[i] = top_k; @@ -333,7 +321,7 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) { // /* method-3 parallel */ // std::vector> id_vec_3(index_file_num); // std::vector> dist_vec_3(index_file_num); -// std::vector k_vec_3(index_file_num); +// std::vector k_vec_3(index_file_num); // for (i = 0; i < index_file_num; i++) { // CopyResult(id_vec_3[i], dist_vec_3[i], top_k, id_vec[i], dist_vec[i], TOPK, nq); // k_vec_3[i] = top_k; diff --git a/core/unittest/db/utils.cpp b/core/unittest/db/utils.cpp index 16e195079c..a78aecd046 100644 --- a/core/unittest/db/utils.cpp +++ b/core/unittest/db/utils.cpp @@ -27,7 +27,9 @@ #include "cache/CpuCacheMgr.h" #include "db/DBFactory.h" #include "db/Options.h" +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif #include "utils/CommonUtil.h" INITIALIZE_EASYLOGGINGPP @@ -68,9 +70,15 @@ static const char " use_blas_threshold: 20\n" "\n" "resource_config:\n" +#ifdef MILVUS_CPU_VERSION + " search_resources:\n" + " - cpu\n" + " index_build_device: cpu # CPU used for building index"; +#else " search_resources:\n" " - gpu0\n" " index_build_device: gpu0 # GPU used for building index"; +#endif void WriteToFile(const std::string& file_path, const char* content) { @@ -118,15 +126,18 @@ BaseTest::InitLog() { void BaseTest::SetUp() { InitLog(); - +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); +#endif } void BaseTest::TearDown() { milvus::cache::CpuCacheMgr::GetInstance()->ClearCache(); milvus::cache::GpuCacheMgr::GetInstance(0)->ClearCache(); +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); +#endif } milvus::engine::DBOptions diff --git a/core/unittest/metrics/test_metrics.cpp b/core/unittest/metrics/test_metrics.cpp index 1b26ad097b..c0d1044bb4 100644 --- a/core/unittest/metrics/test_metrics.cpp +++ b/core/unittest/metrics/test_metrics.cpp @@ -75,8 +75,7 @@ TEST_F(MetricTest, METRIC_TEST) { } std::thread search([&]() { -// milvus::engine::ResultIds result_ids; -// milvus::engine::ResultDistances result_distances; + milvus::engine::QueryResults results; int k = 10; std::this_thread::sleep_for(std::chrono::seconds(2)); @@ -91,7 +90,7 @@ TEST_F(MetricTest, METRIC_TEST) { prev_count = count; START_TIMER; -// stat = db_->Query(group_name, k, qb, qxb, result_ids, result_distances); +// stat = db_->Query(group_name, k, qb, qxb, results); ss << "Search " << j << " With Size " << (float) (count * group_dim * sizeof(float)) / (1024 * 1024) << " M"; diff --git a/core/unittest/scheduler/CMakeLists.txt b/core/unittest/scheduler/CMakeLists.txt index 087f93f017..77404f3cba 100644 --- a/core/unittest/scheduler/CMakeLists.txt +++ b/core/unittest/scheduler/CMakeLists.txt @@ -19,7 +19,7 @@ aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} test_files) -cuda_add_executable(test_scheduler +add_executable(test_scheduler ${common_files} ${entry_file} ${test_files} diff --git a/core/unittest/server/CMakeLists.txt b/core/unittest/server/CMakeLists.txt index 1f89de8d3f..38309611da 100644 --- a/core/unittest/server/CMakeLists.txt +++ b/core/unittest/server/CMakeLists.txt @@ -47,7 +47,7 @@ set(server_test_files ${test_files} ) -cuda_add_executable(test_server ${server_test_files}) +add_executable(test_server ${server_test_files}) set(client_grpc_lib grpcpp_channelz diff --git a/core/unittest/server/test_config.cpp b/core/unittest/server/test_config.cpp index 123ddf5265..76230cbcc3 100644 --- a/core/unittest/server/test_config.cpp +++ b/core/unittest/server/test_config.cpp @@ -112,123 +112,3 @@ TEST_F(ConfigTest, SERVER_CONFIG_TEST) { s = config.ResetDefaultConfig(); ASSERT_TRUE(s.ok()); } - -TEST_F(ConfigTest, SERVER_CONFIG_INVALID_TEST) { - std::string config_path(CONFIG_PATH); - milvus::server::Config& config = milvus::server::Config::GetInstance(); - milvus::Status s; - - s = config.LoadConfigFile(""); - ASSERT_FALSE(s.ok()); - - s = config.LoadConfigFile(config_path + INVALID_CONFIG_FILE); - ASSERT_FALSE(s.ok()); - s = config.LoadConfigFile(config_path + "dummy.yaml"); - ASSERT_FALSE(s.ok()); - - /* server config */ - s = config.SetServerConfigAddress("0.0.0"); - ASSERT_FALSE(s.ok()); - s = config.SetServerConfigAddress("0.0.0.256"); - ASSERT_FALSE(s.ok()); - - s = config.SetServerConfigPort("a"); - ASSERT_FALSE(s.ok()); - s = config.SetServerConfigPort("99999"); - ASSERT_FALSE(s.ok()); - - s = config.SetServerConfigDeployMode("cluster"); - ASSERT_FALSE(s.ok()); - - s = config.SetServerConfigTimeZone("GM"); - ASSERT_FALSE(s.ok()); - s = config.SetServerConfigTimeZone("GMT8"); - ASSERT_FALSE(s.ok()); - s = config.SetServerConfigTimeZone("UTCA"); - ASSERT_FALSE(s.ok()); - - /* db config */ - s = config.SetDBConfigPrimaryPath(""); - ASSERT_FALSE(s.ok()); - -// s = config.SetDBConfigSecondaryPath(""); -// ASSERT_FALSE(s.ok()); - - s = config.SetDBConfigBackendUrl("http://www.google.com"); - ASSERT_FALSE(s.ok()); - s = config.SetDBConfigBackendUrl("sqlite://:@:"); - ASSERT_FALSE(s.ok()); - s = config.SetDBConfigBackendUrl("mysql://root:123456@127.0.0.1/milvus"); - ASSERT_FALSE(s.ok()); - - s = config.SetDBConfigArchiveDiskThreshold("0x10"); - ASSERT_FALSE(s.ok()); - - s = config.SetDBConfigArchiveDaysThreshold("0x10"); - ASSERT_FALSE(s.ok()); - - s = config.SetDBConfigInsertBufferSize("a"); - ASSERT_FALSE(s.ok()); - s = config.SetDBConfigInsertBufferSize("0"); - ASSERT_FALSE(s.ok()); - s = config.SetDBConfigInsertBufferSize("2048"); - ASSERT_FALSE(s.ok()); - - /* metric config */ - s = config.SetMetricConfigEnableMonitor("Y"); - ASSERT_FALSE(s.ok()); - - s = config.SetMetricConfigCollector("zilliz"); - ASSERT_FALSE(s.ok()); - - s = config.SetMetricConfigPrometheusPort("0xff"); - ASSERT_FALSE(s.ok()); - - /* cache config */ - s = config.SetCacheConfigCpuCacheCapacity("a"); - ASSERT_FALSE(s.ok()); - s = config.SetCacheConfigCpuCacheCapacity("0"); - ASSERT_FALSE(s.ok()); - s = config.SetCacheConfigCpuCacheCapacity("2048"); - ASSERT_FALSE(s.ok()); - - s = config.SetCacheConfigCpuCacheThreshold("a"); - ASSERT_FALSE(s.ok()); - s = config.SetCacheConfigCpuCacheThreshold("1.0"); - ASSERT_FALSE(s.ok()); - - s = config.SetCacheConfigGpuCacheCapacity("a"); - ASSERT_FALSE(s.ok()); - s = config.SetCacheConfigGpuCacheCapacity("128"); - ASSERT_FALSE(s.ok()); - - s = config.SetCacheConfigGpuCacheThreshold("a"); - ASSERT_FALSE(s.ok()); - s = config.SetCacheConfigGpuCacheThreshold("1.0"); - ASSERT_FALSE(s.ok()); - - s = config.SetCacheConfigCacheInsertData("N"); - ASSERT_FALSE(s.ok()); - - /* engine config */ - s = config.SetEngineConfigUseBlasThreshold("0xff"); - ASSERT_FALSE(s.ok()); - - s = config.SetEngineConfigOmpThreadNum("a"); - ASSERT_FALSE(s.ok()); - s = config.SetEngineConfigOmpThreadNum("10000"); - ASSERT_FALSE(s.ok()); - - s = config.SetEngineConfigGpuSearchThreshold("-1"); - ASSERT_FALSE(s.ok()); - - /* resource config */ - s = config.SetResourceConfigMode("default"); - ASSERT_FALSE(s.ok()); - - s = config.SetResourceConfigIndexBuildDevice("gup2"); - ASSERT_FALSE(s.ok()); - s = config.SetResourceConfigIndexBuildDevice("gpu16"); - ASSERT_FALSE(s.ok()); -} - diff --git a/core/unittest/server/test_rpc.cpp b/core/unittest/server/test_rpc.cpp index ebbcd810c1..a1de370650 100644 --- a/core/unittest/server/test_rpc.cpp +++ b/core/unittest/server/test_rpc.cpp @@ -23,7 +23,7 @@ #include "server/grpc_impl/GrpcRequestHandler.h" #include "server/grpc_impl/GrpcRequestScheduler.h" #include "server/grpc_impl/GrpcRequestTask.h" -#include "src/version.h" +#include "src/config.h" #include "grpc/gen-milvus/milvus.grpc.pb.h" #include "grpc/gen-status/status.pb.h" diff --git a/core/unittest/server/test_util.cpp b/core/unittest/server/test_util.cpp index 24482740bc..87c2d2f5f2 100644 --- a/core/unittest/server/test_util.cpp +++ b/core/unittest/server/test_util.cpp @@ -314,6 +314,7 @@ TEST(ValidationUtilTest, VALIDATE_NPROBE_TEST) { ASSERT_NE(milvus::server::ValidationUtil::ValidateSearchNprobe(101, schema).code(), milvus::SERVER_SUCCESS); } +#ifdef MILVUS_GPU_VERSION TEST(ValidationUtilTest, VALIDATE_GPU_TEST) { ASSERT_EQ(milvus::server::ValidationUtil::ValidateGpuIndex(0).code(), milvus::SERVER_SUCCESS); ASSERT_NE(milvus::server::ValidationUtil::ValidateGpuIndex(100).code(), milvus::SERVER_SUCCESS); @@ -322,6 +323,7 @@ TEST(ValidationUtilTest, VALIDATE_GPU_TEST) { ASSERT_EQ(milvus::server::ValidationUtil::GetGpuMemory(0, memory).code(), milvus::SERVER_SUCCESS); ASSERT_NE(milvus::server::ValidationUtil::GetGpuMemory(100, memory).code(), milvus::SERVER_SUCCESS); } +#endif TEST(ValidationUtilTest, VALIDATE_IPADDRESS_TEST) { ASSERT_EQ(milvus::server::ValidationUtil::ValidateIpAddress("127.0.0.1").code(), milvus::SERVER_SUCCESS); diff --git a/core/unittest/server/utils.cpp b/core/unittest/server/utils.cpp index 4c03da6ad9..1c99736173 100644 --- a/core/unittest/server/utils.cpp +++ b/core/unittest/server/utils.cpp @@ -60,9 +60,15 @@ static const char " use_blas_threshold: 20 \n" "\n" "resource_config:\n" - " search_resources: \n" +#ifdef MILVUS_CPU_VERSION + " search_resources:\n" + " - cpu\n" + " index_build_device: cpu # CPU used for building index"; +#else + " search_resources:\n" " - gpu0\n" " index_build_device: gpu0 # GPU used for building index"; +#endif static const char* INVALID_CONFIG_STR = "*INVALID*"; diff --git a/core/unittest/wrapper/CMakeLists.txt b/core/unittest/wrapper/CMakeLists.txt index ef145a9f50..232abf6e1a 100644 --- a/core/unittest/wrapper/CMakeLists.txt +++ b/core/unittest/wrapper/CMakeLists.txt @@ -17,7 +17,13 @@ # under the License. #------------------------------------------------------------------------------- -aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} test_files) +set(test_files + test_knowhere.cpp + test_wrapper.cpp) +if (MILVUS_GPU_VERSION) + set(test_files ${test_files} + test_hybrid_index.cpp) +endif () set(wrapper_files ${MILVUS_ENGINE_SRC}/wrapper/DataTransfer.cpp diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index 4ceb07ddb4..97eb3a3a68 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -17,7 +17,11 @@ #include "external/easyloggingpp/easylogging++.h" #include "wrapper/VecIndex.h" + +#ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif + #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "wrapper/utils.h" @@ -30,11 +34,13 @@ using ::testing::Values; using ::testing::Combine; class KnowhereWrapperTest - : public DataGenBase, - public TestWithParam<::std::tuple> { - protected: + : public DataGenBase, + public TestWithParam<::std::tuple> { +protected: void SetUp() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); +#endif std::string generator_type; std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); @@ -48,10 +54,12 @@ class KnowhereWrapperTest } void TearDown() override { +#ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); +#endif } - protected: +protected: milvus::engine::IndexType index_type; milvus::engine::VecIndexPtr index_ = nullptr; knowhere::Config conf; @@ -59,27 +67,30 @@ class KnowhereWrapperTest INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, Values( - //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, - "Default", - 64, - 100000, - 10, - 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, - "Default", - 64, - 100000, - 10, - 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), + //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] +#ifdef MILVUS_GPU_VERSION + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, + "Default", + 64, + 100000, + 10, + 10), +// std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), // std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), +#endif // std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10) - ) + std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, + "Default", + 64, + 100000, + 10, + 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10) +) ); TEST_P(KnowhereWrapperTest, BASE_TEST) { @@ -94,6 +105,8 @@ TEST_P(KnowhereWrapperTest, BASE_TEST) { AssertResult(res_ids, res_dis); } +#ifdef MILVUS_GPU_VERSION + TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -125,6 +138,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { AssertResult(res_ids, res_dis); } } +#endif TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -166,6 +180,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { } #include "wrapper/ConfAdapter.h" + TEST(whatever, test_config) { milvus::engine::TempMetaConf conf; auto nsg_conf = std::make_shared(); diff --git a/core/unittest/wrapper/utils.cpp b/core/unittest/wrapper/utils.cpp index b397a35d7c..7caa662343 100644 --- a/core/unittest/wrapper/utils.cpp +++ b/core/unittest/wrapper/utils.cpp @@ -58,9 +58,15 @@ static const char " blas_threshold: 20\n" "\n" "resource_config:\n" +#ifdef MILVUS_CPU_VERSION + " search_resources:\n" + " - cpu\n" + " index_build_device: cpu # CPU used for building index"; +#else " search_resources:\n" " - gpu0\n" " index_build_device: gpu0 # GPU used for building index"; +#endif void WriteToFile(const std::string& file_path, const char* content) {