From cf1e03980edc2eadc9c3b4ad01be34f9c114af51 Mon Sep 17 00:00:00 2001
From: Heisenberg <zhuangbo.ye@zilliz.com>
Date: Mon, 25 Nov 2019 19:48:47 +0800
Subject: [PATCH] fix issue 527

---
 CHANGELOG.md                                  |  1 +
 core/src/index/unittest/CMakeLists.txt        |  2 +-
 .../faiss_benchmark/faiss_benchmark_test.cpp  | 47 ++++++++++++-------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 53e427dcb9..d6db411cf9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - \#440 - Server cannot startup with gpu_resource_config.enable=false in GPU version
 - \#458 - Index data is not compatible between 0.5 and 0.6
 - \#486 - gpu no usage during index building
+- \#527 - faiss benchmark not compatible with faiss 1.6.0
 
 ## Feature
 - \#12 - Pure CPU version for Milvus
diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt
index 93ae63a9ec..af21d20b48 100644
--- a/core/src/index/unittest/CMakeLists.txt
+++ b/core/src/index/unittest/CMakeLists.txt
@@ -112,6 +112,6 @@ if (KNOWHERE_GPU_VERSION)
     install(TARGETS test_customized_index DESTINATION unittest)
 endif ()
 #add_subdirectory(faiss_ori)
-#add_subdirectory(faiss_benchmark)
+add_subdirectory(faiss_benchmark)
 add_subdirectory(test_nsg)
 
diff --git a/core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp b/core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp
index ce5aa4bafa..0d49004621 100644
--- a/core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp
+++ b/core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp
@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#define USE_FAISS_V0_2_1 0
+
 #include <gtest/gtest.h>
 
 #include <hdf5.h>
@@ -26,21 +28,6 @@
 #include <cstdio>
 #include <vector>
 
-#define USE_FAISS_V1_5_3 0
-
-#if USE_FAISS_V1_5_3
-#include <faiss/gpu/GpuAutoTune.h>
-#include <faiss/utils.h>
-#include <sys/stat.h>
-#include <cstdlib>
-#include <cstring>
-
-#else
-#include <faiss/gpu/GpuCloner.h>
-#include <faiss/index_factory.h>
-#include <faiss/utils/distances.h>
-#endif
-
 #include <faiss/AutoTune.h>
 #include <faiss/Index.h>
 #include <faiss/IndexIVF.h>
@@ -48,8 +35,22 @@
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/index_io.h>
 
+#if USE_FAISS_V0_2_1
+#include <faiss/gpu/GpuAutoTune.h>
+#include <faiss/utils.h>
+#include <sys/stat.h>
+#include <cstdlib>
+#include <cstring>
+#else
+#include <faiss/gpu/GpuCloner.h>
+#include <faiss/index_factory.h>
+#include <faiss/utils/distances.h>
+#endif
+
 #ifdef CUSTOMIZATION
 #include <faiss/gpu/GpuIndexIVFSQHybrid.h>
+#else
+#include <faiss/gpu/GpuIndexIVF.h>
 #endif
 
 /*****************************************************
@@ -295,10 +296,12 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
         cpu_index = faiss::gpu::index_gpu_to_cpu(gpu_index);
         delete gpu_index;
 
+#ifdef CUSTOMIZATION
         faiss::IndexIVF* cpu_ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
         if (cpu_ivf_index != nullptr) {
             cpu_ivf_index->to_readonly();
         }
+#endif
 
         printf("[%.3f s] Writing index file: %s\n", elapsed() - t0, index_file_name.c_str());
         faiss::write_index(cpu_index, index_file_name.c_str());
@@ -374,13 +377,15 @@ test_with_nprobes(const std::string& ann_test_name, const std::string& index_key
     faiss::Index *gpu_index, *index;
     if (query_mode != MODE_CPU) {
         faiss::gpu::GpuClonerOptions option;
+#ifdef CUSTOMIZATION
         option.allInGpu = true;
 
         faiss::IndexComposition index_composition;
         index_composition.index = cpu_index;
         index_composition.quantizer = nullptr;
-
+#endif
         switch (query_mode) {
+#ifdef CUSTOMIZATION
             case MODE_MIX: {
                 index_composition.mode = 1;  // 0: all data, 1: copy quantizer, 2: copy data
 
@@ -403,7 +408,9 @@ test_with_nprobes(const std::string& ann_test_name, const std::string& index_key
                 index = cpu_index;
                 break;
             }
+#endif
             case MODE_GPU:
+#ifdef CUSTOMIZATION
                 index_composition.mode = 0;  // 0: all data, 1: copy quantizer, 2: copy data
 
                 // warm up the transmission
@@ -412,6 +419,14 @@ test_with_nprobes(const std::string& ann_test_name, const std::string& index_key
 
                 copy_time = elapsed();
                 gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
+#else
+                // warm up the transmission
+                gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, cpu_index, &option);
+                delete gpu_index;
+
+                copy_time = elapsed();
+                gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, cpu_index, &option);
+#endif
                 copy_time = elapsed() - copy_time;
                 printf("[%.3f s] Copy data completed, cost %f s\n", elapsed() - t0, copy_time);