MS-538 1. update unittest

Former-commit-id: 38fefd469b3b69ccdcbfb3c741002d8e4703c702

cpp/src/core/include/knowhere/index/vector_index/idmap.h

@@ -32,6 +32,12 @@ class IDMAP : public VectorIndex, public BasicIndex {
     virtual int64_t *GetRawIds();
 
  protected:
+    virtual void search_impl(int64_t n,
+                             const float *data,
+                             int64_t k,
+                             float *distances,
+                             int64_t *labels,
+                             const Config &cfg);
     std::mutex mutex_;
 };
 
@@ -49,6 +55,12 @@ class GPUIDMAP : public IDMAP, public GPUIndex {
     VectorIndexPtr CopyGpuToGpu(const int64_t &device_id, 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;
     BinarySet SerializeImpl() override;
     void LoadImpl(const BinarySet &index_binary) override;
 };

cpp/src/core/src/knowhere/index/preprocessor/normalize.cpp

@@ -9,32 +9,32 @@ namespace knowhere {
 
 DatasetPtr
 NormalizePreprocessor::Preprocess(const DatasetPtr &dataset) {
-    // TODO: wrap dataset->tensor
-    auto tensor = dataset->tensor()[0];
-    auto p_data = (float *)tensor->raw_mutable_data();
-    auto dimension = tensor->shape()[1];
-    auto rows = tensor->shape()[0];
-
-#pragma omp parallel for
-    for (auto i = 0; i < rows; ++i) {
-        Normalize(&(p_data[i * dimension]), dimension);
-    }
+//    // TODO: wrap dataset->tensor
+//    auto tensor = dataset->tensor()[0];
+//    auto p_data = (float *)tensor->raw_mutable_data();
+//    auto dimension = tensor->shape()[1];
+//    auto rows = tensor->shape()[0];
+//
+//#pragma omp parallel for
+//    for (auto i = 0; i < rows; ++i) {
+//        Normalize(&(p_data[i * dimension]), dimension);
+//    }
 }
 
 void
 NormalizePreprocessor::Normalize(float *arr, int64_t dimension) {
-    double vector_length = 0;
-    for (auto j = 0; j < dimension; j++) {
-        double val = arr[j];
-        vector_length += val * val;
-    }
-    vector_length = std::sqrt(vector_length);
-    if (vector_length < 1e-6) {
-        auto val = (float) (1.0 / std::sqrt((double) dimension));
-        for (int j = 0; j < dimension; j++) arr[j] = val;
-    } else {
-        for (int j = 0; j < dimension; j++) arr[j] = (float) (arr[j] / vector_length);
-    }
+    //double vector_length = 0;
+    //for (auto j = 0; j < dimension; j++) {
+    //    double val = arr[j];
+    //    vector_length += val * val;
+    //}
+    //vector_length = std::sqrt(vector_length);
+    //if (vector_length < 1e-6) {
+    //    auto val = (float) (1.0 / std::sqrt((double) dimension));
+    //    for (int j = 0; j < dimension; j++) arr[j] = val;
+    //} else {
+    //    for (int j = 0; j < dimension; j++) arr[j] = (float) (arr[j] / vector_length);
+    //}
 }
 
 } // namespace knowhere

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

@@ -70,10 +70,10 @@ CPUKDTRNG::Train(const DatasetPtr &origin, const Config &train_config) {
     SetParameters(train_config);
     DatasetPtr dataset = origin->Clone();
 
-    if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
-        && preprocessor_) {
+    //if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
+    //    && preprocessor_) {
         preprocessor_->Preprocess(dataset);
-    }
+    //}
 
     auto vectorset = ConvertToVectorSet(dataset);
     auto metaset = ConvertToMetadataSet(dataset);
@@ -88,10 +88,10 @@ CPUKDTRNG::Add(const DatasetPtr &origin, const Config &add_config) {
     SetParameters(add_config);
     DatasetPtr dataset = origin->Clone();
 
-    if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
-        && preprocessor_) {
+    //if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
+    //    && preprocessor_) {
         preprocessor_->Preprocess(dataset);
-    }
+    //}
 
     auto vectorset = ConvertToVectorSet(dataset);
     auto metaset = ConvertToMetadataSet(dataset);

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

@@ -280,15 +280,15 @@ void FaissGpuResourceMgr::InitResource() {
 
     is_init = true;
 
-    std::cout << "InitResource" << std::endl;
+    //std::cout << "InitResource" << std::endl;
     for(auto& device : devices_params_) {
         auto& device_id = device.first;
-        std::cout << "Device Id: " << device_id << std::endl;
+        //std::cout << "Device Id: " << device_id << std::endl;
         auto& device_param = device.second;
         auto& bq = idle_map[device_id];
 
         for (int64_t i = 0; i < device_param.resource_num; ++i) {
-            std::cout << "Resource Id: " << i << std::endl;
+            //std::cout << "Resource Id: " << i << std::endl;
             auto raw_resource = std::make_shared<faiss::gpu::StandardGpuResources>();
 
             // TODO(linxj): enable set pinned memory
@@ -298,7 +298,7 @@ void FaissGpuResourceMgr::InitResource() {
             bq.Put(res_wrapper);
         }
     }
-    std::cout << "End initResource" << std::endl;
+    //std::cout << "End initResource" << std::endl;
 }
 
 ResPtr FaissGpuResourceMgr::GetRes(const int64_t &device_id,
@@ -315,16 +315,6 @@ ResPtr FaissGpuResourceMgr::GetRes(const int64_t &device_id,
     return nullptr;
 }
 
-//bool FaissGpuResourceMgr::GetRes(const int64_t &device_id,
-//                                 ResPtr &res,
-//                                 const int64_t &alloc_size) {
-//    InitResource();
-//
-//    std::lock_guard<std::mutex> lk(res->mutex);
-//    AllocateTempMem(res, device_id, alloc_size);
-//    return true;
-//}
-
 void FaissGpuResourceMgr::MoveToIdle(const int64_t &device_id, const ResPtr &res) {
     auto finder = idle_map.find(device_id);
     if (finder != idle_map.end()) {

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

@@ -50,7 +50,7 @@ DatasetPtr IDMAP::Search(const DatasetPtr &dataset, const Config &config) {
     auto res_ids = (int64_t *) malloc(sizeof(int64_t) * elems);
     auto res_dis = (float *) malloc(sizeof(float) * elems);
 
-    index_->search(rows, (float *) p_data, k, res_dis, res_ids);
+    search_impl(rows, (float *) p_data, k, res_dis, res_ids, Config());
 
     auto id_buf = MakeMutableBufferSmart((uint8_t *) res_ids, sizeof(int64_t) * elems);
     auto dist_buf = MakeMutableBufferSmart((uint8_t *) res_dis, sizeof(float) * elems);
@@ -72,6 +72,11 @@ DatasetPtr IDMAP::Search(const DatasetPtr &dataset, const Config &config) {
     return std::make_shared<Dataset>(array, nullptr);
 }
 
+void IDMAP::search_impl(int64_t n, const float *data, int64_t k, float *distances, int64_t *labels, const Config &cfg) {
+    index_->search(n, (float *) data, k, distances, labels);
+
+}
+
 void IDMAP::Add(const DatasetPtr &dataset, const Config &config) {
     if (!index_) {
         KNOWHERE_THROW_MSG("index not initialize");
@@ -207,6 +212,7 @@ void GPUIDMAP::LoadImpl(const BinarySet &index_binary) {
 
         if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_) ){
             ResScope rs(gpu_id_, res);
+            res_ = res;
             auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index);
             index_.reset(device_index);
         } else {
@@ -230,5 +236,15 @@ int64_t *GPUIDMAP::GetRawIds() {
     KNOWHERE_THROW_MSG("Not support");
 }
 
+void GPUIDMAP::search_impl(int64_t n,
+                           const float *data,
+                           int64_t k,
+                           float *distances,
+                           int64_t *labels,
+                           const Config &cfg) {
+    ResScope rs(res_);
+    index_->search(n, (float *) data, k, distances, labels);
+}
+
 }
 }

cpp/src/core/test/test_idmap.cpp

@@ -12,6 +12,7 @@
 #include "knowhere/index/vector_index/idmap.h"
 #include "knowhere/adapter/structure.h"
 #include "knowhere/index/vector_index/cloner.h"
+#include "knowhere/common/exception.h"
 
 #include "utils.h"
 
@@ -65,19 +66,20 @@ void PrintResult(const DatasetPtr &result,
 }
 
 TEST_F(IDMAPTest, idmap_basic) {
-    assert(!xb.empty());
+    ASSERT_TRUE(!xb.empty());
     Config Default_cfg;
 
     index_->Train(Config::object{{"dim", dim}, {"metric_type", "L2"}});
     index_->Add(base_dataset, Default_cfg);
     EXPECT_EQ(index_->Count(), nb);
     EXPECT_EQ(index_->Dimension(), dim);
-    assert(index_->GetRawVectors() != nullptr);
-    assert(index_->GetRawIds() != nullptr);
+    ASSERT_TRUE(index_->GetRawVectors() != nullptr);
+    ASSERT_TRUE(index_->GetRawIds() != nullptr);
     auto result = index_->Search(query_dataset, Config::object{{"k", k}});
     AssertAnns(result, nq, k);
     PrintResult(result, nq, k);
 
+    index_->Seal();
     auto binaryset = index_->Serialize();
     auto new_index = std::make_shared<IDMAP>();
     new_index->Load(binaryset);
@@ -126,15 +128,15 @@ TEST_F(IDMAPTest, idmap_serialize) {
 }
 
 TEST_F(IDMAPTest, copy_test) {
-    assert(!xb.empty());
+    ASSERT_TRUE(!xb.empty());
     Config Default_cfg;
 
     index_->Train(Config::object{{"dim", dim}, {"metric_type", "L2"}});
     index_->Add(base_dataset, Default_cfg);
     EXPECT_EQ(index_->Count(), nb);
     EXPECT_EQ(index_->Dimension(), dim);
-    assert(index_->GetRawVectors() != nullptr);
-    assert(index_->GetRawIds() != nullptr);
+    ASSERT_TRUE(index_->GetRawVectors() != nullptr);
+    ASSERT_TRUE(index_->GetRawIds() != nullptr);
     auto result = index_->Search(query_dataset, Config::object{{"k", k}});
     AssertAnns(result, nq, k);
     //PrintResult(result, nq, k);
@@ -151,8 +153,16 @@ TEST_F(IDMAPTest, copy_test) {
         auto clone_index = CopyCpuToGpu(index_, device_id, Config());
         auto clone_result = clone_index->Search(query_dataset, Config::object{{"k", k}});
         AssertAnns(clone_result, nq, k);
-        //assert(std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawVectors() != nullptr);
-        //assert(std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawIds() != nullptr);
+        ASSERT_THROW({ std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawVectors(); },
+                     zilliz::knowhere::KnowhereException);
+        ASSERT_THROW({ std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawIds(); },
+                     zilliz::knowhere::KnowhereException);
+
+        auto binary = clone_index->Serialize();
+        clone_index->Load(binary);
+        auto new_result = clone_index->Search(query_dataset, Config::object{{"k", k}});
+        AssertAnns(new_result, nq, k);
+
         auto clone_gpu_idx = clone_index->Clone();
         auto clone_gpu_res = clone_gpu_idx->Search(query_dataset, Config::object{{"k", k}});
         AssertAnns(clone_gpu_res, nq, k);
@@ -161,14 +171,13 @@ TEST_F(IDMAPTest, copy_test) {
         auto host_index = CopyGpuToCpu(clone_index, Config());
         auto host_result = host_index->Search(query_dataset, Config::object{{"k", k}});
         AssertAnns(host_result, nq, k);
-        assert(std::static_pointer_cast<IDMAP>(host_index)->GetRawVectors() != nullptr);
-        assert(std::static_pointer_cast<IDMAP>(host_index)->GetRawIds() != nullptr);
+        ASSERT_TRUE(std::static_pointer_cast<IDMAP>(host_index)->GetRawVectors() != nullptr);
+        ASSERT_TRUE(std::static_pointer_cast<IDMAP>(host_index)->GetRawIds() != nullptr);
 
         // gpu to gpu
         auto device_index = CopyCpuToGpu(index_, device_id, Config());
-        auto device_result = device_index->Search(query_dataset, Config::object{{"k", k}});
+        auto new_device_index = std::static_pointer_cast<GPUIDMAP>(device_index)->CopyGpuToGpu(device_id, Config());
+        auto device_result = new_device_index->Search(query_dataset, Config::object{{"k", k}});
         AssertAnns(device_result, nq, k);
-        //assert(std::static_pointer_cast<GPUIDMAP>(device_index)->GetRawVectors() != nullptr);
-        //assert(std::static_pointer_cast<GPUIDMAP>(device_index)->GetRawIds() != nullptr);
     }
 }

cpp/src/core/test/test_ivf.cpp

@@ -394,8 +394,11 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
 
 
     {
-        index_type = "GPUIVF";
-        index_ = IndexFactory(index_type);
+        index_ =  std::make_shared<GPUIVF>(-1);
+        ASSERT_EQ(std::dynamic_pointer_cast<GPUIVF>(index_)->GetGpuDevice(), -1);
+        std::dynamic_pointer_cast<GPUIVF>(index_)->SetGpuDevice(device_id);
+        ASSERT_EQ(std::dynamic_pointer_cast<GPUIVF>(index_)->GetGpuDevice(), device_id);
+
         auto preprocessor = index_->BuildPreprocessor(base_dataset, preprocess_cfg);
         index_->set_preprocessor(preprocessor);
         train_cfg = Config::object{{"nlist", 1638}, {"gpu_id", device_id}, {"metric_type", "L2"}};
@@ -412,8 +415,9 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
             if (i > search_count - 6 || i < 5)
                 tc.RecordSection("search once");
         }
-        tc.RecordSection("search all");
+        tc.ElapseFromBegin("search all");
     }
+    FaissGpuResourceMgr::GetInstance().Dump();
 
     {
         // IVF-Search
@@ -430,7 +434,7 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
             if (i > search_count - 6 || i < 5)
                 tc.RecordSection("search once");
         }
-        tc.RecordSection("search all");
+        tc.ElapseFromBegin("search all");
     }
 
 }
@@ -461,7 +465,7 @@ TEST_F(GPURESTEST, gpuivfsq) {
             if (i > search_count - 6 || i < 5)
                 tc.RecordSection("search once");
         }
-        tc.RecordSection("search all");
+        tc.ElapseFromBegin("search all");
     }
 
     {
@@ -493,7 +497,7 @@ TEST_F(GPURESTEST, gpuivfsq) {
             if (i > search_count - 6 || i < 5)
                 tc.RecordSection("search once");
         }
-        tc.RecordSection("search all");
+        tc.ElapseFromBegin("search all");
         delete cpu_index;
         delete search_idx;
     }

cpp/src/core/test/test_nsg/test_nsg.cpp

@@ -7,6 +7,7 @@
 #include <gtest/gtest.h>
 #include <memory>
 
+#include "knowhere/common/exception.h"
 #include "knowhere/index/vector_index/gpu_ivf.h"
 #include "knowhere/index/vector_index/nsg_index.h"
 #include "knowhere/index/vector_index/nsg/nsg_io.h"
@@ -71,6 +72,14 @@ TEST_P(NSGInterfaceTest, basic_test) {
     auto new_result = new_index->Search(query_dataset, Config::object{{"k", k}});
     AssertAnns(result, nq, k);
 
+    ASSERT_EQ(index_->Count(), nb);
+    ASSERT_EQ(index_->Dimension(), dim);
+    ASSERT_THROW({index_->Clone();}, zilliz::knowhere::KnowhereException);
+    ASSERT_NO_THROW({
+        index_->Add(base_dataset, Config());
+        index_->Seal();
+    });
+
     {
         //std::cout << "k = 1" << std::endl;
         //new_index->Search(GenQuery(1), Config::object{{"k", 1}});