faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/utils/distances.h

@@ -18,6 +18,8 @@
 
 namespace faiss {
 
+struct IDSelector;
+
 /*********************************************************
  * Optimized distance/norm/inner prod computations
  *********************************************************/
@@ -40,7 +42,7 @@ float fvec_Linf(const float* x, const float* y, size_t d);
  * @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 xb    database vectors (size nb * d)
  * @param dis   output distances (size nq * nb)
  * @param ldq,ldb, ldd strides for the matrices
  */
@@ -63,7 +65,7 @@ void fvec_inner_products_ny(
         size_t d,
         size_t ny);
 
-/* compute ny square L2 distance bewteen x and a set of contiguous y vectors */
+/* compute ny square L2 distance between x and a set of contiguous y vectors */
 void fvec_L2sqr_ny(
         float* dis,
         const float* x,
@@ -71,6 +73,29 @@ void fvec_L2sqr_ny(
         size_t d,
         size_t ny);
 
+/* compute ny square L2 distance between x and a set of contiguous y vectors
+   and return the index of the nearest vector.
+   return 0 if ny == 0. */
+size_t fvec_L2sqr_ny_nearest(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t ny);
+
+/* compute ny square L2 distance between x and a set of transposed contiguous
+   y vectors and return the index of the nearest vector.
+   squared lengths of y should be provided as well
+   return 0 if ny == 0. */
+size_t fvec_L2sqr_ny_nearest_y_transposed(
+        float* distances_tmp_buffer,
+        const float* x,
+        const float* y,
+        const float* y_sqlen,
+        size_t d,
+        size_t d_offset,
+        size_t ny);
+
 /** squared norm of a vector */
 float fvec_norm_L2sqr(const float* x, size_t d);
 
@@ -87,7 +112,7 @@ void fvec_norms_L2sqr(float* norms, 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
+/* This function exists because the Torch counterpart is extremely 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(
@@ -97,6 +122,39 @@ void inner_product_to_L2sqr(
         size_t n1,
         size_t n2);
 
+/*********************************************************
+ * Vector to vector functions
+ *********************************************************/
+
+/** compute c := a + b for vectors
+ *
+ * c and a can overlap, c and b can overlap
+ *
+ * @param a size d
+ * @param b size d
+ * @param c size d
+ */
+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
+ *
+ * @param a size d
+ * @param c size d
+ */
+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
+ *
+ * @param a size d
+ * @param b size d
+ * @param c size d
+ */
+void fvec_sub(size_t d, const float* a, const float* b, float* c);
+
 /***************************************************************************
  * Compute a subset of  distances
  ***************************************************************************/
@@ -165,11 +223,28 @@ FAISS_API extern int distance_compute_blas_database_bs;
 FAISS_API extern int distance_compute_min_k_reservoir;
 
 /** Return the k nearest neighors of each of the nx vectors x among the ny
- *  vector y, w.r.t to max inner product
+ *  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 heap structure, 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,
+        const IDSelector* sel = nullptr);
+
+/**  Return the k nearest neighors of each of the nx vectors x among the ny
+ *  vector y, for the inner product metric.
  *
  * @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
+ * @param distances  output distances, size nq * k
+ * @param indexes    output vector ids, size nq * k
  */
 void knn_inner_product(
         const float* x,
@@ -177,10 +252,18 @@ void knn_inner_product(
         size_t d,
         size_t nx,
         size_t ny,
-        float_minheap_array_t* res);
+        size_t k,
+        float* distances,
+        int64_t* indexes,
+        const IDSelector* sel = nullptr);
 
-/** Same as knn_inner_product, for the L2 distance
- *  @param y_norm2    norms for the y vectors (nullptr or size ny)
+/** Return the k nearest neighors of each of the nx vectors x among the ny
+ *  vector y, for the L2 distance
+ * @param x    query vectors, size nx * d
+ * @param y    database vectors, size ny * d
+ * @param res  result heap strcture, which also provides k. Sorted on output
+ * @param y_norm2    (optional) norms for the y vectors (nullptr or size ny)
+ * @param sel  search in this subset of vectors
  */
 void knn_L2sqr(
         const float* x,
@@ -189,28 +272,74 @@ void knn_L2sqr(
         size_t nx,
         size_t ny,
         float_maxheap_array_t* res,
-        const float* y_norm2 = nullptr);
+        const float* y_norm2 = nullptr,
+        const IDSelector* sel = nullptr);
 
-/* Find the nearest neighbors for nx queries in a set of ny vectors
+/**  Return the k nearest neighors of each of the nx vectors x among the ny
+ *  vector y, for the L2 distance
+ *
+ * @param x    query vectors, size nx * d
+ * @param y    database vectors, size ny * d
+ * @param distances  output distances, size nq * k
+ * @param indexes    output vector ids, size nq * k
+ * @param y_norm2    (optional) norms for the y vectors (nullptr or size ny)
+ * @param sel  search in this subset of vectors
+ */
+void knn_L2sqr(
+        const float* x,
+        const float* y,
+        size_t d,
+        size_t nx,
+        size_t ny,
+        size_t k,
+        float* distances,
+        int64_t* indexes,
+        const float* y_norm2 = nullptr,
+        const IDSelector* sel = nullptr);
+
+/** Find the max inner product neighbors for nx queries in a set of ny vectors
  * indexed by ids. May be useful for re-ranking a pre-selected vector list
+ *
+ * @param x    query vectors, size nx * d
+ * @param y    database vectors, size (max(ids) + 1) * d
+ * @param ids  subset of database vectors to consider, size (nx, nsubset)
+ * @param res  result structure
+ * @param ld_ids stride for the ids array. -1: use nsubset, 0: all queries
+ * process the same subset
  */
 void knn_inner_products_by_idx(
         const float* x,
         const float* y,
-        const int64_t* ids,
+        const int64_t* subset,
         size_t d,
         size_t nx,
-        size_t ny,
-        float_minheap_array_t* res);
+        size_t nsubset,
+        size_t k,
+        float* vals,
+        int64_t* ids,
+        int64_t ld_ids = -1);
 
+/** 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
+ *
+ * @param x    query vectors, size nx * d
+ * @param y    database vectors, size (max(ids) + 1) * d
+ * @param subset subset of database vectors to consider, size (nx, nsubset)
+ * @param res  rIDesult structure
+ * @param ld_subset stride for the subset array. -1: use nsubset, 0: all queries
+ * process the same subset
+ */
 void knn_L2sqr_by_idx(
         const float* x,
         const float* y,
-        const int64_t* ids,
+        const int64_t* subset,
         size_t d,
         size_t nx,
-        size_t ny,
-        float_maxheap_array_t* res);
+        size_t nsubset,
+        size_t k,
+        float* vals,
+        int64_t* ids,
+        int64_t ld_subset = -1);
 
 /***************************************************************************
  * Range search
@@ -234,7 +363,8 @@ void range_search_L2sqr(
         size_t nx,
         size_t ny,
         float radius,
-        RangeSearchResult* result);
+        RangeSearchResult* result,
+        const IDSelector* sel = nullptr);
 
 /// same as range_search_L2sqr for the inner product similarity
 void range_search_inner_product(
@@ -244,7 +374,8 @@ void range_search_inner_product(
         size_t nx,
         size_t ny,
         float radius,
-        RangeSearchResult* result);
+        RangeSearchResult* result,
+        const IDSelector* sel = nullptr);
 
 /***************************************************************************
  * PQ tables computations