Commit fc787705 authored by xiaojun.lin's avatar xiaojun.lin
Browse files

update v1 


Former-commit-id: c012da78a4fd1074660e0c83bc0a594c77a923cc
parent 7c244ec8

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

+9 −0
Original line number Diff line number Diff line
@@ -71,4 +71,13 @@ GPUIVFSQ::CopyGpuToCpu(const Config& config) { 
    return std::make_shared<IVFSQ>(new_index);

}



void

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

#ifdef CUSTOMIZATION

    GPUIVF::search_impl(n, data, k, distances, labels, cfg);

#else

    IVF::search_impl(n, data, k, distances, labels, cfg);

#endif

}



}  // namespace knowhere

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

+4 −0
Original line number Diff line number Diff line
@@ -38,6 +38,10 @@ class GPUIVFSQ : public GPUIVF { 


    VectorIndexPtr

    CopyGpuToCpu(const Config& config) override;



 protected:

    void

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

};



}  // namespace knowhere

cpp/src/core/unittest/test_ivf.cpp

+23 −68
Original line number Diff line number Diff line
@@ -154,8 +154,8 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P 
INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,

                        Values(std::make_tuple("IVF", ParameterType::ivf),

                               std::make_tuple("GPUIVF", ParameterType::ivf),

                               //                            std::make_tuple("IVFPQ", ParameterType::ivfpq),

                               //                            std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),

                               std::make_tuple("IVFPQ", ParameterType::ivfpq),

                               std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),

                               std::make_tuple("IVFSQ", ParameterType::ivfsq),

#ifdef CUSTOMIZATION

                               std::make_tuple("IVFSQHybrid", ParameterType::ivfsq),
@@ -240,25 +240,26 @@ TEST_P(IVFTest, hybrid) { 
        auto result = hybrid_1_idx->Search(query_dataset, conf);

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

        PrintResult(result, nq, k);

        hybrid_1_idx->UnsetQuantizer();

    }



    {

        auto hybrid_2_idx = std::make_shared<knowhere::IVFSQHybrid>(device_id);



        auto binaryset = index_->Serialize();

        hybrid_2_idx->Load(binaryset);



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

        quantizer_conf->mode = 1;

        quantizer_conf->gpu_id = device_id;

        auto q = hybrid_2_idx->LoadQuantizer(quantizer_conf);

        quantizer_conf->mode = 2;

        hybrid_2_idx->LoadData(q, quantizer_conf);



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

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

        PrintResult(result, nq, k);

    }

//    {

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

//

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

//        hybrid_2_idx->Load(binaryset);

//

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

//        quantizer_conf->mode = 1;

//        quantizer_conf->gpu_id = device_id;

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

//        quantizer_conf->mode = 2;

//        hybrid_2_idx->LoadData(q, quantizer_conf);

//

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

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

//        PrintResult(result, nq, k);

//    }

}



// TEST_P(IVFTest, gpu_to_cpu) {
@@ -438,6 +439,7 @@ TEST_P(IVFTest, clone_test) { 
    }

}



#ifdef CUSTOMIZATION

TEST_P(IVFTest, seal_test) {

    // FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
@@ -472,6 +474,7 @@ TEST_P(IVFTest, seal_test) { 
    auto with_seal = tc.RecordSection("With seal");

    ASSERT_GE(without_seal, with_seal);

}

#endif



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

 protected:
@@ -637,7 +640,7 @@ TEST_F(GPURESTEST, copyandsearch) { 
    // search and copy at the same time

    printf("==================\n");



    index_type = "GPUIVFSQ";

    index_type = "GPUIVF";

    index_ = IndexFactory(index_type);



    auto conf = std::make_shared<knowhere::IVFSQCfg>();
@@ -693,54 +696,6 @@ TEST_F(GPURESTEST, copyandsearch) { 


    std::thread search_thread(search_func);

    std::thread load_thread(load_func);

    search_thread.join();

    load_thread.join();

    tc.RecordSection("Copy&search total");

}



TEST_F(GPURESTEST, TrainAndSearch) {

    index_type = "GPUIVFSQ";

    index_ = IndexFactory(index_type);



    auto conf = std::make_shared<knowhere::IVFSQCfg>();

    conf->nlist = 1638;

    conf->d = dim;

    conf->gpu_id = device_id;

    conf->metric_type = knowhere::METRICTYPE::L2;

    conf->k = k;

    conf->nbits = 8;

    conf->nprobe = 1;



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

    index_->set_preprocessor(preprocessor);

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

    auto new_index = IndexFactory(index_type);

    new_index->set_index_model(model);

    new_index->Add(base_dataset, conf);

    auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config());

    cpu_idx->Seal();

    auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());



    constexpr int train_count = 1;

    constexpr int search_count = 5000;

    auto train_stage = [&] {

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

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

            auto test_idx = IndexFactory(index_type);

            test_idx->set_index_model(model);

            test_idx->Add(base_dataset, conf);

        }

    };

    auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) {

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

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

            AssertAnns(result, nq, k);

        }

    };



    // TimeRecorder tc("record");

    // train_stage();

    // tc.RecordSection("train cost");

    // search_stage(search_idx);

    // tc.RecordSection("search cost");

cpp/src/core/unittest/test_kdt.cpp

+1 −0
Original line number Diff line number Diff line
@@ -36,6 +36,7 @@ class KDTTest : public DataGen, public ::testing::Test { 
 protected:

    void

    SetUp() override {

        Generate(96, 1000, 10);

        index_ = std::make_shared<knowhere::CPUKDTRNG>();



        auto tempconf = std::make_shared<knowhere::KDTCfg>();

cpp/src/core/unittest/test_nsg/test_nsg.cpp

+7 −7
Original line number Diff line number Diff line
@@ -38,17 +38,17 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test { 
    SetUp() override {

        // Init_with_default();

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

        Generate(256, 1000000, 1);

        Generate(256, 1000000 / 100, 1);

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



        auto tmp_conf = std::make_shared<knowhere::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->knng = 20;

        tmp_conf->nprobe = 8;

        tmp_conf->nlist = 163;

        tmp_conf->search_length = 40;

        tmp_conf->out_degree = 30;

        tmp_conf->candidate_pool_size = 100;

        tmp_conf->metric_type = knowhere::METRICTYPE::L2;

        train_conf = tmp_conf;