Commit ab47dee4 authored by jinhai's avatar jinhai
Browse files

Merge branch 'knowhere_ut' into '0.5.0' 

MS-648 Knowhere ut

See merge request megasearch/milvus!757

Former-commit-id: 4cf9c1ae0671bf9d7fbf227ebd1dc90ec3ebf508
parents 7c3e5bfe 876d7da1

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

+2 −0
Original line number Diff line number Diff line
@@ -189,6 +189,8 @@ IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { 
        if (quantizer_conf->mode != 2) {

            KNOWHERE_THROW_MSG("mode only support 2 in this func");

        }

    } else {

        KNOWHERE_THROW_MSG("conf error");

    }

    //    if (quantizer_conf->gpu_id != gpu_id_) {

    //        KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card");

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

+2 −2
Original line number Diff line number Diff line
@@ -63,7 +63,7 @@ FaissGpuResourceMgr::InitResource() { 


        mutex_cache_.emplace(device_id, std::make_unique<std::mutex>());



        // std::cout << "Device Id: " << device_id << std::endl;

        // std::cout << "Device Id: " << DEVICEID << std::endl;

        auto& device_param = device.second;

        auto& bq = idle_map_[device_id];
@@ -119,7 +119,7 @@ void 
FaissGpuResourceMgr::Dump() {

    for (auto& item : idle_map_) {

        auto& bq = item.second;

        std::cout << "device_id: " << item.first << ", resource count:" << bq.Size();

        std::cout << "DEVICEID: " << item.first << ", resource count:" << bq.Size();

    }

}

core/src/index/unittest/CMakeLists.txt

+8 −0
Original line number Diff line number Diff line
@@ -73,9 +73,17 @@ target_link_libraries(test_kdt 
        SPTAGLibStatic

        ${depend_libs} ${unittest_libs} ${basic_libs})



add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})

target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})



add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})

target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})



install(TARGETS test_ivf DESTINATION unittest)

install(TARGETS test_idmap DESTINATION unittest)

install(TARGETS test_kdt DESTINATION unittest)

install(TARGETS test_gpuresource DESTINATION unittest)

install(TARGETS test_customized_index DESTINATION unittest)



#add_subdirectory(faiss_ori)

add_subdirectory(test_nsg)

core/src/index/unittest/Helper.h

0 → 100644
+120 −0
Original line number Diff line number Diff line 
// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.



#include <memory>

#include <string>



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

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

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

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

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

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

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



constexpr int DEVICEID = 0;

constexpr int64_t DIM = 128;

constexpr int64_t NB = 10000;

constexpr int64_t NQ = 10;

constexpr int64_t K = 10;

constexpr int64_t PINMEM = 1024 * 1024 * 200;

constexpr int64_t TEMPMEM = 1024 * 1024 * 300;

constexpr int64_t RESNUM = 2;



knowhere::IVFIndexPtr

IndexFactory(const std::string& type) {

    if (type == "IVF") {

        return std::make_shared<knowhere::IVF>();

    } else if (type == "IVFPQ") {

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

    } else if (type == "GPUIVF") {

        return std::make_shared<knowhere::GPUIVF>(DEVICEID);

    } else if (type == "GPUIVFPQ") {

        return std::make_shared<knowhere::GPUIVFPQ>(DEVICEID);

    } else if (type == "IVFSQ") {

        return std::make_shared<knowhere::IVFSQ>();

    } else if (type == "GPUIVFSQ") {

        return std::make_shared<knowhere::GPUIVFSQ>(DEVICEID);

    } else if (type == "IVFSQHybrid") {

        return std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

    }

}



enum class ParameterType {

    ivf,

    ivfpq,

    ivfsq,

};



class ParamGenerator {

 public:

    static ParamGenerator&

    GetInstance() {

        static ParamGenerator instance;

        return instance;

    }



    knowhere::Config

    Gen(const ParameterType& type) {

        if (type == ParameterType::ivf) {

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

            tempconf->d = DIM;

            tempconf->gpu_id = DEVICEID;

            tempconf->nlist = 100;

            tempconf->nprobe = 4;

            tempconf->k = K;

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

            return tempconf;

        } else if (type == ParameterType::ivfpq) {

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

            tempconf->d = DIM;

            tempconf->gpu_id = DEVICEID;

            tempconf->nlist = 100;

            tempconf->nprobe = 4;

            tempconf->k = K;

            tempconf->m = 4;

            tempconf->nbits = 8;

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

            return tempconf;

        } else if (type == ParameterType::ivfsq) {

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

            tempconf->d = DIM;

            tempconf->gpu_id = DEVICEID;

            tempconf->nlist = 100;

            tempconf->nprobe = 4;

            tempconf->k = K;

            tempconf->nbits = 8;

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

            return tempconf;

        }

    }

};



#include <gtest/gtest.h>



class TestGpuIndexBase : public ::testing::Test {

 protected:

    void

    SetUp() override {

        knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);

    }



    void

    TearDown() override {

        knowhere::FaissGpuResourceMgr::GetInstance().Free();

    }

};

core/src/index/unittest/test_customized_index.cpp

0 → 100644
+122 −0
Original line number Diff line number Diff line 
// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.



#include <gtest/gtest.h>



#include "unittest/Helper.h"

#include "unittest/utils.h"



class SingleIndexTest : public DataGen, public TestGpuIndexBase {

 protected:

    void

    SetUp() override {

        TestGpuIndexBase::SetUp();

        Generate(DIM, NB, NQ);

        k = K;

    }



    void

    TearDown() override {

        TestGpuIndexBase::TearDown();

    }



 protected:

    std::string index_type;

    knowhere::IVFIndexPtr index_ = nullptr;

};



#ifdef CUSTOMIZATION

TEST_F(SingleIndexTest, IVFSQHybrid) {

    assert(!xb.empty());



    index_type = "IVFSQHybrid";

    index_ = IndexFactory(index_type);

    auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq);

    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 binaryset = index_->Serialize();

    {

        // copy cpu to gpu

        auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

        cpu_idx->Load(binaryset);



        {

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

                auto gpu_idx = cpu_idx->CopyCpuToGpu(DEVICEID, conf);

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

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

                // PrintResult(result, nq, k);

            }

        }

    }



    {

        // quantization already in gpu, only copy data

        auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

        cpu_idx->Load(binaryset);



        auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);

        auto gpu_idx = pair.first;

        auto quantization = pair.second;



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

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

        //        PrintResult(result, nq, k);



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

        quantizer_conf->mode = 2;  // only copy data

        quantizer_conf->gpu_id = DEVICEID;

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

            auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

            hybrid_idx->Load(binaryset);



            auto new_idx = hybrid_idx->LoadData(quantization, quantizer_conf);

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

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

            //            PrintResult(result, nq, k);

        }

    }



    {

        // quantization already in gpu, only set quantization

        auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

        cpu_idx->Load(binaryset);



        auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);

        auto quantization = pair.second;



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

            auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);

            hybrid_idx->Load(binaryset);



            hybrid_idx->SetQuantizer(quantization);

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

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

            //            PrintResult(result, nq, k);

            hybrid_idx->UnsetQuantizer();

        }

    }

}



#endif