admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-01-06 10:52:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'jinhai' of http://191.168.1.105:6060/jinhai/vecwise_engine into jinhai

Browse Source 
Former-commit-id: 4d8bb348e42f7d791daf9dc5c59e66d45d395214
This commit is contained in:
jinhai 2019-05-10 23:44:19 -04:00
commit 1390398cdb
16 changed files with 85 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
cpp/CMakeLists.txt
View File

@ -14,6 +14,10 @@ set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
message("CUDA_TOOLKIT_ROOT_DIR=${CUDA_TOOLKIT_ROOT_DIR}")
message("CUDA_NVCC_FLAGS=${CUDA_NVCC_FLAGS}")
if (GPU_VERSION STREQUAL "ON")
add_definitions("-DGPU_VERSION")
endif ()
set(CMAKE_CXX_STANDARD 14)
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")

7
cpp/build.sh
View File

@ -2,9 +2,10 @@
BUILD_TYPE="Debug"
BUILD_UNITTEST="off"
BUILD_GPU="OFF"
INSTALL_PREFIX=$(pwd)/megasearch
while getopts "p:t:uh" arg
while getopts "p:t:uhg" arg
do
case $arg in
t)
@ -17,6 +18,9 @@ do
p)
INSTALL_PREFIX=$OPTARG
;;
g)
BUILD_GPU="ON"
;;
h) # help
echo "
@ -51,6 +55,7 @@ CMAKE_CMD="cmake -DBUILD_UNIT_TEST=${BUILD_UNITTEST} \
-DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
-DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
-DCMAKE_CUDA_COMPILER=${CUDA_COMPILER} \
-DGPU_VERSION=${BUILD_GPU} \
$@ ../"
echo ${CMAKE_CMD}

41
cpp/src/CMakeLists.txt
View File

@ -30,23 +30,35 @@ set(vecwise_engine_src
include_directories(/usr/include)
include_directories(/usr/local/cuda/include)
find_library(cuda_library cudart cublas HINTS /usr/local/cuda/lib64)
if (GPU_VERSION STREQUAL "ON")
link_directories(/usr/local/cuda/lib64)
set(engine_libs
pthread
libfaiss.a
libgomp.a
libopenblas.a
libgfortran.a
libquadmath.a
cudart
cublas
libsqlite3.a
)
else()
set(engine_libs
pthread
libfaiss.a
libgomp.a
libopenblas.a
libgfortran.a
libquadmath.a
libsqlite3.a
)
endif ()
cuda_add_library(vecwise_engine STATIC ${vecwise_engine_src})
set(engine_libs
pthread
libfaiss.a
libgomp.a
libopenblas.a
libgfortran.a
libquadmath.a
cudart
cublas
libsqlite3.a
)
target_link_libraries(vecwise_engine ${engine_libs} ${cuda_library})
target_link_libraries(vecwise_engine ${engine_libs})
add_executable(vecwise_server
${config_files}
@ -70,6 +82,7 @@ set(server_libs
libz.a
libzstd.a
liblz4.a
dl
)
target_link_libraries(vecwise_server ${server_libs})

9
cpp/src/db/DB.cpp
View File

@ -19,7 +19,14 @@ DB::~DB() {}
void DB::Open(const Options& options, DB** dbptr) {
*dbptr = nullptr;
*dbptr = DBFactory::Build(options);
#ifdef GPU_VERSION
std::string default_index_type{"Faiss,IVF"};
#else
std::string default_index_type{"Faiss,IDMap"};
#endif
*dbptr = DBFactory::Build(options, default_index_type);
return;
}

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

@ -11,7 +11,6 @@
#include <thread>
#include <iostream>
#include <cstring>
#include <wrapper/Topk.h>
#include <easylogging++.h>
#include <cache/CpuCacheMgr.h>
@ -142,12 +141,27 @@ Status DBImpl<EngineT>::search(const std::string& group_id, size_t k, size_t nq,
}
};
auto topk_cpu = [](const std::vector<float> &input_data,
const int &k,
float *output_distence,
long *output_ids) -> void {
std::map<float, int> inverted_table;
for (int i = 0; i < input_data.size(); ++i) {
inverted_table[input_data[i]] = i;
}
int count = 0;
for (auto it = inverted_table.begin(); it != inverted_table.end() && count < k; ++it, ++count) {
output_distence[count] = it->first;
output_ids[count] = it->second;
}
};
auto cluster_topk = [&]() -> void {
QueryResult res;
for (auto &result_pair : batchresult) {
auto &dis = result_pair.second;
auto &nns = result_pair.first;
TopK(dis.data(), dis.size(), k, output_distence, output_ids);
topk_cpu(dis, k, output_distence, output_ids);
for (int i = 0; i < k; ++i) {
res.emplace_back(nns[output_ids[i]]); // mapping
}

3
cpp/src/db/MemManager.cpp
View File

@ -33,9 +33,6 @@ template<typename EngineT>
void MemVectors<EngineT>::add(size_t n_, const float* vectors_, IDNumbers& vector_ids_) {
_pIdGenerator->getNextIDNumbers(n_, vector_ids_);
pEE_->AddWithIds(n_, vectors_, vector_ids_.data());
for(auto i=0 ; i<n_; i++) {
vector_ids_.push_back(i);
}
}
template<typename EngineT>

3
cpp/src/wrapper/Index.cpp
View File

@ -4,7 +4,8 @@
// Proprietary and confidential.
////////////////////////////////////////////////////////////////////////////////
#ifdef CUDA_VERSION
// TODO: maybe support static search
#ifdef GPU_VERSION
#include "faiss/gpu/GpuAutoTune.h"
#include "faiss/gpu/StandardGpuResources.h"
#include "faiss/gpu/utils/DeviceUtils.h"

2
cpp/src/wrapper/Index.h
View File

@ -13,8 +13,6 @@
#include <fstream>
#include "faiss/AutoTune.h"
#include "faiss/AuxIndexStructures.h"
#include "faiss/gpu/GpuAutoTune.h"
#include "faiss/index_io.h"
#include "Operand.h"

15
cpp/src/wrapper/IndexBuilder.cpp
View File

@ -6,9 +6,12 @@
#include "mutex"
#ifdef GPU_VERSION
#include <faiss/gpu/StandardGpuResources.h>
#include "faiss/gpu/GpuIndexIVFFlat.h"
#include "faiss/gpu/GpuAutoTune.h"
#endif
#include "faiss/IndexFlat.h"
#include "IndexBuilder.h"
@ -34,6 +37,7 @@ Index_ptr IndexBuilder::build_all(const long &nb,
const long &nt,
const float *xt) {
std::shared_ptr<faiss::Index> host_index = nullptr;
#ifdef GPU_VERSION
{
// TODO: list support index-type.
faiss::Index *ori_index = faiss::index_factory(opd_->d, opd_->get_index_type(nb).c_str());
@ -52,6 +56,17 @@ Index_ptr IndexBuilder::build_all(const long &nb,
delete device_index;
delete ori_index;
}
#else
{
faiss::Index *index = faiss::index_factory(opd_->d, opd_->get_index_type(nb).c_str());
if (!index->is_trained) {
nt == 0 || xt == nullptr ? index->train(nb, xb)
: index->train(nt, xt);
}
index->add_with_ids(nb, xb, ids);
host_index.reset(index);
}
#endif
return std::make_shared<Index>(host_index);
}

0
cpp/src/wrapper/Arithmetic.h → cpp/src/wrapper/gpu/Arithmetic.h
View File

0
cpp/src/wrapper/Topk.cu → cpp/src/wrapper/gpu/Topk.cu
View File

0
cpp/src/wrapper/Topk.h → cpp/src/wrapper/gpu/Topk.h
View File

3
cpp/unittest/db/db_tests.cpp
View File

@ -74,7 +74,7 @@ TEST_F(DBTest, DB_TEST) {
for (auto result : results[k]) {
ss << result << " ";
}
LOG(DEBUG) << ss.str();
/* LOG(DEBUG) << ss.str(); */
}
ASSERT_TRUE(count >= prev_count);
std::this_thread::sleep_for(std::chrono::seconds(1));
@ -86,6 +86,7 @@ TEST_F(DBTest, DB_TEST) {
for (auto i=0; i<loop; ++i) {
if (i==40) {
db_->add_vectors(group_name, qb, qxb, target_ids);
ASSERT_EQ(target_ids.size(), qb);
} else {
db_->add_vectors(group_name, nb, xb, vector_ids);
}

6
cpp/unittest/faiss_wrapper/CMakeLists.txt
View File

@ -48,7 +48,7 @@ target_link_libraries(wrapper_test ${unittest_libs} ${wrapper_libs})
set(topk_test_src
topk_test.cpp
${CMAKE_SOURCE_DIR}/src/wrapper/Topk.cu)
${CMAKE_SOURCE_DIR}/src/wrapper/gpu/Topk.cu)
cuda_add_executable(topk_test ${topk_test_src})
target_link_libraries(topk_test ${unittest_libs} ${faiss_libs})
#cuda_add_executable(topk_test ${topk_test_src})
#target_link_libraries(topk_test ${unittest_libs} ${faiss_libs})

2
cpp/unittest/faiss_wrapper/topk_test.cpp
View File

@ -6,7 +6,7 @@
#include <gtest/gtest.h>
#include "wrapper/Topk.h"
#include "wrapper/gpu/Topk.h"
using namespace zilliz::vecwise::engine;

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

@ -7,7 +7,7 @@ include_directories(../../src)
include_directories(/usr/include)
include_directories(/usr/local/cuda/include)
find_library(cuda_library cudart cublas HINTS /usr/local/cuda/lib64)
link_directories(/usr/local/cuda/lib64)
aux_source_directory(../../src/config config_files)
aux_source_directory(../../src/cache cache_srcs)