Commit 81c37afc authored by jinhai's avatar jinhai
Browse files

Merge branch 'refactor_knowhere' into 'branch-0.5.0' 

MS-565 Refactor knowhere

See merge request megasearch/milvus!632

Former-commit-id: bc42aef029e853c270141d5a5727fed0c4f624fe
parents 4d72034f 2652b3bc

cpp/src/core/CMakeLists.txt

+1 −1
Original line number Diff line number Diff line
@@ -93,7 +93,7 @@ endif() 
set(CORE_INCLUDE_DIRS ${CORE_INCLUDE_DIRS} PARENT_SCOPE)



if(BUILD_UNIT_TEST STREQUAL "ON")

#    add_subdirectory(test)

    add_subdirectory(test)

endif()



config_summary()

cpp/src/core/knowhere/CMakeLists.txt

+1 −4
Original line number Diff line number Diff line
@@ -8,7 +8,6 @@ link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64) 
include_directories(${CORE_SOURCE_DIR}/knowhere)

include_directories(${CORE_SOURCE_DIR}/thirdparty)

include_directories(${CORE_SOURCE_DIR}/thirdparty/SPTAG/AnnService)

include_directories(${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include)



set(SPTAG_SOURCE_DIR ${CORE_SOURCE_DIR}/thirdparty/SPTAG)

file(GLOB HDR_FILES
@@ -55,6 +54,7 @@ set(index_srcs 
        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

        )



set(depend_libs
@@ -117,7 +117,6 @@ set(CORE_INCLUDE_DIRS 
        ${CORE_SOURCE_DIR}/knowhere

        ${CORE_SOURCE_DIR}/thirdparty

        ${CORE_SOURCE_DIR}/thirdparty/SPTAG/AnnService

        ${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include

        ${ARROW_INCLUDE_DIR}

        ${FAISS_INCLUDE_DIR}

        ${OPENBLAS_INCLUDE_DIR}
@@ -129,8 +128,6 @@ set(CORE_INCLUDE_DIRS ${CORE_INCLUDE_DIRS} PARENT_SCOPE) 


#INSTALL(DIRECTORY

#        ${CORE_SOURCE_DIR}/include/knowhere

#        ${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include/jsoncons

#        ${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include/jsoncons_ext

#        ${ARROW_INCLUDE_DIR}/arrow

#        ${FAISS_PREFIX}/include/faiss

#        ${OPENBLAS_INCLUDE_DIR}/

cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp

+1 −2
Original line number Diff line number Diff line
@@ -69,8 +69,7 @@ ConvertToQueryResult(const DatasetPtr &dataset, const Config &config) { 
    auto dimension = tensor->shape()[1];

    auto rows = tensor->shape()[0];



    auto k = config[META_K].as<int64_t>();

    std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, k, true));

    std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, config->k, true));

    for (auto i = 0; i < rows; ++i) {

        query_results[i].SetTarget(&p_data[i * dimension]);

    }

cpp/src/core/knowhere/knowhere/common/Config.h

+32 −5
Original line number Diff line number Diff line
@@ -18,15 +18,42 @@ 


#pragma once



#include <jsoncons/json.hpp>



#include <memory>



namespace zilliz {

namespace knowhere {





using Config = jsoncons::json;



enum class METRICTYPE {

    INVALID = 0,

    L2 = 1,

    IP = 2,

};



// General Config

constexpr int64_t INVALID_VALUE = -1;

constexpr int64_t DEFAULT_K = INVALID_VALUE;

constexpr int64_t DEFAULT_DIM = INVALID_VALUE;

constexpr int64_t DEFAULT_GPUID = INVALID_VALUE;

constexpr METRICTYPE DEFAULT_TYPE = METRICTYPE::INVALID;



struct Cfg {

    METRICTYPE metric_type = DEFAULT_TYPE;

    int64_t k = DEFAULT_K;

    int64_t gpu_id = DEFAULT_GPUID;

    int64_t d = DEFAULT_DIM;



    Cfg(const int64_t &dim,

        const int64_t &k,

        const int64_t &gpu_id,

        METRICTYPE type)

        : d(dim), k(k), gpu_id(gpu_id), metric_type(type) {}



    Cfg() = default;



    virtual bool

    CheckValid(){};

};

using Config = std::shared_ptr<Cfg>;



} // namespace knowhere

} // namespace zilliz

cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp

+11 −7
Original line number Diff line number Diff line
@@ -36,10 +36,11 @@ namespace zilliz { 
namespace knowhere {



IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) {

    auto nlist = config["nlist"].as<size_t>();

    gpu_id_ = config.get_with_default("gpu_id", gpu_id_);

    auto metric_type = config["metric_type"].as_string() == "L2" ?

                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;

    auto build_cfg = std::dynamic_pointer_cast<IVFCfg>(config);

    if (build_cfg != nullptr) {

        build_cfg->CheckValid(); // throw exception

    }

    gpu_id_ = build_cfg->gpu_id;



    GETTENSOR(dataset)
@@ -48,7 +49,9 @@ IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) { 
        ResScope rs(temp_resource, gpu_id_, true);

        faiss::gpu::GpuIndexIVFFlatConfig idx_config;

        idx_config.device = gpu_id_;

        faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, nlist, metric_type, idx_config);

        faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim,

                                                 build_cfg->nlist, GetMetricType(build_cfg->metric_type),

                                                 idx_config);

        device_index.train(rows, (float *) p_data);



        std::shared_ptr<faiss::Index> host_index = nullptr;
@@ -143,9 +146,10 @@ void GPUIVF::search_impl(int64_t n, 
                         const Config &cfg) {

    std::lock_guard<std::mutex> lk(mutex_);



    // TODO(linxj): gpu index support GenParams

    if (auto device_index = std::static_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {

        auto nprobe = cfg.get_with_default("nprobe", size_t(1));

        device_index->setNumProbes(nprobe);

        auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(cfg);

        device_index->setNumProbes(search_cfg->nprobe);



        {

            // TODO(linxj): allocate mem