milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add CPU version

pull/192/head
youny626 2019-11-05 10:23:59 +08:00
parent af319887d4
commit 4ca6b59b91
79 changed files with 2103 additions and 1610 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

242
core/CMakeLists.txt
View File

@ -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 ()

18
core/build.sh
View File

@ -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
fi

5
core/cmake/DefineOptions.cmake
View File

@ -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")

13
core/conf/server_config.template
View File

@ -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
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

42
core/src/CMakeLists.txt
View File

@ -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
)

3
core/src/config.h Normal file
View File

@ -0,0 +1,3 @@
#define MILVUS_VERSION "0.5.0"
#define BUILD_TYPE "Debug"
#define BUILD_TIME "2019-11-05 10:23.18"

3
core/src/config.h.in Normal file
View File

@ -0,0 +1,3 @@
#cmakedefine MILVUS_VERSION "@MILVUS_VERSION@"
#cmakedefine BUILD_TYPE "@BUILD_TYPE@"
#cmakedefine BUILD_TIME @BUILD_TIME@

7
core/src/db/DB.h
View File

@ -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<std::string>& 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;

18
core/src/db/DBImpl.cpp
View File

@ -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<std::string>& 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<std::string>& 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();

9
core/src/db/DBImpl.h
View File

@ -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<std::string>& 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();

8
core/src/db/Types.h
View File

@ -19,7 +19,6 @@
#include "db/engine/ExecutionEngine.h"
#include <faiss/Index.h>
#include <stdint.h>
#include <utility>
#include <vector>
@ -27,13 +26,12 @@
namespace milvus {
namespace engine {
using IDNumber = faiss::Index::idx_t;
typedef int64_t IDNumber;
typedef IDNumber* IDNumberPtr;
typedef std::vector<IDNumber> IDNumbers;
typedef std::vector<faiss::Index::idx_t> ResultIds;
typedef std::vector<faiss::Index::distance_t> ResultDistances;
typedef std::vector<std::pair<IDNumber, double>> QueryResult;
typedef std::vector<QueryResult> QueryResults;
struct TableIndex {
int32_t engine_type_ = (int)EngineType::FAISS_IDMAP;

38
core/src/db/engine/ExecutionEngineImpl.cpp
View File

@ -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<VecIndex>(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();
}

47
core/src/index/CMakeLists.txt
View File

@ -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()

11
core/src/index/cmake/DefineOptionsCore.cmake
View File

@ -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)

322
core/src/index/cmake/ThirdPartyPackagesCore.cmake
View File

@ -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 ()

57
core/src/index/knowhere/CMakeLists.txt
View File

@ -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

113
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp Normal file
View File

@ -0,0 +1,113 @@
#include "IndexGPUIDMAP.h"
#include <faiss/AutoTune.h>
#include <faiss/IndexFlat.h>
#include <faiss/MetaIndexes.h>
#include <faiss/index_io.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuCloner.h>
#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<std::mutex> lk(mutex_);
faiss::Index *device_index = index_.get();
faiss::Index *host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IDMAP>(new_index);
}
VectorIndexPtr
GPUIDMAP::Clone() {
auto cpu_idx = CopyGpuToCpu(Config());
if (auto idmap = std::dynamic_pointer_cast<IDMAP>(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<uint8_t>();
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<IDMAP>(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

47
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h Normal file
View File

@ -0,0 +1,47 @@
#pragma once
#include "IndexGPUIVF.h"
#include "IndexIVF.h"
#include "IndexIDMAP.h"
#include <memory>
#include <utility>
namespace knowhere {
class GPUIDMAP : public IDMAP, public GPUIndex {
public:
explicit GPUIDMAP(std::shared_ptr<faiss::Index> 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<GPUIDMAP>;
} // knowhere

385
core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
View File

@ -17,10 +17,18 @@
#include <faiss/IndexFlat.h>
#include <faiss/MetaIndexes.h>
#include <faiss/gpu/GpuCloner.h>
#include <faiss/index_factory.h>
#include <faiss/clone_index.h>
#include <faiss/AutoTune.h>
#include <faiss/index_io.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuCloner.h>
#endif
#include <vector>
#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<std::mutex> lk(mutex_);
return SerializeImpl();
}
void
IDMAP::Load(const BinarySet& index_binary) {
std::lock_guard<std::mutex> 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<BufferPtr> id_bufs{nullptr, id_buf};
std::vector<BufferPtr> dist_bufs{nullptr, dist_buf};
auto int64_type = std::make_shared<arrow::Int64Type>();
auto float_type = std::make_shared<arrow::FloatType>();
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<NumericArray<arrow::Int64Type>>(id_array_data);
auto dists = std::make_shared<NumericArray<arrow::FloatType>>(dist_array_data);
std::vector<ArrayPtr> array{ids, dists};
return std::make_shared<Dataset>(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<std::mutex> lk(mutex_);
GETTENSOR(dataset)
// TODO: magic here.
auto array = dataset->array()[0];
auto p_ids = array->data()->GetValues<int64_t>(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<faiss::IndexIDMap*>(index_.get());
auto flat_index = dynamic_cast<faiss::IndexFlat*>(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<faiss::IndexIDMap*>(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<std::mutex> lk(mutex_);
auto clone_index = faiss::clone_index(index_.get());
std::shared_ptr<faiss::Index> new_index;
new_index.reset(clone_index);
return std::make_shared<IDMAP>(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<faiss::Index> device_index;
device_index.reset(gpu_index);
return std::make_shared<GPUIDMAP>(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<std::mutex> lk(mutex_);
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IDMAP>(new_index);
}
VectorIndexPtr
GPUIDMAP::Clone() {
auto cpu_idx = CopyGpuToCpu(Config());
if (auto idmap = std::dynamic_pointer_cast<IDMAP>(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<uint8_t>();
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<std::mutex> 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<std::mutex> 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<BufferPtr> id_bufs{nullptr, id_buf};
std::vector<BufferPtr> dist_bufs{nullptr, dist_buf};
auto int64_type = std::make_shared<arrow::Int64Type>();
auto float_type = std::make_shared<arrow::FloatType>();
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<NumericArray<arrow::Int64Type>>(id_array_data);
auto dists = std::make_shared<NumericArray<arrow::FloatType>>(dist_array_data);
std::vector<ArrayPtr> array{ids, dists};
return std::make_shared<Dataset>(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<std::mutex> lk(mutex_);
GETTENSOR(dataset)
// TODO: magic here.
auto array = dataset->array()[0];
auto p_ids = array->data()->GetValues<int64_t>(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<faiss::IndexIDMap *>(index_.get());
auto flat_index = dynamic_cast<faiss::IndexFlat *>(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<faiss::IndexIDMap *>(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<std::mutex> lk(mutex_);
auto clone_index = faiss::clone_index(index_.get());
std::shared_ptr<faiss::Index> new_index;
new_index.reset(clone_index);
return std::make_shared<IDMAP>(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<faiss::Index> device_index;
device_index.reset(gpu_index);
return std::make_shared<GPUIDMAP>(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<IDMAP>(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

29
core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h
View File

@ -17,7 +17,6 @@
#pragma once
#include "IndexGPUIVF.h"
#include "IndexIVF.h"
#include <memory>
@ -67,32 +66,4 @@ class IDMAP : public VectorIndex, public FaissBaseIndex {
using IDMAPPtr = std::shared_ptr<IDMAP>;
class GPUIDMAP : public IDMAP, public GPUIndex {
public:
explicit GPUIDMAP(std::shared_ptr<faiss::Index> 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<GPUIDMAP>;
} // namespace knowhere

24
core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
View File

@ -15,11 +15,19 @@
// specific language governing permissions and limitations
// under the License.
#include <faiss/AutoTune.h>
#include <faiss/IVFlib.h>
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/index_io.h>
#include <faiss/index_factory.h>
#include <faiss/clone_index.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuCloner.h>
#endif
#include <chrono>
#include <memory>
@ -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<double, std::micro>(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

17
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp
View File

@ -15,15 +15,23 @@
// specific language governing permissions and limitations
// under the License.
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuCloner.h>
#endif
#include <faiss/index_factory.h>
#include <memory>
#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<faiss::Index>& 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

8
core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp
View File

@ -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<GPUIVF>(build_cfg->gpu_id);
#else
auto preprocess_index = std::make_shared<IVF>();
#endif
auto model = preprocess_index->Train(dataset, config);
preprocess_index->set_index_model(model);
preprocess_index->AddWithoutIds(dataset, config);

86
core/src/index/unittest/CMakeLists.txt
View File

@ -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 ()
#<IVF-TEST>
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})
#<IDMAP-TEST>
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})
#<KDT-TEST>
@ -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)

32
core/src/index/unittest/Helper.h
View File

@ -18,13 +18,16 @@
#include <memory>
#include <string>
#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<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "GPUIVF") {
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
}
#ifdef MILVUS_GPU_VERSION
else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(DEVICEID);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(DEVICEID);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(DEVICEID);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(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<knowhere::IVFCfg>();
tempconf->d = DIM;
@ -107,14 +113,18 @@ class ParamGenerator {
#include <gtest/gtest.h>
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
}
};

64
core/src/index/unittest/faiss_benchmark/CMakeLists.txt
View File

@ -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 ()

63
core/src/index/unittest/faiss_ori/CMakeLists.txt
View File

@ -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
)
#<GPU-TEST>
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)
#<GPU-TEST>
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 ()

6
core/src/index/unittest/test_idmap.cpp
View File

@ -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

86
core/src/index/unittest/test_ivf.cpp
View File

@ -20,13 +20,24 @@
#include <iostream>
#include <thread>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuIndexIVFFlat.h>
#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<std::string, ParameterType>> {
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<std::string, P
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
protected:
@ -64,15 +78,18 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
};
INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
Values(std::make_tuple("IVF", ParameterType::ivf),
std::make_tuple("GPUIVF", ParameterType::ivf),
std::make_tuple("IVFPQ", ParameterType::ivfpq),
std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std::make_tuple("IVFSQ", ParameterType::ivfsq),
Values(
#ifdef MILVUS_GPU_VERSION
std::make_tuple("GPUIVF", ParameterType::ivf),
std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq),
#ifdef CUSTOMIZATION
std::make_tuple("IVFSQHybrid", ParameterType::ivfsq),
std::make_tuple("IVFSQHybrid", ParameterType::ivfsq),
#endif
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)));
#endif
std::make_tuple("IVF", ParameterType::ivf),
std::make_tuple("IVFPQ", ParameterType::ivfpq),
std::make_tuple("IVFSQ", ParameterType::ivfsq)));
TEST_P(IVFTest, ivf_basic) {
assert(!xb.empty());
@ -92,9 +109,9 @@ TEST_P(IVFTest, ivf_basic) {
}
TEST_P(IVFTest, ivf_serialize) {
auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) {
auto serialize = [](const std::string &filename, knowhere::BinaryPtr &bin, uint8_t *ret) {
FileIOWriter writer(filename);
writer(static_cast<void*>(bin->data.get()), bin->size);
writer(static_cast<void *>(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<std::string> 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

3
core/src/index/unittest/test_nsg/CMakeLists.txt
View File

@ -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)

6
core/src/index/unittest/test_nsg/test_nsg.cpp
View File

@ -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<knowhere::NSG>();
@ -60,7 +64,9 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test {
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
protected:

33
core/src/main.cpp
View File

@ -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;
}

25
core/src/metrics/SystemInfo.cpp
View File

@ -19,7 +19,6 @@
#include "utils/Log.h"
#include <dirent.h>
#include <nvml.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
@ -29,6 +28,10 @@
#include <string>
#include <utility>
#ifdef MILVUS_GPU_VERSION
#include <nvml.h>
#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<uint64_t, uint64_t> in_and_out_octets = Octets();
@ -209,10 +214,14 @@ SystemInfo::CPUPercent() {
std::vector<uint64_t>
SystemInfo::GPUMemoryTotal() {
// get GPU usage percent
if (!initialized_)
Init();
std::vector<uint64_t> 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<uint64_t> 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<uint64_t> 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;
}

31
core/src/scheduler/JobMgr.cpp
View File

@ -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<std::string> path;
auto spec_label = std::static_pointer_cast<SpecResLabel>(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<SpecResLabel>(task->label());
auto src = res_mgr_->GetDiskResources()[0];
auto dest = spec_label->resource();
std::vector<std::string> 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<std::string> path;
auto spec_label = std::static_pointer_cast<SpecResLabel>(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

13
core/src/scheduler/SchedInst.cpp
View File

@ -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 <set>
@ -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);
}
}

1
core/src/scheduler/SchedInst.h
View File

@ -106,6 +106,7 @@ class OptimizerInst {
has_cpu = true;
}
}
std::vector<PassPtr> pass_list;
pass_list.push_back(std::make_shared<LargeSQ8HPass>());
pass_list.push_back(std::make_shared<HybridPass>());

11
core/src/scheduler/TaskCreator.cpp
View File

@ -70,15 +70,8 @@ TaskCreator::Create(const DeleteJobPtr& job) {
std::vector<TaskPtr>
TaskCreator::Create(const BuildIndexJobPtr& job) {
std::vector<TaskPtr> 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<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));

4
core/src/scheduler/Utils.cpp
View File

@ -19,7 +19,9 @@
#include "server/Config.h"
#include "utils/Log.h"
#ifdef MILVUS_GPU_VERSION
#include <cuda_runtime.h>
#endif
#include <chrono>
#include <set>
#include <string>
@ -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;
}

102
core/src/scheduler/action/PushTaskToNeighbour.cpp
View File

@ -138,41 +138,73 @@ Action::SpecifiedResourceLabelTaskScheduler(const ResourceMgrPtr& res_mgr, Resou
std::shared_ptr<LoadCompletedEvent> 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<std::vector<std::string>> paths;
// std::vector<uint64_t> transport_costs;
// for (auto& res : compute_resources) {
// std::vector<std::string> 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<std::vector<std::string>> paths;
std::vector<uint64_t> transport_costs;
for (auto& res : compute_resources) {
std::vector<std::string> 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<SpecResLabel>(label);
// if (spec_label->resource().lock()->type() == ResourceType::CPU) {
// std::vector<std::string> 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<uint64_t>::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();

5
core/src/scheduler/job/BuildIndexJob.cpp
View File

@ -50,10 +50,7 @@ void
BuildIndexJob::BuildIndexDone(size_t to_index_id) {
std::unique_lock<std::mutex> 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;
}

16
core/src/scheduler/job/SearchJob.cpp
View File

@ -49,21 +49,13 @@ void
SearchJob::SearchDone(size_t index_id) {
std::unique_lock<std::mutex> 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&

22
core/src/scheduler/job/SearchJob.h
View File

@ -29,7 +29,6 @@
#include <vector>
#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<size_t, TableFileSchemaPtr>;
using ResultIds = engine::ResultIds;
using ResultDistances = engine::ResultDistances;
using IdDistPair = std::pair<int64_t, double>;
using Id2DistVec = std::vector<IdDistPair>;
using ResultSet = std::vector<Id2DistVec>;
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_;

2
core/src/scheduler/optimizer/OnlyGPUPass.cpp
View File

@ -46,7 +46,7 @@ OnlyGPUPass::Run(const TaskPtr& task) {
auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(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;
}

115
core/src/scheduler/task/SearchTask.cpp
View File

@ -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<std::mutex> 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<int64_t>& input_ids, const std::vector<float>& 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

5
core/src/scheduler/task/SearchTask.h
View File

@ -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<int64_t>& input_ids, const std::vector<float>& input_distance,
uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending, scheduler::ResultSet& result);
// static void
// MergeTopkArray(std::vector<int64_t>& tar_ids, std::vector<float>& tar_distance, uint64_t& tar_input_k,

2
core/src/sdk/grpc/ClientProxy.cpp
View File

@ -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 <memory>
#include <string>

75
core/src/server/Config.cpp
View File

@ -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<std::string>& 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<std::string>& 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();
}

4
core/src/server/Config.h
View File

@ -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:

2
core/src/server/Server.cpp
View File

@ -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"

28
core/src/server/grpc_impl/GrpcRequestTask.cpp
View File

@ -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);
}
}

10
core/src/utils/ValidationUtil.cpp
View File

@ -20,7 +20,9 @@
#include "db/engine/ExecutionEngine.h"
#include <arpa/inet.h>
#ifdef MILVUS_GPU_VERSION
#include <cuda_runtime.h>
#endif
#include <algorithm>
#include <cmath>
#include <regex>
@ -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();
}

2
core/src/wrapper/ConfAdapter.cpp
View File

@ -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;

9
core/src/wrapper/KnowhereResource.cpp
View File

@ -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 <map>
@ -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();
}

501
core/src/wrapper/VecImpl.cpp
View File

@ -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 <src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h>
#include <src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.h>
#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<int64_t>(1, i * k + j)) << " ";
// ss_dist << *(dists->data()->GetValues<float>(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<int64_t>(1, i * k + j)) << " ";
// ss_dist << *(dists->data()->GetValues<float>(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<int64_t>(1, 0);
auto p_dist = dis_array->data()->GetValues<float>(1, 0);
auto p_ids = ids_array->data()->GetValues<int64_t>(1, 0);
auto p_dist = dis_array->data()->GetValues<float>(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<VecIndexImpl>(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<VecIndexImpl>(cpu_index, ConvertToCpuIndexType(type));
new_index->dim = dim;
return new_index;
}
VecIndexPtr
VecIndexImpl::Clone() {
// TODO(linxj): exception handle
auto clone_index = std::make_shared<VecIndexImpl>(index_->Clone(), type);
clone_index->dim = dim;
return clone_index;
}
int64_t
VecIndexImpl::GetDeviceId() {
if (auto device_idx = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
return device_idx->GetGpuDevice();
}
// else
return -1; // -1 == cpu
}
float*
BFIndex::GetRawVectors() {
auto raw_index = std::dynamic_pointer_cast<knowhere::IDMAP>(index_);
if (raw_index) {
return raw_index->GetRawVectors();
}
return nullptr;
}
int64_t*
BFIndex::GetRawIds() {
return std::static_pointer_cast<knowhere::IDMAP>(index_)->GetRawIds();
}
ErrorCode
BFIndex::Build(const Config& cfg) {
try {
dim = cfg->d;
std::static_pointer_cast<knowhere::IDMAP>(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<knowhere::IDMAP>(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<knowhere::GPUIndex>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(index_)) {
return std::make_shared<IVFHybridIndex>(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<VecIndexPtr, knowhere::QuantizerPtr>
IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) {
try {
// TODO(linxj): Hardcode here
if (auto hybrid_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg);
auto new_idx = std::make_shared<IVFHybridIndex>(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<VecIndexImpl>(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<VecIndexImpl>(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<VecIndexImpl>(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<knowhere::GPUIndex>(index_)) {
return device_idx->GetGpuDevice();
}
#else
// else
return -1; // -1 == cpu
#endif
}
float *
BFIndex::GetRawVectors() {
auto raw_index = std::dynamic_pointer_cast<knowhere::IDMAP>(index_);
if (raw_index) {
return raw_index->GetRawVectors();
}
return nullptr;
}
int64_t *
BFIndex::GetRawIds() {
return std::static_pointer_cast<knowhere::IDMAP>(index_)->GetRawIds();
}
ErrorCode
BFIndex::Build(const Config &cfg) {
try {
dim = cfg->d;
std::static_pointer_cast<knowhere::IDMAP>(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<knowhere::IDMAP>(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

35
core/src/wrapper/VecImpl.h
View File

@ -77,41 +77,6 @@ class VecIndexImpl : public VecIndex {
std::shared_ptr<knowhere::VectorIndex> index_ = nullptr;
};
class IVFMixIndex : public VecIndexImpl {
public:
explicit IVFMixIndex(std::shared_ptr<knowhere::VectorIndex> 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<knowhere::VectorIndex> 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<VecIndexPtr, knowhere::QuantizerPtr>
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<knowhere::VectorIndex> index)

47
core/src/wrapper/VecIndex.cpp
View File

@ -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 <cuda.h>
#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<knowhere::IVF>();
break;
}
case IndexType::FAISS_IVFFLAT_GPU: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
break;
}
case IndexType::FAISS_IVFFLAT_MIX: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFFLAT_MIX);
}
case IndexType::FAISS_IVFPQ_CPU: {
index = std::make_shared<knowhere::IVFPQ>();
break;
}
case IndexType::FAISS_IVFPQ_GPU: {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
break;
}
case IndexType::SPTAG_KDT_RNT_CPU: {
index = std::make_shared<knowhere::CPUKDTRNG>();
break;
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
}
case IndexType::FAISS_IVFSQ8_CPU: {
index = std::make_shared<knowhere::IVFSQ>();
break;
}
#ifdef MILVUS_GPU_VERSION
case IndexType::FAISS_IVFFLAT_GPU: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
break;
}
case IndexType::FAISS_IVFPQ_GPU: {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
break;
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
}
case IndexType::FAISS_IVFSQ8_GPU: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
break;
}
case IndexType::FAISS_IVFFLAT_MIX: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFFLAT_MIX);
}
#ifdef CUSTOMIZATION
case IndexType::FAISS_IVFSQ8_HYBRID: {
index = std::make_shared<knowhere::IVFSQHybrid>(gpu_device);
return std::make_shared<IVFHybridIndex>(index, IndexType::FAISS_IVFSQ8_HYBRID);
}
#endif
#endif
case IndexType::NSG_MIX: { // TODO(linxj): bug.
index = std::make_shared<knowhere::NSG>(gpu_device);

16
core/src/version.h.macro → core/src/wrapper/WrapperException.cpp
View File

@ -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@
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

36
core/src/wrapper/WrapperException.h Normal file
View File

@ -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 <exception>
#include <string>
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

164
core/src/wrapper/gpu/GPUVecImpl.cpp Normal file
View File

@ -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<knowhere::GPUIndex>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(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<knowhere::IVFSQHybrid>(index_)) {
return std::make_shared<IVFHybridIndex>(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<VecIndexPtr, knowhere::QuantizerPtr>
IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) {
try {
// TODO(linxj): Hardcode here
if (auto hybrid_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg);
auto new_idx = std::make_shared<IVFHybridIndex>(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

66
core/src/wrapper/gpu/GPUVecImpl.h Normal file
View File

@ -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 <memory>
#include <utility>
#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<knowhere::VectorIndex> 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<knowhere::VectorIndex> 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<VecIndexPtr, knowhere::QuantizerPtr>
CopyToGpuWithQuantizer(const int64_t &device_id, const Config &cfg) override;
VecIndexPtr
LoadData(const knowhere::QuantizerPtr &q, const Config &conf) override;
};
} // namespace engine
} // namespace milvus

1
core/thirdparty/versions.txt vendored
View File

@ -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:

6
core/ubuntu_build_deps.sh
View File

@ -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

20
core/unittest/CMakeLists.txt
View File

@ -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)

2
core/unittest/db/CMakeLists.txt
View File

@ -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}
)

47
core/unittest/db/test_db.cpp
View File

@ -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<std::string> 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<std::string> 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<std::string> 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);

22
core/unittest/db/test_db_mysql.cpp
View File

@ -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());
}

24
core/unittest/db/test_mem.cpp
View File

@ -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(); */
}

168
core/unittest/db/test_search.cpp
View File

@ -19,7 +19,6 @@
#include <cmath>
#include <vector>
#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<int64_t>& output_ids,
std::vector<float>& 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<int64_t>& output_ids,
std::vector<float>& output_distance,
uint64_t output_topk,
std::vector<int64_t>& input_ids,
std::vector<float>& 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<int64_t>& input_ids_1,
const std::vector<float>& input_distance_1,
const std::vector<int64_t>& input_ids_2,
const std::vector<float>& 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<float>
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<int64_t> ids1, ids2;
std::vector<float> 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<int64_t> ids1, ids2;
// std::vector<float> 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<size_t> thread_vec = {4};
std::vector<size_t> nq_vec = {1000};
std::vector<size_t> topk_vec = {64};
size_t NQ = nq_vec[nq_vec.size() - 1];
size_t TOPK = topk_vec[topk_vec.size() - 1];
std::vector<int32_t> thread_vec = {4, 8};
std::vector<int32_t> nq_vec = {1, 10, 100};
std::vector<int32_t> 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<ms::ResultIds> id_vec;
std::vector<ms::ResultDistances> dist_vec;
ms::ResultIds input_ids;
ms::ResultDistances input_distances;
std::vector<std::vector<int64_t>> id_vec;
std::vector<std::vector<float>> dist_vec;
std::vector<int64_t> input_ids;
std::vector<float> 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<ms::ResultIds> id_vec_1(index_file_num);
std::vector<ms::ResultDistances> dist_vec_1(index_file_num);
std::vector<std::vector<int64_t>> id_vec_1(index_file_num);
std::vector<std::vector<float>> 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<std::vector<int64_t>> id_vec_2(index_file_num);
// std::vector<std::vector<float>> dist_vec_2(index_file_num);
// std::vector<size_t> k_vec_2(index_file_num);
// std::vector<uint64_t> 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<std::vector<int64_t>> id_vec_3(index_file_num);
// std::vector<std::vector<float>> dist_vec_3(index_file_num);
// std::vector<size_t> k_vec_3(index_file_num);
// std::vector<uint64_t> 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;

13
core/unittest/db/utils.cpp
View File

@ -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

5
core/unittest/metrics/test_metrics.cpp
View File

@ -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";

2
core/unittest/scheduler/CMakeLists.txt
View File

@ -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}

2
core/unittest/server/CMakeLists.txt
View File

@ -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

120
core/unittest/server/test_config.cpp
View File

@ -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());
}

2
core/unittest/server/test_rpc.cpp
View File

@ -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"

2
core/unittest/server/test_util.cpp
View File

@ -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);

8
core/unittest/server/utils.cpp
View File

@ -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*";

8
core/unittest/wrapper/CMakeLists.txt
View File

@ -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

61
core/unittest/wrapper/test_wrapper.cpp
View File

@ -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<milvus::engine::IndexType, std::string, int, int, int, int>> {
protected:
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
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<milvus::engine::NSGConfAdapter>();

6
core/unittest/wrapper/utils.cpp
View File

@ -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) {