MS-565 split class version 1

Former-commit-id: 1a60045de13e055d4b0a7e9da389b7e377c141af

cpp/src/core/knowhere/CMakeLists.txt

@@ -42,13 +42,19 @@ set(index_srcs
         knowhere/index/vector_index/IndexIDMAP.cpp
         knowhere/index/vector_index/IndexIVF.cpp
         knowhere/index/vector_index/IndexGPUIVF.cpp
-        knowhere/index/vector_index/utils/KDTParameterMgr.cpp
+        knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
         knowhere/index/vector_index/IndexNSG.cpp
         knowhere/index/vector_index/nsg/NSG.cpp
         knowhere/index/vector_index/nsg/NSGIO.cpp
         knowhere/index/vector_index/nsg/NSGHelper.cpp
-        knowhere/index/vector_index/utils/Cloner.cpp
-        knowhere/index/vector_index/utils/FaissGpuResourceMgr.cpp
+        knowhere/index/vector_index/helpers/Cloner.cpp
+        knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
+        knowhere/index/vector_index/IndexIVFSQ.cpp
+        knowhere/index/vector_index/IndexGPUIVFSQ.cpp
+        knowhere/index/vector_index/IndexIVFPQ.cpp
+        knowhere/index/vector_index/IndexGPUIVFPQ.cpp
+        knowhere/index/vector_index/FaissBaseIndex.cpp
+        knowhere/index/vector_index/helpers/FaissIO.cpp
         )
 
 set(depend_libs

cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "knowhere/index/vector_index/utils/Definitions.h"
+#include "knowhere/index/vector_index/helpers/Definitions.h"
 #include "SptagAdapter.h"
 #include "Structure.h"
 

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

+// 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 <faiss/index_io.h>
+#include <faiss/IndexIVF.h>
+
+#include "knowhere/common/Exception.h"
+#include "knowhere/index/vector_index/helpers/FaissIO.h"
+#include "FaissBaseIndex.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+FaissBaseIndex::FaissBaseIndex(std::shared_ptr<faiss::Index> index) : index_(std::move(index)) {}
+
+BinarySet FaissBaseIndex::SerializeImpl() {
+    try {
+        faiss::Index *index = index_.get();
+
+        SealImpl();
+
+        MemoryIOWriter writer;
+        faiss::write_index(index, &writer);
+        auto data = std::make_shared<uint8_t>();
+        data.reset(writer.data_);
+
+        BinarySet res_set;
+        // TODO(linxj): use virtual func Name() instead of raw string.
+        res_set.Append("IVF", data, writer.rp);
+        return res_set;
+    } catch (std::exception &e) {
+        KNOWHERE_THROW_MSG(e.what());
+    }
+}
+
+void FaissBaseIndex::LoadImpl(const BinarySet &index_binary) {
+    auto binary = index_binary.GetByName("IVF");
+
+    MemoryIOReader reader;
+    reader.total = binary->size;
+    reader.data_ = binary->data.get();
+
+    faiss::Index *index = faiss::read_index(&reader);
+
+    index_.reset(index);
+}
+
+void FaissBaseIndex::SealImpl() {
+// TODO(linxj): enable
+//#ifdef ZILLIZ_FAISS
+    faiss::Index *index = index_.get();
+    auto idx = dynamic_cast<faiss::IndexIVF *>(index);
+    if (idx != nullptr) {
+        idx->to_readonly();
+    }
+    //else {
+    //    KNOHWERE_ERROR_MSG("Seal failed");
+    //}
+//#endif
+}
+
+} // knowhere
+} // zilliz
+

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

+// 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.
+
+
+#pragma once
+
+#include <memory>
+
+#include <faiss/Index.h>
+
+#include "knowhere/common/BinarySet.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+class FaissBaseIndex {
+ protected:
+    explicit FaissBaseIndex(std::shared_ptr<faiss::Index> index);
+
+    virtual BinarySet
+    SerializeImpl();
+
+    virtual void
+    LoadImpl(const BinarySet &index_binary);
+
+    virtual void
+    SealImpl();
+
+ protected:
+    std::shared_ptr<faiss::Index> index_ = nullptr;
+};
+
+} // knowhere
+} // zilliz
+
+
+

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

@@ -20,18 +20,17 @@
 #include <faiss/gpu/GpuIndexFlat.h>
 #include <faiss/gpu/GpuIndexIVF.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
-#include <faiss/gpu/GpuIndexIVFPQ.h>
 #include <faiss/gpu/GpuAutoTune.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/index_io.h>
 
 
 #include "knowhere/common/Exception.h"
-#include "knowhere/index/vector_index/utils/Cloner.h"
+#include "knowhere/index/vector_index/helpers/Cloner.h"
 #include "knowhere/adapter/VectorAdapter.h"
 #include "IndexGPUIVF.h"
+#include "knowhere/index/vector_index/helpers/FaissIO.h"
 
-#include <algorithm>
 
 namespace zilliz {
 namespace knowhere {
@@ -187,84 +186,6 @@ void GPUIVF::Add(const DatasetPtr &dataset, const Config &config) {
     }
 }
 
-IndexModelPtr GPUIVFPQ::Train(const DatasetPtr &dataset, const Config &config) {
-    auto nlist = config["nlist"].as<size_t>();
-    auto M = config["M"].as<size_t>();        // number of subquantizers(subvectors)
-    auto nbits = config["nbits"].as<size_t>();// number of bit per subvector index
-    auto gpu_num = config.get_with_default("gpu_id", gpu_id_);
-    auto metric_type = config["metric_type"].as_string() == "L2" ?
-                       faiss::METRIC_L2 : faiss::METRIC_L2; // IP not support.
-
-    GETTENSOR(dataset)
-
-    // TODO(linxj): set device here.
-    // TODO(linxj): set gpu resource here.
-    faiss::gpu::StandardGpuResources res;
-    faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, nlist, M, nbits, metric_type);
-    device_index.train(rows, (float *) p_data);
-
-    std::shared_ptr<faiss::Index> host_index = nullptr;
-    host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index));
-
-    return std::make_shared<IVFIndexModel>(host_index);
-}
-
-std::shared_ptr<faiss::IVFSearchParameters> GPUIVFPQ::GenParams(const Config &config) {
-    auto params = std::make_shared<faiss::IVFPQSearchParameters>();
-    params->nprobe = config.get_with_default("nprobe", size_t(1));
-    //params->scan_table_threshold = 0;
-    //params->polysemous_ht = 0;
-    //params->max_codes = 0;
-
-    return params;
-}
-
-VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config &config) {
-    KNOWHERE_THROW_MSG("not support yet");
-}
-
-IndexModelPtr GPUIVFSQ::Train(const DatasetPtr &dataset, const Config &config) {
-    auto nlist = config["nlist"].as<size_t>();
-    auto nbits = config["nbits"].as<size_t>(); // TODO(linxj):  gpu only support SQ4 SQ8 SQ16
-    gpu_id_ = config.get_with_default("gpu_id", gpu_id_);
-    auto metric_type = config["metric_type"].as_string() == "L2" ?
-                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
-
-    GETTENSOR(dataset)
-
-    std::stringstream index_type;
-    index_type << "IVF" << nlist << "," << "SQ" << nbits;
-    auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type);
-
-    auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
-    if (temp_resource != nullptr) {
-        ResScope rs(temp_resource, gpu_id_, true);
-        auto device_index = faiss::gpu::index_cpu_to_gpu(temp_resource->faiss_res.get(), gpu_id_, build_index);
-        device_index->train(rows, (float *) p_data);
-
-        std::shared_ptr<faiss::Index> host_index = nullptr;
-        host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));
-
-        delete device_index;
-        delete build_index;
-
-        return std::make_shared<IVFIndexModel>(host_index);
-    } else {
-        KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");
-    }
-}
-
-VectorIndexPtr GPUIVFSQ::CopyGpuToCpu(const Config &config) {
-    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<IVFSQ>(new_index);
-}
-
 void GPUIndex::SetGpuDevice(const int &gpu_id) {
     gpu_id_ = gpu_id;
 }

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

@@ -20,79 +20,79 @@
 
 
 #include "IndexIVF.h"
-#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h"
+#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
 
 
 namespace zilliz {
 namespace knowhere {
 
 class GPUIndex {
- public:
-    explicit GPUIndex(const int &device_id) : gpu_id_(device_id) {}
-    GPUIndex(const int& device_id, const ResPtr& resource): gpu_id_(device_id), res_(resource){}
+public:
+     explicit GPUIndex(const int &device_id) : gpu_id_(device_id) {}
 
-    virtual VectorIndexPtr CopyGpuToCpu(const Config &config) = 0;
-    virtual VectorIndexPtr CopyGpuToGpu(const int64_t &device_id, const Config &config) = 0;
+     GPUIndex(const int& device_id, const ResPtr& resource): gpu_id_(device_id), res_(resource) {}
 
-    void SetGpuDevice(const int &gpu_id);
-    const int64_t &GetGpuDevice();
+     virtual VectorIndexPtr
+     CopyGpuToCpu(const Config &config) = 0;
 
- protected:
-    int64_t gpu_id_;
-    ResWPtr res_;
+     virtual VectorIndexPtr
+     CopyGpuToGpu(const int64_t &device_id, const Config &config) = 0;
+
+     void
+     SetGpuDevice(const int &gpu_id);
+
+     const int64_t &
+     GetGpuDevice();
+
+protected:
+     int64_t gpu_id_;
+     ResWPtr res_;
 };
 
 class GPUIVF : public IVF, public GPUIndex {
- public:
-    explicit GPUIVF(const int &device_id) : IVF(), GPUIndex(device_id) {}
-    explicit GPUIVF(std::shared_ptr<faiss::Index> index, const int64_t &device_id, ResPtr &resource)
-        : IVF(std::move(index)), GPUIndex(device_id, resource) {};
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
-    void Add(const DatasetPtr &dataset, const Config &config) override;
-    void set_index_model(IndexModelPtr model) override;
-    //DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override;
-    VectorIndexPtr CopyGpuToCpu(const Config &config) override;
-    VectorIndexPtr CopyGpuToGpu(const int64_t &device_id, const Config &config) override;
-    VectorIndexPtr Clone() final;
-
-    // TODO(linxj): Deprecated
-    virtual IVFIndexPtr Copy_index_gpu_to_cpu();
-
- protected:
-    void search_impl(int64_t n,
-                     const float *data,
-                     int64_t k,
-                     float *distances,
-                     int64_t *labels,
-                     const Config &cfg) override;
-    BinarySet SerializeImpl() override;
-    void LoadImpl(const BinarySet &index_binary) override;
-};
+public:
+     explicit GPUIVF(const int &device_id) : IVF(), GPUIndex(device_id) {}
 
-class GPUIVFSQ : public GPUIVF {
- public:
-    explicit GPUIVFSQ(const int &device_id) : GPUIVF(device_id) {}
-    explicit GPUIVFSQ(std::shared_ptr<faiss::Index> index, const int64_t &device_id, ResPtr &resource)
-        : GPUIVF(std::move(index), device_id, resource) {};
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
+     explicit GPUIVF(std::shared_ptr<faiss::Index> index, const int64_t &device_id, ResPtr &resource)
+          : IVF(std::move(index)), GPUIndex(device_id, resource) {};
 
- public:
-    VectorIndexPtr CopyGpuToCpu(const Config &config) override;
-};
+     IndexModelPtr
+     Train(const DatasetPtr &dataset, const Config &config) override;
 
-class GPUIVFPQ : public GPUIVF {
- public:
-    explicit GPUIVFPQ(const int &device_id) : GPUIVF(device_id) {}
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
+     void
+     Add(const DatasetPtr &dataset, const Config &config) override;
 
- public:
-    VectorIndexPtr CopyGpuToCpu(const Config &config) override;
+     void
+     set_index_model(IndexModelPtr model) override;
 
- protected:
-    // TODO(linxj): remove GenParams.
-    std::shared_ptr<faiss::IVFSearchParameters> GenParams(const Config &config) override;
-};
+     //DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override;
+     VectorIndexPtr
+     CopyGpuToCpu(const Config &config) override;
 
+     VectorIndexPtr
+     CopyGpuToGpu(const int64_t &device_id, const Config &config) override;
+
+     VectorIndexPtr
+     Clone() final;
+
+     // TODO(linxj): Deprecated
+     virtual IVFIndexPtr Copy_index_gpu_to_cpu();
+
+protected:
+     void
+     search_impl(int64_t n,
+                 const float *data,
+                 int64_t k,
+                 float *distances,
+                 int64_t *labels,
+                 const Config &cfg) override;
+
+     BinarySet
+     SerializeImpl() override;
+
+     void
+     LoadImpl(const BinarySet &index_binary) override;
+};
 
 } // namespace knowhere
 } // namespace zilliz

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

+// 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 <faiss/gpu/GpuIndexIVFPQ.h>
+#include <faiss/gpu/GpuAutoTune.h>
+#include <faiss/IndexIVFPQ.h>
+
+#include "IndexGPUIVFPQ.h"
+#include "knowhere/common/Exception.h"
+#include "knowhere/adapter/VectorAdapter.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+IndexModelPtr GPUIVFPQ::Train(const DatasetPtr &dataset, const Config &config) {
+    auto nlist = config["nlist"].as<size_t>();
+    auto M = config["M"].as<size_t>();        // number of subquantizers(subvectors)
+    auto nbits = config["nbits"].as<size_t>();// number of bit per subvector index
+    auto gpu_num = config.get_with_default("gpu_id", gpu_id_);
+    auto metric_type = config["metric_type"].as_string() == "L2" ?
+                       faiss::METRIC_L2 : faiss::METRIC_L2; // IP not support.
+
+    GETTENSOR(dataset)
+
+    // TODO(linxj): set device here.
+    // TODO(linxj): set gpu resource here.
+    faiss::gpu::StandardGpuResources res;
+    faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, nlist, M, nbits, metric_type);
+    device_index.train(rows, (float *) p_data);
+
+    std::shared_ptr<faiss::Index> host_index = nullptr;
+    host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index));
+
+    return std::make_shared<IVFIndexModel>(host_index);
+}
+
+std::shared_ptr<faiss::IVFSearchParameters> GPUIVFPQ::GenParams(const Config &config) {
+    auto params = std::make_shared<faiss::IVFPQSearchParameters>();
+    params->nprobe = config.get_with_default("nprobe", size_t(1));
+    //params->scan_table_threshold = 0;
+    //params->polysemous_ht = 0;
+    //params->max_codes = 0;
+
+    return params;
+}
+
+VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config &config) {
+    KNOWHERE_THROW_MSG("not support yet");
+}
+
+} // knowhere
+} // zilliz

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

+// 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.
+
+
+#pragma once
+
+#include "IndexGPUIVF.h"
+
+namespace zilliz {
+namespace knowhere {
+
+class GPUIVFPQ : public GPUIVF {
+public:
+    explicit GPUIVFPQ(const int &device_id) : GPUIVF(device_id) {}
+
+    IndexModelPtr
+    Train(const DatasetPtr &dataset, const Config &config) override;
+
+public:
+    VectorIndexPtr
+    CopyGpuToCpu(const Config &config) override;
+
+protected:
+    // TODO(linxj): remove GenParams.
+    std::shared_ptr<faiss::IVFSearchParameters>
+    GenParams(const Config &config) override;
+};
+
+} // knowhere
+} // zilliz
+
+
+

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

+// 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 <faiss/gpu/GpuAutoTune.h>
+
+#include "knowhere/adapter/VectorAdapter.h"
+#include "knowhere/common/Exception.h"
+#include "IndexGPUIVFSQ.h"
+#include "IndexIVFSQ.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+    IndexModelPtr GPUIVFSQ::Train(const DatasetPtr &dataset, const Config &config) {
+        auto nlist = config["nlist"].as<size_t>();
+        auto nbits = config["nbits"].as<size_t>(); // TODO(linxj):  gpu only support SQ4 SQ8 SQ16
+        gpu_id_ = config.get_with_default("gpu_id", gpu_id_);
+        auto metric_type = config["metric_type"].as_string() == "L2" ?
+                           faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
+
+        GETTENSOR(dataset)
+
+        std::stringstream index_type;
+        index_type << "IVF" << nlist << "," << "SQ" << nbits;
+        auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type);
+
+        auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
+        if (temp_resource != nullptr) {
+            ResScope rs(temp_resource, gpu_id_, true);
+            auto device_index = faiss::gpu::index_cpu_to_gpu(temp_resource->faiss_res.get(), gpu_id_, build_index);
+            device_index->train(rows, (float *) p_data);
+
+            std::shared_ptr<faiss::Index> host_index = nullptr;
+            host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));
+
+            delete device_index;
+            delete build_index;
+
+            return std::make_shared<IVFIndexModel>(host_index);
+        } else {
+            KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");
+        }
+    }
+
+    VectorIndexPtr GPUIVFSQ::CopyGpuToCpu(const Config &config) {
+        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<IVFSQ>(new_index);
+    }
+
+} // knowhere
+} // zilliz
+

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

+// 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.
+
+
+#pragma once
+
+#include "IndexGPUIVF.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+class GPUIVFSQ : public GPUIVF {
+public:
+    explicit GPUIVFSQ(const int &device_id) : GPUIVF(device_id) {}
+
+    explicit GPUIVFSQ(std::shared_ptr<faiss::Index> index, const int64_t &device_id, ResPtr &resource)
+            : GPUIVF(std::move(index), device_id, resource) {};
+
+    IndexModelPtr
+    Train(const DatasetPtr &dataset, const Config &config) override;
+
+    VectorIndexPtr
+    CopyGpuToCpu(const Config &config) override;
+};
+
+} // knowhere
+} // zilliz
+

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

@@ -25,6 +25,7 @@
 
 #include "knowhere/common/Exception.h"
 #include "knowhere/adapter/VectorAdapter.h"
+#include "knowhere/index/vector_index/helpers/FaissIO.h"
 #include "IndexIDMAP.h"
 
 

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

@@ -25,10 +25,10 @@
 namespace zilliz {
 namespace knowhere {
 
-class IDMAP : public VectorIndex, public BasicIndex {
+class IDMAP : public VectorIndex, public FaissBaseIndex {
  public:
-    IDMAP() : BasicIndex(nullptr) {};
-    explicit IDMAP(std::shared_ptr<faiss::Index> index) : BasicIndex(std::move(index)) {};
+    IDMAP() : FaissBaseIndex(nullptr) {};
+    explicit IDMAP(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {};
     BinarySet Serialize() override;
     void Load(const BinarySet &index_binary) override;
     void Train(const Config &config);

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

@@ -238,7 +238,7 @@ void IVF::Seal() {
 }
 
 
-IVFIndexModel::IVFIndexModel(std::shared_ptr<faiss::Index> index) : BasicIndex(std::move(index)) {}
+IVFIndexModel::IVFIndexModel(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {}
 
 BinarySet IVFIndexModel::Serialize() {
     if (!index_ || !index_->is_trained) {
@@ -257,159 +257,8 @@ void IVFIndexModel::SealImpl() {
     // do nothing
 }
 
-IndexModelPtr IVFSQ::Train(const DatasetPtr &dataset, const Config &config) {
-    auto nlist = config["nlist"].as<size_t>();
-    auto nbits = config["nbits"].as<size_t>(); // TODO(linxj): only support SQ4 SQ6 SQ8 SQ16
-    auto metric_type = config["metric_type"].as_string() == "L2" ?
-                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
-
-    GETTENSOR(dataset)
-
-    std::stringstream index_type;
-    index_type << "IVF" << nlist << "," << "SQ" << nbits;
-    auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type);
-    build_index->train(rows, (float *) p_data);
-
-    std::shared_ptr<faiss::Index> ret_index;
-    ret_index.reset(build_index);
-    return std::make_shared<IVFIndexModel>(ret_index);
-}
-
-VectorIndexPtr IVFSQ::Clone_impl(const std::shared_ptr<faiss::Index> &index) {
-    return std::make_shared<IVFSQ>(index);
-}
-
-VectorIndexPtr IVFSQ::CopyCpuToGpu(const int64_t &device_id, const Config &config) {
-    if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)){
-        ResScope rs(res, device_id, false);
-        faiss::gpu::GpuClonerOptions option;
-        option.allInGpu = true;
-
-        auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get(), &option);
-
-        std::shared_ptr<faiss::Index> device_index;
-        device_index.reset(gpu_index);
-        return std::make_shared<GPUIVFSQ>(device_index, device_id, res);
-    } else {
-        KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
-    }
-}
-
-IndexModelPtr IVFPQ::Train(const DatasetPtr &dataset, const Config &config) {
-    auto nlist = config["nlist"].as<size_t>();
-    auto M = config["M"].as<size_t>();        // number of subquantizers(subvector)
-    auto nbits = config["nbits"].as<size_t>();// number of bit per subvector index
-    auto metric_type = config["metric_type"].as_string() == "L2" ?
-                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
-
-    GETTENSOR(dataset)
-
-    faiss::Index *coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
-    auto index = std::make_shared<faiss::IndexIVFPQ>(coarse_quantizer, dim, nlist, M, nbits);
-    index->train(rows, (float *) p_data);
-
-    return std::make_shared<IVFIndexModel>(index);
-}
 
-std::shared_ptr<faiss::IVFSearchParameters> IVFPQ::GenParams(const Config &config) {
-    auto params = std::make_shared<faiss::IVFPQSearchParameters>();
-    params->nprobe = config.get_with_default("nprobe", size_t(1));
-    //params->scan_table_threshold = 0;
-    //params->polysemous_ht = 0;
-    //params->max_codes = 0;
-
-    return params;
-}
-
-VectorIndexPtr IVFPQ::Clone_impl(const std::shared_ptr<faiss::Index> &index) {
-    return std::make_shared<IVFPQ>(index);
-}
-
-BasicIndex::BasicIndex(std::shared_ptr<faiss::Index> index) : index_(std::move(index)) {}
-
-BinarySet BasicIndex::SerializeImpl() {
-    try {
-        faiss::Index *index = index_.get();
-
-        SealImpl();
-
-        MemoryIOWriter writer;
-        faiss::write_index(index, &writer);
-        auto data = std::make_shared<uint8_t>();
-        data.reset(writer.data_);
-
-        BinarySet res_set;
-        // TODO(linxj): use virtual func Name() instead of raw string.
-        res_set.Append("IVF", data, writer.rp);
-        return res_set;
-    } catch (std::exception &e) {
-        KNOWHERE_THROW_MSG(e.what());
-    }
-}
-
-void BasicIndex::LoadImpl(const BinarySet &index_binary) {
-    auto binary = index_binary.GetByName("IVF");
-
-    MemoryIOReader reader;
-    reader.total = binary->size;
-    reader.data_ = binary->data.get();
 
-    faiss::Index *index = faiss::read_index(&reader);
-
-    index_.reset(index);
-}
-
-void BasicIndex::SealImpl() {
-// TODO(linxj): enable
-//#ifdef ZILLIZ_FAISS
-    faiss::Index *index = index_.get();
-    auto idx = dynamic_cast<faiss::IndexIVF *>(index);
-    if (idx != nullptr) {
-        idx->to_readonly();
-    }
-    //else {
-    //    KNOHWERE_ERROR_MSG("Seal failed");
-    //}
-//#endif
-}
-
-// TODO(linxj): Get From Config File
-static size_t magic_num = 2;
-size_t MemoryIOWriter::operator()(const void *ptr, size_t size, size_t nitems) {
-    auto total_need = size * nitems + rp;
-
-    if (!data_) { // data == nullptr
-        total = total_need * magic_num;
-        rp = size * nitems;
-        data_ = new uint8_t[total];
-        memcpy((void *) (data_), ptr, rp);
-    }
-
-    if (total_need > total) {
-        total = total_need * magic_num;
-        auto new_data = new uint8_t[total];
-        memcpy((void *) new_data, (void *) data_, rp);
-        delete data_;
-        data_ = new_data;
-
-        memcpy((void *) (data_ + rp), ptr, size * nitems);
-        rp = total_need;
-    } else {
-        memcpy((void *) (data_ + rp), ptr, size * nitems);
-        rp = total_need;
-    }
-
-    return nitems;
-}
-
-size_t MemoryIOReader::operator()(void *ptr, size_t size, size_t nitems) {
-    if (rp >= total) return 0;
-    size_t nremain = (total - rp) / size;
-    if (nremain < nitems) nitems = nremain;
-    memcpy(ptr, (void *) (data_ + rp), size * nitems);
-    rp += size * nitems;
-    return nitems;
-}
 
 
 }

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

@@ -21,136 +21,106 @@
 #include <memory>
 #include <mutex>
 
-#include <faiss/IndexIVF.h>
-#include <faiss/AuxIndexStructures.h>
-#include <faiss/Index.h>
-
 #include "VectorIndex.h"
+#include "FaissBaseIndex.h"
+#include "faiss/IndexIVF.h"
 
 
 namespace zilliz {
 namespace knowhere {
 
-class BasicIndex {
- protected:
-    explicit BasicIndex(std::shared_ptr<faiss::Index> index);
-    virtual BinarySet SerializeImpl();
-    virtual void LoadImpl(const BinarySet &index_binary);
-    virtual void
-    SealImpl();
-
- protected:
-    std::shared_ptr<faiss::Index> index_ = nullptr;
-};
-
 using Graph = std::vector<std::vector<int64_t>>;
 
-class IVF : public VectorIndex, protected BasicIndex {
+class IVF : public VectorIndex, protected FaissBaseIndex {
  public:
-    IVF() : BasicIndex(nullptr) {};
-    explicit IVF(std::shared_ptr<faiss::Index> index) : BasicIndex(std::move(index)) {}
-    VectorIndexPtr Clone() override;;
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
-    void set_index_model(IndexModelPtr model) override;
-    void Add(const DatasetPtr &dataset, const Config &config) override;
-    void AddWithoutIds(const DatasetPtr &dataset, const Config &config);
-    DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override;
-    void GenGraph(const int64_t &k, Graph &graph, const DatasetPtr &dataset, const Config &config);
-    BinarySet Serialize() override;
-    void Load(const BinarySet &index_binary) override;
-    int64_t Count() override;
-    int64_t Dimension() override;
-
-    void
-    Seal() override;
-
-    virtual VectorIndexPtr CopyCpuToGpu(const int64_t &device_id, const Config &config);
+    IVF() : FaissBaseIndex(nullptr) {};
 
+    explicit IVF(std::shared_ptr<faiss::Index> index) : FaissBaseIndex(std::move(index)) {}
 
- protected:
-    virtual std::shared_ptr<faiss::IVFSearchParameters> GenParams(const Config &config);
+    VectorIndexPtr
+    Clone() override;;
 
-    virtual VectorIndexPtr Clone_impl(const std::shared_ptr<faiss::Index> &index);
+    IndexModelPtr
+    Train(const DatasetPtr &dataset, const Config &config) override;
 
-    virtual void search_impl(int64_t n,
-                             const float *data,
-                             int64_t k,
-                             float *distances,
-                             int64_t *labels,
-                             const Config &cfg);
+    void
+    set_index_model(IndexModelPtr model) override;
 
- protected:
-    std::mutex mutex_;
-};
+    void
+    Add(const DatasetPtr &dataset, const Config &config) override;
 
-using IVFIndexPtr = std::shared_ptr<IVF>;
+    void
+    AddWithoutIds(const DatasetPtr &dataset, const Config &config);
 
-class IVFSQ : public IVF {
- public:
-    explicit IVFSQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {}
-    IVFSQ() = default;
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
-    VectorIndexPtr CopyCpuToGpu(const int64_t &device_id, const Config &config) override;
- protected:
-    VectorIndexPtr Clone_impl(const std::shared_ptr<faiss::Index> &index) override;
-};
+    DatasetPtr
+    Search(const DatasetPtr &dataset, const Config &config) override;
 
-class IVFPQ : public IVF {
- public:
-    explicit IVFPQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {}
-    IVFPQ() = default;
-    IndexModelPtr Train(const DatasetPtr &dataset, const Config &config) override;
- protected:
-    std::shared_ptr<faiss::IVFSearchParameters> GenParams(const Config &config) override;
-    VectorIndexPtr Clone_impl(const std::shared_ptr<faiss::Index> &index) override;
-};
+    void
+    GenGraph(const int64_t &k, Graph &graph, const DatasetPtr &dataset, const Config &config);
 
+    BinarySet
+    Serialize() override;
 
-//class OPQIVFPQ : public IVFPQ {
-// public:
-//    PreprocessorPtr BuildPreprocessor(const Dataset &dataset, const Config &config) override;
-//};
+    void
+    Load(const BinarySet &index_binary) override;
 
+    int64_t
+    Count() override;
 
-class GPUIVF;
+    int64_t
+    Dimension() override;
 
+    void
+    Seal() override;
 
-struct MemoryIOWriter : public faiss::IOWriter {
-    uint8_t *data_ = nullptr;
-    size_t total = 0;
-    size_t rp = 0;
+    virtual VectorIndexPtr
+    CopyCpuToGpu(const int64_t &device_id, const Config &config);
 
-    size_t operator()(const void *ptr, size_t size, size_t nitems) override;
-};
 
+ protected:
+    virtual std::shared_ptr<faiss::IVFSearchParameters>
+    GenParams(const Config &config);
 
-struct MemoryIOReader : public faiss::IOReader {
-    uint8_t *data_;
-    size_t rp = 0;
-    size_t total = 0;
+    virtual VectorIndexPtr
+    Clone_impl(const std::shared_ptr<faiss::Index> &index);
 
-    size_t operator()(void *ptr, size_t size, size_t nitems) override;
+    virtual void
+    search_impl(int64_t n,
+                const float *data,
+                int64_t k,
+                float *distances,
+                int64_t *labels,
+                const Config &cfg);
 
+ protected:
+    std::mutex mutex_;
 };
 
+using IVFIndexPtr = std::shared_ptr<IVF>;
 
-class IVFIndexModel : public IndexModel, public BasicIndex {
+class GPUIVF;
+class IVFIndexModel : public IndexModel, public FaissBaseIndex {
     friend IVF;
     friend GPUIVF;
 
  public:
     explicit IVFIndexModel(std::shared_ptr<faiss::Index> index);
-    IVFIndexModel() : BasicIndex(nullptr) {};
-    BinarySet Serialize() override;
+
+    IVFIndexModel() : FaissBaseIndex(nullptr) {};
+
+    BinarySet
+    Serialize() override;
+
+    void
+    Load(const BinarySet &binary) override;
+
  protected:
-    void SealImpl() override;
- public:
-    void Load(const BinarySet &binary) override;
+    void
+    SealImpl() override;
 
  protected:
     std::mutex mutex_;
 };
-
 using IVFIndexModelPtr = std::shared_ptr<IVFIndexModel>;
 
 }

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

+// 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 <faiss/IndexFlat.h>
+#include <faiss/IndexIVFPQ.h>
+
+#include "IndexIVFPQ.h"
+#include "knowhere/common/Exception.h"
+#include "knowhere/adapter/VectorAdapter.h"
+
+namespace zilliz {
+namespace knowhere {
+
+IndexModelPtr IVFPQ::Train(const DatasetPtr &dataset, const Config &config) {
+    auto nlist = config["nlist"].as<size_t>();
+    auto M = config["M"].as<size_t>();        // number of subquantizers(subvector)
+    auto nbits = config["nbits"].as<size_t>();// number of bit per subvector index
+    auto metric_type = config["metric_type"].as_string() == "L2" ?
+                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
+
+    GETTENSOR(dataset)
+
+    faiss::Index *coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
+    auto index = std::make_shared<faiss::IndexIVFPQ>(coarse_quantizer, dim, nlist, M, nbits);
+    index->train(rows, (float *) p_data);
+
+    return std::make_shared<IVFIndexModel>(index);
+}
+
+std::shared_ptr<faiss::IVFSearchParameters> IVFPQ::GenParams(const Config &config) {
+    auto params = std::make_shared<faiss::IVFPQSearchParameters>();
+    params->nprobe = config.get_with_default("nprobe", size_t(1));
+    //params->scan_table_threshold = 0;
+    //params->polysemous_ht = 0;
+    //params->max_codes = 0;
+
+    return params;
+}
+
+VectorIndexPtr IVFPQ::Clone_impl(const std::shared_ptr<faiss::Index> &index) {
+    return std::make_shared<IVFPQ>(index);
+}
+
+} // knowhere
+} // zilliz

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

+// 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.
+
+
+#pragma once
+
+#include "IndexIVF.h"
+
+namespace zilliz {
+namespace knowhere {
+
+class IVFPQ : public IVF {
+public:
+    explicit IVFPQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {}
+
+    IVFPQ() = default;
+
+    IndexModelPtr
+    Train(const DatasetPtr &dataset, const Config &config) override;
+
+protected:
+    std::shared_ptr<faiss::IVFSearchParameters>
+    GenParams(const Config &config) override;
+
+    VectorIndexPtr
+    Clone_impl(const std::shared_ptr<faiss::Index> &index) override;
+};
+
+} // knowhere
+} // zilliz
+
+
+

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

+// 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 <faiss/gpu/GpuAutoTune.h>
+
+#include "knowhere/common/Exception.h"
+#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
+#include "knowhere/adapter/VectorAdapter.h"
+#include "IndexIVFSQ.h"
+#include "IndexGPUIVFSQ.h"
+
+
+namespace zilliz {
+namespace knowhere {
+
+IndexModelPtr IVFSQ::Train(const DatasetPtr &dataset, const Config &config) {
+    auto nlist = config["nlist"].as<size_t>();
+    auto nbits = config["nbits"].as<size_t>(); // TODO(linxj): only support SQ4 SQ6 SQ8 SQ16
+    auto metric_type = config["metric_type"].as_string() == "L2" ?
+                       faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT;
+
+    GETTENSOR(dataset)
+
+    std::stringstream index_type;
+    index_type << "IVF" << nlist << "," << "SQ" << nbits;
+    auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type);
+    build_index->train(rows, (float *) p_data);
+
+    std::shared_ptr<faiss::Index> ret_index;
+    ret_index.reset(build_index);
+    return std::make_shared<IVFIndexModel>(ret_index);
+}
+
+VectorIndexPtr IVFSQ::Clone_impl(const std::shared_ptr<faiss::Index> &index) {
+    return std::make_shared<IVFSQ>(index);
+}
+
+VectorIndexPtr IVFSQ::CopyCpuToGpu(const int64_t &device_id, const Config &config) {
+    if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)){
+        ResScope rs(res, device_id, false);
+        faiss::gpu::GpuClonerOptions option;
+        option.allInGpu = true;
+
+        auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get(), &option);
+
+        std::shared_ptr<faiss::Index> device_index;
+        device_index.reset(gpu_index);
+        return std::make_shared<GPUIVFSQ>(device_index, device_id, res);
+    } else {
+        KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
+    }
+}
+
+} // knowhere
+} // zilliz

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

+// 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.
+
+
+#pragma once
+
+#include "IndexIVF.h"
+
+namespace zilliz {
+namespace knowhere {
+
+class IVFSQ : public IVF {
+public:
+    explicit IVFSQ(std::shared_ptr<faiss::Index> index) : IVF(std::move(index)) {}
+
+    IVFSQ() = default;
+
+    IndexModelPtr
+    Train(const DatasetPtr &dataset, const Config &config) override;
+
+    VectorIndexPtr
+    CopyCpuToGpu(const int64_t &device_id, const Config &config) override;
+
+protected:
+    VectorIndexPtr
+    Clone_impl(const std::shared_ptr<faiss::Index> &index) override;
+};
+
+} // knowhere
+} // zilliz
+

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

@@ -25,9 +25,9 @@
 #undef mkdir
 
 #include "IndexKDT.h"
-#include "knowhere/index/vector_index/utils/Definitions.h"
+#include "knowhere/index/vector_index/helpers/Definitions.h"
 //#include "knowhere/index/preprocessor/normalize.h"
-#include "knowhere/index/vector_index/utils/KDTParameterMgr.h"
+#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h"
 #include "knowhere/adapter/SptagAdapter.h"
 #include "knowhere/common/Exception.h"
 

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

@@ -18,6 +18,8 @@
 
 #include "knowhere/common/Exception.h"
 #include "knowhere/index/vector_index/IndexIVF.h"
+#include "knowhere/index/vector_index/IndexIVFSQ.h"
+#include "knowhere/index/vector_index/IndexIVFPQ.h"
 #include "knowhere/index/vector_index/IndexGPUIVF.h"
 #include "knowhere/index/vector_index/IndexIDMAP.h"
 #include "Cloner.h"
@@ -43,7 +45,7 @@ VectorIndexPtr CopyCpuToGpu(const VectorIndexPtr &index, const int64_t &device_i
     if (auto cpu_index = std::dynamic_pointer_cast<IVFSQ>(index)) {
         return cpu_index->CopyCpuToGpu(device_id, config);
     } else if (auto cpu_index = std::dynamic_pointer_cast<IVFPQ>(index)) {
-        KNOWHERE_THROW_MSG("IVFPQ not support tranfer to gpu");
+        KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu");
     } else if (auto cpu_index = std::dynamic_pointer_cast<IVF>(index)) {
         return cpu_index->CopyCpuToGpu(device_id, config);
     } else if (auto cpu_index = std::dynamic_pointer_cast<IDMAP>(index)) {

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

+// 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 <cstring>
+
+#include "FaissIO.h"
+
+namespace zilliz {
+namespace knowhere {
+
+// TODO(linxj): Get From Config File
+static size_t magic_num = 2;
+size_t MemoryIOWriter::operator()(const void *ptr, size_t size, size_t nitems) {
+    auto total_need = size * nitems + rp;
+
+    if (!data_) { // data == nullptr
+        total = total_need * magic_num;
+        rp = size * nitems;
+        data_ = new uint8_t[total];
+        memcpy((void *) (data_), ptr, rp);
+    }
+
+    if (total_need > total) {
+        total = total_need * magic_num;
+        auto new_data = new uint8_t[total];
+        memcpy((void *) new_data, (void *) data_, rp);
+        delete data_;
+        data_ = new_data;
+
+        memcpy((void *) (data_ + rp), ptr, size * nitems);
+        rp = total_need;
+    } else {
+        memcpy((void *) (data_ + rp), ptr, size * nitems);
+        rp = total_need;
+    }
+
+    return nitems;
+}
+
+size_t MemoryIOReader::operator()(void *ptr, size_t size, size_t nitems) {
+    if (rp >= total) return 0;
+    size_t nremain = (total - rp) / size;
+    if (nremain < nitems) nitems = nremain;
+    memcpy(ptr, (void *) (data_ + rp), size * nitems);
+    rp += size * nitems;
+    return nitems;
+}
+
+} // knowhere
+} // zilliz

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

+// 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.
+
+
+#pragma once
+
+#include <faiss/AuxIndexStructures.h>
+
+namespace zilliz {
+namespace knowhere {
+
+struct MemoryIOWriter : public faiss::IOWriter {
+    uint8_t *data_ = nullptr;
+    size_t total = 0;
+    size_t rp = 0;
+
+    size_t
+    operator()(const void *ptr, size_t size, size_t nitems) override;
+};
+
+struct MemoryIOReader : public faiss::IOReader {
+    uint8_t *data_;
+    size_t rp = 0;
+    size_t total = 0;
+
+    size_t
+    operator()(void *ptr, size_t size, size_t nitems) override;
+};
+
+} // knowhere
+} // zilliz
+
+
+

cpp/src/core/knowhere/knowhere/index/vector_index/nsg/NSGIO.h

@@ -18,6 +18,7 @@
 
 #pragma once
 
+#include "knowhere/index/vector_index/helpers/FaissIO.h"
 #include "NSG.h"
 #include "knowhere/index/vector_index/IndexIVF.h"
 

cpp/src/wrapper/KnowhereResource.cpp

@@ -17,7 +17,7 @@
 
 
 #include "KnowhereResource.h"
-#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h"
+#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
 #include "server/ServerConfig.h"
 
 #include <map>

cpp/src/wrapper/vec_impl.cpp

@@ -20,7 +20,7 @@
 #include "knowhere/index/vector_index/IndexIDMAP.h"
 #include "knowhere/index/vector_index/IndexGPUIVF.h"
 #include "knowhere/common/Exception.h"
-#include "knowhere/index/vector_index/utils/Cloner.h"
+#include "knowhere/index/vector_index/helpers/Cloner.h"
 
 #include "vec_impl.h"
 #include "data_transfer.h"

cpp/src/wrapper/vec_index.cpp

@@ -16,8 +16,12 @@
 // under the License.
 
 #include "knowhere/index/vector_index/IndexIVF.h"
-#include "knowhere/index/vector_index/IndexIDMAP.h"
 #include "knowhere/index/vector_index/IndexGPUIVF.h"
+#include "knowhere/index/vector_index/IndexIVFSQ.h"
+#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
+#include "knowhere/index/vector_index/IndexIVFPQ.h"
+#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
+#include "knowhere/index/vector_index/IndexIDMAP.h"
 #include "knowhere/index/vector_index/IndexKDT.h"
 #include "knowhere/index/vector_index/IndexNSG.h"
 #include "knowhere/common/Exception.h"

cpp/unittest/db/utils.cpp

@@ -25,9 +25,7 @@
 #include "cache/CpuCacheMgr.h"
 #include "db/DBFactory.h"
 #include "db/Options.h"
-#include "server/ServerConfig.h"
-#include "knowhere/index/vector_index/IndexGPUIVF.h"
-#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h"
+#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
 
 INITIALIZE_EASYLOGGINGPP
 

cpp/unittest/wrapper/wrapper_test.cpp

@@ -20,8 +20,7 @@
 #include "utils/easylogging++.h"
 
 #include "src/wrapper/vec_index.h"
-#include "knowhere/index/vector_index/IndexGPUIVF.h"
-#include "knowhere/index/vector_index/utils/FaissGpuResourceMgr.h"
+#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
 
 #include "utils.h"