enable hybrid IVFSQ8

Former-commit-id: 9bbf835e63818ef6cc5f792d0ab11247794d0c3a

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

@@ -63,7 +63,7 @@ VectorIndexPtr
 IVFSQHybrid::CopyGpuToCpu(const Config &config) {
     std::lock_guard<std::mutex> lk(mutex_);
 
-    if (auto device_idx = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) {
+    if (auto device_idx = std::dynamic_pointer_cast<faiss::IndexIVF>(index_)) {
         faiss::Index *device_index = index_.get();
         faiss::Index *host_index = faiss::gpu::index_gpu_to_cpu(device_index);
 
@@ -113,6 +113,7 @@ IVFSQHybrid::search_impl(int64_t n,
     if (gpu_mode) {
         GPUIVF::search_impl(n, data, k, distances, labels, cfg);
     } else {
+        ResScope rs(res_, gpu_id_);
         IVF::search_impl(n, data, k, distances, labels, cfg);
     }
 }
@@ -121,7 +122,9 @@ QuantizerPtr
 IVFSQHybrid::LoadQuantizer(const Config &conf) {
     auto quantizer_conf = std::dynamic_pointer_cast<QuantizerCfg>(conf);
     if (quantizer_conf != nullptr) {
-        quantizer_conf->CheckValid(); // throw exception
+        if(quantizer_conf->mode != 1) {
+            KNOWHERE_THROW_MSG("mode only support 1 in this func");
+        }
     }
     gpu_id_ = quantizer_conf->gpu_id;
 
@@ -133,13 +136,14 @@ IVFSQHybrid::LoadQuantizer(const Config &conf) {
         auto index_composition = new faiss::IndexComposition;
         index_composition->index = index_.get();
         index_composition->quantizer = nullptr;
-        index_composition->mode = quantizer_conf->mode; // 1 or 2
+        index_composition->mode = quantizer_conf->mode; // only 1
 
         auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index_composition, &option);
         delete gpu_index;
 
-        std::shared_ptr<FaissIVFQuantizer> q;
-        q->quantizer = index_composition;
+        auto q = std::make_shared<FaissIVFQuantizer>();
+        q->quantizer = index_composition->quantizer;
+        res_ = res;
         return q;
     } else {
         KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
@@ -153,15 +157,13 @@ IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) {
         KNOWHERE_THROW_MSG("Quantizer type error");
     }
 
-    if (ivf_quantizer->quantizer->mode == 2) gpu_mode = true; // all in gpu
-
     faiss::IndexIVF *ivf_index =
         dynamic_cast<faiss::IndexIVF *>(index_.get());
 
     faiss::gpu::GpuIndexFlat *is_gpu_flat_index = dynamic_cast<faiss::gpu::GpuIndexFlat *>(ivf_index->quantizer);
     if (is_gpu_flat_index == nullptr) {
         delete ivf_index->quantizer;
-        ivf_index->quantizer = ivf_quantizer->quantizer->quantizer;
+        ivf_index->quantizer = ivf_quantizer->quantizer;
     }
 }
 
@@ -175,5 +177,38 @@ IVFSQHybrid::UnsetQuantizer() {
     ivf_index->quantizer = nullptr;
 }
 
+void
+IVFSQHybrid::LoadData(const knowhere::QuantizerPtr &q, const Config &conf) {
+    auto quantizer_conf = std::dynamic_pointer_cast<QuantizerCfg>(conf);
+    if (quantizer_conf != nullptr) {
+        if(quantizer_conf->mode != 2) {
+            KNOWHERE_THROW_MSG("mode only support 2 in this func");
+        }
+    }
+    if (quantizer_conf->gpu_id != gpu_id_) {
+        KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card");
+    }
+
+    if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) {
+        ResScope rs(res, gpu_id_, false);
+        faiss::gpu::GpuClonerOptions option;
+        option.allInGpu = true;
+
+        auto ivf_quantizer = std::dynamic_pointer_cast<FaissIVFQuantizer>(q);
+        if (ivf_quantizer == nullptr) KNOWHERE_THROW_MSG("quantizer type not faissivfquantizer");
+
+        auto index_composition = new faiss::IndexComposition;
+        index_composition->index = index_.get();
+        index_composition->quantizer = ivf_quantizer->quantizer;
+        index_composition->mode = quantizer_conf->mode; // only 2
+
+        auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index_composition, &option);
+        index_.reset(gpu_index);
+        gpu_mode = true; // all in gpu
+    } else {
+        KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
+    }
+}
+
 }
 }

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

@@ -29,15 +29,20 @@ namespace zilliz {
 namespace knowhere {
 
 struct FaissIVFQuantizer : public Quantizer {
-    faiss::IndexComposition *quantizer = nullptr;
+    faiss::gpu::GpuIndexFlat *quantizer = nullptr;
 };
 using FaissIVFQuantizerPtr = std::shared_ptr<FaissIVFQuantizer>;
 
 class IVFSQHybrid : public GPUIVFSQ {
  public:
-    explicit IVFSQHybrid(const int &device_id) : GPUIVFSQ(device_id) {}
+    explicit IVFSQHybrid(const int &device_id) : GPUIVFSQ(device_id) {
+        gpu_mode = false;
+    }
 
-    explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {gpu_mode = false;}
+    explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index) : GPUIVFSQ(-1) {
+        index_ = index;
+        gpu_mode = false;
+    }
 
     explicit IVFSQHybrid(std::shared_ptr<faiss::Index> index, const int64_t &device_id, ResPtr &resource)
         : GPUIVFSQ(index, device_id, resource) {
@@ -54,6 +59,9 @@ class IVFSQHybrid : public GPUIVFSQ {
     void
     UnsetQuantizer();
 
+    void
+    LoadData(const knowhere::QuantizerPtr &q, const Config& conf);
+
     IndexModelPtr
     Train(const DatasetPtr &dataset, const Config &config) override;
 

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

@@ -100,6 +100,11 @@ public:
         Lock();
     }
 
+    ResScope(ResWPtr &res, const int64_t& device_id, const bool& isown)
+        : resource(res), device_id(device_id), move(true), own(isown) {
+        Lock();
+    }
+
     // specif for search
     // get the ownership of gpuresource and gpu
     ResScope(ResWPtr &res, const int64_t &device_id)

cpp/src/core/test/CMakeLists.txt

@@ -41,6 +41,7 @@ set(ivf_srcs
         ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp
         ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
         ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
+        ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
         ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
         ${CORE_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp
         )

cpp/src/core/test/test_ivf.cpp

@@ -65,6 +65,8 @@ IVFIndexPtr IndexFactory(const std::string &type) {
         return std::make_shared<IVFSQ>();
     } else if (type == "GPUIVFSQ") {
         return std::make_shared<GPUIVFSQ>(device_id);
+    } else if (type == "IVFSQHybrid") {
+        return std::make_shared<IVFSQHybrid>(device_id);
     }
 }
 
@@ -72,6 +74,7 @@ enum class ParameterType {
     ivf,
     ivfpq,
     ivfsq,
+    ivfsqhybrid,
     nsg,
 };
 
@@ -105,7 +108,7 @@ class ParamGenerator {
             tempconf->metric_type = METRICTYPE::L2;
             return tempconf;
         }
-        else if (type == ParameterType::ivfsq) {
+        else if (type == ParameterType::ivfsq || type == ParameterType::ivfsqhybrid) {
             auto tempconf = std::make_shared<IVFSQCfg>();
             tempconf->d = DIM;
             tempconf->gpu_id = device_id;
@@ -160,6 +163,7 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
 //                            std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
                             std::make_tuple("IVFSQ", ParameterType::ivfsq),
                             std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)
+                            std::make_tuple("IVFSQHybrid", ParameterType::ivfsqhybrid)
                         )
 );
 
@@ -210,36 +214,60 @@ 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);
+TEST_P(IVFTest, hybrid) {
+    if (index_type != "IVFSQHybrid") {
+        return;
+    }
+    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);
-//    //PrintResult(result, nq, k);
-//}
+
+    {
+        auto hybrid_1_idx = std::make_shared<IVFSQHybrid>(device_id);
+
+        auto binaryset = index_->Serialize();
+        hybrid_1_idx->Load(binaryset);
+
+        auto quantizer_conf = std::make_shared<QuantizerCfg>();
+        quantizer_conf->mode = 1;
+        quantizer_conf->gpu_id = device_id;
+        auto q = hybrid_1_idx->LoadQuantizer(quantizer_conf);
+        hybrid_1_idx->SetQuantizer(q);
+        auto result = hybrid_1_idx->Search(query_dataset, conf);
+        AssertAnns(result, nq, conf->k);
+        PrintResult(result, nq, k);
+    }
+
+    {
+        auto hybrid_2_idx = std::make_shared<IVFSQHybrid>(device_id);
+
+        auto binaryset = index_->Serialize();
+        hybrid_2_idx->Load(binaryset);
+
+        auto quantizer_conf = std::make_shared<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) {
 //    if (index_type.find("GPU") == std::string::npos) { return; }
@@ -350,29 +378,35 @@ TEST_P(IVFTest, clone_test) {
         }
     };
 
+//    {
+//        // clone in place
+//        std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFPQ", "IVFSQ", "GPUIVFSQ"};
+//        auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
+//        if (finder != support_idx_vec.cend()) {
+//            EXPECT_NO_THROW({
+//                                auto clone_index = index_->Clone();
+//                                auto clone_result = clone_index->Search(query_dataset, conf);
+//                                //AssertAnns(result, nq, conf->k);
+//                                AssertEqual(result, clone_result);
+//                                std::cout << "inplace clone [" << index_type << "] success" << std::endl;
+//                            });
+//        } else {
+//            EXPECT_THROW({
+//                             std::cout << "inplace clone [" << index_type << "] failed" << std::endl;
+//                             auto clone_index = index_->Clone();
+//                         }, KnowhereException);
+//        }
+//    }
+
     {
-        // clone in place
-        std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFPQ", "IVFSQ", "GPUIVFSQ"};
-        auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
-        if (finder != support_idx_vec.cend()) {
-            EXPECT_NO_THROW({
-                                auto clone_index = index_->Clone();
-                                auto clone_result = clone_index->Search(query_dataset, conf);
-                                //AssertAnns(result, nq, conf->k);
-                                AssertEqual(result, clone_result);
-                                std::cout << "inplace clone [" << index_type << "] success" << std::endl;
-                            });
-        } else {
-            EXPECT_THROW({
-                             std::cout << "inplace clone [" << index_type << "] failed" << std::endl;
-                             auto clone_index = index_->Clone();
-                         }, KnowhereException);
+        if (index_type == "IVFSQHybrid") {
+            return;
         }
     }
 
     {
         // copy from gpu to cpu
-        std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ"};
+        std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
         auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
         if (finder != support_idx_vec.cend()) {
             EXPECT_NO_THROW({
@@ -412,7 +446,7 @@ TEST_P(IVFTest, clone_test) {
 TEST_P(IVFTest, seal_test) {
     //FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
 
-    std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ"};
+    std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
     auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
     if (finder == support_idx_vec.cend()) {
         return;

cpp/src/wrapper/VecImpl.cpp

@@ -304,6 +304,7 @@ IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
         WRAPPER_LOG_ERROR << e.what();
         return Status(KNOWHERE_ERROR, e.what());
     }
+    return Status::OK();
 }
 
 Status
@@ -323,6 +324,26 @@ IVFHybridIndex::UnsetQuantizer() {
         WRAPPER_LOG_ERROR << e.what();
         return Status(KNOWHERE_ERROR, e.what());
     }
+    return Status::OK();
+}
+
+Status IVFHybridIndex::LoadData(const knowhere::QuantizerPtr &q, const Config &conf) {
+    try {
+        // TODO(linxj): Hardcode here
+        if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
+            new_idx->LoadData(q, conf);
+        } else {
+            WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
+            return Status(KNOWHERE_ERROR, "not support");
+        }
+    } catch (knowhere::KnowhereException &e) {
+        WRAPPER_LOG_ERROR << e.what();
+        return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
+    } catch (std::exception &e) {
+        WRAPPER_LOG_ERROR << e.what();
+        return Status(KNOWHERE_ERROR, e.what());
+    }
+    return Status::OK();
 }
 
 } // namespace engine

cpp/src/wrapper/VecImpl.h

@@ -111,6 +111,9 @@ class IVFHybridIndex : public IVFMixIndex {
 
     Status
     UnsetQuantizer() override;
+
+    Status
+    LoadData(const knowhere::QuantizerPtr &q, const Config &conf) override;
 };
 
 class BFIndex : public VecIndexImpl {

cpp/src/wrapper/VecIndex.h

@@ -104,15 +104,19 @@ class VecIndex {
     Load(const zilliz::knowhere::BinarySet &index_binary) = 0;
 
     // TODO(linxj): refactor later
+    ////////////////
     virtual knowhere::QuantizerPtr
-    LoadQuantizer(const Config& conf) { return Status::OK(); }
+    LoadQuantizer(const Config& conf) { return nullptr; }
+
+    virtual Status
+    LoadData(const knowhere::QuantizerPtr &q, const Config &conf) { return Status::OK(); }
 
-    // TODO(linxj): refactor later
     virtual Status
     SetQuantizer(const knowhere::QuantizerPtr& q) { return Status::OK(); }
 
     virtual Status
     UnsetQuantizer() { return Status::OK(); }
+    ////////////////
 };
 
 extern Status