add hybrid mode

Former-commit-id: feae4c8df728c1ae8f96e2d19c0b115a5a420575

cpp/src/core/knowhere/CMakeLists.txt

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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; }
 //

cpp/src/wrapper/VecImpl.cpp

@@ -19,6 +19,7 @@
 #include "wrapper/VecImpl.h"
 #include "utils/Log.h"
 #include "knowhere/index/vector_index/IndexIDMAP.h"
+#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
 #include "knowhere/index/vector_index/IndexGPUIVF.h"
 #include "knowhere/common/Exception.h"
 #include "knowhere/index/vector_index/helpers/Cloner.h"
@@ -251,7 +252,7 @@ IVFMixIndex::BuildAll(const int64_t &nb,
         index_->set_index_model(model);
         index_->Add(dataset, cfg);
 
-        if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)) {
+        if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
             auto host_index = device_index->CopyGpuToCpu(Config());
             index_ = host_index;
             type = ConvertToCpuIndexType(type);
@@ -276,6 +277,33 @@ IVFMixIndex::Load(const zilliz::knowhere::BinarySet &index_binary) {
     return Status::OK();
 }
 
+knowhere::QuantizerPtr IVFHybridIndex::LoadQuantizer(const Config& conf) {
+    // TODO(linxj): Hardcode here
+    if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)){
+        return new_idx->LoadQuantizer(conf);
+    } else {
+        WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
+    }
+}
+
+Status IVFHybridIndex::SetQuantizer(knowhere::QuantizerPtr q) {
+    try {
+        // TODO(linxj): Hardcode here
+        if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
+            new_idx->SetQuantizer(q);
+        } 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());
+    }
+}
+
 } // namespace engine
 } // namespace milvus
 } // namespace zilliz

cpp/src/wrapper/VecImpl.h

@@ -101,6 +101,13 @@ class IVFMixIndex : public VecIndexImpl {
     Load(const zilliz::knowhere::BinarySet &index_binary) override;
 };
 
+class IVFHybridIndex : public IVFMixIndex {
+ public:
+    knowhere::QuantizerPtr LoadQuantizer(const Config& conf) override;
+
+    Status SetQuantizer(knowhere::QuantizerPtr q) override;
+};
+
 class BFIndex : public VecIndexImpl {
  public:
     explicit BFIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index) : VecIndexImpl(std::move(index),

cpp/src/wrapper/VecIndex.cpp

@@ -25,6 +25,7 @@
 #include "knowhere/index/vector_index/IndexIDMAP.h"
 #include "knowhere/index/vector_index/IndexKDT.h"
 #include "knowhere/index/vector_index/IndexNSG.h"
+#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
 #include "knowhere/common/Exception.h"
 #include "VecImpl.h"
 #include "utils/Log.h"
@@ -137,6 +138,10 @@ GetVecIndexFactory(const IndexType &type, const Config &cfg) {
             index = std::make_shared<zilliz::knowhere::GPUIVFSQ>(gpu_device);
             break;
         }
+        case IndexType::FAISS_IVFSQ8_HYBRID: {
+            index = std::make_shared<zilliz::knowhere::IVFSQHybrid>(gpu_device);
+            break;
+        }
         case IndexType::NSG_MIX: { // TODO(linxj): bug.
             index = std::make_shared<zilliz::knowhere::NSG>(gpu_device);
             break;

cpp/src/wrapper/VecIndex.h

@@ -24,6 +24,7 @@
 #include "utils/Status.h"
 #include "knowhere/common/Config.h"
 #include "knowhere/common/BinarySet.h"
+#include "knowhere/index/vector_index/Quantizer.h"
 
 namespace zilliz {
 namespace milvus {
@@ -43,6 +44,7 @@ enum class IndexType {
     FAISS_IVFSQ8_MIX,
     FAISS_IVFSQ8_CPU,
     FAISS_IVFSQ8_GPU,
+    FAISS_IVFSQ8_HYBRID, // only support build on gpu.
     NSG_MIX,
 };
 
@@ -100,6 +102,14 @@ class VecIndex {
 
     virtual Status
     Load(const zilliz::knowhere::BinarySet &index_binary) = 0;
+
+    // TODO(linxj): refactor later
+    virtual knowhere::QuantizerPtr
+    LoadQuantizer(const Config& conf) {}
+
+    // TODO(linxj): refactor later
+    virtual Status
+    SetQuantizer(knowhere::QuantizerPtr q) {}
 };
 
 extern Status