convert simdlib in distances_simd.cpp (#4884)

Summary:
Pull Request resolved: #4884

Templatize fvec_sub, fvec_add, and compute_PQ_dis_tables_dsub2 on SIMDLevel.

Move these 256-bit simdlib-based functions from distances_simd.cpp into a new header simd_impl/distances_simdlib256.h, templatized on THE_SIMDLEVEL. Add with_simd_level_256bit() to simd_dispatch.h that maps AVX512->AVX2 and ARM_SVE->ARM_NEON using simd256_level_selector. Dispatch wrappers use with_simd_level_256bit with lambdas.

Reviewed By: algoriddle

Differential Revision: D95570445

fbshipit-source-id: 557b6546e9526c40a44f7165584e96a9e39a3b89

Author: mdouze

faiss/CMakeLists.txt

@@ -322,6 +322,7 @@ set(FAISS_HEADERS
   utils/hamming_distance/avx2-inl.h
   utils/hamming_distance/avx512-inl.h
   utils/simd_impl/distances_autovec-inl.h
+  utils/simd_impl/distances_simdlib256.h
   utils/simd_impl/distances_sse-inl.h
 )
 

faiss/impl/simd_dispatch.h

@@ -136,4 +136,21 @@ inline auto with_simd_level(LambdaType&& action) {
     DISPATCH_SIMDLevel(action.template operator());
 }
 
+/**
+ * Like with_simd_level, but maps to the 256-bit SIMD equivalent:
+ *   AVX512, AVX512_SPR -> AVX2
+ *   ARM_SVE -> ARM_NEON
+ *   AVX2, ARM_NEON, NONE -> unchanged
+ *
+ * Use for functions implemented with simd8float32 (256-bit) operations
+ * that don't have dedicated AVX512 or SVE implementations.
+ */
+template <typename LambdaType>
+inline auto with_simd_level_256bit(LambdaType&& action) {
+    return with_simd_level([&]<SIMDLevel level>() {
+        constexpr SIMDLevel level256 = simd256_level_selector<level>::value;
+        return action.template operator()<level256>();
+    });
+}
+
 } // namespace faiss

faiss/utils/distances.cpp

@@ -172,6 +172,30 @@ int fvec_madd_and_argmin(
     return fvec_madd_and_argmin_dispatch(n, a, bf, b, c);
 }
 
+void fvec_sub(size_t d, const float* a, const float* b, float* c) {
+    fvec_sub_dispatch(d, a, b, c);
+}
+
+void fvec_add(size_t d, const float* a, const float* b, float* c) {
+    fvec_add_dispatch(d, a, b, c);
+}
+
+void fvec_add(size_t d, const float* a, float b, float* c) {
+    fvec_add_scalar_dispatch(d, a, b, c);
+}
+
+void compute_PQ_dis_tables_dsub2(
+        size_t d,
+        size_t ksub,
+        const float* all_centroids,
+        size_t nx,
+        const float* x,
+        bool is_inner_product,
+        float* dis_tables) {
+    compute_PQ_dis_tables_dsub2_dispatch(
+            d, ksub, all_centroids, nx, x, is_inner_product, dis_tables);
+}
+
 /***************************************************************************
  * Matrix/vector ops
  ***************************************************************************/

faiss/utils/distances.h

@@ -261,6 +261,9 @@ void inner_product_to_L2sqr(
  */
 void fvec_add(size_t d, const float* a, const float* b, float* c);
 
+template <SIMDLevel>
+void fvec_add(size_t d, const float* a, const float* b, float* c);
+
 /** compute c := a + b for a, c vectors and b a scalar
  *
  * c and a can overlap
@@ -270,6 +273,9 @@ void fvec_add(size_t d, const float* a, const float* b, float* c);
  */
 void fvec_add(size_t d, const float* a, float b, float* c);
 
+template <SIMDLevel>
+void fvec_add(size_t d, const float* a, float b, float* c);
+
 /** compute c := a - b for vectors
  *
  * c and a can overlap, c and b can overlap
@@ -280,6 +286,9 @@ void fvec_add(size_t d, const float* a, float b, float* c);
  */
 void fvec_sub(size_t d, const float* a, const float* b, float* c);
 
+template <SIMDLevel>
+void fvec_sub(size_t d, const float* a, const float* b, float* c);
+
 /***************************************************************************
  * Compute a subset of  distances
  ***************************************************************************/
@@ -542,6 +551,16 @@ void compute_PQ_dis_tables_dsub2(
         bool is_inner_product,
         float* dis_tables);
 
+template <SIMDLevel>
+void compute_PQ_dis_tables_dsub2(
+        size_t d,
+        size_t ksub,
+        const float* centroids,
+        size_t nx,
+        const float* x,
+        bool is_inner_product,
+        float* dis_tables);
+
 /***************************************************************************
  * Templatized versions of distance functions
  ***************************************************************************/

faiss/utils/distances_dispatch.h

@@ -167,4 +167,45 @@ inline int fvec_madd_and_argmin_dispatch(
     DISPATCH_SIMDLevel(fvec_madd_and_argmin, n, a, bf, b, c);
 }
 
+inline void fvec_sub_dispatch(
+        size_t d,
+        const float* a,
+        const float* b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_sub<level>(d, a, b, c); });
+}
+
+inline void fvec_add_dispatch(
+        size_t d,
+        const float* a,
+        const float* b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
+}
+
+inline void fvec_add_scalar_dispatch(
+        size_t d,
+        const float* a,
+        float b,
+        float* c) {
+    with_simd_level_256bit(
+            [&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
+}
+
+inline void compute_PQ_dis_tables_dsub2_dispatch(
+        size_t d,
+        size_t ksub,
+        const float* centroids,
+        size_t nx,
+        const float* x,
+        bool is_inner_product,
+        float* dis_tables) {
+    with_simd_level_256bit([&]<SIMDLevel level>() {
+        compute_PQ_dis_tables_dsub2<level>(
+                d, ksub, centroids, nx, x, is_inner_product, dis_tables);
+    });
+}
+
 } // namespace faiss

faiss/utils/distances_simd.cpp

@@ -22,6 +22,10 @@
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 
+#define THE_SIMDLEVEL SIMDLevel::NONE
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/distances_simdlib256.h>
+
 namespace faiss {
 
 /*******
@@ -168,177 +172,3 @@ int fvec_madd_and_argmin<SIMDLevel::NONE>(
 }
 
 } // namespace faiss
-
-namespace faiss {
-
-/***************************************************************************
- * PQ tables computations
- ***************************************************************************/
-
-namespace {
-
-/// compute the IP for dsub = 2 for 8 centroids and 4 sub-vectors at a time
-template <bool is_inner_product>
-void pq2_8cents_table(
-        const simd8float32 centroids[8],
-        const simd8float32 x,
-        float* out,
-        size_t ldo,
-        size_t nout = 4) {
-    simd8float32 ips[4];
-
-    for (int i = 0; i < 4; i++) {
-        simd8float32 p1, p2;
-        if (is_inner_product) {
-            p1 = x * centroids[2 * i];
-            p2 = x * centroids[2 * i + 1];
-        } else {
-            p1 = (x - centroids[2 * i]);
-            p1 = p1 * p1;
-            p2 = (x - centroids[2 * i + 1]);
-            p2 = p2 * p2;
-        }
-        ips[i] = hadd(p1, p2);
-    }
-
-    simd8float32 ip02a = geteven(ips[0], ips[1]);
-    simd8float32 ip02b = geteven(ips[2], ips[3]);
-    simd8float32 ip0 = getlow128(ip02a, ip02b);
-    simd8float32 ip2 = gethigh128(ip02a, ip02b);
-
-    simd8float32 ip13a = getodd(ips[0], ips[1]);
-    simd8float32 ip13b = getodd(ips[2], ips[3]);
-    simd8float32 ip1 = getlow128(ip13a, ip13b);
-    simd8float32 ip3 = gethigh128(ip13a, ip13b);
-
-    switch (nout) {
-        case 4:
-            ip3.storeu(out + 3 * ldo);
-            [[fallthrough]];
-        case 3:
-            ip2.storeu(out + 2 * ldo);
-            [[fallthrough]];
-        case 2:
-            ip1.storeu(out + 1 * ldo);
-            [[fallthrough]];
-        case 1:
-            ip0.storeu(out);
-    }
-}
-
-simd8float32 load_simd8float32_partial(const float* x, int n) {
-    ALIGNED(32) float tmp[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-    float* wp = tmp;
-    for (int i = 0; i < n; i++) {
-        *wp++ = *x++;
-    }
-    return simd8float32(tmp);
-}
-
-} // anonymous namespace
-
-void compute_PQ_dis_tables_dsub2(
-        size_t d,
-        size_t ksub,
-        const float* all_centroids,
-        size_t nx,
-        const float* x,
-        bool is_inner_product,
-        float* dis_tables) {
-    size_t M = d / 2;
-    FAISS_THROW_IF_NOT(ksub % 8 == 0);
-
-    for (size_t m0 = 0; m0 < M; m0 += 4) {
-        int m1 = std::min(M, m0 + 4);
-        for (int k0 = 0; k0 < ksub; k0 += 8) {
-            simd8float32 centroids[8];
-            for (int k = 0; k < 8; k++) {
-                ALIGNED(32) float centroid[8];
-                size_t wp = 0;
-                size_t rp = (m0 * ksub + k + k0) * 2;
-                for (int m = m0; m < m1; m++) {
-                    centroid[wp++] = all_centroids[rp];
-                    centroid[wp++] = all_centroids[rp + 1];
-                    rp += 2 * ksub;
-                }
-                centroids[k] = simd8float32(centroid);
-            }
-            for (size_t i = 0; i < nx; i++) {
-                simd8float32 xi;
-                if (m1 == m0 + 4) {
-                    xi.loadu(x + i * d + m0 * 2);
-                } else {
-                    xi = load_simd8float32_partial(
-                            x + i * d + m0 * 2, 2 * (m1 - m0));
-                }
-
-                if (is_inner_product) {
-                    pq2_8cents_table<true>(
-                            centroids,
-                            xi,
-                            dis_tables + (i * M + m0) * ksub + k0,
-                            ksub,
-                            m1 - m0);
-                } else {
-                    pq2_8cents_table<false>(
-                            centroids,
-                            xi,
-                            dis_tables + (i * M + m0) * ksub + k0,
-                            ksub,
-                            m1 - m0);
-                }
-            }
-        }
-    }
-}
-
-/*********************************************************
- * Vector to vector functions
- *********************************************************/
-
-void fvec_sub(size_t d, const float* a, const float* b, float* c) {
-    size_t i;
-    for (i = 0; i + 7 < d; i += 8) {
-        simd8float32 ci, ai, bi;
-        ai.loadu(a + i);
-        bi.loadu(b + i);
-        ci = ai - bi;
-        ci.storeu(c + i);
-    }
-    // finish non-multiple of 8 remainder
-    for (; i < d; i++) {
-        c[i] = a[i] - b[i];
-    }
-}
-
-void fvec_add(size_t d, const float* a, const float* b, float* c) {
-    size_t i;
-    for (i = 0; i + 7 < d; i += 8) {
-        simd8float32 ci, ai, bi;
-        ai.loadu(a + i);
-        bi.loadu(b + i);
-        ci = ai + bi;
-        ci.storeu(c + i);
-    }
-    // finish non-multiple of 8 remainder
-    for (; i < d; i++) {
-        c[i] = a[i] + b[i];
-    }
-}
-
-void fvec_add(size_t d, const float* a, float b, float* c) {
-    size_t i;
-    simd8float32 bv(b);
-    for (i = 0; i + 7 < d; i += 8) {
-        simd8float32 ci, ai;
-        ai.loadu(a + i);
-        ci = ai + bv;
-        ci.storeu(c + i);
-    }
-    // finish non-multiple of 8 remainder
-    for (; i < d; i++) {
-        c[i] = a[i] + b;
-    }
-}
-
-} // namespace faiss

faiss/utils/simd_impl/distances_aarch64.cpp

@@ -15,6 +15,9 @@
 #define AUTOVEC_LEVEL SIMDLevel::ARM_NEON
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 
+#define THE_SIMDLEVEL SIMDLevel::ARM_NEON
+#include <faiss/utils/simd_impl/distances_simdlib256.h>
+
 namespace faiss {
 
 template <>

faiss/utils/simd_impl/distances_avx2.cpp

@@ -13,6 +13,10 @@
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_autovec-inl.h>
 
+#define THE_SIMDLEVEL SIMDLevel::AVX2
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/simd_impl/distances_simdlib256.h>
+
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)
 #include <faiss/utils/simd_impl/distances_sse-inl.h>
 // NOLINTNEXTLINE(facebook-hte-InlineHeader)