Unverified Commit 41f3604e authored by Jin Hai's avatar Jin Hai Committed by GitHub
Browse files

Merge pull request #629 from tinkerlin/issue-548 

Issue 548
parents 9b389264 3a87f16b

CHANGELOG.md

+1 −0
Original line number Diff line number Diff line
@@ -36,6 +36,7 @@ Please mark all change in change log and use the ticket from JIRA. 
- \#543 - client raise exception in shards when search results is empty

- \#545 - Avoid dead circle of build index thread when error occurs

- \#547 - NSG build failed using GPU-edition if set gpu_enable false

- \#548 - NSG search accuracy is too low

- \#552 - Server down during building index_type: IVF_PQ using GPU-edition

- \#561 - Milvus server should report exception/error message or terminate on mysql metadata backend error

- \#579 - Build index hang in GPU version when gpu_resources disabled

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

+34 −0
Original line number Diff line number Diff line
@@ -126,4 +126,38 @@ GPUIDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, 
    index_->search(n, (float*)data, k, distances, labels);

}



void

GPUIDMAP::GenGraph(float* data, const int64_t& k, Graph& graph, const Config& config) {

    int64_t K = k + 1;

    auto ntotal = Count();



    size_t dim = config->d;

    auto batch_size = 1000;

    auto tail_batch_size = ntotal % batch_size;

    auto batch_search_count = ntotal / batch_size;

    auto total_search_count = tail_batch_size == 0 ? batch_search_count : batch_search_count + 1;



    std::vector<float> res_dis(K * batch_size);

    graph.resize(ntotal);

    Graph res_vec(total_search_count);

    for (int i = 0; i < total_search_count; ++i) {

        auto b_size = (i == (total_search_count - 1)) && tail_batch_size != 0 ? tail_batch_size : batch_size;



        auto& res = res_vec[i];

        res.resize(K * b_size);



        auto xq = data + batch_size * dim * i;

        search_impl(b_size, (float*)xq, K, res_dis.data(), res.data(), config);



        for (int j = 0; j < b_size; ++j) {

            auto& node = graph[batch_size * i + j];

            node.resize(k);

            auto start_pos = j * K + 1;

            for (int m = 0, cursor = start_pos; m < k && cursor < start_pos + k; ++m, ++cursor) {

                node[m] = res[cursor];

            }

        }

    }

}



}  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h

+4 −0
Original line number Diff line number Diff line
@@ -23,6 +23,7 @@ 


#include <memory>

#include <utility>

#include <vector>



namespace knowhere {
@@ -47,6 +48,9 @@ class GPUIDMAP : public IDMAP, public GPUIndex { 
    VectorIndexPtr

    CopyGpuToGpu(const int64_t& device_id, const Config& config) override;



    void

    GenGraph(float* data, const int64_t& k, Graph& graph, const Config& config);



 protected:

    void

    search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override;

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

+20 −0
Original line number Diff line number Diff line
@@ -121,6 +121,26 @@ IDMAP::Add(const DatasetPtr& dataset, const Config& config) { 
    index_->add_with_ids(rows, (float*)p_data, p_ids);

}



void

IDMAP::AddWithoutId(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];



    std::vector<int64_t> new_ids(rows);

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

        new_ids[i] = i;

    }



    index_->add_with_ids(rows, (float*)p_data, new_ids.data());

}



int64_t

IDMAP::Count() {

    return index_->ntotal;

core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h

+14 −0
Original line number Diff line number Diff line
@@ -34,20 +34,31 @@ class IDMAP : public VectorIndex, public FaissBaseIndex { 


    BinarySet

    Serialize() override;



    void

    Load(const BinarySet& index_binary) override;



    void

    Train(const Config& config);



    DatasetPtr

    Search(const DatasetPtr& dataset, const Config& config) override;



    int64_t

    Count() override;



    //    VectorIndexPtr

    //    Clone() override;



    int64_t

    Dimension() override;



    void

    Add(const DatasetPtr& dataset, const Config& config) override;



    void

    AddWithoutId(const DatasetPtr& dataset, const Config& config);



    VectorIndexPtr

    CopyCpuToGpu(const int64_t& device_id, const Config& config);

    void
@@ -55,12 +66,15 @@ class IDMAP : public VectorIndex, public FaissBaseIndex { 


    virtual float*

    GetRawVectors();



    virtual int64_t*

    GetRawIds();



 protected:

    virtual void

    search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg);



 protected:

    std::mutex mutex_;

};