faiss 0.6.4 → 0.6.5

data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/utils/distances.h>
+#include <faiss/utils/simd_impl/masked_sum-avx2.h>
 
 #include <immintrin.h>
 
@@ -105,6 +106,19 @@ void fvec_madd<SIMDLevel::AVX2>(
     }
 }
 
+template <>
+float fvec_masked_sum<SIMDLevel::AVX2>(
+        const float* x,
+        const uint8_t* bits,
+        size_t d) {
+    return masked_sum<SIMDLevel::AVX2>(x, bits, d);
+}
+
+// Internal linkage: the same template is compiled with other -m flags in
+// the other per-SIMD distances TU, and a shared symbol would let the linker
+// hand one level the other's code.
+namespace {
+
 template <size_t DIM>
 void fvec_L2sqr_ny_y_transposed_D(
         float* distances,
@@ -181,6 +195,8 @@ void fvec_L2sqr_ny_y_transposed_D(
     }
 }
 
+} // namespace
+
 template <>
 void fvec_L2sqr_ny_transposed<SIMDLevel::AVX2>(
         float* dis,
@@ -199,8 +215,12 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::AVX2>(
     switch (d) {
         DISPATCH(1)
         DISPATCH(2)
+        DISPATCH(3)
         DISPATCH(4)
+        DISPATCH(6)
         DISPATCH(8)
+        DISPATCH(12)
+        DISPATCH(16)
         default:
             return fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
                     dis, x, y, y_sqlen, d, d_offset, ny);
@@ -1060,6 +1080,11 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::AVX2>(
             &fvec_L2sqr_ny_nearest_D8<SIMDLevel::AVX2>);
 }
 
+// Internal linkage: the same template is compiled with other -m flags in
+// the other per-SIMD distances TU, and a shared symbol would let the linker
+// hand one level the other's code.
+namespace {
+
 template <size_t DIM>
 size_t fvec_L2sqr_ny_nearest_y_transposed_D(
         float* /*distances_tmp_buffer*/,
@@ -1176,6 +1201,8 @@ size_t fvec_L2sqr_ny_nearest_y_transposed_D(
     return current_min_index;
 }
 
+} // namespace
+
 template <>
 size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX2>(
         float* distances_tmp_buffer,
@@ -1213,9 +1240,11 @@ int fvec_madd_and_argmin<SIMDLevel::AVX2>(
     return fvec_madd_and_argmin_sse(n, a, bf, b, c);
 }
 
-template <>
-void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
-        const float* x,
+namespace {
+
+template <class QueryTiles>
+void exhaustive_L2sqr_blas_cmax_avx2(
+        QueryTiles& queries,
         const float* y,
         size_t d,
         size_t nx,
@@ -1235,7 +1264,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
     std::unique_ptr<float[]> x_norms(new float[nx]);
     std::unique_ptr<float[]> del2;
 
-    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+    queries.prepare_norms(nx, x_norms.get());
 
     if (!y_norms) {
         float* y_norms2 = new float[ny];
@@ -1250,6 +1279,8 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
             i1 = nx;
         }
 
+        const float* xt = queries.tile(i0, i1, x_norms.get());
+
         res.begin_multiple(i0, i1);
 
         for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
@@ -1269,7 +1300,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
                        &one,
                        y + j0 * d,
                        &di,
-                       x + i0 * d,
+                       xt,
                        &di,
                        &zero,
                        ip_block.get(),
@@ -1427,4 +1458,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
     }
 }
 
+} // namespace
+
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX2>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    Fp32QueryTiles queries{x, d};
+    exhaustive_L2sqr_blas_cmax_avx2(queries, y, d, nx, ny, res, y_norms);
+}
+
+template <>
+void exhaustive_L2sqr_blas_cmax_fp16<SIMDLevel::AVX2>(
+        const uint16_t* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    Fp16QueryTiles queries(x, d, nx);
+    exhaustive_L2sqr_blas_cmax_avx2(queries, y, d, nx, ny, res, y_norms);
+}
+
 } // namespace faiss

data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp

@@ -6,6 +6,7 @@
  */
 
 #include <faiss/utils/distances.h>
+#include <faiss/utils/simd_impl/masked_sum-avx512.h>
 
 #include <immintrin.h>
 
@@ -78,6 +79,19 @@ void fvec_madd<SIMDLevel::AVX512>(
     }
 }
 
+template <>
+float fvec_masked_sum<SIMDLevel::AVX512>(
+        const float* x,
+        const uint8_t* bits,
+        size_t d) {
+    return masked_sum<SIMDLevel::AVX512>(x, bits, d);
+}
+
+// Internal linkage: the same template is compiled with other -m flags in
+// the other per-SIMD distances TU, and a shared symbol would let the linker
+// hand one level the other's code.
+namespace {
+
 template <size_t DIM>
 void fvec_L2sqr_ny_y_transposed_D(
         float* distances,
@@ -151,6 +165,8 @@ void fvec_L2sqr_ny_y_transposed_D(
     }
 }
 
+} // namespace
+
 template <>
 void fvec_L2sqr_ny_transposed<SIMDLevel::AVX512>(
         float* dis,
@@ -169,8 +185,12 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::AVX512>(
     switch (d) {
         DISPATCH(1)
         DISPATCH(2)
+        DISPATCH(3)
         DISPATCH(4)
+        DISPATCH(6)
         DISPATCH(8)
+        DISPATCH(12)
+        DISPATCH(16)
         default:
             return fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
                     dis, x, y, y_sqlen, d, d_offset, ny);
@@ -983,6 +1003,11 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::AVX512>(
             &fvec_L2sqr_ny_nearest_D8<SIMDLevel::AVX512>);
 }
 
+// Internal linkage: the same template is compiled with other -m flags in
+// the other per-SIMD distances TU, and a shared symbol would let the linker
+// hand one level the other's code.
+namespace {
+
 template <size_t DIM>
 size_t fvec_L2sqr_ny_nearest_y_transposed_D(
         float* /* distances_tmp_buffer */,
@@ -1096,6 +1121,8 @@ size_t fvec_L2sqr_ny_nearest_y_transposed_D(
     return current_min_index;
 }
 
+} // namespace
+
 template <>
 size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX512>(
         float* distances_tmp_buffer,
@@ -1133,9 +1160,11 @@ int fvec_madd_and_argmin<SIMDLevel::AVX512>(
     return fvec_madd_and_argmin_sse(n, a, bf, b, c);
 }
 
-template <>
-void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
-        const float* x,
+namespace {
+
+template <class QueryTiles>
+void exhaustive_L2sqr_blas_cmax_avx512(
+        QueryTiles& queries,
         const float* y,
         size_t d,
         size_t nx,
@@ -1154,7 +1183,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
     std::unique_ptr<float[]> x_norms(new float[nx]);
     std::unique_ptr<float[]> del2;
 
-    fvec_norms_L2sqr(x_norms.get(), x, d, nx);
+    queries.prepare_norms(nx, x_norms.get());
 
     if (!y_norms) {
         float* y_norms2 = new float[ny];
@@ -1169,6 +1198,8 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
             i1 = nx;
         }
 
+        const float* xt = queries.tile(i0, i1, x_norms.get());
+
         res.begin_multiple(i0, i1);
 
         for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
@@ -1188,7 +1219,7 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
                        &one,
                        y + j0 * d,
                        &di,
-                       x + i0 * d,
+                       xt,
                        &di,
                        &zero,
                        ip_block.get(),
@@ -1340,4 +1371,32 @@ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
     }
 }
 
+} // namespace
+
+template <>
+void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    Fp32QueryTiles queries{x, d};
+    exhaustive_L2sqr_blas_cmax_avx512(queries, y, d, nx, ny, res, y_norms);
+}
+
+template <>
+void exhaustive_L2sqr_blas_cmax_fp16<SIMDLevel::AVX512>(
+        const uint16_t* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        Top1BlockResultHandler<CMax<float, int64_t>>& res,
+        const float* y_norms) {
+    Fp16QueryTiles queries(x, d, nx);
+    exhaustive_L2sqr_blas_cmax_avx512(queries, y, d, nx, ny, res, y_norms);
+}
+
 } // namespace faiss