Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/matthijs/github_faiss/faiss/MetaIndexes.cpp
1 
2 /**
3  * Copyright (c) 2015-present, Facebook, Inc.
4  * All rights reserved.
5  *
6  * This source code is licensed under the CC-by-NC license found in the
7  * LICENSE file in the root directory of this source tree.
8  */
9 
10 // Copyright 2004-present Facebook. All Rights Reserved
11 // -*- c++ -*-
12 
13 #include "MetaIndexes.h"
14 
15 #include <pthread.h>
16 
17 #include <cstdio>
18 
19 #include "FaissAssert.h"
20 #include "Heap.h"
21 
22 namespace faiss {
23 
24 /*****************************************************
25  * IndexIDMap implementation
26  *******************************************************/
27 
28 IndexIDMap::IndexIDMap (Index *index):
29  index (index),
30  own_fields (false)
31 {
32  FAISS_ASSERT (index->ntotal == 0 || !"index must be empty on input");
33  is_trained = index->is_trained;
34  metric_type = index->metric_type;
35  verbose = index->verbose;
36  d = index->d;
37  set_typename ();
38 }
39 
40 void IndexIDMap::add (idx_t, const float *)
41 {
42  FAISS_ASSERT (!"add does not make sense with IndexIDMap, "
43  "use add_with_ids");
44 }
45 
46 
47 void IndexIDMap::train (idx_t n, const float *x)
48 {
49  index->train (n, x);
50  is_trained = index->is_trained;
51 }
52 
54 {
55  index->reset ();
56  ntotal = 0;
57 }
58 
59 
60 void IndexIDMap::add_with_ids (idx_t n, const float * x, const long *xids)
61 {
62  index->add (n, x);
63  for (idx_t i = 0; i < n; i++)
64  id_map.push_back (xids[i]);
65  ntotal = index->ntotal;
66 }
67 
68 
69 void IndexIDMap::search (idx_t n, const float *x, idx_t k,
70  float *distances, idx_t *labels) const
71 {
72  index->search (n, x, k, distances, labels);
73  idx_t *li = labels;
74  for (idx_t i = 0; i < n * k; i++) {
75  li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
76  }
77 }
78 
79 
80 
81 IndexIDMap::~IndexIDMap ()
82 {
83  if (own_fields) delete index;
84 }
85 
86 void IndexIDMap::set_typename ()
87 {
88  index_typename = "IDMap[" + index->index_typename + "]";
89 }
90 
91 
92 /*****************************************************
93  * IndexShards implementation
94  *******************************************************/
95 
96 // subroutines
97 namespace {
98 
99 
100 typedef Index::idx_t idx_t;
101 
102 
103 template<class Job>
104 struct Thread {
105  Job job;
106  pthread_t thread;
107 
108  Thread () {}
109 
110  explicit Thread (const Job & job): job(job) {}
111 
112  void start () {
113  pthread_create (&thread, nullptr, run, this);
114  }
115 
116  void wait () {
117  pthread_join (thread, nullptr);
118  }
119 
120  static void * run (void *arg) {
121  static_cast<Thread*> (arg)->job.run();
122  return nullptr;
123  }
124 
125 };
126 
127 
128 /// callback + thread management to train 1 shard
129 struct TrainJob {
130  IndexShards *index; // the relevant index
131  int no; // shard number
132  idx_t n; // train points
133  const float *x;
134 
135  void run ()
136  {
137  if (index->verbose)
138  printf ("begin train shard %d on %ld points\n", no, n);
139  index->shard_indexes [no]->train(n, x);
140  if (index->verbose)
141  printf ("end train shard %d\n", no);
142  }
143 
144 };
145 
146 struct AddJob {
147  IndexShards *index; // the relevant index
148  int no; // shard number
149 
150  idx_t n;
151  const float *x;
152  const idx_t *ids;
153 
154  void run ()
155  {
156  if (index->verbose)
157  printf ("begin add shard %d on %ld points\n", no, n);
158  if (ids)
159  index->shard_indexes[no]->add_with_ids (n, x, ids);
160  else
161  index->shard_indexes[no]->add (n, x);
162  if (index->verbose)
163  printf ("end add shard %d on %ld points\n", no, n);
164  }
165 };
166 
167 
168 
169 /// callback + thread management to query in 1 shard
170 struct QueryJob {
171  const IndexShards *index; // the relevant index
172  int no; // shard number
173 
174  // query params
175  idx_t n;
176  const float *x;
177  idx_t k;
178  float *distances;
179  idx_t *labels;
180 
181 
182  void run ()
183  {
184  if (index->verbose)
185  printf ("begin query shard %d on %ld points\n", no, n);
186  index->shard_indexes [no]->search (n, x, k,
187  distances, labels);
188  if (index->verbose)
189  printf ("end query shard %d\n", no);
190  }
191 
192 
193 };
194 
195 
196 
197 
198 // add translation to all valid labels
199 void translate_labels (long n, idx_t *labels, long translation)
200 {
201  if (translation == 0) return;
202  for (long i = 0; i < n; i++) {
203  if(labels[i] < 0) return;
204  labels[i] += translation;
205  }
206 }
207 
208 
209 /** merge result tables from several shards.
210  * @param all_distances size nshard * n * k
211  * @param all_labels idem
212  * @param translartions label translations to apply, size nshard
213  */
214 
215 template <class C>
216 void merge_tables (long n, long k, long nshard,
217  float *distances, idx_t *labels,
218  const float *all_distances,
219  idx_t *all_labels,
220  const long *translations)
221 {
222  if(k == 0) {
223  return;
224  }
225 
226  long stride = n * k;
227 #pragma omp parallel
228  {
229  std::vector<int> buf (2 * nshard);
230  int * pointer = buf.data();
231  int * shard_ids = pointer + nshard;
232  std::vector<float> buf2 (nshard);
233  float * heap_vals = buf2.data();
234 #pragma omp for
235  for (long i = 0; i < n; i++) {
236  // the heap maps values to the shard where they are
237  // produced.
238  const float *D_in = all_distances + i * k;
239  const idx_t *I_in = all_labels + i * k;
240  int heap_size = 0;
241 
242  for (long s = 0; s < nshard; s++) {
243  pointer[s] = 0;
244  if (I_in[stride * s] >= 0)
245  heap_push<C> (++heap_size, heap_vals, shard_ids,
246  D_in[stride * s], s);
247  }
248 
249  float *D = distances + i * k;
250  idx_t *I = labels + i * k;
251 
252  for (int j = 0; j < k; j++) {
253  if (heap_size == 0) {
254  I[j] = -1;
255  D[j] = C::neutral();
256  } else {
257  // pop best element
258  int s = shard_ids[0];
259  int & p = pointer[s];
260  D[j] = heap_vals[0];
261  I[j] = I_in[stride * s + p] + translations[s];
262 
263  heap_pop<C> (heap_size--, heap_vals, shard_ids);
264  p++;
265  if (p < k && I_in[stride * s + p] >= 0)
266  heap_push<C> (++heap_size, heap_vals, shard_ids,
267  D_in[stride * s + p], s);
268  }
269  }
270  }
271  }
272 }
273 
274 
275 };
276 
277 
278 
279 
280 IndexShards::IndexShards (idx_t d, bool threaded, bool successive_ids):
281  Index (d), own_fields (false),
282  threaded (threaded), successive_ids (successive_ids)
283 {
284 
285 }
286 
287 void IndexShards::add_shard (Index *idx)
288 {
289  shard_indexes.push_back (idx);
290  sync_with_shard_indexes ();
291 }
292 
293 void IndexShards::sync_with_shard_indexes ()
294 {
295  if (shard_indexes.empty()) return;
296  Index * index0 = shard_indexes[0];
297  d = index0->d;
298  metric_type = index0->metric_type;
299  is_trained = index0->is_trained;
300  ntotal = index0->ntotal;
301  for (int i = 1; i < shard_indexes.size(); i++) {
302  Index * index = shard_indexes[i];
303  FAISS_ASSERT (metric_type == index->metric_type);
304  FAISS_ASSERT (d == index->d);
305  ntotal += index->ntotal;
306  }
307 }
308 
309 
310 void IndexShards::train (idx_t n, const float *x)
311 {
312 
313  // pre-alloc because we don't want reallocs
314  std::vector<Thread<TrainJob > > tss (shard_indexes.size());
315  int nt = 0;
316  for (int i = 0; i < shard_indexes.size(); i++) {
317  if(!shard_indexes[i]->is_trained) {
318  TrainJob ts = {this, i, n, x};
319  if (threaded) {
320  tss[nt] = Thread<TrainJob> (ts);
321  tss[nt++].start();
322  } else {
323  ts.run();
324  }
325  }
326  }
327  for (int i = 0; i < nt; i++) {
328  tss[i].wait();
329  }
330  sync_with_shard_indexes ();
331 }
332 
333 void IndexShards::add (idx_t n, const float *x)
334 {
335  add_with_ids (n, x, nullptr);
336 }
337 
338  /**
339  * Cases (successive_ids, xids):
340  * - true, non-NULL ERROR: it makes no sense to pass in ids and
341  * request them to be shifted
342  * - true, NULL OK, but should be called only once (calls add()
343  * on sub-indexes).
344  * - false, non-NULL OK: will call add_with_ids with passed in xids
345  * distributed evenly over shards
346  * - false, NULL OK: will call add_with_ids on each sub-index,
347  * starting at ntotal
348  */
349 
350 void IndexShards::add_with_ids (idx_t n, const float * x, const long *xids)
351 {
352 
353  FAISS_ASSERT(!(successive_ids && xids) ||
354  !"It makes no sense to pass in ids and request them to be shifted");
355 
356  if (successive_ids) {
357  FAISS_ASSERT(!xids ||
358  !"It makes no sense to pass in ids and request them to be shifted");
359  FAISS_ASSERT(ntotal == 0 ||
360  !"when adding to IndexShards with sucessive_ids, only add() "
361  "in a single pass is supported");
362  }
363 
364  long nshard = shard_indexes.size();
365  const long *ids = xids;
366  if (!ids && !successive_ids) {
367  long *aids = new long[n];
368  for (long i = 0; i < n; i++)
369  aids[i] = ntotal + i;
370  ids = aids;
371  }
372 
373  std::vector<Thread<AddJob > > asa (shard_indexes.size());
374  int nt = 0;
375  for (int i = 0; i < nshard; i++) {
376  long i0 = i * n / nshard;
377  long i1 = (i + 1) * n / nshard;
378 
379  AddJob as = {this, i,
380  i1 - i0, x + i0 * d,
381  ids ? ids + i0 : nullptr};
382  if (threaded) {
383  asa[nt] = Thread<AddJob>(as);
384  asa[nt++].start();
385  } else {
386  as.run();
387  }
388  }
389  for (int i = 0; i < nt; i++) {
390  asa[i].wait();
391  }
392  if (ids != xids) delete [] ids;
393  ntotal += n;
394 }
395 
396 
397 
398 
399 
401 {
402  for (int i = 0; i < shard_indexes.size(); i++) {
403  shard_indexes[i]->reset ();
404  }
405  sync_with_shard_indexes ();
406 }
407 
409  idx_t n, const float *x, idx_t k,
410  float *distances, idx_t *labels) const
411 {
412  long nshard = shard_indexes.size();
413  float *all_distances = new float [nshard * k * n];
414  idx_t *all_labels = new idx_t [nshard * k * n];
415 
416 #if 1
417 
418  // pre-alloc because we don't want reallocs
419  std::vector<Thread<QueryJob> > qss (nshard);
420  for (int i = 0; i < nshard; i++) {
421  QueryJob qs = {
422  this, i, n, x, k,
423  all_distances + i * k * n,
424  all_labels + i * k * n
425  };
426  if (threaded) {
427  qss[i] = Thread<QueryJob> (qs);
428  qss[i].start();
429  } else {
430  qs.run();
431  }
432  }
433 
434  if (threaded) {
435  for (int i = 0; i < qss.size(); i++) {
436  qss[i].wait();
437  }
438  }
439 #else
440 
441  // pre-alloc because we don't want reallocs
442  std::vector<QueryJob> qss (nshard);
443  for (int i = 0; i < nshard; i++) {
444  QueryJob qs = {
445  this, i, n, x, k,
446  all_distances + i * k * n,
447  all_labels + i * k * n
448  };
449  if (threaded) {
450  qss[i] = qs;
451  } else {
452  qs.run();
453  }
454  }
455 
456  if (threaded) {
457 #pragma omp parallel for
458  for (int i = 0; i < qss.size(); i++) {
459  qss[i].run();
460  }
461  }
462 
463 #endif
464  std::vector<long> translations (nshard, 0);
465  if (successive_ids) {
466  translations[0] = 0;
467  for (int s = 0; s + 1 < nshard; s++)
468  translations [s + 1] = translations [s] +
469  shard_indexes [s]->ntotal;
470  }
471 
472  if (metric_type == METRIC_L2) {
473  merge_tables< CMin<float, int> > (
474  n, k, nshard, distances, labels,
475  all_distances, all_labels, translations.data ());
476  } else {
477  merge_tables< CMax<float, int> > (
478  n, k, nshard, distances, labels,
479  all_distances, all_labels, translations.data ());
480  }
481  delete [] all_distances;
482  delete [] all_labels;
483 }
484 
485 
486 void IndexShards::set_typename ()
487 {
488 
489 }
490 
491 IndexShards::~IndexShards ()
492 {
493  if (own_fields) {
494  for (int s = 0; s < shard_indexes.size(); s++)
495  delete shard_indexes [s];
496  }
497 }
498 
499 
500 /*****************************************************
501  * IndexSplitVectors implementation
502  *******************************************************/
503 
504 
505 IndexSplitVectors::IndexSplitVectors (idx_t d, bool threaded):
506  Index (d), own_fields (false),
507  threaded (threaded), sum_d (0)
508 {
509 
510 }
511 
512 void IndexSplitVectors::add_sub_index (Index *index)
513 {
514  sub_indexes.push_back (index);
515  sync_with_sub_indexes ();
516 }
517 
518 void IndexSplitVectors::sync_with_sub_indexes ()
519 {
520  if (sub_indexes.empty()) return;
521  Index * index0 = sub_indexes[0];
522  sum_d = index0->d;
523  metric_type = index0->metric_type;
524  is_trained = index0->is_trained;
525  ntotal = index0->ntotal;
526  for (int i = 1; i < sub_indexes.size(); i++) {
527  Index * index = sub_indexes[i];
528  FAISS_ASSERT (metric_type == index->metric_type);
529  FAISS_ASSERT (ntotal == index->ntotal);
530  sum_d += index->d;
531  }
532 
533 }
534 
535 void IndexSplitVectors::add (idx_t n, const float *x)
536 {
537  FAISS_ASSERT (!"not implemented");
538 }
539 
540 namespace {
541 
542 /// callback + thread management to query in 1 shard
543 struct SplitQueryJob {
544  const IndexSplitVectors *index; // the relevant index
545  int no; // shard number
546 
547  // query params
548  idx_t n;
549  const float *x;
550  idx_t k;
551  float *distances;
552  idx_t *labels;
553 
554 
555  void run ()
556  {
557  if (index->verbose)
558  printf ("begin query shard %d on %ld points\n", no, n);
559  const Index * sub_index = index->sub_indexes[no];
560  long sub_d = sub_index->d, d = index->d;
561  idx_t ofs = 0;
562  for (int i = 0; i < no; i++) ofs += index->sub_indexes[i]->d;
563  float *sub_x = new float [sub_d * n];
564  for (idx_t i = 0; i < n; i++)
565  memcpy (sub_x + i * sub_d, x + ofs + i * d, sub_d * sizeof (sub_x));
566  sub_index->search (n, sub_x, k, distances, labels);
567  delete [] sub_x;
568  if (index->verbose)
569  printf ("end query shard %d\n", no);
570  }
571 
572 };
573 
574 
575 
576 }
577 
578 
579 
581  idx_t n, const float *x, idx_t k,
582  float *distances, idx_t *labels) const
583 {
584  FAISS_ASSERT (k == 1 || !"search implemented only for k=1");
585  FAISS_ASSERT (sum_d == d || !"not enough indexes compared to # dimensions");
586 
587  long nshard = sub_indexes.size();
588  float *all_distances = new float [nshard * k * n];
589  idx_t *all_labels = new idx_t [nshard * k * n];
590 
591  // pre-alloc because we don't want reallocs
592  std::vector<Thread<SplitQueryJob> > qss (nshard);
593  for (int i = 0; i < nshard; i++) {
594  SplitQueryJob qs = {
595  this, i, n, x, k,
596  i == 0 ? distances : all_distances + i * k * n,
597  i == 0 ? labels : all_labels + i * k * n
598  };
599  if (threaded) {
600  qss[i] = Thread<SplitQueryJob> (qs);
601  qss[i].start();
602  } else {
603  qs.run();
604  }
605  }
606 
607  if (threaded) {
608  for (int i = 0; i < qss.size(); i++) {
609  qss[i].wait();
610  }
611  }
612 
613  long factor = 1;
614  for (int i = 0; i < nshard; i++) {
615  if (i > 0) { // results of 0 are already in the table
616  const float *distances_i = all_distances + i * k * n;
617  const idx_t *labels_i = all_labels + i * k * n;
618  for (long j = 0; j < n; j++) {
619  if (labels[j] >= 0 && labels_i[j] >= 0) {
620  labels[j] += labels_i[j] * factor;
621  distances[j] += distances_i[j];
622  } else {
623  labels[j] = -1;
624  distances[j] = 0.0 / 0.0;
625  }
626  }
627  }
628  factor *= sub_indexes[i]->ntotal;
629  }
630  delete [] all_labels;
631  delete [] all_distances;
632 }
633 
634 
635 void IndexSplitVectors::train (idx_t n, const float *x)
636 {
637  FAISS_ASSERT (!"not implemented");
638 }
639 
641 {
642  FAISS_ASSERT (!"not implemented");
643 }
644 
645 void IndexSplitVectors::set_typename ()
646 {}
647 
648 IndexSplitVectors::~IndexSplitVectors ()
649 {
650  if (own_fields) {
651  for (int s = 0; s < sub_indexes.size(); s++)
652  delete sub_indexes [s];
653  }
654 }
655 
656 
657 
658 
659 
660 
661 }; // namespace faiss
virtual void train(idx_t n, const float *x) override
IndexShards(idx_t d, bool threaded=false, bool successive_ids=true)
virtual void reset()=0
removes all elements from the database.
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: MetaIndexes.cpp:69
virtual void add_with_ids(idx_t n, const float *x, const long *xids) override
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
virtual void reset() override
removes all elements from the database.
virtual void add(idx_t n, const float *x) override
this will fail. Use add_with_ids
Definition: MetaIndexes.cpp:40
int d
vector dimension
Definition: Index.h:66
std::vector< long > id_map
! whether pointers are deleted in destructo
Definition: MetaIndexes.h:29
virtual void add(idx_t n, const float *x)=0
virtual void train(idx_t n, const float *x) override
virtual void add(idx_t n, const float *x) override
supported only for sub-indices that implement add_with_ids
long idx_t
all indices are this type
Definition: Index.h:64
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
bool verbose
verbosity level
Definition: Index.h:68
virtual void add(idx_t n, const float *x) override
bool threaded
should the sub-indexes be deleted along with this?
Definition: MetaIndexes.h:62
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
virtual void reset() override
removes all elements from the database.
virtual void reset() override
removes all elements from the database.
Definition: MetaIndexes.cpp:53
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
virtual void train(idx_t n, const float *x) override
Definition: MetaIndexes.cpp:47
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:71
virtual void train(idx_t n, const float *x)
Definition: Index.h:92
virtual void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: MetaIndexes.cpp:60
IndexSplitVectors(idx_t d, bool threaded=false)
sum of dimensions seen so far
bool own_fields
! the sub-index
Definition: MetaIndexes.h:28