faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/utils/simd_impl/distances_sse-inl.h

@@ -0,0 +1,392 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/utils/distances.h>
+
+#include <immintrin.h>
+
+namespace faiss {
+
+[[maybe_unused]] inline void fvec_madd_sse(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    n >>= 2;
+    __m128 bf4 = _mm_set_ps1(bf);
+    __m128* a4 = (__m128*)a;
+    __m128* b4 = (__m128*)b;
+    __m128* c4 = (__m128*)c;
+
+    while (n--) {
+        *c4 = _mm_add_ps(*a4, _mm_mul_ps(bf4, *b4));
+        b4++;
+        a4++;
+        c4++;
+    }
+}
+
+/// helper function
+inline float horizontal_sum(const __m128 v) {
+    // say, v is [x0, x1, x2, x3]
+
+    // v0 is [x2, x3, ..., ...]
+    const __m128 v0 = _mm_shuffle_ps(v, v, _MM_SHUFFLE(0, 0, 3, 2));
+    // v1 is [x0 + x2, x1 + x3, ..., ...]
+    const __m128 v1 = _mm_add_ps(v, v0);
+    // v2 is [x1 + x3, ..., .... ,...]
+    __m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
+    // v3 is [x0 + x1 + x2 + x3, ..., ..., ...]
+    const __m128 v3 = _mm_add_ps(v1, v2);
+    // return v3[0]
+    return _mm_cvtss_f32(v3);
+}
+
+/// Function that does a component-wise operation between x and y
+/// to compute inner products
+struct ElementOpIP {
+    static float op(float x, float y) {
+        return x * y;
+    }
+
+    static __m128 op(__m128 x, __m128 y) {
+        return _mm_mul_ps(x, y);
+    }
+};
+
+/// Function that does a component-wise operation between x and y
+/// to compute L2 distances. ElementOp can then be used in the fvec_op_ny
+/// functions below
+struct ElementOpL2 {
+    static float op(float x, float y) {
+        float tmp = x - y;
+        return tmp * tmp;
+    }
+
+    static __m128 op(__m128 x, __m128 y) {
+        __m128 tmp = _mm_sub_ps(x, y);
+        return _mm_mul_ps(tmp, tmp);
+    }
+};
+
+template <class ElementOp>
+void fvec_op_ny_D1(float* dis, const float* x, const float* y, size_t ny) {
+    float x0s = x[0];
+    __m128 x0 = _mm_set_ps(x0s, x0s, x0s, x0s);
+
+    size_t i;
+    for (i = 0; i + 3 < ny; i += 4) {
+        __m128 accu = ElementOp::op(x0, _mm_loadu_ps(y));
+        y += 4;
+        dis[i] = _mm_cvtss_f32(accu);
+        __m128 tmp = _mm_shuffle_ps(accu, accu, 1);
+        dis[i + 1] = _mm_cvtss_f32(tmp);
+        tmp = _mm_shuffle_ps(accu, accu, 2);
+        dis[i + 2] = _mm_cvtss_f32(tmp);
+        tmp = _mm_shuffle_ps(accu, accu, 3);
+        dis[i + 3] = _mm_cvtss_f32(tmp);
+    }
+    while (i < ny) { // handle non-multiple-of-4 case
+        dis[i++] = ElementOp::op(x0s, *y++);
+    }
+}
+
+template <class ElementOp>
+void fvec_op_ny_D2(float* dis, const float* x, const float* y, size_t ny) {
+    __m128 x0 = _mm_set_ps(x[1], x[0], x[1], x[0]);
+
+    size_t i;
+    for (i = 0; i + 1 < ny; i += 2) {
+        __m128 accu = ElementOp::op(x0, _mm_loadu_ps(y));
+        y += 4;
+        accu = _mm_hadd_ps(accu, accu);
+        dis[i] = _mm_cvtss_f32(accu);
+        accu = _mm_shuffle_ps(accu, accu, 3);
+        dis[i + 1] = _mm_cvtss_f32(accu);
+    }
+    if (i < ny) { // handle odd case
+        dis[i] = ElementOp::op(x[0], y[0]) + ElementOp::op(x[1], y[1]);
+    }
+}
+
+template <class ElementOp>
+void fvec_op_ny_D4(float* dis, const float* x, const float* y, size_t ny) {
+    __m128 x0 = _mm_loadu_ps(x);
+
+    for (size_t i = 0; i < ny; i++) {
+        __m128 accu = ElementOp::op(x0, _mm_loadu_ps(y));
+        y += 4;
+        dis[i] = horizontal_sum(accu);
+    }
+}
+
+template <class ElementOp>
+void fvec_op_ny_D8(float* dis, const float* x, const float* y, size_t ny) {
+    __m128 x0 = _mm_loadu_ps(x);
+    __m128 x1 = _mm_loadu_ps(x + 4);
+
+    for (size_t i = 0; i < ny; i++) {
+        __m128 accu = ElementOp::op(x0, _mm_loadu_ps(y));
+        y += 4;
+        accu = _mm_add_ps(accu, ElementOp::op(x1, _mm_loadu_ps(y)));
+        y += 4;
+        accu = _mm_hadd_ps(accu, accu);
+        accu = _mm_hadd_ps(accu, accu);
+        dis[i] = _mm_cvtss_f32(accu);
+    }
+}
+
+template <class ElementOp>
+void fvec_op_ny_D12(float* dis, const float* x, const float* y, size_t ny) {
+    __m128 x0 = _mm_loadu_ps(x);
+    __m128 x1 = _mm_loadu_ps(x + 4);
+    __m128 x2 = _mm_loadu_ps(x + 8);
+
+    for (size_t i = 0; i < ny; i++) {
+        __m128 accu = ElementOp::op(x0, _mm_loadu_ps(y));
+        y += 4;
+        accu = _mm_add_ps(accu, ElementOp::op(x1, _mm_loadu_ps(y)));
+        y += 4;
+        accu = _mm_add_ps(accu, ElementOp::op(x2, _mm_loadu_ps(y)));
+        y += 4;
+        dis[i] = horizontal_sum(accu);
+    }
+}
+
+template <class ElementOpIP>
+void fvec_inner_products_ny_ref(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+#define DISPATCH(dval)                                  \
+    case dval:                                          \
+        fvec_op_ny_D##dval<ElementOpIP>(dis, x, y, ny); \
+        return;
+
+    switch (d) {
+        DISPATCH(1)
+        DISPATCH(2)
+        DISPATCH(4)
+        DISPATCH(8)
+        DISPATCH(12)
+        default:
+            fvec_inner_products_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
+            return;
+    }
+#undef DISPATCH
+}
+
+template <class ElementOpL2>
+void fvec_L2sqr_ny_ref(
+        float* dis,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny) {
+    // optimized for a few special cases
+
+#define DISPATCH(dval)                                  \
+    case dval:                                          \
+        fvec_op_ny_D##dval<ElementOpL2>(dis, x, y, ny); \
+        return;
+
+    switch (d) {
+        DISPATCH(1)
+        DISPATCH(2)
+        DISPATCH(4)
+        DISPATCH(8)
+        DISPATCH(12)
+        default:
+            fvec_L2sqr_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
+            return;
+    }
+#undef DISPATCH
+}
+
+template <SIMDLevel>
+size_t fvec_L2sqr_ny_nearest_D2(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t ny);
+
+template <SIMDLevel>
+size_t fvec_L2sqr_ny_nearest_D4(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t ny);
+
+template <SIMDLevel>
+size_t fvec_L2sqr_ny_nearest_D8(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t ny);
+
+template <SIMDLevel SIMD>
+size_t fvec_L2sqr_ny_nearest_x86(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny,
+        size_t (*fvec_L2sqr_ny_nearest_D2_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t) = &fvec_L2sqr_ny_nearest_D2<SIMD>,
+        size_t (*fvec_L2sqr_ny_nearest_D4_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t) = &fvec_L2sqr_ny_nearest_D4<SIMD>,
+        size_t (*fvec_L2sqr_ny_nearest_D8_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t) = &fvec_L2sqr_ny_nearest_D8<SIMD>);
+
+template <SIMDLevel SIMD>
+size_t fvec_L2sqr_ny_nearest_x86(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny,
+        size_t (*fvec_L2sqr_ny_nearest_D2_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t),
+        size_t (*fvec_L2sqr_ny_nearest_D4_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t),
+        size_t (*fvec_L2sqr_ny_nearest_D8_func)(
+                float*,
+                const float*,
+                const float*,
+                size_t)) {
+    switch (d) {
+        case 2:
+            return fvec_L2sqr_ny_nearest_D2_func(
+                    distances_tmp_buffer, x, y, ny);
+        case 4:
+            return fvec_L2sqr_ny_nearest_D4_func(
+                    distances_tmp_buffer, x, y, ny);
+        case 8:
+            return fvec_L2sqr_ny_nearest_D8_func(
+                    distances_tmp_buffer, x, y, ny);
+    }
+
+    return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
+            distances_tmp_buffer, x, y, d, ny);
+}
+
+template <SIMDLevel SIMD>
+inline size_t fvec_L2sqr_ny_nearest(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny);
+
+inline int fvec_madd_and_argmin_sse_ref(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    n >>= 2;
+    __m128 bf4 = _mm_set_ps1(bf);
+    __m128 vmin4 = _mm_set_ps1(1e20);
+    __m128i imin4 = _mm_set1_epi32(-1);
+    __m128i idx4 = _mm_set_epi32(3, 2, 1, 0);
+    __m128i inc4 = _mm_set1_epi32(4);
+    __m128* a4 = (__m128*)a;
+    __m128* b4 = (__m128*)b;
+    __m128* c4 = (__m128*)c;
+
+    while (n--) {
+        __m128 vc4 = _mm_add_ps(*a4, _mm_mul_ps(bf4, *b4));
+        *c4 = vc4;
+        __m128i mask = _mm_castps_si128(_mm_cmpgt_ps(vmin4, vc4));
+        // imin4 = _mm_blendv_epi8 (imin4, idx4, mask); // slower!
+
+        imin4 = _mm_or_si128(
+                _mm_and_si128(mask, idx4), _mm_andnot_si128(mask, imin4));
+        vmin4 = _mm_min_ps(vmin4, vc4);
+        b4++;
+        a4++;
+        c4++;
+        idx4 = _mm_add_epi32(idx4, inc4);
+    }
+
+    // 4 values -> 2
+    {
+        idx4 = _mm_shuffle_epi32(imin4, 3 << 2 | 2);
+        __m128 vc4 = _mm_shuffle_ps(vmin4, vmin4, 3 << 2 | 2);
+        __m128i mask = _mm_castps_si128(_mm_cmpgt_ps(vmin4, vc4));
+        imin4 = _mm_or_si128(
+                _mm_and_si128(mask, idx4), _mm_andnot_si128(mask, imin4));
+        vmin4 = _mm_min_ps(vmin4, vc4);
+    }
+    // 2 values -> 1
+    {
+        idx4 = _mm_shuffle_epi32(imin4, 1);
+        __m128 vc4 = _mm_shuffle_ps(vmin4, vmin4, 1);
+        __m128i mask = _mm_castps_si128(_mm_cmpgt_ps(vmin4, vc4));
+        imin4 = _mm_or_si128(
+                _mm_and_si128(mask, idx4), _mm_andnot_si128(mask, imin4));
+        // vmin4 = _mm_min_ps (vmin4, vc4);
+    }
+    return _mm_cvtsi128_si32(imin4);
+}
+
+inline int fvec_madd_and_argmin_sse(
+        size_t n,
+        const float* a,
+        float bf,
+        const float* b,
+        float* c) {
+    if ((n & 3) == 0 &&
+        ((((uintptr_t)a) | ((uintptr_t)b) | ((uintptr_t)c)) & 15) == 0) {
+        return fvec_madd_and_argmin_sse_ref(n, a, bf, b, c);
+    } else {
+        return fvec_madd_and_argmin<SIMDLevel::NONE>(n, a, bf, b, c);
+    }
+}
+
+// reads 0 <= d < 4 floats as __m128
+inline __m128 masked_read(int d, const float* x) {
+    assert(0 <= d && d < 4);
+    ALIGNED(16) float buf[4] = {0, 0, 0, 0};
+    switch (d) {
+        case 3:
+            buf[2] = x[2];
+            [[fallthrough]];
+        case 2:
+            buf[1] = x[1];
+            [[fallthrough]];
+        case 1:
+            buf[0] = x[0];
+            break;
+        default:
+            break;
+    }
+    return _mm_load_ps(buf);
+    // cannot use AVX2 _mm_mask_set1_epi32
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/utils/{distances_fused/simdlib_based.h → simd_impl/exhaustive_L2sqr_blas_cmax.h}

@@ -8,17 +8,14 @@
 #pragma once
 
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/impl/platform_macros.h>
-
-#include <faiss/utils/Heap.h>
-
-#if defined(__AVX2__) || defined(__aarch64__)
+#include <faiss/utils/simd_levels.h>
 
 namespace faiss {
 
-// Returns true if the fused kernel is available and the data was processed.
-// Returns false if the fused kernel is not available.
-bool exhaustive_L2sqr_fused_cmax_simdlib(
+/// BLAS-accelerated exhaustive L2 search for the k=1 (top-1) case.
+/// Specializations live in the per-SIMD translation units under simd_impl/.
+template <SIMDLevel>
+void exhaustive_L2sqr_blas_cmax(
         const float* x,
         const float* y,
         size_t d,
@@ -28,5 +25,3 @@ bool exhaustive_L2sqr_fused_cmax_simdlib(
         const float* y_norms);
 
 } // namespace faiss
-
-#endif