Commit c181e405 authored by jinhai's avatar jinhai
Browse files

Merge branch 'enable_nsg' into 'branch-0.5.0' 

MS-596 enable knowhere and fix compile error

See merge request megasearch/milvus!636

Former-commit-id: 39892263eaf578f745a608a4a03305c27f9fa0f3
parents afc55cca c99b3b4a

cpp/src/core/CMakeLists.txt

+4 −2
Original line number Diff line number Diff line
@@ -46,9 +46,11 @@ if(NOT CMAKE_BUILD_TYPE) 
endif(NOT CMAKE_BUILD_TYPE)



if(CMAKE_BUILD_TYPE STREQUAL "Release")

    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -fopenmp")

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

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

cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSG.cpp

+4 −2
Original line number Diff line number Diff line
@@ -24,6 +24,7 @@ 


#include "NSG.h"

#include "knowhere/common/Exception.h"

#include "knowhere/common/Log.h"

#include "knowhere/common/Timer.h"

#include "NSGHelper.h"
@@ -83,8 +84,9 @@ void NsgIndex::Build_with_ids(size_t nb, const float *data, const long *ids, con 
    for (int i = 0; i < ntotal; ++i) {

        total_degree += nsg[i].size();

    }

    std::cout << "graph physical size: " << total_degree * sizeof(node_t) / 1024 / 1024;

    std::cout << "average degree: " << total_degree / ntotal;



    KNOWHERE_LOG_DEBUG << "Graph physical size: " << total_degree * sizeof(node_t) / 1024 / 1024 << "m";

    KNOWHERE_LOG_DEBUG << "Average degree: " << total_degree / ntotal;

    /////



    is_trained = true;

cpp/src/core/test/CMakeLists.txt

+3 −11
Original line number Diff line number Diff line
@@ -29,7 +29,7 @@ set(util_srcs 
        ${CORE_SOURCE_DIR}/knowhere/knowhere/adapter/ArrowAdapter.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/common/Exception.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/common/Timer.cpp

        utils.cpp

        ${CORE_SOURCE_DIR}/test/utils.cpp

        )



#<IVF-TEST>
@@ -52,18 +52,10 @@ target_link_libraries(test_ivf ${depend_libs} ${unittest_libs} ${basic_libs}) 


#<IDMAP-TEST>

set(idmap_srcs

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVF.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp

        )

if(NOT TARGET test_idmap)

    add_executable(test_idmap test_idmap.cpp ${idmap_srcs} ${util_srcs})

    add_executable(test_idmap test_idmap.cpp ${idmap_srcs} ${ivf_srcs} ${util_srcs})

endif()

target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs})
@@ -86,5 +78,5 @@ install(TARGETS test_idmap DESTINATION unittest) 
install(TARGETS test_kdt DESTINATION unittest)



#add_subdirectory(faiss_ori)

#add_subdirectory(test_nsg)

add_subdirectory(test_nsg)

cpp/src/core/test/test_nsg/CMakeLists.txt

+6 −20
Original line number Diff line number Diff line 
##############################

include_directories(/usr/local/include/gperftools)

link_directories(/usr/local/lib)

#include_directories(/usr/local/include/gperftools)

#link_directories(/usr/local/lib)



add_definitions(-std=c++11 -O3 -lboost -march=native -Wall -DINFO)
@@ -13,29 +13,15 @@ else () 
endif ()

message(${OpenMP_CXX_FLAGS})



include_directories(${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/nsg)



aux_source_directory(${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/nsg nsg_src)

include_directories(${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg)

aux_source_directory(${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg nsg_src)



set(interface_src

        ${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/ivf.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/gpu_ivf.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/cloner.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/idmap.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/index/vector_index/nsg_index.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/adapter/structure.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/common/exception.cpp

        ${CORE_SOURCE_DIR}/src/knowhere/common/timer.cpp

        ../utils.cpp

        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexNSG.cpp

        )



if(NOT TARGET test_nsg)

    add_executable(test_nsg

            test_nsg.cpp

            ${interface_src}

            ${nsg_src}

            ${util_srcs}

            )

    add_executable(test_nsg test_nsg.cpp ${interface_src} ${nsg_src} ${util_srcs} ${ivf_srcs})

endif()



target_link_libraries(test_nsg ${depend_libs} ${unittest_libs} ${basic_libs})

cpp/src/core/test/test_nsg/test_nsg.cpp

+30 −22
Original line number Diff line number Diff line
@@ -19,10 +19,11 @@ 
#include <gtest/gtest.h>

#include <memory>



#include "knowhere/common/exception.h"

#include "knowhere/index/vector_index/gpu_ivf.h"

#include "knowhere/index/vector_index/nsg_index.h"

#include "knowhere/index/vector_index/nsg/nsg_io.h"

#include "knowhere/common/Exception.h"

#include "knowhere/index/vector_index/FaissBaseIndex.h"

#include "knowhere/index/vector_index/IndexNSG.h"

#include "knowhere/index/vector_index/nsg/NSGIO.h"

#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"



#include "../utils.h"
@@ -32,16 +33,31 @@ using ::testing::TestWithParam; 
using ::testing::Values;

using ::testing::Combine;



constexpr int64_t DEVICE_ID = 0;

constexpr int64_t DEVICE_ID = 1;



class NSGInterfaceTest : public DataGen, public TestWithParam<::std::tuple<Config, Config>> {

 class NSGInterfaceTest : public DataGen, public ::testing::Test {

 protected:

    void SetUp() override {

        //Init_with_default();

        FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, 1024*1024*200, 1024*1024*600, 2);

        Generate(256, 10000, 1);

        Generate(256, 1000000, 1);

        index_ = std::make_shared<NSG>();

        std::tie(train_cfg, search_cfg) = GetParam();



        auto tmp_conf = std::make_shared<NSGCfg>();

        tmp_conf->gpu_id = DEVICE_ID;

        tmp_conf->knng = 100;

        tmp_conf->nprobe = 32;

        tmp_conf->nlist = 16384;

        tmp_conf->search_length = 60;

        tmp_conf->out_degree = 70;

        tmp_conf->candidate_pool_size = 500;

        tmp_conf->metric_type = METRICTYPE::L2;

        train_conf = tmp_conf;



        auto tmp2_conf = std::make_shared<NSGCfg>();

        tmp2_conf->k = k;

        tmp2_conf->search_length = 30;

        search_conf = tmp2_conf;

    }



    void TearDown() override {
@@ -50,18 +66,10 @@ class NSGInterfaceTest : public DataGen, public TestWithParam<::std::tuple<Confi 


 protected:

    std::shared_ptr<NSG> index_;

    Config train_cfg;

    Config search_cfg;

    Config train_conf;

    Config search_conf;

};



INSTANTIATE_TEST_CASE_P(NSGparameters, NSGInterfaceTest,

                        Values(std::make_tuple(

                            // search length > out_degree

                            Config::object{{"nlist", 128}, {"nprobe", 50}, {"knng", 100}, {"metric_type", "L2"},

                                           {"search_length", 60}, {"out_degree", 70}, {"candidate_pool_size", 500}},

                            Config::object{{"k", 20}, {"search_length", 30}}))

);



void AssertAnns(const DatasetPtr &result,

                const int &nq,

                const int &k) {
@@ -71,17 +79,17 @@ void AssertAnns(const DatasetPtr &result, 
    }

}



TEST_P(NSGInterfaceTest, basic_test) {

TEST_F(NSGInterfaceTest, basic_test) {

    assert(!xb.empty());



    auto model = index_->Train(base_dataset, train_cfg);

    auto result = index_->Search(query_dataset, search_cfg);

    auto model = index_->Train(base_dataset, train_conf);

    auto result = index_->Search(query_dataset, search_conf);

    AssertAnns(result, nq, k);



    auto binaryset = index_->Serialize();

    auto new_index = std::make_shared<NSG>();

    new_index->Load(binaryset);

    auto new_result = new_index->Search(query_dataset, Config::object{{"k", k}});

    auto new_result = new_index->Search(query_dataset, search_conf);

    AssertAnns(result, nq, k);



    ASSERT_EQ(index_->Count(), nb);