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

Merge pull request #324 from Yukikaze-CZR/0.6.0 

Support new Index type IVFPQ #127
parents 6c996b5a 4d1345b4

CHANGELOG.md

+2 −1
Original line number Diff line number Diff line
@@ -14,6 +14,7 @@ Please mark all change in change log and use the ticket from JIRA. 
- \#12 - Pure CPU version for Milvus

- \#77 - Support table partition

- \#226 - Experimental shards middleware for Milvus

- \#127 - Support new Index type IVFPQ



## Improvement

- \#275 - Rename C++ SDK IndexType

core/src/db/engine/ExecutionEngine.h

+3 −1
Original line number Diff line number Diff line
@@ -26,6 +26,7 @@ 
namespace milvus {

namespace engine {



// TODO(linxj): replace with VecIndex::IndexType

enum class EngineType {

    INVALID = 0,

    FAISS_IDMAP = 1,
@@ -33,7 +34,8 @@ enum class EngineType { 
    FAISS_IVFSQ8,

    NSG_MIX,

    FAISS_IVFSQ8H,

    MAX_VALUE = FAISS_IVFSQ8H,

    FAISS_PQ,

    MAX_VALUE = FAISS_PQ,

};



enum class MetricType {

core/src/db/engine/ExecutionEngineImpl.cpp

+8 −0
Original line number Diff line number Diff line
@@ -116,6 +116,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) { 
            index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID);

            break;

        }

        case EngineType::FAISS_PQ: {

#ifdef MILVUS_CPU_VERSION

            index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_CPU);

#else

            index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX);

#endif

            break;

        }

        default: {

            ENGINE_LOG_ERROR << "Unsupported index type";

            return nullptr;

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

+21 −12
Original line number Diff line number Diff line
@@ -39,17 +39,19 @@ GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) { 


    GETTENSOR(dataset)



    // TODO(linxj): set device here.

    // TODO(linxj): set gpu resource here.

    faiss::gpu::StandardGpuResources res;

    faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, build_cfg->nlist, build_cfg->m, build_cfg->nbits,

    auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);

    if (temp_resource != nullptr) {

        ResScope rs(temp_resource, gpu_id_, true);

        auto device_index = new faiss::gpu::GpuIndexIVFPQ(temp_resource->faiss_res.get(), dim, build_cfg->nlist,

                                                          build_cfg->m, build_cfg->nbits,

                                                          GetMetricType(build_cfg->metric_type));  // IP not support

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



        device_index->train(rows, (float*)p_data);

        std::shared_ptr<faiss::Index> host_index = nullptr;

    host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index));



        host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));

        return std::make_shared<IVFIndexModel>(host_index);

    } else {

        KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");

    }

}



std::shared_ptr<faiss::IVFSearchParameters>
@@ -66,7 +68,14 @@ GPUIVFPQ::GenParams(const Config& config) { 


VectorIndexPtr

GPUIVFPQ::CopyGpuToCpu(const Config& config) {

    KNOWHERE_THROW_MSG("not support yet");

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



    faiss::Index* device_index = index_.get();

    faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);



    std::shared_ptr<faiss::Index> new_index;

    new_index.reset(host_index);

    return std::make_shared<IVFPQ>(new_index);

}



}  // namespace knowhere

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

+5 −0
Original line number Diff line number Diff line
@@ -18,6 +18,7 @@ 
#pragma once



#include <memory>

#include <utility>



#include "IndexGPUIVF.h"
@@ -28,6 +29,10 @@ class GPUIVFPQ : public GPUIVF { 
    explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) {

    }



    GPUIVFPQ(std::shared_ptr<faiss::Index> index, const int64_t& device_id, ResPtr& resource)

        : GPUIVF(std::move(index), device_id, resource) {

    }



    IndexModelPtr

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