Merge branch 'cpu-version' into jinhai

Former-commit-id: 17a38d07e55b49888cb2abdd6aed3eca85e7a5f9

cpp/CMakeLists.txt

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

cpp/build.sh

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

cpp/src/CMakeLists.txt

@@ -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)
 
-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
+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})
 
-target_link_libraries(vecwise_engine ${engine_libs} ${cuda_library})
+target_link_libraries(vecwise_engine ${engine_libs})
 
 add_executable(vecwise_server
         ${config_files}
@@ -71,6 +83,7 @@ set(server_libs
         libz.a
         libzstd.a
         liblz4.a
+        dl
         )
 
 target_link_libraries(vecwise_server ${server_libs})

cpp/src/db/DB.cpp

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

cpp/src/db/DBImpl.cpp

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

cpp/src/wrapper/Index.cpp

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

cpp/src/wrapper/Index.h

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

cpp/src/wrapper/IndexBuilder.cpp

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

cpp/unittest/faiss_wrapper/CMakeLists.txt

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

cpp/unittest/faiss_wrapper/topk_test.cpp

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

cpp/unittest/server/CMakeLists.txt

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