mirrors
/
faiss
mirror of https://github.com/facebookresearch/faiss.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Parallelize add_with_id of IndexIVFFlat and IndexIVFFlatDedup (#1805) 

Summary:
This PR parallelized the `add_with_ids` methods of `IndexIVFFlat` and `IndexIVFFlatDedup`. Related to https://github.com/facebookresearch/faiss/issues/1617.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1805

Reviewed By: wickedfoo

Differential Revision: D27619557

Pulled By: mdouze

fbshipit-source-id: 74e0d74c7c51870534372a7ddf6fa0badba2686c
pull/1810/head
Chengqi Deng 2021-04-07 08:08:44 -07:00 committed by Facebook GitHub Bot
parent 267edb120b
commit 213ab22b71
1 changed files with 67 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
faiss/IndexIVFFlat.cpp
View File

@ -9,6 +9,8 @@
#include <faiss/IndexIVFFlat.h> #include <faiss/IndexIVFFlat.h>
#include <omp.h>
#include <cinttypes> #include <cinttypes>
#include <cstdio> #include <cstdio>
@ -47,9 +49,21 @@ void IndexIVFFlat::add_core(
direct_map.check_can_add(xids); direct_map.check_can_add(xids);
int64_t n_add = 0; int64_t n_add = 0;
#pragma omp parallel reduction(+ : n_add)
{
int nt = omp_get_num_threads();
int rank = omp_get_thread_num();
// each thread takes care of a subset of lists
for (size_t i = 0; i < n; i++) { for (size_t i = 0; i < n; i++) {
idx_t id = xids ? xids[i] : ntotal + i;
idx_t list_no = coarse_idx[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; size_t offset;
if (list_no >= 0) { if (list_no >= 0) {
@ -59,8 +73,12 @@ void IndexIVFFlat::add_core(
} else { } else {
offset = 0; offset = 0;
} }
#pragma omp critical
// executed by one thread at a time
direct_map.add_single_id(id, list_no, offset); direct_map.add_single_id(id, list_no, offset);
} }
}
if (verbose) { if (verbose) {
printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64 printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64
@ -249,14 +267,21 @@ void IndexIVFFlatDedup::add_with_ids(
quantizer->assign(na, x, idx); quantizer->assign(na, x, idx);
int64_t n_add = 0, n_dup = 0; int64_t n_add = 0, n_dup = 0;
// TODO make a omp loop with this
#pragma omp parallel reduction(+ : n_add, n_dup)
{
int nt = omp_get_num_threads();
int rank = omp_get_thread_num();
// each thread takes care of a subset of lists
for (size_t i = 0; i < na; i++) { for (size_t i = 0; i < na; i++) {
idx_t id = xids ? xids[i] : ntotal + i;
int64_t list_no = idx[i]; int64_t list_no = idx[i];
if (list_no < 0) { if (list_no < 0 || list_no % nt != rank) {
continue; continue;
} }
idx_t id = xids ? xids[i] : ntotal + i;
const float* xi = x + i * d; const float* xi = x + i * d;
// search if there is already an entry with that id // search if there is already an entry with that id
@ -277,11 +302,16 @@ void IndexIVFFlatDedup::add_with_ids(
// mark equivalence // mark equivalence
idx_t id2 = invlists->get_single_id(list_no, offset); idx_t id2 = invlists->get_single_id(list_no, offset);
std::pair<idx_t, idx_t> pair(id2, id); std::pair<idx_t, idx_t> pair(id2, id);
#pragma omp critical
// executed by one thread at a time
instances.insert(pair); instances.insert(pair);
n_dup++; n_dup++;
} }
n_add++; n_add++;
} }
}
if (verbose) { if (verbose) {
printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64 printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64
" vectors" " vectors"