Range search bug

Summary:
This diff fixes a serious bug in the range search implementation.

During range search in a flat index, (exhaustive_L2sqr_seq and exhaustive_inner_product_seq) when running in multiple threads, the per-thread results are collected into RangeSearchPartialResult structures.

When the computation is finished, they are aggregated into a RangeSearchResult. In the previous version of the code, this loop was nested into a second loop that is used to check for KeyboardInterrupts. Thus, at each iteration, the results were overwritten.

The fix removes the outer loop. It is most likely useless anyways because the sequential code is called only for a small number of queries, for a larger number the BLAS version is used.

Reviewed By: wickedfoo

Differential Revision: D28486415

fbshipit-source-id: 89a52b17f6ca1ef68fc5e758f0e5a44d0df9fe38

faiss/impl/AuxIndexStructures.cpp

@@ -35,6 +35,9 @@ RangeSearchResult::RangeSearchResult(idx_t nq, bool alloc_lims) : nq(nq) {
 /// called when lims contains the nb of elements result entries
 /// for each query
 void RangeSearchResult::do_allocation() {
+    // works only if all the partial results are aggregated
+    // simulatenously
+    FAISS_THROW_IF_NOT(labels == nullptr && distances == nullptr);
     size_t ofs = 0;
     for (int i = 0; i < nq; i++) {
         size_t n = lims[i];

faiss/utils/distances.cpp

@@ -104,34 +104,25 @@ void exhaustive_inner_product_seq(
         size_t nx,
         size_t ny,
         ResultHandler& res) {
-    size_t check_period = InterruptCallback::get_period_hint(ny * d);
-
-    check_period *= omp_get_max_threads();
-
     using SingleResultHandler = typename ResultHandler::SingleResultHandler;
 
-    for (size_t i0 = 0; i0 < nx; i0 += check_period) {
-        size_t i1 = std::min(i0 + check_period, nx);
-
 #pragma omp parallel
-        {
-            SingleResultHandler resi(res);
+    {
+        SingleResultHandler resi(res);
 #pragma omp for
-            for (int64_t i = i0; i < i1; i++) {
-                const float* x_i = x + i * d;
-                const float* y_j = y;
+        for (int64_t i = 0; i < nx; i++) {
+            const float* x_i = x + i * d;
+            const float* y_j = y;
 
-                resi.begin(i);
+            resi.begin(i);
 
-                for (size_t j = 0; j < ny; j++) {
-                    float ip = fvec_inner_product(x_i, y_j, d);
-                    resi.add_result(ip, j);
-                    y_j += d;
-                }
-                resi.end();
+            for (size_t j = 0; j < ny; j++) {
+                float ip = fvec_inner_product(x_i, y_j, d);
+                resi.add_result(ip, j);
+                y_j += d;
             }
+            resi.end();
         }
-        InterruptCallback::check();
     }
 }
 
@@ -143,32 +134,25 @@ void exhaustive_L2sqr_seq(
         size_t nx,
         size_t ny,
         ResultHandler& res) {
-    size_t check_period = InterruptCallback::get_period_hint(ny * d);
-    check_period *= omp_get_max_threads();
     using SingleResultHandler = typename ResultHandler::SingleResultHandler;
 
-    for (size_t i0 = 0; i0 < nx; i0 += check_period) {
-        size_t i1 = std::min(i0 + check_period, nx);
-
 #pragma omp parallel
-        {
-            SingleResultHandler resi(res);
+    {
+        SingleResultHandler resi(res);
 #pragma omp for
-            for (int64_t i = i0; i < i1; i++) {
-                const float* x_i = x + i * d;
-                const float* y_j = y;
-                resi.begin(i);
-                for (size_t j = 0; j < ny; j++) {
-                    float disij = fvec_L2sqr(x_i, y_j, d);
-                    resi.add_result(disij, j);
-                    y_j += d;
-                }
-                resi.end();
+        for (int64_t i = 0; i < nx; i++) {
+            const float* x_i = x + i * d;
+            const float* y_j = y;
+            resi.begin(i);
+            for (size_t j = 0; j < ny; j++) {
+                float disij = fvec_L2sqr(x_i, y_j, d);
+                resi.add_result(disij, j);
+                y_j += d;
             }
+            resi.end();
         }
-        InterruptCallback::check();
     }
-};
+}
 
 /** Find the nearest neighbors for nx queries in a set of ny vectors */
 template <class ResultHandler>

tests/test_index.py

@@ -1002,3 +1002,24 @@ class TestValidIndexParams(unittest.TestCase):
 
         self.assertEqual(D.shape[0], nq)
         self.assertEqual(D.shape[1], k)
+
+
+class TestLargeRangeSearch(unittest.TestCase):
+
+    def test_range_search(self):
+        # test for https://github.com/facebookresearch/faiss/issues/1889
+        d = 256
+        nq = 16
+        nb = 1000000
+
+        # faiss.cvar.distance_compute_blas_threshold = 10
+        faiss.omp_set_num_threads(1)
+
+        index = faiss.IndexFlatL2(d)
+        xb = np.zeros((nb, d), dtype="float32")
+        index.add(xb)
+
+        xq = np.zeros((nq, d), dtype="float32")
+        lims, D, I = index.range_search(xq, 1.0)
+
+        assert len(D) == len(xb) * len(xq)