faiss 0.1.0 → 0.1.1

data/vendor/faiss/utils/distances.h

@@ -0,0 +1,243 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+/* All distance functions for L2 and IP distances.
+ * The actual functions are implemented in distances.cpp and distances_simd.cpp */
+
+#pragma once
+
+#include <stdint.h>
+
+#include <faiss/utils/Heap.h>
+
+
+namespace faiss {
+
+ /*********************************************************
+ * Optimized distance/norm/inner prod computations
+ *********************************************************/
+
+
+/// Squared L2 distance between two vectors
+float fvec_L2sqr (
+        const float * x,
+        const float * y,
+        size_t d);
+
+/// inner product
+float  fvec_inner_product (
+        const float * x,
+        const float * y,
+        size_t d);
+
+/// L1 distance
+float fvec_L1 (
+        const float * x,
+        const float * y,
+        size_t d);
+
+float fvec_Linf (
+        const float * x,
+        const float * y,
+        size_t d);
+
+
+/** Compute pairwise distances between sets of vectors
+ *
+ * @param d     dimension of the vectors
+ * @param nq    nb of query vectors
+ * @param nb    nb of database vectors
+ * @param xq    query vectors (size nq * d)
+ * @param xb    database vectros (size nb * d)
+ * @param dis   output distances (size nq * nb)
+ * @param ldq,ldb, ldd strides for the matrices
+ */
+void pairwise_L2sqr (int64_t d,
+                     int64_t nq, const float *xq,
+                     int64_t nb, const float *xb,
+                     float *dis,
+                     int64_t ldq = -1, int64_t ldb = -1, int64_t ldd = -1);
+
+/* compute the inner product between nx vectors x and one y */
+void fvec_inner_products_ny (
+        float * ip,         /* output inner product */
+        const float * x,
+        const float * y,
+        size_t d, size_t ny);
+
+/* compute ny square L2 distance bewteen x and a set of contiguous y vectors */
+void fvec_L2sqr_ny (
+        float * dis,
+        const float * x,
+        const float * y,
+        size_t d, size_t ny);
+
+
+/** squared norm of a vector */
+float fvec_norm_L2sqr (const float * x,
+                       size_t d);
+
+/** compute the L2 norms for a set of vectors
+ *
+ * @param  ip       output norms, size nx
+ * @param  x        set of vectors, size nx * d
+ */
+void fvec_norms_L2 (float * ip, const float * x, size_t d, size_t nx);
+
+/// same as fvec_norms_L2, but computes square norms
+void fvec_norms_L2sqr (float * ip, const float * x, size_t d, size_t nx);
+
+/* L2-renormalize a set of vector. Nothing done if the vector is 0-normed */
+void fvec_renorm_L2 (size_t d, size_t nx, float * x);
+
+
+/* This function exists because the Torch counterpart is extremly slow
+   (not multi-threaded + unexpected overhead even in single thread).
+   It is here to implement the usual property |x-y|^2=|x|^2+|y|^2-2<x|y>  */
+void inner_product_to_L2sqr (float * dis,
+                             const float * nr1,
+                             const float * nr2,
+                             size_t n1, size_t n2);
+
+/***************************************************************************
+ * Compute a subset of  distances
+ ***************************************************************************/
+
+ /* compute the inner product between x and a subset y of ny vectors,
+   whose indices are given by idy.  */
+void fvec_inner_products_by_idx (
+        float * ip,
+        const float * x,
+        const float * y,
+        const int64_t *ids,
+        size_t d, size_t nx, size_t ny);
+
+/* same but for a subset in y indexed by idsy (ny vectors in total) */
+void fvec_L2sqr_by_idx (
+        float * dis,
+        const float * x,
+        const float * y,
+        const int64_t *ids, /* ids of y vecs */
+        size_t d, size_t nx, size_t ny);
+
+
+/** compute dis[j] = L2sqr(x[ix[j]], y[iy[j]]) forall j=0..n-1
+ *
+ * @param x  size (max(ix) + 1, d)
+ * @param y  size (max(iy) + 1, d)
+ * @param ix size n
+ * @param iy size n
+ * @param dis size n
+ */
+void pairwise_indexed_L2sqr (
+        size_t d, size_t n,
+        const float * x, const int64_t *ix,
+        const float * y, const int64_t *iy,
+        float *dis);
+
+/* same for inner product */
+void pairwise_indexed_inner_product (
+        size_t d, size_t n,
+        const float * x, const int64_t *ix,
+        const float * y, const int64_t *iy,
+        float *dis);
+
+/***************************************************************************
+ * KNN functions
+ ***************************************************************************/
+
+// threshold on nx above which we switch to BLAS to compute distances
+extern int distance_compute_blas_threshold;
+
+/** Return the k nearest neighors of each of the nx vectors x among the ny
+ *  vector y, w.r.t to max inner product
+ *
+ * @param x    query vectors, size nx * d
+ * @param y    database vectors, size ny * d
+ * @param res  result array, which also provides k. Sorted on output
+ */
+void knn_inner_product (
+        const float * x,
+        const float * y,
+        size_t d, size_t nx, size_t ny,
+        float_minheap_array_t * res);
+
+/** Same as knn_inner_product, for the L2 distance */
+void knn_L2sqr (
+        const float * x,
+        const float * y,
+        size_t d, size_t nx, size_t ny,
+        float_maxheap_array_t * res);
+
+
+
+/** same as knn_L2sqr, but base_shift[bno] is subtracted to all
+ * computed distances.
+ *
+ * @param base_shift   size ny
+ */
+void knn_L2sqr_base_shift (
+         const float * x,
+         const float * y,
+         size_t d, size_t nx, size_t ny,
+         float_maxheap_array_t * res,
+         const float *base_shift);
+
+/* Find the nearest neighbors for nx queries in a set of ny vectors
+ * indexed by ids. May be useful for re-ranking a pre-selected vector list
+ */
+void knn_inner_products_by_idx (
+        const float * x,
+        const float * y,
+        const int64_t *  ids,
+        size_t d, size_t nx, size_t ny,
+        float_minheap_array_t * res);
+
+void knn_L2sqr_by_idx (const float * x,
+                       const float * y,
+                       const int64_t * ids,
+                       size_t d, size_t nx, size_t ny,
+                       float_maxheap_array_t * res);
+
+/***************************************************************************
+ * Range search
+ ***************************************************************************/
+
+
+
+/// Forward declaration, see AuxIndexStructures.h
+struct RangeSearchResult;
+
+/** Return the k nearest neighors of each of the nx vectors x among the ny
+ *  vector y, w.r.t to max inner product
+ *
+ * @param x      query vectors, size nx * d
+ * @param y      database vectors, size ny * d
+ * @param radius search radius around the x vectors
+ * @param result result structure
+ */
+void range_search_L2sqr (
+        const float * x,
+        const float * y,
+        size_t d, size_t nx, size_t ny,
+        float radius,
+        RangeSearchResult *result);
+
+/// same as range_search_L2sqr for the inner product similarity
+void range_search_inner_product (
+        const float * x,
+        const float * y,
+        size_t d, size_t nx, size_t ny,
+        float radius,
+        RangeSearchResult *result);
+
+
+
+
+} // namespace faiss

data/vendor/faiss/utils/distances_simd.cpp

@@ -0,0 +1,809 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/utils/distances.h>
+
+#include <cstdio>
+#include <cassert>
+#include <cstring>
+#include <cmath>
+
+#ifdef __SSE__
+#include <immintrin.h>
+#endif
+
+#ifdef __aarch64__
+#include <arm_neon.h>
+#endif
+
+#include <omp.h>
+
+namespace faiss {
+
+#ifdef __AVX__
+#define USE_AVX
+#endif
+
+
+/*********************************************************
+ * Optimized distance computations
+ *********************************************************/
+
+
+/* Functions to compute:
+   - L2 distance between 2 vectors
+   - inner product between 2 vectors
+   - L2 norm of a vector
+
+   The functions should probably not be invoked when a large number of
+   vectors are be processed in batch (in which case Matrix multiply
+   is faster), but may be useful for comparing vectors isolated in
+   memory.
+
+   Works with any vectors of any dimension, even unaligned (in which
+   case they are slower).
+
+*/
+
+
+/*********************************************************
+ * Reference implementations
+ */
+
+
+float fvec_L2sqr_ref (const float * x,
+                     const float * y,
+                     size_t d)
+{
+    size_t i;
+    float res = 0;
+    for (i = 0; i < d; i++) {
+        const float tmp = x[i] - y[i];
+       res += tmp * tmp;
+    }
+    return res;
+}
+
+float fvec_L1_ref (const float * x,
+                   const float * y,
+                   size_t d)
+{
+    size_t i;
+    float res = 0;
+    for (i = 0; i < d; i++) {
+        const float tmp = x[i] - y[i];
+        res += fabs(tmp);
+    }
+    return res;
+}
+
+float fvec_Linf_ref (const float * x,
+                     const float * y,
+                     size_t d)
+{
+    size_t i;
+    float res = 0;
+    for (i = 0; i < d; i++) {
+      res = fmax(res, fabs(x[i] - y[i]));
+    }
+    return res;
+}
+
+float fvec_inner_product_ref (const float * x,
+                             const float * y,
+                             size_t d)
+{
+    size_t i;
+    float res = 0;
+    for (i = 0; i < d; i++)
+       res += x[i] * y[i];
+    return res;
+}
+
+float fvec_norm_L2sqr_ref (const float *x, size_t d)
+{
+    size_t i;
+    double res = 0;
+    for (i = 0; i < d; i++)
+       res += x[i] * x[i];
+    return res;
+}
+
+
+void fvec_L2sqr_ny_ref (float * dis,
+                    const float * x,
+                    const float * y,
+                    size_t d, size_t ny)
+{
+    for (size_t i = 0; i < ny; i++) {
+        dis[i] = fvec_L2sqr (x, y, d);
+        y += d;
+    }
+}
+
+
+
+
+/*********************************************************
+ * SSE and AVX implementations
+ */
+
+#ifdef __SSE__
+
+// reads 0 <= d < 4 floats as __m128
+static inline __m128 masked_read (int d, const float *x)
+{
+    assert (0 <= d && d < 4);
+    __attribute__((__aligned__(16))) float buf[4] = {0, 0, 0, 0};
+    switch (d) {
+      case 3:
+        buf[2] = x[2];
+      case 2:
+        buf[1] = x[1];
+      case 1:
+        buf[0] = x[0];
+    }
+    return _mm_load_ps (buf);
+    // cannot use AVX2 _mm_mask_set1_epi32
+}
+
+float fvec_norm_L2sqr (const float *  x,
+                      size_t d)
+{
+    __m128 mx;
+    __m128 msum1 = _mm_setzero_ps();
+
+    while (d >= 4) {
+        mx = _mm_loadu_ps (x); x += 4;
+        msum1 = _mm_add_ps (msum1, _mm_mul_ps (mx, mx));
+        d -= 4;
+    }
+
+    mx = masked_read (d, x);
+    msum1 = _mm_add_ps (msum1, _mm_mul_ps (mx, mx));
+
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    return  _mm_cvtss_f32 (msum1);
+}
+
+namespace {
+
+float sqr (float x) {
+    return x * x;
+}
+
+
+void fvec_L2sqr_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 tmp, accu;
+        tmp = x0 - _mm_loadu_ps (y); y += 4;
+        accu = tmp * tmp;
+        dis[i] = _mm_cvtss_f32 (accu);
+        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++] = sqr(x0s - *y++);
+    }
+}
+
+
+void fvec_L2sqr_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 tmp, accu;
+        tmp = x0 - _mm_loadu_ps (y); y += 4;
+        accu = tmp * tmp;
+        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] = sqr(x[0] - y[0]) + sqr(x[1] - y[1]);
+    }
+}
+
+
+
+void fvec_L2sqr_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 tmp, accu;
+        tmp = x0 - _mm_loadu_ps (y); y += 4;
+        accu = tmp * tmp;
+        accu = _mm_hadd_ps (accu, accu);
+        accu = _mm_hadd_ps (accu, accu);
+        dis[i] = _mm_cvtss_f32 (accu);
+    }
+}
+
+
+void fvec_L2sqr_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 tmp, accu;
+        tmp = x0 - _mm_loadu_ps (y); y += 4;
+        accu = tmp * tmp;
+        tmp = x1 - _mm_loadu_ps (y); y += 4;
+        accu += tmp * tmp;
+        accu = _mm_hadd_ps (accu, accu);
+        accu = _mm_hadd_ps (accu, accu);
+        dis[i] = _mm_cvtss_f32 (accu);
+    }
+}
+
+
+void fvec_L2sqr_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 tmp, accu;
+        tmp = x0 - _mm_loadu_ps (y); y += 4;
+        accu = tmp * tmp;
+        tmp = x1 - _mm_loadu_ps (y); y += 4;
+        accu += tmp * tmp;
+        tmp = x2 - _mm_loadu_ps (y); y += 4;
+        accu += tmp * tmp;
+        accu = _mm_hadd_ps (accu, accu);
+        accu = _mm_hadd_ps (accu, accu);
+        dis[i] = _mm_cvtss_f32 (accu);
+    }
+}
+
+
+} // anonymous namespace
+
+void fvec_L2sqr_ny (float * dis, const float * x,
+                        const float * y, size_t d, size_t ny) {
+    // optimized for a few special cases
+    switch(d) {
+    case 1:
+        fvec_L2sqr_ny_D1 (dis, x, y, ny);
+        return;
+    case 2:
+        fvec_L2sqr_ny_D2 (dis, x, y, ny);
+        return;
+    case 4:
+        fvec_L2sqr_ny_D4 (dis, x, y, ny);
+        return;
+    case 8:
+        fvec_L2sqr_ny_D8 (dis, x, y, ny);
+        return;
+    case 12:
+        fvec_L2sqr_ny_D12 (dis, x, y, ny);
+        return;
+    default:
+        fvec_L2sqr_ny_ref (dis, x, y, d, ny);
+        return;
+    }
+}
+
+
+
+#endif
+
+#ifdef USE_AVX
+
+// reads 0 <= d < 8 floats as __m256
+static inline __m256 masked_read_8 (int d, const float *x)
+{
+    assert (0 <= d && d < 8);
+    if (d < 4) {
+        __m256 res = _mm256_setzero_ps ();
+        res = _mm256_insertf128_ps (res, masked_read (d, x), 0);
+        return res;
+    } else {
+        __m256 res = _mm256_setzero_ps ();
+        res = _mm256_insertf128_ps (res, _mm_loadu_ps (x), 0);
+        res = _mm256_insertf128_ps (res, masked_read (d - 4, x + 4), 1);
+        return res;
+    }
+}
+
+float fvec_inner_product (const float * x,
+                          const float * y,
+                          size_t d)
+{
+    __m256 msum1 = _mm256_setzero_ps();
+
+    while (d >= 8) {
+        __m256 mx = _mm256_loadu_ps (x); x += 8;
+        __m256 my = _mm256_loadu_ps (y); y += 8;
+        msum1 = _mm256_add_ps (msum1, _mm256_mul_ps (mx, my));
+        d -= 8;
+    }
+
+    __m128 msum2 = _mm256_extractf128_ps(msum1, 1);
+    msum2 +=       _mm256_extractf128_ps(msum1, 0);
+
+    if (d >= 4) {
+        __m128 mx = _mm_loadu_ps (x); x += 4;
+        __m128 my = _mm_loadu_ps (y); y += 4;
+        msum2 = _mm_add_ps (msum2, _mm_mul_ps (mx, my));
+        d -= 4;
+    }
+
+    if (d > 0) {
+        __m128 mx = masked_read (d, x);
+        __m128 my = masked_read (d, y);
+        msum2 = _mm_add_ps (msum2, _mm_mul_ps (mx, my));
+    }
+
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    return  _mm_cvtss_f32 (msum2);
+}
+
+float fvec_L2sqr (const float * x,
+                 const float * y,
+                 size_t d)
+{
+    __m256 msum1 = _mm256_setzero_ps();
+
+    while (d >= 8) {
+        __m256 mx = _mm256_loadu_ps (x); x += 8;
+        __m256 my = _mm256_loadu_ps (y); y += 8;
+        const __m256 a_m_b1 = mx - my;
+        msum1 += a_m_b1 * a_m_b1;
+        d -= 8;
+    }
+
+    __m128 msum2 = _mm256_extractf128_ps(msum1, 1);
+    msum2 +=       _mm256_extractf128_ps(msum1, 0);
+
+    if (d >= 4) {
+        __m128 mx = _mm_loadu_ps (x); x += 4;
+        __m128 my = _mm_loadu_ps (y); y += 4;
+        const __m128 a_m_b1 = mx - my;
+        msum2 += a_m_b1 * a_m_b1;
+        d -= 4;
+    }
+
+    if (d > 0) {
+        __m128 mx = masked_read (d, x);
+        __m128 my = masked_read (d, y);
+        __m128 a_m_b1 = mx - my;
+        msum2 += a_m_b1 * a_m_b1;
+    }
+
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    return  _mm_cvtss_f32 (msum2);
+}
+
+float fvec_L1 (const float * x, const float * y, size_t d)
+{
+    __m256 msum1 = _mm256_setzero_ps();
+    __m256 signmask = __m256(_mm256_set1_epi32 (0x7fffffffUL));
+
+    while (d >= 8) {
+        __m256 mx = _mm256_loadu_ps (x); x += 8;
+        __m256 my = _mm256_loadu_ps (y); y += 8;
+        const __m256 a_m_b = mx - my;
+        msum1 += _mm256_and_ps(signmask, a_m_b);
+        d -= 8;
+    }
+
+    __m128 msum2 = _mm256_extractf128_ps(msum1, 1);
+    msum2 +=       _mm256_extractf128_ps(msum1, 0);
+    __m128 signmask2 = __m128(_mm_set1_epi32 (0x7fffffffUL));
+
+    if (d >= 4) {
+        __m128 mx = _mm_loadu_ps (x); x += 4;
+        __m128 my = _mm_loadu_ps (y); y += 4;
+        const __m128 a_m_b = mx - my;
+        msum2 += _mm_and_ps(signmask2, a_m_b);
+        d -= 4;
+    }
+
+    if (d > 0) {
+        __m128 mx = masked_read (d, x);
+        __m128 my = masked_read (d, y);
+        __m128 a_m_b = mx - my;
+        msum2 += _mm_and_ps(signmask2, a_m_b);
+    }
+
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    msum2 = _mm_hadd_ps (msum2, msum2);
+    return  _mm_cvtss_f32 (msum2);
+}
+
+float fvec_Linf (const float * x, const float * y, size_t d)
+{
+    __m256 msum1 = _mm256_setzero_ps();
+    __m256 signmask = __m256(_mm256_set1_epi32 (0x7fffffffUL));
+
+    while (d >= 8) {
+        __m256 mx = _mm256_loadu_ps (x); x += 8;
+        __m256 my = _mm256_loadu_ps (y); y += 8;
+        const __m256 a_m_b = mx - my;
+        msum1 = _mm256_max_ps(msum1, _mm256_and_ps(signmask, a_m_b));
+        d -= 8;
+    }
+
+    __m128 msum2 = _mm256_extractf128_ps(msum1, 1);
+    msum2 = _mm_max_ps (msum2, _mm256_extractf128_ps(msum1, 0));
+    __m128 signmask2 = __m128(_mm_set1_epi32 (0x7fffffffUL));
+
+    if (d >= 4) {
+        __m128 mx = _mm_loadu_ps (x); x += 4;
+        __m128 my = _mm_loadu_ps (y); y += 4;
+        const __m128 a_m_b = mx - my;
+        msum2 = _mm_max_ps(msum2, _mm_and_ps(signmask2, a_m_b));
+        d -= 4;
+    }
+
+    if (d > 0) {
+        __m128 mx = masked_read (d, x);
+        __m128 my = masked_read (d, y);
+        __m128 a_m_b = mx - my;
+        msum2 = _mm_max_ps(msum2, _mm_and_ps(signmask2, a_m_b));
+    }
+
+    msum2 = _mm_max_ps(_mm_movehl_ps(msum2, msum2), msum2);
+    msum2 = _mm_max_ps(msum2, _mm_shuffle_ps (msum2, msum2, 1));
+    return  _mm_cvtss_f32 (msum2);
+}
+
+#elif defined(__SSE__) // But not AVX
+
+float fvec_L1 (const float * x, const float * y, size_t d)
+{
+    return fvec_L1_ref (x, y, d);
+}
+
+float fvec_Linf (const float * x, const float * y, size_t d)
+{
+    return fvec_Linf_ref (x, y, d);
+}
+
+
+float fvec_L2sqr (const float * x,
+                 const float * y,
+                 size_t d)
+{
+    __m128 msum1 = _mm_setzero_ps();
+
+    while (d >= 4) {
+        __m128 mx = _mm_loadu_ps (x); x += 4;
+        __m128 my = _mm_loadu_ps (y); y += 4;
+        const __m128 a_m_b1 = mx - my;
+        msum1 += a_m_b1 * a_m_b1;
+        d -= 4;
+    }
+
+    if (d > 0) {
+        // add the last 1, 2 or 3 values
+        __m128 mx = masked_read (d, x);
+        __m128 my = masked_read (d, y);
+        __m128 a_m_b1 = mx - my;
+        msum1 += a_m_b1 * a_m_b1;
+    }
+
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    return  _mm_cvtss_f32 (msum1);
+}
+
+
+float fvec_inner_product (const float * x,
+                         const float * y,
+                         size_t d)
+{
+    __m128 mx, my;
+    __m128 msum1 = _mm_setzero_ps();
+
+    while (d >= 4) {
+        mx = _mm_loadu_ps (x); x += 4;
+        my = _mm_loadu_ps (y); y += 4;
+        msum1 = _mm_add_ps (msum1, _mm_mul_ps (mx, my));
+        d -= 4;
+    }
+
+    // add the last 1, 2, or 3 values
+    mx = masked_read (d, x);
+    my = masked_read (d, y);
+    __m128 prod = _mm_mul_ps (mx, my);
+
+    msum1 = _mm_add_ps (msum1, prod);
+
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    msum1 = _mm_hadd_ps (msum1, msum1);
+    return  _mm_cvtss_f32 (msum1);
+}
+
+#elif defined(__aarch64__)
+
+
+float fvec_L2sqr (const float * x,
+                  const float * y,
+                  size_t d)
+{
+    if (d & 3) return fvec_L2sqr_ref (x, y, d);
+    float32x4_t accu = vdupq_n_f32 (0);
+    for (size_t i = 0; i < d; i += 4) {
+        float32x4_t xi = vld1q_f32 (x + i);
+        float32x4_t yi = vld1q_f32 (y + i);
+        float32x4_t sq = vsubq_f32 (xi, yi);
+        accu = vfmaq_f32 (accu, sq, sq);
+    }
+    float32x4_t a2 = vpaddq_f32 (accu, accu);
+    return vdups_laneq_f32 (a2, 0) + vdups_laneq_f32 (a2, 1);
+}
+
+float fvec_inner_product (const float * x,
+                          const float * y,
+                          size_t d)
+{
+    if (d & 3) return fvec_inner_product_ref (x, y, d);
+    float32x4_t accu = vdupq_n_f32 (0);
+    for (size_t i = 0; i < d; i += 4) {
+        float32x4_t xi = vld1q_f32 (x + i);
+        float32x4_t yi = vld1q_f32 (y + i);
+        accu = vfmaq_f32 (accu, xi, yi);
+    }
+    float32x4_t a2 = vpaddq_f32 (accu, accu);
+    return vdups_laneq_f32 (a2, 0) + vdups_laneq_f32 (a2, 1);
+}
+
+float fvec_norm_L2sqr (const float *x, size_t d)
+{
+    if (d & 3) return fvec_norm_L2sqr_ref (x, d);
+    float32x4_t accu = vdupq_n_f32 (0);
+    for (size_t i = 0; i < d; i += 4) {
+        float32x4_t xi = vld1q_f32 (x + i);
+        accu = vfmaq_f32 (accu, xi, xi);
+    }
+    float32x4_t a2 = vpaddq_f32 (accu, accu);
+    return vdups_laneq_f32 (a2, 0) + vdups_laneq_f32 (a2, 1);
+}
+
+// not optimized for ARM
+void fvec_L2sqr_ny (float * dis, const float * x,
+                        const float * y, size_t d, size_t ny) {
+    fvec_L2sqr_ny_ref (dis, x, y, d, ny);
+}
+
+float fvec_L1 (const float * x, const float * y, size_t d)
+{
+    return fvec_L1_ref (x, y, d);
+}
+
+float fvec_Linf (const float * x, const float * y, size_t d)
+{
+    return fvec_Linf_ref (x, y, d);
+}
+
+
+#else
+// scalar implementation
+
+float fvec_L2sqr (const float * x,
+                  const float * y,
+                  size_t d)
+{
+    return fvec_L2sqr_ref (x, y, d);
+}
+
+float fvec_L1 (const float * x, const float * y, size_t d)
+{
+    return fvec_L1_ref (x, y, d);
+}
+
+float fvec_Linf (const float * x, const float * y, size_t d)
+{
+    return fvec_Linf_ref (x, y, d);
+}
+
+float fvec_inner_product (const float * x,
+                             const float * y,
+                             size_t d)
+{
+    return fvec_inner_product_ref (x, y, d);
+}
+
+float fvec_norm_L2sqr (const float *x, size_t d)
+{
+    return fvec_norm_L2sqr_ref (x, d);
+}
+
+void fvec_L2sqr_ny (float * dis, const float * x,
+                        const float * y, size_t d, size_t ny) {
+    fvec_L2sqr_ny_ref (dis, x, y, d, ny);
+}
+
+
+#endif
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+/***************************************************************************
+ * heavily optimized table computations
+ ***************************************************************************/
+
+
+static inline void fvec_madd_ref (size_t n, const float *a,
+                           float bf, const float *b, float *c) {
+    for (size_t i = 0; i < n; i++)
+        c[i] = a[i] + bf * b[i];
+}
+
+#ifdef __SSE__
+
+static 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++;
+    }
+}
+
+void fvec_madd (size_t n, const float *a,
+                float bf, const float *b, float *c)
+{
+    if ((n & 3) == 0 &&
+        ((((long)a) | ((long)b) | ((long)c)) & 15) == 0)
+        fvec_madd_sse (n, a, bf, b, c);
+    else
+        fvec_madd_ref (n, a, bf, b, c);
+}
+
+#else
+
+void fvec_madd (size_t n, const float *a,
+                float bf, const float *b, float *c)
+{
+    fvec_madd_ref (n, a, bf, b, c);
+}
+
+#endif
+
+static inline int fvec_madd_and_argmin_ref (size_t n, const float *a,
+                                         float bf, const float *b, float *c) {
+    float vmin = 1e20;
+    int imin = -1;
+
+    for (size_t i = 0; i < n; i++) {
+        c[i] = a[i] + bf * b[i];
+        if (c[i] < vmin) {
+            vmin = c[i];
+            imin = i;
+        }
+    }
+    return imin;
+}
+
+#ifdef __SSE__
+
+static inline int fvec_madd_and_argmin_sse (
+        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 = (__m128i)_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 = (__m128i)_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 = (__m128i)_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);
+}
+
+
+int fvec_madd_and_argmin (size_t n, const float *a,
+                          float bf, const float *b, float *c)
+{
+    if ((n & 3) == 0 &&
+        ((((long)a) | ((long)b) | ((long)c)) & 15) == 0)
+        return fvec_madd_and_argmin_sse (n, a, bf, b, c);
+    else
+        return fvec_madd_and_argmin_ref (n, a, bf, b, c);
+}
+
+#else
+
+int fvec_madd_and_argmin (size_t n, const float *a,
+                          float bf, const float *b, float *c)
+{
+  return fvec_madd_and_argmin_ref (n, a, bf, b, c);
+}
+
+#endif
+
+
+
+
+} // namespace faiss