Add partitioning_avx2.cpp for PQFastScan (#159)

Signed-off-by: chasingegg <[email protected]>
zilliztech · Oct 19, 2023 · ce65e5e · ce65e5e
1 parent 0548a61
commit ce65e5e
Show file tree

Hide file tree

Showing 4 changed files with 453 additions and 2 deletions.
diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake
@@ -10,6 +10,7 @@ knowhere_file_glob(GLOB FAISS_AVX2_SRCS
  thirdparty/faiss/faiss/impl/*avx.cpp
  thirdparty/faiss/faiss/impl/pq4_fast_scan_search_1.cpp
  thirdparty/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp
+ thirdparty/faiss/faiss/utils/partitioning_avx2.cpp
  thirdparty/faiss/faiss/IndexPQFastScan.cpp
  thirdparty/faiss/faiss/IndexIVFPQFastScan.cpp)
 

diff --git a/thirdparty/faiss/faiss/impl/simd_result_handlers.h b/thirdparty/faiss/faiss/impl/simd_result_handlers.h
@@ -427,15 +427,15 @@ struct ReservoirTopN {
 
  void shrink() {
  uint64_t t0 = get_cy();
- threshold = partition<C>(vals, ids, i, n);
+ threshold = partition_avx2<C>(vals, ids, i, n);
  i = n;
  cycles += get_cy() - t0;
  }
 
  void shrink_fuzzy() {
  uint64_t t0 = get_cy();
  assert(i == capacity);
- threshold = partition_fuzzy<C>(
+ threshold = partition_fuzzy_avx2<C>(
  vals, ids, capacity, n, (capacity + n) / 2, &i);
  cycles += get_cy() - t0;
  }

diff --git a/thirdparty/faiss/faiss/utils/partitioning.h b/thirdparty/faiss/faiss/utils/partitioning.h
@@ -38,6 +38,27 @@ inline typename C::T partition(
  return partition_fuzzy<C>(vals, ids, n, q, q, nullptr);
 }
 
+// avx2 version
+template <class C>
+typename C::T partition_fuzzy_avx2(
+ typename C::T* vals,
+ typename C::TI* ids,
+ size_t n,
+ size_t q_min,
+ size_t q_max,
+ size_t* q_out);
+
+// avx2 version
+/** simplified interface for when the parition is not fuzzy */
+template <class C>
+inline typename C::T partition_avx2(
+ typename C::T* vals,
+ typename C::TI* ids,
+ size_t n,
+ size_t q) {
+ return partition_fuzzy_avx2<C>(vals, ids, n, q, q, nullptr);
+}
+
 /** low level SIMD histogramming functions */
 
 /** 8-bin histogram of (x - min) >> shift