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

104
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,19 +49,35 @@ 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;
for (size_t i = 0; i < n; i++) {
idx_t id = xids ? xids[i] : ntotal + i;
idx_t list_no = coarse_idx[i];
size_t offset;
if (list_no >= 0) { #pragma omp parallel reduction(+ : n_add)
const float* xi = x + i * d; {
offset = invlists->add_entry(list_no, id, (const uint8_t*)xi); int nt = omp_get_num_threads();
n_add++; int rank = omp_get_thread_num();
} else {
offset = 0; // each thread takes care of a subset of lists
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) {
const float* xi = x + i * d;
offset = invlists->add_entry(list_no, id, (const uint8_t*)xi);
n_add++;
} else {
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) {
@ -249,38 +267,50 @@ 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
for (size_t i = 0; i < na; i++) {
idx_t id = xids ? xids[i] : ntotal + i;
int64_t list_no = idx[i];
if (list_no < 0) { #pragma omp parallel reduction(+ : n_add, n_dup)
continue; {
} int nt = omp_get_num_threads();
const float* xi = x + i * d; int rank = omp_get_thread_num();
// search if there is already an entry with that id // each thread takes care of a subset of lists
InvertedLists::ScopedCodes codes(invlists, list_no); for (size_t i = 0; i < na; i++) {
int64_t list_no = idx[i];
int64_t n = invlists->list_size(list_no); if (list_no < 0 || list_no % nt != rank) {
int64_t offset = -1; continue;
for (int64_t o = 0; o < n; o++) {
if (!memcmp(codes.get() + o * code_size, xi, code_size)) {
offset = o;
break;
} }
}
if (offset == -1) { // not found idx_t id = xids ? xids[i] : ntotal + i;
invlists->add_entry(list_no, id, (const uint8_t*)xi); const float* xi = x + i * d;
} else {
// mark equivalence // search if there is already an entry with that id
idx_t id2 = invlists->get_single_id(list_no, offset); InvertedLists::ScopedCodes codes(invlists, list_no);
std::pair<idx_t, idx_t> pair(id2, id);
instances.insert(pair); int64_t n = invlists->list_size(list_no);
n_dup++; int64_t offset = -1;
for (int64_t o = 0; o < n; o++) {
if (!memcmp(codes.get() + o * code_size, xi, code_size)) {
offset = o;
break;
}
}
if (offset == -1) { // not found
invlists->add_entry(list_no, id, (const uint8_t*)xi);
} else {
// mark equivalence
idx_t id2 = invlists->get_single_id(list_no, offset);
std::pair<idx_t, idx_t> pair(id2, id);
#pragma omp critical
// executed by one thread at a time
instances.insert(pair);
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