FAISS HNSW

benchmark/CMakeLists.txt

@@ -51,6 +51,7 @@ benchmark_test(benchmark_float_bitset          hdf5/benchmark_float_bitset.cpp)
 benchmark_test(benchmark_float_qps             hdf5/benchmark_float_qps.cpp)
 benchmark_test(benchmark_float_range           hdf5/benchmark_float_range.cpp)
 benchmark_test(benchmark_float_range_bitset    hdf5/benchmark_float_range_bitset.cpp)
+benchmark_test(benchmark_faiss_hnsw            hdf5/benchmark_faiss_hnsw.cpp)
 
 benchmark_test(gen_hdf5_file hdf5/gen_hdf5_file.cpp)
 benchmark_test(gen_fbin_file hdf5/gen_fbin_file.cpp)

benchmark/benchmark_base.h

@@ -14,7 +14,10 @@
 #include <math.h>
 #include <sys/time.h>
 
+#include <cstdint>
 #include <string>
+#include <unordered_map>
+#include <unordered_set>
 
 #define CALC_TIME_SPAN(X)       \
     double t_start = elapsed(); \
@@ -39,7 +42,7 @@ class Benchmark_base {
         }
     }
 
-    inline double
+    static inline double
     elapsed() {
         struct timeval tv;
         gettimeofday(&tv, nullptr);
@@ -86,7 +89,7 @@ class Benchmark_base {
         return (hit * 1.0f / (nq * min_k));
     }
 
-    float
+    static float
     CalcRecall(const int64_t* g_ids, const int64_t* ids, int32_t nq, int32_t k) {
         int32_t hit = 0;
         for (int32_t i = 0; i < nq; i++) {

benchmark/hdf5/benchmark_knowhere.h

@@ -11,7 +11,12 @@
 
 #pragma once
 
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
 #include <exception>
+#include <optional>
+#include <string>
 #include <vector>
 
 #include "benchmark/utils.h"
@@ -32,7 +37,7 @@ std::string kIPIndexPrefix = kIPIndexDir + "/ip";
 
 class Benchmark_knowhere : public Benchmark_hdf5 {
  public:
-    void
+    static void
     write_index(knowhere::Index<knowhere::IndexNode>& index, const std::string& filename, const knowhere::Json& conf) {
         FileIOWriter writer(filename);
 
@@ -53,7 +58,7 @@ class Benchmark_knowhere : public Benchmark_hdf5 {
         }
     }
 
-    void
+    static void
     read_index(knowhere::Index<knowhere::IndexNode>& index, const std::string& filename, const knowhere::Json& conf) {
         FileIOReader reader(filename);
         int64_t file_size = reader.size();
@@ -86,66 +91,80 @@ class Benchmark_knowhere : public Benchmark_hdf5 {
     }
 
     template <typename T>
-    std::string
-    get_index_name(const std::vector<int32_t>& params) {
+    static std::string
+    get_index_name(const std::string& ann_test_name, const std::string& index_type,
+                   const std::vector<int32_t>& params) {
         std::string params_str = "";
         for (size_t i = 0; i < params.size(); i++) {
             params_str += "_" + std::to_string(params[i]);
         }
         if constexpr (std::is_same_v<T, knowhere::fp32>) {
-            return ann_test_name_ + "_" + index_type_ + params_str + "_fp32" + ".index";
+            return ann_test_name + "_" + index_type + params_str + "_fp32" + ".index";
         } else if constexpr (std::is_same_v<T, knowhere::fp16>) {
-            return ann_test_name_ + "_" + index_type_ + params_str + "_fp16" + ".index";
+            return ann_test_name + "_" + index_type + params_str + "_fp16" + ".index";
         } else if constexpr (std::is_same_v<T, knowhere::bf16>) {
-            return ann_test_name_ + "_" + index_type_ + params_str + "_bf16" + ".index";
+            return ann_test_name + "_" + index_type + params_str + "_bf16" + ".index";
         } else {
-            return ann_test_name_ + "_" + index_type_ + params_str + ".index";
+            return ann_test_name + "_" + index_type + params_str + ".index";
         }
     }
 
+    template <typename T>
+    std::string
+    get_index_name(const std::vector<int32_t>& params) {
+        return this->get_index_name<T>(ann_test_name_, index_type_, params);
+    }
+
     template <typename T>
     knowhere::Index<knowhere::IndexNode>
-    create_index(const std::string& index_file_name, const knowhere::Json& conf) {
+    create_index(const std::string& index_type, const std::string& index_file_name,
+                 const knowhere::DataSetPtr& default_ds_ptr, const knowhere::Json& conf,
+                 const std::optional<std::string>& additional_name = std::nullopt) {
+        std::string additional_name_s = additional_name.value_or("");
+
+        printf("[%.3f s] Creating %sindex \"%s\"\n", get_time_diff(), additional_name_s.c_str(), index_type.c_str());
+
         auto version = knowhere::Version::GetCurrentVersion().VersionNumber();
-        printf("[%.3f s] Creating index \"%s\"\n", get_time_diff(), index_type_.c_str());
-        index_ = knowhere::IndexFactory::Instance().Create<T>(index_type_, version);
+        auto index = knowhere::IndexFactory::Instance().Create<T>(index_type, version);
 
         try {
-            printf("[%.3f s] Reading index file: %s\n", get_time_diff(), index_file_name.c_str());
-            read_index(index_.value(), index_file_name, conf);
+            printf("[%.3f s] Reading %sindex file: %s\n", get_time_diff(), additional_name_s.c_str(),
+                   index_file_name.c_str());
+
+            read_index(index.value(), index_file_name, conf);
         } catch (...) {
-            printf("[%.3f s] Building all on %d vectors\n", get_time_diff(), nb_);
-            auto ds_ptr = knowhere::GenDataSet(nb_, dim_, xb_);
-            auto base = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
-            index_.value().Build(base, conf);
+            printf("[%.3f s] Building %sindex all on %ld vectors\n", get_time_diff(), additional_name_s.c_str(),
+                   default_ds_ptr->GetRows());
+
+            auto base = knowhere::ConvertToDataTypeIfNeeded<T>(default_ds_ptr);
+            index.value().Build(base, conf);
 
-            printf("[%.3f s] Writing index file: %s\n", get_time_diff(), index_file_name.c_str());
-            write_index(index_.value(), index_file_name, conf);
+            printf("[%.3f s] Writing %sindex file: %s\n", get_time_diff(), additional_name_s.c_str(),
+                   index_file_name.c_str());
+
+            write_index(index.value(), index_file_name, conf);
         }
-        return index_.value();
+
+        return index.value();
+    }
+
+    template <typename T>
+    knowhere::Index<knowhere::IndexNode>
+    create_index(const std::string& index_file_name, const knowhere::Json& conf) {
+        auto idx = this->create_index<T>(index_type_, index_file_name, knowhere::GenDataSet(nb_, dim_, xb_), conf);
+        index_ = idx;
+        return idx;
     }
 
     knowhere::Index<knowhere::IndexNode>
     create_golden_index(const knowhere::Json& conf) {
-        auto version = knowhere::Version::GetCurrentVersion().VersionNumber();
         golden_index_type_ = knowhere::IndexEnum::INDEX_FAISS_IDMAP;
-
         std::string golden_index_file_name = ann_test_name_ + "_" + golden_index_type_ + "_GOLDEN" + ".index";
-        printf("[%.3f s] Creating golden index \"%s\"\n", get_time_diff(), golden_index_type_.c_str());
-        golden_index_ = knowhere::IndexFactory::Instance().Create<knowhere::fp32>(golden_index_type_, version);
-
-        try {
-            printf("[%.3f s] Reading golden index file: %s\n", get_time_diff(), golden_index_file_name.c_str());
-            read_index(golden_index_.value(), golden_index_file_name, conf);
-        } catch (...) {
-            printf("[%.3f s] Building golden index on %d vectors\n", get_time_diff(), nb_);
-            knowhere::DataSetPtr ds_ptr = knowhere::GenDataSet(nb_, dim_, xb_);
-            golden_index_.value().Build(ds_ptr, conf);
 
-            printf("[%.3f s] Writing golden index file: %s\n", get_time_diff(), golden_index_file_name.c_str());
-            write_index(golden_index_.value(), golden_index_file_name, conf);
-        }
-        return golden_index_.value();
+        auto idx = this->create_index<knowhere::fp32>(golden_index_type_, golden_index_file_name,
+                                                      knowhere::GenDataSet(nb_, dim_, xb_), conf, "golden ");
+        golden_index_ = idx;
+        return idx;
     }
 
     void

benchmark/utils.h

@@ -11,6 +11,9 @@
 
 #pragma once
 
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
 #include <cstdio>
 #include <cstdlib>
 #include <fstream>

cmake/libs/libfaiss.cmake

@@ -1,7 +1,8 @@
 knowhere_file_glob(
   GLOB FAISS_SRCS thirdparty/faiss/faiss/*.cpp
   thirdparty/faiss/faiss/impl/*.cpp thirdparty/faiss/faiss/invlists/*.cpp
-  thirdparty/faiss/faiss/utils/*.cpp)
+  thirdparty/faiss/faiss/utils/*.cpp
+  thirdparty/faiss/faiss/cppcontrib/knowhere/*.cpp)
 
 knowhere_file_glob(GLOB FAISS_AVX512_SRCS
                    thirdparty/faiss/faiss/impl/*avx512.cpp)

include/knowhere/comp/index_param.h

@@ -53,6 +53,11 @@ constexpr const char* INDEX_HNSW_SQ8 = "HNSW_SQ8";
 constexpr const char* INDEX_HNSW_SQ8_REFINE = "HNSW_SQ8_REFINE";
 constexpr const char* INDEX_DISKANN = "DISKANN";
 
+constexpr const char* INDEX_FAISS_HNSW_FLAT = "FAISS_HNSW_FLAT";
+constexpr const char* INDEX_FAISS_HNSW_SQ = "FAISS_HNSW_SQ";
+constexpr const char* INDEX_FAISS_HNSW_PQ = "FAISS_HNSW_PQ";
+constexpr const char* INDEX_FAISS_HNSW_PRQ = "FAISS_HNSW_PRQ";
+
 constexpr const char* INDEX_SPARSE_INVERTED_INDEX = "SPARSE_INVERTED_INDEX";
 constexpr const char* INDEX_SPARSE_WAND = "SPARSE_WAND";
 }  // namespace IndexEnum
@@ -150,6 +155,10 @@ constexpr const char* HNSW_M = "M";
 constexpr const char* EF = "ef";
 constexpr const char* OVERVIEW_LEVELS = "overview_levels";
 
+// FAISS additional Params
+constexpr const char* SQ_TYPE = "sq_type";  // for IVF_SQ and HNSW_SQ
+constexpr const char* PRQ_NUM = "nrq";      // for PRQ, number of redisual quantizers
+
 // Sparse Params
 constexpr const char* DROP_RATIO_BUILD = "drop_ratio_build";
 constexpr const char* DROP_RATIO_SEARCH = "drop_ratio_search";

include/knowhere/index/index_table.h

@@ -66,6 +66,22 @@ static std::set<std::pair<std::string, VecType>> legal_knowhere_index = {
     {IndexEnum::INDEX_HNSW_SQ8_REFINE, VecType::VECTOR_FLOAT},
     {IndexEnum::INDEX_HNSW_SQ8_REFINE, VecType::VECTOR_FLOAT16},
     {IndexEnum::INDEX_HNSW_SQ8_REFINE, VecType::VECTOR_BFLOAT16},
+    // faiss hnsw
+    {IndexEnum::INDEX_FAISS_HNSW_FLAT, VecType::VECTOR_FLOAT},
+    {IndexEnum::INDEX_FAISS_HNSW_FLAT, VecType::VECTOR_FLOAT16},
+    {IndexEnum::INDEX_FAISS_HNSW_FLAT, VecType::VECTOR_BFLOAT16},
+
+    {IndexEnum::INDEX_FAISS_HNSW_SQ, VecType::VECTOR_FLOAT},
+    {IndexEnum::INDEX_FAISS_HNSW_SQ, VecType::VECTOR_FLOAT16},
+    {IndexEnum::INDEX_FAISS_HNSW_SQ, VecType::VECTOR_BFLOAT16},
+
+    {IndexEnum::INDEX_FAISS_HNSW_PQ, VecType::VECTOR_FLOAT},
+    {IndexEnum::INDEX_FAISS_HNSW_PQ, VecType::VECTOR_FLOAT16},
+    {IndexEnum::INDEX_FAISS_HNSW_PQ, VecType::VECTOR_BFLOAT16},
+
+    {IndexEnum::INDEX_FAISS_HNSW_PRQ, VecType::VECTOR_FLOAT},
+    {IndexEnum::INDEX_FAISS_HNSW_PRQ, VecType::VECTOR_FLOAT16},
+    {IndexEnum::INDEX_FAISS_HNSW_PRQ, VecType::VECTOR_BFLOAT16},
     // diskann
     {IndexEnum::INDEX_DISKANN, VecType::VECTOR_FLOAT},
     {IndexEnum::INDEX_DISKANN, VecType::VECTOR_FLOAT16},
@@ -116,6 +132,16 @@ static std::set<std::string> legal_support_mmap_knowhere_index = {
     IndexEnum::INDEX_HNSW_SQ8_REFINE,
     IndexEnum::INDEX_HNSW_SQ8_REFINE,
     IndexEnum::INDEX_HNSW_SQ8_REFINE,
+
+    // faiss hnsw
+    IndexEnum::INDEX_FAISS_HNSW_FLAT,
+
+    IndexEnum::INDEX_FAISS_HNSW_SQ,
+
+    IndexEnum::INDEX_FAISS_HNSW_PQ,
+
+    IndexEnum::INDEX_FAISS_HNSW_PRQ,
+
     // sparse index
     IndexEnum::INDEX_SPARSE_INVERTED_INDEX,
     IndexEnum::INDEX_SPARSE_WAND,

include/knowhere/utils.h

@@ -109,44 +109,69 @@ GetKey(const std::string& name) {
 
 template <typename InType, typename OutType>
 inline DataSetPtr
-data_type_conversion(const DataSet& src) {
+data_type_conversion(const DataSet& src, const std::optional<int64_t> start = std::nullopt,
+                     const std::optional<int64_t> count = std::nullopt) {
     auto dim = src.GetDim();
     auto rows = src.GetRows();
 
-    auto des_data = new OutType[dim * rows];
-    auto src_data = (InType*)src.GetTensor();
-    for (auto i = 0; i < dim * rows; i++) {
-        des_data[i] = (OutType)src_data[i];
+    // check the acceptable range
+    int64_t start_row = start.value_or(0);
+    if (start_row < 0 || start_row >= rows) {
+        return nullptr;
+    }
+
+    int64_t count_rows = count.value_or(rows - start_row);
+    if (count_rows < 0 || start_row + count_rows > rows) {
+        return nullptr;
+    }
+
+    // map
+    auto* des_data = new OutType[dim * count_rows];
+    auto* src_data = (const InType*)src.GetTensor();
+    for (auto i = 0; i < dim * count_rows; i++) {
+        des_data[i] = (OutType)src_data[i + start_row * dim];
     }
 
     auto des = std::make_shared<DataSet>();
-    des->SetRows(rows);
+    des->SetRows(count_rows);
     des->SetDim(dim);
     des->SetTensor(des_data);
     des->SetIsOwner(true);
     return des;
 }
 
 // Convert DataSet from DataType to float
+// * no start, no count, float -> returns the source without cloning
+// * no start, no count, no float -> returns a clone with a different type
+// * start, no count -> returns a clone that starts from a given row 'start'
+// * no start, count -> returns a clone that starts from a row 0 and has 'count' rows
+// * start, count -> returns a clone that start from a given row 'start' and has 'count' rows
+// * invalid start, count values -> returns nullptr
 template <typename DataType>
 inline DataSetPtr
-ConvertFromDataTypeIfNeeded(const DataSetPtr ds) {
+ConvertFromDataTypeIfNeeded(const DataSetPtr& ds, const std::optional<int64_t> start = std::nullopt,
+                            const std::optional<int64_t> count = std::nullopt) {
     if constexpr (std::is_same_v<DataType, typename MockData<DataType>::type>) {
-        return ds;
-    } else {
-        return data_type_conversion<DataType, typename MockData<DataType>::type>(*ds);
+        if (!start.has_value() && !count.has_value()) {
+            return ds;
+        }
     }
+
+    return data_type_conversion<DataType, typename MockData<DataType>::type>(*ds, start, count);
 }
 
 // Convert DataSet from float to DataType
 template <typename DataType>
 inline DataSetPtr
-ConvertToDataTypeIfNeeded(const DataSetPtr ds) {
+ConvertToDataTypeIfNeeded(const DataSetPtr& ds, const std::optional<int64_t> start = std::nullopt,
+                          const std::optional<int64_t> count = std::nullopt) {
     if constexpr (std::is_same_v<DataType, typename MockData<DataType>::type>) {
-        return ds;
-    } else {
-        return data_type_conversion<typename MockData<DataType>::type, DataType>(*ds);
+        if (!start.has_value() && !count.has_value()) {
+            return ds;
+        }
     }
+
+    return data_type_conversion<typename MockData<DataType>::type, DataType>(*ds, start, count);
 }
 
 template <typename T>