Using DirectMapAdd to fix IVFFLAT parallel adding (#1842)

Summary: Signed-off-by: shengjun.li <[email protected]> When `direct_map.type` is `DirectMap::Type::Array`, IVFFLAT parallel adding may fail here. ```C++ void DirectMap::add_single_id(idx_t id, idx_t list_no, size_t offset) { ... if (type == Array) { assert(id == array.size()); ``` However, DirectMapAdd has solved it. Pull Request resolved: #1842 Reviewed By: beauby Differential Revision: D28025735 Pulled By: mdouze fbshipit-source-id: 74c423eacd226c9f7b1882532dcde517f6956409
facebookresearch · Apr 27, 2021 · c3842ae · c3842ae
1 parent 8bc7fb6
commit c3842ae
Showing 1 changed file with 9 additions and 15 deletions.
diff --git a/faiss/IndexIVFFlat.cpp b/faiss/IndexIVFFlat.cpp
@@ -50,6 +50,8 @@ void IndexIVFFlat::add_core(
 
     int64_t n_add = 0;
 
+    DirectMapAdd dm_adder(direct_map, n, xids);
+
 #pragma omp parallel reduction(+ : n_add)
     {
         int nt = omp_get_num_threads();
@@ -59,24 +61,16 @@ void IndexIVFFlat::add_core(
         for (size_t i = 0; i < n; i++) {
             idx_t list_no = coarse_idx[i];
 
-            if (list_no % nt != rank) {
-                continue;
-            }
-
-            idx_t id = xids ? xids[i] : ntotal + i;
-            size_t offset;
-
-            if (list_no >= 0) {
+            if (list_no >= 0 && list_no % nt == rank) {
+                idx_t id = xids ? xids[i] : ntotal + i;
                 const float* xi = x + i * d;
-                offset = invlists->add_entry(list_no, id, (const uint8_t*)xi);
+                size_t offset =
+                        invlists->add_entry(list_no, id, (const uint8_t*)xi);
+                dm_adder.add(i, list_no, offset);
                 n_add++;
-            } else {
-                offset = 0;
+            } else if (rank == 0 && list_no == -1) {
+                dm_adder.add(i, -1, 0);
             }
-
-#pragma omp critical
-            // executed by one thread at a time
-            direct_map.add_single_id(id, list_no, offset);
         }
     }