Commit 9cc0c598 authored by xiaojun.lin's avatar xiaojun.lin
Browse files

add hybrid mode 


Former-commit-id: 436deb945e42a7205fa241150cc4ab23edb05913
parent a42dd429

cpp/src/core/knowhere/CMakeLists.txt

+1 −0
Original line number Diff line number Diff line
@@ -50,6 +50,7 @@ set(index_srcs 
        knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp

        knowhere/index/vector_index/IndexIVFSQ.cpp

        knowhere/index/vector_index/IndexGPUIVFSQ.cpp

        knowhere/index/vector_index/IndexIVFSQHybrid.cpp

        knowhere/index/vector_index/IndexIVFPQ.cpp

        knowhere/index/vector_index/IndexGPUIVFPQ.cpp

        knowhere/index/vector_index/FaissBaseIndex.cpp

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

+0 −11
Original line number Diff line number Diff line
@@ -127,17 +127,6 @@ void GPUIVF::LoadImpl(const BinarySet &index_binary) { 
    }

}



IVFIndexPtr GPUIVF::Copy_index_gpu_to_cpu() {

    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<IVF>(new_index);

}



void GPUIVF::search_impl(int64_t n,

                         const float *data,

                         int64_t k,

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

+0 −3
Original line number Diff line number Diff line
@@ -75,9 +75,6 @@ public: 
     VectorIndexPtr

     Clone() final;



     // TODO(linxj): Deprecated

     virtual IVFIndexPtr Copy_index_gpu_to_cpu();



protected:

     void

     search_impl(int64_t n,

cpp/src/core/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp

+5 −0
Original line number Diff line number Diff line
@@ -22,6 +22,7 @@ 
#include "knowhere/index/vector_index/IndexIVFPQ.h"

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

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

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

#include "Cloner.h"
@@ -38,6 +39,10 @@ VectorIndexPtr CopyGpuToCpu(const VectorIndexPtr &index, const Config &config) { 
}



VectorIndexPtr CopyCpuToGpu(const VectorIndexPtr &index, const int64_t &device_id, const Config &config) {

    if (auto device_index = std::dynamic_pointer_cast<IVFSQHybrid>(index)) {

        return device_index->CopyCpuToGpu(device_id, config);

    }



    if (auto device_index = std::dynamic_pointer_cast<GPUIndex>(index)) {

        return device_index->CopyGpuToGpu(device_id, config);

    }

cpp/src/core/test/test_ivf.cpp

+32 −0
Original line number Diff line number Diff line
@@ -25,6 +25,7 @@ 
#include <faiss/gpu/GpuAutoTune.h>

#include <faiss/gpu/GpuIndexIVFFlat.h>



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

#include "knowhere/common/Exception.h"

#include "knowhere/common/Timer.h"

#include "knowhere/adapter/Structure.h"
@@ -209,6 +210,37 @@ TEST_P(IVFTest, ivf_basic) { 
    //PrintResult(result, nq, k);

}



//TEST_P(IVFTest, hybrid) {

//    assert(!xb.empty());

//

//    auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);

//    index_->set_preprocessor(preprocessor);

//

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

//    index_->set_index_model(model);

//    index_->Add(base_dataset, conf);

//    EXPECT_EQ(index_->Count(), nb);

//    EXPECT_EQ(index_->Dimension(), dim);

//

////    auto new_idx = ChooseTodo();

////    auto result = new_idx->Search(query_dataset, conf);

////    AssertAnns(result, nq, conf->k);

//

//    auto iss_idx = std::make_shared<IVFSQHybrid>(device_id);

//

//    auto binaryset = index_->Serialize();

//    iss_idx->Load(binaryset);

//

//    auto quantizer_conf = std::make_shared<QuantizerCfg>();

//    quantizer_conf->mode = 1;

//    quantizer_conf->gpu_id = 1;

//    auto q = iss_idx->LoadQuantizer(quantizer_conf);

//    iss_idx->SetQuantizer(q);

//    auto result = iss_idx->Search(query_dataset, conf);

//    AssertAnns(result, nq, conf->k);

//    //PrintResult(result, nq, k);

//}



//TEST_P(IVFTest, gpu_to_cpu) {

//    if (index_type.find("GPU") == std::string::npos) { return; }

//