Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/matthijs/github_faiss/faiss/IndexIVFPQ.cpp
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 /* Copyright 2004-present Facebook. All Rights Reserved.
10 Inverted list structure.
11 */
12
13 #include "IndexIVFPQ.h"
14
15 #include <cmath>
16 #include <cstdio>
17 #include <cassert>
18
19 #include <sys/mman.h>
20
21 #include <algorithm>
22
23 #include "Heap.h"
24 #include "utils.h"
25
26 #include "Clustering.h"
27 #include "IndexFlat.h"
28
29 #include "hamming.h"
30
31 #include "FaissAssert.h"
32
33 #include "AuxIndexStructures.h"
34
35 namespace faiss {
36
37
38
39
40
41 /*****************************************
42 * IndexIVFPQ implementation
43 ******************************************/
44
45 IndexIVFPQ::IndexIVFPQ (Index * quantizer, size_t d, size_t nlist,
46 size_t M, size_t nbits_per_idx):
47 IndexIVF (quantizer, d, nlist, METRIC_L2),
48 pq (d, M, nbits_per_idx)
49 {
50 FAISS_THROW_IF_NOT (nbits_per_idx <= 8);
51 code_size = pq.code_size;
52 is_trained = false;
53 by_residual = true;
54 use_precomputed_table = 0;
55 scan_table_threshold = 0;
56 max_codes = 0; // means unlimited
57
58 polysemous_training = nullptr;
59 do_polysemous_training = false;
60 polysemous_ht = 0;
61
62 }
63
64
65 void IndexIVFPQ::train_residual (idx_t n, const float *x)
66 {
67 train_residual_o (n, x, nullptr);
68 }
69
70
71 void IndexIVFPQ::train_residual_o (idx_t n, const float *x, float *residuals_2)
72 {
73 const float * x_in = x;
74
75 x = fvecs_maybe_subsample (
76 d, (size_t*)&n, pq.cp.max_points_per_centroid * pq.ksub,
77 x, verbose, pq.cp.seed);
78
79 ScopeDeleter<float> del_x (x_in == x ? nullptr : x);
80
81 const float *trainset;
82 ScopeDeleter<float> del_residuals;
83 if (by_residual) {
84 if(verbose) printf("computing residuals\n");
85 idx_t * assign = new idx_t [n]; // assignement to coarse centroids
86 ScopeDeleter<idx_t> del (assign);
87 quantizer->assign (n, x, assign);
88 float *residuals = new float [n * d];
89 del_residuals.set (residuals);
90 for (idx_t i = 0; i < n; i++)
91 quantizer->compute_residual (x + i * d, residuals+i*d, assign[i]);
92
93 trainset = residuals;
94 } else {
95 trainset = x;
96 }
97 if (verbose)
98 printf ("training %zdx%zd product quantizer on %ld vectors in %dD\n",
99 pq.M, pq.ksub, n, d);
100 pq.verbose = verbose;
101 pq.train (n, trainset);
102
103 if (do_polysemous_training) {
104 if (verbose)
105 printf("doing polysemous training for PQ\n");
106 PolysemousTraining default_pt;
107 PolysemousTraining *pt = polysemous_training;
108 if (!pt) pt = &default_pt;
109 pt->optimize_pq_for_hamming (pq, n, trainset);
110 }
111
112 // prepare second-level residuals for refine PQ
113 if (residuals_2) {
114 uint8_t *train_codes = new uint8_t [pq.code_size * n];
115 ScopeDeleter<uint8_t> del (train_codes);
116 pq.compute_codes (trainset, train_codes, n);
117
118 for (idx_t i = 0; i < n; i++) {
119 const float *xx = trainset + i * d;
120 float * res = residuals_2 + i * d;
121 pq.decode (train_codes + i * pq.code_size, res);
122 for (int j = 0; j < d; j++)
123 res[j] = xx[j] - res[j];
124 }
125
126 }
127
128 if (by_residual) {
129 precompute_table ();
130 }
131
132 }
133
134
135 /* produce a binary signature based on the residual vector */
136 void IndexIVFPQ::encode (long key, const float * x, uint8_t * code) const
137 {
138 if (by_residual) {
139 float residual_vec[d];
140 quantizer->compute_residual (x, residual_vec, key);
141 pq.compute_code (residual_vec, code);
142 }
143 else pq.compute_code (x, code);
144 }
145
146
147
148
149
150 void IndexIVFPQ::encode_multiple (size_t n, long *keys,
151 const float * x, uint8_t * xcodes,
152 bool compute_keys) const
153 {
154 if (compute_keys)
155 quantizer->assign (n, x, keys);
156
157 if (by_residual) {
158 float *residuals = new float [n * d];
159 ScopeDeleter<float> del (residuals);
160 // TODO: parallelize?
161 for (size_t i = 0; i < n; i++)
162 quantizer->compute_residual (x + i * d, residuals + i * d, keys[i]);
163 pq.compute_codes (residuals, xcodes, n);
164 } else {
165 pq.compute_codes (x, xcodes, n);
166 }
167 }
168
169 void IndexIVFPQ::decode_multiple (size_t n, const long *keys,
170 const uint8_t * xcodes, float * x) const
171 {
172 pq.decode (xcodes, x, n);
173 if (by_residual) {
174 std::vector<float> centroid (d);
175 for (size_t i = 0; i < n; i++) {
176 quantizer->reconstruct (keys[i], centroid.data());
177 float *xi = x + i * d;
178 for (size_t j = 0; j < d; j++) {
179 xi [j] += centroid [j];
180 }
181 }
182 }
183 }
184
185
186 void IndexIVFPQ::add_with_ids (idx_t n, const float * x, const long *xids)
187 {
188 add_core_o (n, x, xids, nullptr);
189 }
190
191
192 void IndexIVFPQ::add_core_o (idx_t n, const float * x, const long *xids,
193 float *residuals_2, const long *precomputed_idx)
194 {
195 FAISS_THROW_IF_NOT (is_trained);
196 double t0 = getmillisecs ();
197 const long * idx;
198 ScopeDeleter<long> del_idx;
199
200 if (precomputed_idx) {
201 idx = precomputed_idx;
202 } else {
203 long * idx0 = new long [n];
204 del_idx.set (idx0);
205 quantizer->assign (n, x, idx0);
206 idx = idx0;
207 }
208
209 double t1 = getmillisecs ();
210 uint8_t * xcodes = new uint8_t [n * code_size];
211 ScopeDeleter<uint8_t> del_xcodes (xcodes);
212
213 const float *to_encode = nullptr;
214 ScopeDeleter<float> del_to_encode;
215
216 if (by_residual) {
217 float *residuals = new float [n * d];
218 // TODO: parallelize?
219 for (size_t i = 0; i < n; i++) {
220 if (idx[i] < 0)
221 memset (residuals + i * d, 0, sizeof(*residuals) * d);
222 else
223 quantizer->compute_residual (
224 x + i * d, residuals + i * d, idx[i]);
225 }
226 to_encode = residuals;
227 del_to_encode.set (to_encode);
228 } else {
229 to_encode = x;
230 }
231 pq.compute_codes (to_encode, xcodes, n);
232
233 double t2 = getmillisecs ();
234 // TODO: parallelize?
235 size_t n_ignore = 0;
236 for (size_t i = 0; i < n; i++) {
237 idx_t key = idx[i];
238 if (key < 0) {
239 n_ignore ++;
240 if (residuals_2)
241 memset (residuals_2, 0, sizeof(*residuals_2) * d);
242 continue;
243 }
244 idx_t id = xids ? xids[i] : ntotal + i;
245 ids[key].push_back (id);
246 uint8_t *code = xcodes + i * code_size;
247 for (size_t j = 0; j < code_size; j++)
248 codes[key].push_back (code[j]);
249
250 if (residuals_2) {
251 float *res2 = residuals_2 + i * d;
252 const float *xi = to_encode + i * d;
253 pq.decode (code, res2);
254 for (int j = 0; j < d; j++)
255 res2[j] = xi[j] - res2[j];
256 }
257
258 if (maintain_direct_map)
259 direct_map.push_back (key << 32 | (ids[key].size() - 1));
260 }
261
262
263 double t3 = getmillisecs ();
264 if(verbose) {
265 char comment[100] = {0};
266 if (n_ignore > 0)
267 snprintf (comment, 100, "(%ld vectors ignored)", n_ignore);
268 printf(" add_core times: %.3f %.3f %.3f %s\n",
269 t1 - t0, t2 - t1, t3 - t2, comment);
270 }
271 ntotal += n;
272 }
273
274 void IndexIVFPQ::reconstruct_n (idx_t i0, idx_t ni, float *recons) const
275 {
276 FAISS_THROW_IF_NOT (ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
277
278 std::vector<float> centroid (d);
279
280 for (int key = 0; key < nlist; key++) {
281 const std::vector<long> & idlist = ids[key];
282 const uint8_t * code_line = codes[key].data();
283
284 for (long ofs = 0; ofs < idlist.size(); ofs++) {
285 long id = idlist[ofs];
286 if (!(id >= i0 && id < i0 + ni)) continue;
287 float *r = recons + d * (id - i0);
288 if (by_residual) {
289 quantizer->reconstruct (key, centroid.data());
290 pq.decode (code_line + ofs * pq.code_size, r);
291 for (int j = 0; j < d; j++) {
292 r[j] += centroid[j];
293 }
294 } else {
295 pq.decode (code_line + ofs * pq.code_size, r);
296 }
297 }
298 }
299 }
300
301
302 void IndexIVFPQ::reconstruct (idx_t key, float * recons) const
303 {
304 FAISS_THROW_IF_NOT (direct_map.size() == ntotal);
305
306 int list_no = direct_map[key] >> 32;
307 int ofs = direct_map[key] & 0xffffffff;
308
309 quantizer->reconstruct (list_no, recons);
310 const uint8_t * code = &(codes[list_no][ofs * pq.code_size]);
311
312 for (size_t m = 0; m < pq.M; m++) {
313 float * out = recons + m * pq.dsub;
314 const float * cent = pq.get_centroids (m, code[m]);
315 for (size_t i = 0; i < pq.dsub; i++) {
316 out[i] += cent[i];
317 }
318 }
319 }
320
321
322
323
324
325
326 /** Precomputed tables for residuals
327 *
328 * During IVFPQ search with by_residual, we compute
329 *
330 * d = || x - y_C - y_R ||^2
331 *
332 * where x is the query vector, y_C the coarse centroid, y_R the
333 * refined PQ centroid. The expression can be decomposed as:
334 *
335 * d = || x - y_C ||^2 + || y_R ||^2 + 2 * (y_C|y_R) - 2 * (x|y_R)
336 * --------------- --------------------------- -------
337 * term 1 term 2 term 3
338 *
339 * When using multiprobe, we use the following decomposition:
340 * - term 1 is the distance to the coarse centroid, that is computed
341 * during the 1st stage search.
342 * - term 2 can be precomputed, as it does not involve x. However,
343 * because of the PQ, it needs nlist * M * ksub storage. This is why
344 * use_precomputed_table is off by default
345 * - term 3 is the classical non-residual distance table.
346 *
347 * Since y_R defined by a product quantizer, it is split across
348 * subvectors and stored separately for each subvector. If the coarse
349 * quantizer is a MultiIndexQuantizer then the table can be stored
350 * more compactly.
351 *
352 * At search time, the tables for term 2 and term 3 are added up. This
353 * is faster when the length of the lists is > ksub * M.
354 */
355
356 void IndexIVFPQ::precompute_table ()
357 {
358 if (use_precomputed_table == 0) { // then choose the type of table
359 if (quantizer->metric_type == METRIC_INNER_PRODUCT) {
360 fprintf(stderr, "IndexIVFPQ::precompute_table: WARN precomputed "
361 "tables not needed for inner product quantizers\n");
362 return;
363 }
364 const MultiIndexQuantizer *miq =
365 dynamic_cast<const MultiIndexQuantizer *> (quantizer);
366 if (miq && pq.M % miq->pq.M == 0)
367 use_precomputed_table = 2;
368 else
369 use_precomputed_table = 1;
370 } // otherwise assume user has set appropriate flag on input
371
372 if (verbose) {
373 printf ("precomputing IVFPQ tables type %d\n",
374 use_precomputed_table);
375 }
376
377 // squared norms of the PQ centroids
378 std::vector<float> r_norms (pq.M * pq.ksub, NAN);
379 for (int m = 0; m < pq.M; m++)
380 for (int j = 0; j < pq.ksub; j++)
381 r_norms [m * pq.ksub + j] =
382 fvec_norm_L2sqr (pq.get_centroids (m, j), pq.dsub);
383
384 if (use_precomputed_table == 1) {
385
386 precomputed_table.resize (nlist * pq.M * pq.ksub);
387 std::vector<float> centroid (d);
388
389 for (size_t i = 0; i < nlist; i++) {
390 quantizer->reconstruct (i, centroid.data());
391
392 float *tab = &precomputed_table[i * pq.M * pq.ksub];
393 pq.compute_inner_prod_table (centroid.data(), tab);
394 fvec_madd (pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
395 }
396 } else if (use_precomputed_table == 2) {
397 const MultiIndexQuantizer *miq =
398 dynamic_cast<const MultiIndexQuantizer *> (quantizer);
399 FAISS_THROW_IF_NOT (miq);
400 const ProductQuantizer &cpq = miq->pq;
401 FAISS_THROW_IF_NOT (pq.M % cpq.M == 0);
402
403 precomputed_table.resize(cpq.ksub * pq.M * pq.ksub);
404
405 // reorder PQ centroid table
406 std::vector<float> centroids (d * cpq.ksub, NAN);
407
408 for (int m = 0; m < cpq.M; m++) {
409 for (size_t i = 0; i < cpq.ksub; i++) {
410 memcpy (centroids.data() + i * d + m * cpq.dsub,
411 cpq.get_centroids (m, i),
412 sizeof (*centroids.data()) * cpq.dsub);
413 }
414 }
415
416 pq.compute_inner_prod_tables (cpq.ksub, centroids.data (),
417 precomputed_table.data ());
418
419 for (size_t i = 0; i < cpq.ksub; i++) {
420 float *tab = &precomputed_table[i * pq.M * pq.ksub];
421 fvec_madd (pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
422 }
423
424 }
425 }
426
427 namespace {
428
429 static uint64_t get_cycles () {
430 uint32_t high, low;
431 asm volatile("rdtsc \n\t"
432 : "=a" (low),
433 "=d" (high));
434 return ((uint64_t)high << 32) | (low);
435 }
436
437 #define TIC t0 = get_cycles()
438 #define TOC get_cycles () - t0
439
440
441
442 /** QueryTables manages the various ways of searching an
443 * IndexIVFPQ. The code contains a lot of branches, depending on:
444 * - metric_type: are we computing L2 or Inner product similarity?
445 * - by_residual: do we encode raw vectors or residuals?
446 * - use_precomputed_table: are x_R|x_C tables precomputed?
447 * - polysemous_ht: are we filtering with polysemous codes?
448 */
449 struct QueryTables {
450
451 /*****************************************************
452 * General data from the IVFPQ
453 *****************************************************/
454
455 const IndexIVFPQ & ivfpq;
456
457 // copied from IndexIVFPQ for easier access
458 int d;
459 const ProductQuantizer & pq;
460 MetricType metric_type;
461 bool by_residual;
462 int use_precomputed_table;
463
464 // pre-allocated data buffers
465 float * sim_table, * sim_table_2;
466 float * residual_vec, *decoded_vec;
467
468 // single data buffer
469 std::vector<float> mem;
470
471 // for table pointers
472 std::vector<const float *> sim_table_ptrs;
473
474 explicit QueryTables (const IndexIVFPQ & ivfpq):
475 ivfpq(ivfpq),
476 d(ivfpq.d),
477 pq (ivfpq.pq),
478 metric_type (ivfpq.metric_type),
479 by_residual (ivfpq.by_residual),
480 use_precomputed_table (ivfpq.use_precomputed_table)
481 {
482 mem.resize (pq.ksub * pq.M * 2 + d *2);
483 sim_table = mem.data();
484 sim_table_2 = sim_table + pq.ksub * pq.M;
485 residual_vec = sim_table_2 + pq.ksub * pq.M;
486 decoded_vec = residual_vec + d;
487
488 // for polysemous
489 if (ivfpq.polysemous_ht != 0) {
490 q_code.resize (pq.code_size);
491 }
492 init_list_cycles = 0;
493 sim_table_ptrs.resize (pq.M);
494 }
495
496 /*****************************************************
497 * What we do when query is known
498 *****************************************************/
499
500 // field specific to query
501 const float * qi;
502
503 // query-specific intialization
504 void init_query (const float * qi) {
505 this->qi = qi;
506 if (metric_type == METRIC_INNER_PRODUCT)
507 init_query_IP ();
508 else
509 init_query_L2 ();
510 if (!by_residual && ivfpq.polysemous_ht != 0)
511 pq.compute_code (qi, q_code.data());
512 }
513
514 void init_query_IP () {
515 // precompute some tables specific to the query qi
516 pq.compute_inner_prod_table (qi, sim_table);
517 // we compute negated inner products for use with the maxheap
518 for (int i = 0; i < pq.ksub * pq.M; i++) {
519 sim_table[i] = - sim_table[i];
520 }
521 }
522
523 void init_query_L2 () {
524 if (!by_residual) {
525 pq.compute_distance_table (qi, sim_table);
526 } else if (use_precomputed_table) {
527 pq.compute_inner_prod_table (qi, sim_table_2);
528 }
529 }
530
531 /*****************************************************
532 * When inverted list is known: prepare computations
533 *****************************************************/
534
535 // fields specific to list
536 Index::idx_t key;
537 float coarse_dis;
538 std::vector<uint8_t> q_code;
539
540 uint64_t init_list_cycles;
541
542 /// once we know the query and the centroid, we can prepare the
543 /// sim_table that will be used for accumulation
544 /// and dis0, the initial value
545 float precompute_list_tables () {
546 float dis0 = 0;
547 uint64_t t0; TIC;
548 if (by_residual) {
549 if (metric_type == METRIC_INNER_PRODUCT)
550 dis0 = precompute_list_tables_IP ();
551 else
552 dis0 = precompute_list_tables_L2 ();
553 }
554 init_list_cycles += TOC;
555 return dis0;
556 }
557
558 float precompute_list_table_pointers () {
559 float dis0 = 0;
560 uint64_t t0; TIC;
561 if (by_residual) {
562 if (metric_type == METRIC_INNER_PRODUCT)
563 FAISS_THROW_MSG ("not implemented");
564 else
565 dis0 = precompute_list_table_pointers_L2 ();
566 }
567 init_list_cycles += TOC;
568 return dis0;
569 }
570
571 /*****************************************************
572 * compute tables for inner prod
573 *****************************************************/
574
575 float precompute_list_tables_IP ()
576 {
577 // prepare the sim_table that will be used for accumulation
578 // and dis0, the initial value
579 ivfpq.quantizer->reconstruct (key, decoded_vec);
580 // decoded_vec = centroid
581 float dis0 = -fvec_inner_product (qi, decoded_vec, d);
582
583 if (ivfpq.polysemous_ht) {
584 for (int i = 0; i < d; i++) {
585 residual_vec [i] = qi[i] - decoded_vec[i];
586 }
587 pq.compute_code (residual_vec, q_code.data());
588 }
589 return dis0;
590 }
591
592
593 /*****************************************************
594 * compute tables for L2 distance
595 *****************************************************/
596
597 float precompute_list_tables_L2 ()
598 {
599 float dis0 = 0;
600
601 if (use_precomputed_table == 0) {
602 ivfpq.quantizer->compute_residual (qi, residual_vec, key);
603 pq.compute_distance_table (residual_vec, sim_table);
604 } else if (use_precomputed_table == 1) {
605 dis0 = coarse_dis;
606
607 fvec_madd (pq.M * pq.ksub,
608 &ivfpq.precomputed_table [key * pq.ksub * pq.M],
609 -2.0, sim_table_2,
610 sim_table);
611 } else if (use_precomputed_table == 2) {
612 dis0 = coarse_dis;
613
614 const MultiIndexQuantizer *miq =
615 dynamic_cast<const MultiIndexQuantizer *> (ivfpq.quantizer);
616 FAISS_THROW_IF_NOT (miq);
617 const ProductQuantizer &cpq = miq->pq;
618 int Mf = pq.M / cpq.M;
619
620 const float *qtab = sim_table_2; // query-specific table
621 float *ltab = sim_table; // (output) list-specific table
622
623 long k = key;
624 for (int cm = 0; cm < cpq.M; cm++) {
625 // compute PQ index
626 int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
627 k >>= cpq.nbits;
628
629 // get corresponding table
630 const float *pc = &ivfpq.precomputed_table
631 [(ki * pq.M + cm * Mf) * pq.ksub];
632
633 if (ivfpq.polysemous_ht == 0) {
634
635 // sum up with query-specific table
636 fvec_madd (Mf * pq.ksub,
637 pc,
638 -2.0, qtab,
639 ltab);
640 ltab += Mf * pq.ksub;
641 qtab += Mf * pq.ksub;
642 } else {
643 for (int m = cm * Mf; m < (cm + 1) * Mf; m++) {
644 q_code[m] = fvec_madd_and_argmin
645 (pq.ksub, pc, -2, qtab, ltab);
646 pc += pq.ksub;
647 ltab += pq.ksub;
648 qtab += pq.ksub;
649 }
650 }
651
652 }
653 }
654
655 return dis0;
656 }
657
658 float precompute_list_table_pointers_L2 ()
659 {
660 float dis0 = 0;
661
662 if (use_precomputed_table == 1) {
663 dis0 = coarse_dis;
664
665 const float * s = &ivfpq.precomputed_table [key * pq.ksub * pq.M];
666 for (int m = 0; m < pq.M; m++) {
667 sim_table_ptrs [m] = s;
668 s += pq.ksub;
669 }
670 } else if (use_precomputed_table == 2) {
671 dis0 = coarse_dis;
672
673 const MultiIndexQuantizer *miq =
674 dynamic_cast<const MultiIndexQuantizer *> (ivfpq.quantizer);
675 FAISS_THROW_IF_NOT (miq);
676 const ProductQuantizer &cpq = miq->pq;
677 int Mf = pq.M / cpq.M;
678
679 long k = key;
680 int m0 = 0;
681 for (int cm = 0; cm < cpq.M; cm++) {
682 int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
683 k >>= cpq.nbits;
684
685 const float *pc = &ivfpq.precomputed_table
686 [(ki * pq.M + cm * Mf) * pq.ksub];
687
688 for (int m = m0; m < m0 + Mf; m++) {
689 sim_table_ptrs [m] = pc;
690 pc += pq.ksub;
691 }
692 m0 += Mf;
693 }
694 } else {
695 FAISS_THROW_MSG ("need precomputed tables");
696 }
697
698 if (ivfpq.polysemous_ht) {
699 FAISS_THROW_MSG ("not implemented");
700 // Not clear that it makes sense to implemente this,
701 // because it costs M * ksub, which is what we wanted to
702 // avoid with the tables pointers.
703 }
704
705 return dis0;
706 }
707
708
709 };
710
711
712 /*****************************************************
713 * Scaning the codes.
714 * The scanning functions call their favorite precompute_*
715 * function to precompute the tables they need.
716 *****************************************************/
717 template <typename IDType>
718 struct InvertedListScanner: QueryTables {
719
720 const uint8_t * __restrict list_codes;
721 const IDType * list_ids;
722 size_t list_size;
723
724 explicit InvertedListScanner (const IndexIVFPQ & ivfpq):
725 QueryTables (ivfpq)
726 {
727 FAISS_THROW_IF_NOT (pq.byte_per_idx == 1);
728 n_hamming_pass = 0;
729 }
730
731 /// list_specific intialization
732 void init_list (Index::idx_t key, float coarse_dis,
733 size_t list_size_in, const IDType *list_ids_in,
734 const uint8_t *list_codes_in) {
735 this->key = key;
736 this->coarse_dis = coarse_dis;
737 list_size = list_size_in;
738 list_codes = list_codes_in;
739 list_ids = list_ids_in;
740 }
741
742 /*****************************************************
743 * Scaning the codes: simple PQ scan.
744 *****************************************************/
745
746 /// version of the scan where we use precomputed tables
747 void scan_list_with_table (
748 size_t k, float * heap_sim, long * heap_ids, bool store_pairs)
749 {
750 float dis0 = precompute_list_tables ();
751
752 for (size_t j = 0; j < list_size; j++) {
753
754 float dis = dis0;
755 const float *tab = sim_table;
756
757 for (size_t m = 0; m < pq.M; m++) {
758 dis += tab[*list_codes++];
759 tab += pq.ksub;
760 }
761
762 if (dis < heap_sim[0]) {
763 maxheap_pop (k, heap_sim, heap_ids);
764 long id = store_pairs ? (key << 32 | j) : list_ids[j];
765 maxheap_push (k, heap_sim, heap_ids, dis, id);
766 }
767 }
768 }
769
770
771 /// tables are not precomputed, but pointers are provided to the
772 /// relevant X_c|x_r tables
773 void scan_list_with_pointer (
774 size_t k, float * heap_sim, long * heap_ids, bool store_pairs)
775 {
776
777 float dis0 = precompute_list_table_pointers ();
778
779 for (size_t j = 0; j < list_size; j++) {
780
781 float dis = dis0;
782 const float *tab = sim_table_2;
783
784 for (size_t m = 0; m < pq.M; m++) {
785 int ci = *list_codes++;
786 dis += sim_table_ptrs [m][ci] - 2 * tab [ci];
787 tab += pq.ksub;
788 }
789
790 if (dis < heap_sim[0]) {
791 maxheap_pop (k, heap_sim, heap_ids);
792 long id = store_pairs ? (key << 32 | j) : list_ids[j];
793 maxheap_push (k, heap_sim, heap_ids, dis, id);
794 }
795 }
796
797 }
798
799 /// nothing is precomputed: access residuals on-the-fly
800 void scan_on_the_fly_dist (
801 size_t k, float * heap_sim, long * heap_ids, bool store_pairs)
802 {
803
804 if (by_residual && use_precomputed_table) {
805 scan_list_with_pointer (k, heap_sim, heap_ids, store_pairs);
806 return;
807 }
808
809 const float *dvec;
810 float dis0 = 0;
811
812 if (by_residual) {
813 if (metric_type == METRIC_INNER_PRODUCT) {
814 ivfpq.quantizer->reconstruct (key, residual_vec);
815 dis0 = fvec_inner_product (residual_vec, qi, d);
816 } else {
817 ivfpq.quantizer->compute_residual (qi, residual_vec, key);
818 }
819 dvec = residual_vec;
820 } else {
821 dvec = qi;
822 dis0 = 0;
823 }
824
825 for (size_t j = 0; j < list_size; j++) {
826
827 pq.decode (list_codes, decoded_vec);
828 list_codes += pq.code_size;
829
830 float dis;
831 if (metric_type == METRIC_INNER_PRODUCT) {
832 dis = -dis0 - fvec_inner_product (decoded_vec, qi, d);
833 } else {
834 dis = fvec_L2sqr (decoded_vec, dvec, d);
835 }
836
837 if (dis < heap_sim[0]) {
838 maxheap_pop (k, heap_sim, heap_ids);
839 long id = store_pairs ? (key << 32 | j) : list_ids[j];
840 maxheap_push (k, heap_sim, heap_ids, dis, id);
841 }
842 }
843 }
844
845 /*****************************************************
846 * Scanning codes with polysemous filtering
847 *****************************************************/
848
849 // code for the query
850 size_t n_hamming_pass;
851
852
853 template <class HammingComputer>
854 void scan_list_polysemous_hc (
855 size_t k, float * heap_sim, long * heap_ids, bool store_pairs)
856 {
857 float dis0 = precompute_list_tables ();
858 int ht = ivfpq.polysemous_ht;
859
860 int code_size = pq.code_size;
861
862 HammingComputer hc (q_code.data(), code_size);
863
864 for (size_t j = 0; j < list_size; j++) {
865 const uint8_t *b_code = list_codes;
866 int hd = hc.hamming (b_code);
867 if (hd < ht) {
868 n_hamming_pass ++;
869
870 float dis = dis0;
871 const float *tab = sim_table;
872
873 for (size_t m = 0; m < pq.M; m++) {
874 dis += tab[*b_code++];
875 tab += pq.ksub;
876 }
877
878 if (dis < heap_sim[0]) {
879 maxheap_pop (k, heap_sim, heap_ids);
880 long id = store_pairs ? (key << 32 | j) : list_ids[j];
881 maxheap_push (k, heap_sim, heap_ids, dis, id);
882 }
883 }
884 list_codes += code_size;
885 }
886 }
887
888 void scan_list_polysemous (
889 size_t k, float * heap_sim, long * heap_ids, bool store_pairs)
890 {
891 switch (pq.code_size) {
892 #define HANDLE_CODE_SIZE(cs) \
893 case cs: \
894 scan_list_polysemous_hc <HammingComputer ## cs> \
895 (k, heap_sim, heap_ids, store_pairs); \
896 break
897 HANDLE_CODE_SIZE(4);
898 HANDLE_CODE_SIZE(8);
899 HANDLE_CODE_SIZE(16);
900 HANDLE_CODE_SIZE(20);
901 HANDLE_CODE_SIZE(32);
902 HANDLE_CODE_SIZE(64);
903 #undef HANDLE_CODE_SIZE
904 default:
905 if (pq.code_size % 8 == 0)
906 scan_list_polysemous_hc <HammingComputerM8>
907 (k, heap_sim, heap_ids, store_pairs);
908 else
909 scan_list_polysemous_hc <HammingComputerM4>
910 (k, heap_sim, heap_ids, store_pairs);
911 break;
912 }
913 }
914
915 };
916
917
918
919
920 } // anonymous namespace
921
922
923 IndexIVFPQStats indexIVFPQ_stats;
924
925 void IndexIVFPQStats::reset () {
926 memset (this, 0, sizeof (*this));
927 }
928
929
930
931 void IndexIVFPQ::search_preassigned (idx_t nx, const float *qx, idx_t k,
932 const idx_t *keys,
933 const float *coarse_dis,
934 float *distances, idx_t *labels,
935 bool store_pairs) const
936 {
937 float_maxheap_array_t res = {
938 size_t(nx), size_t(k),
939 labels, distances
940 };
941
942 #pragma omp parallel
943 {
944 InvertedListScanner<long> qt (*this);
945 size_t stats_nlist = 0;
946 size_t stats_ncode = 0;
947 uint64_t init_query_cycles = 0;
948 uint64_t scan_cycles = 0;
949 uint64_t heap_cycles = 0;
950
951 #pragma omp for
952 for (size_t i = 0; i < nx; i++) {
953 const float *qi = qx + i * d;
954 const long * keysi = keys + i * nprobe;
955 const float *coarse_dis_i = coarse_dis + i * nprobe;
956 float * heap_sim = res.get_val (i);
957 long * heap_ids = res.get_ids (i);
958
959 uint64_t t0;
960 TIC;
961 maxheap_heapify (k, heap_sim, heap_ids);
962 heap_cycles += TOC;
963
964 TIC;
965 qt.init_query (qi);
966 init_query_cycles += TOC;
967
968 size_t nscan = 0;
969
970 for (size_t ik = 0; ik < nprobe; ik++) {
971 long key = keysi[ik]; /* select the list */
972 if (key < 0) {
973 // not enough centroids for multiprobe
974 continue;
975 }
976 FAISS_THROW_IF_NOT_FMT (
977 key < (long) nlist,
978 "Invalid key=%ld at ik=%ld nlist=%ld\n",
979 key, ik, nlist);
980
981 size_t list_size = ids[key].size();
982 stats_nlist ++;
983 nscan += list_size;
984
985 if (list_size == 0) continue;
986
987 qt.init_list (key, coarse_dis_i[ik],
988 list_size, ids[key].data(),
989 codes[key].data());
990
991 TIC;
992 if (polysemous_ht > 0) {
993 qt.scan_list_polysemous
994 (k, heap_sim, heap_ids, store_pairs);
995 } else if (list_size > scan_table_threshold) {
996 qt.scan_list_with_table (k, heap_sim, heap_ids, store_pairs);
997 } else {
998 qt.scan_on_the_fly_dist (k, heap_sim, heap_ids, store_pairs);
999 }
1000 scan_cycles += TOC;
1001
1002 if (max_codes && nscan >= max_codes) break;
1003 }
1004 stats_ncode += nscan;
1005 TIC;
1006 maxheap_reorder (k, heap_sim, heap_ids);
1007
1008 if (metric_type == METRIC_INNER_PRODUCT) {
1009 for (size_t j = 0; j < k; j++)
1010 heap_sim[j] = -heap_sim[j];
1011 }
1012 heap_cycles += TOC;
1013 }
1014
1015 #pragma omp critical
1016 {
1017 indexIVFPQ_stats.n_hamming_pass += qt.n_hamming_pass;
1018 indexIVFPQ_stats.nlist += stats_nlist;
1019 indexIVFPQ_stats.ncode += stats_ncode;
1020
1021 indexIVFPQ_stats.init_query_cycles += init_query_cycles;
1022 indexIVFPQ_stats.init_list_cycles += qt.init_list_cycles;
1023 indexIVFPQ_stats.scan_cycles += scan_cycles - qt.init_list_cycles;
1024 indexIVFPQ_stats.heap_cycles += heap_cycles;
1025 }
1026
1027 }
1028 indexIVFPQ_stats.nq += nx;
1029 }
1030
1031
1032 void IndexIVFPQ::search_and_reconstruct (idx_t n, const float *x, idx_t k,
1033 float *distances, idx_t *labels,
1034 float *reconstructed)
1035 {
1036 long * idx = new long [n * nprobe];
1037 ScopeDeleter<long> del (idx);
1038 float * coarse_dis = new float [n * nprobe];
1039 ScopeDeleter<float> del2 (coarse_dis);
1040
1041 quantizer->search (n, x, nprobe, coarse_dis, idx);
1042
1043 search_preassigned (n, x, k, idx, coarse_dis,
1044 distances, labels, true);
1045
1046 for (long i = 0; i < n; i++) {
1047 for (long j = 0; j < k; j++) {
1048 long ij = i * k + j;
1049 idx_t res = labels[ij];
1050 float *recons = reconstructed + d * (ij);
1051 if (res < 0) {
1052 // fill with NaNs
1053 memset(recons, -1, sizeof(*recons) * d);
1054 } else {
1055 int list_no = res >> 32;
1056 int ofs = res & 0xffffffff;
1057 labels[ij] = ids[list_no][ofs];
1058
1059 quantizer->reconstruct (list_no, recons);
1060 const uint8_t * code = &(codes[list_no][ofs * pq.code_size]);
1061
1062 for (size_t m = 0; m < pq.M; m++) {
1063 float * out = recons + m * pq.dsub;
1064 const float * cent = pq.get_centroids (m, code[m]);
1065 for (size_t l = 0; l < pq.dsub; l++) {
1066 out[l] += cent[l];
1067 }
1068 }
1069 }
1070 }
1071 }
1072
1073 }
1074
1075
1076
1077
1078
1079 IndexIVFPQ::IndexIVFPQ ()
1080 {
1081 // initialize some runtime values
1082 use_precomputed_table = 0;
1083 scan_table_threshold = 0;
1084 do_polysemous_training = false;
1085 polysemous_ht = 0;
1086 max_codes = 0;
1087 polysemous_training = nullptr;
1088 }
1089
1090
1091 struct CodeCmp {
1092 const uint8_t *tab;
1093 size_t code_size;
1094 bool operator () (int a, int b) const {
1095 return cmp (a, b) > 0;
1096 }
1097 int cmp (int a, int b) const {
1098 return memcmp (tab + a * code_size, tab + b * code_size,
1099 code_size);
1100 }
1101 };
1102
1103
1104 size_t IndexIVFPQ::find_duplicates (idx_t *dup_ids, size_t *lims) const
1105 {
1106 size_t ngroup = 0;
1107 lims[0] = 0;
1108 for (size_t list_no = 0; list_no < nlist; list_no++) {
1109 size_t n = ids[list_no].size();
1110 std::vector<int> ord (n);
1111 for (int i = 0; i < n; i++) ord[i] = i;
1112 CodeCmp cs = { codes[list_no].data(), code_size };
1113 std::sort (ord.begin(), ord.end(), cs);
1114
1115 const idx_t *list_ids = ids[list_no].data();
1116 int prev = -1; // all elements from prev to i-1 are equal
1117 for (int i = 0; i < n; i++) {
1118 if (prev >= 0 && cs.cmp (ord [prev], ord [i]) == 0) {
1119 // same as previous => remember
1120 if (prev + 1 == i) { // start new group
1121 ngroup++;
1122 lims[ngroup] = lims[ngroup - 1];
1123 dup_ids [lims [ngroup]++] = list_ids [ord [prev]];
1124 }
1125 dup_ids [lims [ngroup]++] = list_ids [ord [i]];
1126 } else { // not same as previous.
1127 prev = i;
1128 }
1129 }
1130 }
1131 return ngroup;
1132 }
1133
1134
1135
1136
1137 /*****************************************
1138 * IndexIVFPQR implementation
1139 ******************************************/
1140
1141 IndexIVFPQR::IndexIVFPQR (
1142 Index * quantizer, size_t d, size_t nlist,
1143 size_t M, size_t nbits_per_idx,
1144 size_t M_refine, size_t nbits_per_idx_refine):
1145 IndexIVFPQ (quantizer, d, nlist, M, nbits_per_idx),
1146 refine_pq (d, M_refine, nbits_per_idx_refine),
1147 k_factor (4)
1148 {
1149 by_residual = true;
1150 }
1151
1152 IndexIVFPQR::IndexIVFPQR ():
1153 k_factor (1)
1154 {
1155 by_residual = true;
1156 }
1157
1158
1159
1160 void IndexIVFPQR::reset()
1161 {
1162 IndexIVFPQ::reset();
1163 refine_codes.clear();
1164 }
1165
1166
1167
1168
1169 void IndexIVFPQR::train_residual (idx_t n, const float *x)
1170 {
1171
1172 float * residual_2 = new float [n * d];
1173 ScopeDeleter <float> del(residual_2);
1174
1175 train_residual_o (n, x, residual_2);
1176
1177 if (verbose)
1178 printf ("training %zdx%zd 2nd level PQ quantizer on %ld %dD-vectors\n",
1179 refine_pq.M, refine_pq.ksub, n, d);
1180
1181 refine_pq.cp.max_points_per_centroid = 1000;
1182 refine_pq.cp.verbose = verbose;
1183
1184 refine_pq.train (n, residual_2);
1185
1186 }
1187
1188
1189 void IndexIVFPQR::add_with_ids (idx_t n, const float *x, const long *xids) {
1190 add_core (n, x, xids, nullptr);
1191 }
1192
1193 void IndexIVFPQR::add_core (idx_t n, const float *x, const long *xids,
1194 const long *precomputed_idx) {
1195
1196 float * residual_2 = new float [n * d];
1197 ScopeDeleter <float> del(residual_2);
1198
1199 idx_t n0 = ntotal;
1200
1201 add_core_o (n, x, xids, residual_2, precomputed_idx);
1202
1203 refine_codes.resize (ntotal * refine_pq.code_size);
1204
1205 refine_pq.compute_codes (
1206 residual_2, &refine_codes[n0 * refine_pq.code_size], n);
1207
1208
1209 }
1210
1211
1212 void IndexIVFPQR::search (
1213 idx_t n, const float *x, idx_t k,
1214 float *distances, idx_t *labels) const
1215 {
1216 FAISS_THROW_IF_NOT (is_trained);
1217 long * idx = new long [n * nprobe];
1218 ScopeDeleter<long> del (idx);
1219 float * L1_dis = new float [n * nprobe];
1220 ScopeDeleter<float> del2 (L1_dis);
1221 uint64_t t0;
1222 TIC;
1223 quantizer->search (n, x, nprobe, L1_dis, idx);
1224 indexIVFPQ_stats.assign_cycles += TOC;
1225
1226 TIC;
1227 size_t k_coarse = long(k * k_factor);
1228 idx_t *coarse_labels = new idx_t [k_coarse * n];
1229 ScopeDeleter<idx_t> del3 (coarse_labels);
1230 { // query with quantizer levels 1 and 2.
1231 float *coarse_distances = new float [k_coarse * n];
1232 ScopeDeleter<float> del(coarse_distances);
1233
1234 search_preassigned (n, x, k_coarse,
1235 idx, L1_dis, coarse_distances, coarse_labels,
1236 true);
1237 }
1238
1239
1240 indexIVFPQ_stats.search_cycles += TOC;
1241
1242 TIC;
1243
1244 // 3rd level refinement
1245 size_t n_refine = 0;
1246 #pragma omp parallel reduction(+ : n_refine)
1247 {
1248 // tmp buffers
1249 float *residual_1 = new float [2 * d];
1250 ScopeDeleter<float> del (residual_1);
1251 float *residual_2 = residual_1 + d;
1252 #pragma omp for
1253 for (idx_t i = 0; i < n; i++) {
1254 const float *xq = x + i * d;
1255 const long * shortlist = coarse_labels + k_coarse * i;
1256 float * heap_sim = distances + k * i;
1257 long * heap_ids = labels + k * i;
1258 maxheap_heapify (k, heap_sim, heap_ids);
1259
1260 for (int j = 0; j < k_coarse; j++) {
1261 long sl = shortlist[j];
1262
1263 if (sl == -1) continue;
1264
1265 int list_no = sl >> 32;
1266 int ofs = sl & 0xffffffff;
1267
1268 assert (list_no >= 0 && list_no < nlist);
1269 assert (ofs >= 0 && ofs < ids[list_no].size());
1270
1271 // 1st level residual
1272 quantizer->compute_residual (xq, residual_1, list_no);
1273
1274 // 2nd level residual
1275 const uint8_t * l2code = &codes[list_no][ofs * pq.code_size];
1276 pq.decode (l2code, residual_2);
1277 for (int l = 0; l < d; l++)
1278 residual_2[l] = residual_1[l] - residual_2[l];
1279
1280 // 3rd level residual's approximation
1281 idx_t id = ids[list_no][ofs];
1282 assert (0 <= id && id < ntotal);
1283 refine_pq.decode (&refine_codes [id * refine_pq.code_size],
1284 residual_1);
1285
1286 float dis = fvec_L2sqr (residual_1, residual_2, d);
1287
1288 if (dis < heap_sim[0]) {
1289 maxheap_pop (k, heap_sim, heap_ids);
1290 maxheap_push (k, heap_sim, heap_ids, dis, id);
1291 }
1292 n_refine ++;
1293 }
1294 maxheap_reorder (k, heap_sim, heap_ids);
1295 }
1296 }
1297 indexIVFPQ_stats.nrefine += n_refine;
1298 indexIVFPQ_stats.refine_cycles += TOC;
1299 }
1300
1301 void IndexIVFPQR::reconstruct_n (idx_t i0, idx_t ni, float *recons) const
1302 {
1303 std::vector<float> r3 (d);
1304
1305 IndexIVFPQ::reconstruct_n (i0, ni, recons);
1306
1307 for (idx_t i = i0; i < i0 + ni; i++) {
1308 float *r = recons + i * d;
1309 refine_pq.decode (&refine_codes [i * refine_pq.code_size], r3.data());
1310
1311 for (int j = 0; j < d; j++)
1312 r[j] += r3[j];
1313
1314 }
1315
1316 }
1317
1318
1319
1320 void IndexIVFPQR::merge_from (IndexIVF &other_in, idx_t add_id)
1321 {
1322 IndexIVFPQR *other = dynamic_cast<IndexIVFPQR *> (&other_in);
1323 FAISS_THROW_IF_NOT(other);
1324
1325 IndexIVF::merge_from (other_in, add_id);
1326
1327 refine_codes.insert (refine_codes.end(),
1328 other->refine_codes.begin(),
1329 other->refine_codes.end());
1330 other->refine_codes.clear();
1331 }
1332
1333 long IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
1334 FAISS_THROW_MSG("not implemented");
1335 return 0;
1336 }
1337
1338 /*****************************************
1339 * IndexIVFPQCompact implementation
1340 ******************************************/
1341
1342 IndexIVFPQCompact::IndexIVFPQCompact ()
1343 {
1344 alloc_type = Alloc_type_none;
1345 limits = nullptr;
1346 compact_ids = nullptr;
1347 compact_codes = nullptr;
1348 }
1349
1350
1351 IndexIVFPQCompact::IndexIVFPQCompact (const IndexIVFPQ &other)
1352 {
1353 FAISS_THROW_IF_NOT_MSG (other.ntotal < (1UL << 31),
1354 "IndexIVFPQCompact cannot store more than 2G images");
1355
1356 // here it would be more convenient to just use the
1357 // copy-constructor, but it would copy the lists as well: too much
1358 // overhead...
1359
1360 // copy fields from Index
1361 d = other.d;
1362 ntotal = other.ntotal;
1363 verbose = other.verbose;
1364 is_trained = other.is_trained;
1365 metric_type = other.metric_type;
1366
1367 // copy fields from IndexIVF (except ids)
1368 nlist = other.nlist;
1369 nprobe = other.nprobe;
1370 quantizer = other.quantizer;
1371 quantizer_trains_alone = other.quantizer_trains_alone;
1372 own_fields = false;
1373 direct_map = other.direct_map;
1374
1375 // copy fields from IndexIVFPQ (except codes)
1376 by_residual = other.by_residual;
1377 use_precomputed_table = other.use_precomputed_table;
1378 precomputed_table = other.precomputed_table;
1379 code_size = other.code_size;
1380 pq = other.pq;
1381 do_polysemous_training = other.do_polysemous_training;
1382 polysemous_training = nullptr;
1383
1384 scan_table_threshold = other.scan_table_threshold;
1385 max_codes = other.max_codes;
1386 polysemous_ht = other.polysemous_ht;
1387
1388 //allocate
1389 alloc_type = Alloc_type_new;
1390 limits = new uint32_t [nlist + 1];
1391 compact_ids = new uint32_t [ntotal];
1392 compact_codes = new uint8_t [ntotal * code_size];
1393
1394
1395 // copy content from other
1396 size_t ofs = 0;
1397 for (size_t i = 0; i < nlist; i++) {
1398 limits [i] = ofs;
1399 const std::vector<long> &other_ids = other.ids[i];
1400 for (size_t j = 0; j < other_ids.size(); j++) {
1401 long id = other_ids[j];
1402 FAISS_THROW_IF_NOT_MSG (id < (1UL << 31),
1403 "IndexIVFPQCompact cannot store ids > 2G");
1404 compact_ids[ofs + j] = id;
1405 }
1406 memcpy (compact_codes + ofs * code_size,
1407 other.codes[i].data(),
1408 other.codes[i].size());
1409 ofs += other_ids.size();
1410 }
1411 FAISS_THROW_IF_NOT (ofs == ntotal);
1412 limits [nlist] = ofs;
1413
1414 }
1415
1416 void IndexIVFPQCompact::add (idx_t, const float *) {
1417 FAISS_THROW_MSG ("cannot add to an IndexIVFPQCompact");
1418 }
1419
1420 void IndexIVFPQCompact::reset () {
1421 FAISS_THROW_MSG ("cannot reset an IndexIVFPQCompact");
1422 }
1423
1424 void IndexIVFPQCompact::train (idx_t, const float *) {
1425 FAISS_THROW_MSG ("cannot train an IndexIVFPQCompact");
1426 }
1427
1428
1429
1430
1431 IndexIVFPQCompact::~IndexIVFPQCompact ()
1432 {
1433 if (alloc_type == Alloc_type_new) {
1434 delete [] limits;
1435 delete [] compact_codes;
1436 delete [] compact_ids;
1437 } else if (alloc_type == Alloc_type_mmap) {
1438 munmap (mmap_buffer, mmap_length);
1439
1440 }
1441
1442 }
1443
1444
1445
1446 void IndexIVFPQCompact::search_preassigned (idx_t nx, const float *qx, idx_t k,
1447 const idx_t *keys,
1448 const float *coarse_dis,
1449 float *distances, idx_t *labels,
1450 bool store_pairs) const
1451 {
1452 float_maxheap_array_t res = {
1453 size_t(nx), size_t(k),
1454 labels, distances
1455 };
1456
1457 #pragma omp parallel
1458 {
1459 InvertedListScanner<uint32_t> qt (*this);
1460 size_t stats_nlist = 0;
1461 size_t stats_ncode = 0;
1462 uint64_t init_query_cycles = 0;
1463 uint64_t scan_cycles = 0;
1464 uint64_t heap_cycles = 0;
1465
1466 #pragma omp for
1467 for (size_t i = 0; i < nx; i++) {
1468 const float *qi = qx + i * d;
1469 const long * keysi = keys + i * nprobe;
1470 const float *coarse_dis_i = coarse_dis + i * nprobe;
1471 float * heap_sim = res.get_val (i);
1472 long * heap_ids = res.get_ids (i);
1473
1474 uint64_t t0;
1475 TIC;
1476 maxheap_heapify (k, heap_sim, heap_ids);
1477 heap_cycles += TOC;
1478
1479 TIC;
1480 qt.init_query (qi);
1481 init_query_cycles += TOC;
1482
1483 size_t nscan = 0;
1484
1485 for (size_t ik = 0; ik < nprobe; ik++) {
1486 long key = keysi[ik]; /* select the list */
1487 if (key < 0) {
1488 // not enough centroids for multiprobe
1489 continue;
1490 }
1491 if (key >= (long) nlist) {
1492 fprintf (stderr, "Invalid key=%ld nlist=%ld\n", key, nlist);
1493 throw;
1494 }
1495 size_t list_size = limits[key + 1] - limits[key];
1496 stats_nlist ++;
1497 nscan += list_size;
1498
1499 if (list_size == 0) continue;
1500
1501 qt.init_list (key, coarse_dis_i[ik],
1502 list_size, compact_ids + limits[key],
1503 compact_codes + limits[key] * code_size);
1504
1505 TIC;
1506 if (polysemous_ht > 0) {
1507 qt.scan_list_polysemous
1508 (k, heap_sim, heap_ids, store_pairs);
1509 } else if (list_size > scan_table_threshold) {
1510 qt.scan_list_with_table (k, heap_sim, heap_ids, store_pairs);
1511 } else {
1512 qt.scan_on_the_fly_dist (k, heap_sim, heap_ids, store_pairs);
1513 }
1514 scan_cycles += TOC;
1515
1516 if (max_codes && nscan >= max_codes) break;
1517 }
1518 stats_ncode += nscan;
1519 TIC;
1520 maxheap_reorder (k, heap_sim, heap_ids);
1521
1522 if (metric_type == METRIC_INNER_PRODUCT) {
1523 for (size_t j = 0; j < k; j++) {
1524 heap_sim[i] = -heap_sim[i];
1525 }
1526 }
1527 heap_cycles += TOC;
1528 }
1529
1530 #pragma omp critical
1531 {
1532 indexIVFPQ_stats.n_hamming_pass += qt.n_hamming_pass;
1533 indexIVFPQ_stats.nlist += stats_nlist;
1534 indexIVFPQ_stats.ncode += stats_ncode;
1535
1536 indexIVFPQ_stats.init_query_cycles += init_query_cycles;
1537 indexIVFPQ_stats.init_list_cycles += qt.init_list_cycles;
1538 indexIVFPQ_stats.scan_cycles += scan_cycles - qt.init_list_cycles;
1539 indexIVFPQ_stats.heap_cycles += heap_cycles;
1540 }
1541
1542 }
1543 indexIVFPQ_stats.nq += nx;
1544 }
1545
1546
1547
1548 } // namespace faiss
faiss::IndexIVFPQCompact::compact_ids
uint32_t * compact_ids
size ntotal
Definition: IndexIVFPQ.h:227
faiss::IndexIVF
Definition: IndexIVF.h:45
faiss::IndexIVFPQCompact::compact_codes
uint8_t * compact_codes
size ntotal * code_size
Definition: IndexIVFPQ.h:228
faiss::IndexIVFPQ::precompute_table
void precompute_table()
build precomputed table
Definition: IndexIVFPQ.cpp:356
faiss::IndexIVFPQR::merge_from
void merge_from(IndexIVF &other, idx_t add_id) override
Definition: IndexIVFPQ.cpp:1320
faiss::IndexIVFPQ::reconstruct
void reconstruct(idx_t key, float *recons) const override
Definition: IndexIVFPQ.cpp:302
faiss::IndexIVFPQR
Definition: IndexIVFPQ.h:168
faiss::ProductQuantizer::decode
void decode(const uint8_t *code, float *x) const
decode a vector from a given code (or n vectors if third argument)
Definition: ProductQuantizer.cpp:338
faiss::IndexIVFPQR::refine_pq
ProductQuantizer refine_pq
3rd level quantizer
Definition: IndexIVFPQ.h:169
faiss::fvec_L2sqr
float fvec_L2sqr(const float *x, const float *y, size_t d)
Squared L2 distance between two vectors.
Definition: utils.cpp:574
faiss::IndexIVFPQ::polysemous_training
PolysemousTraining * polysemous_training
if NULL, use default
Definition: IndexIVFPQ.h:35
faiss::HeapArray::get_val
T * get_val(size_t key)
Return the list of values for a heap.
Definition: Heap.h:360
faiss::IndexIVFPQCompact::add
void add(idx_t, const float *) override
the three following functions will fail at runtime
Definition: IndexIVFPQ.cpp:1416
faiss::IndexIVFPQR::reconstruct_n
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexIVFPQ.cpp:1301
faiss::IndexIVFPQCompact::search_preassigned
void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs) const override
Definition: IndexIVFPQ.cpp:1446
faiss::IndexIVF::quantizer_trains_alone
char quantizer_trains_alone
Definition: IndexIVF.h:56
faiss::fvecs_maybe_subsample
const float * fvecs_maybe_subsample(size_t d, size_t *n, size_t nmax, const float *x, bool verbose, long seed)
Definition: utils.cpp:1973
faiss::IndexIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:47
faiss::Index::assign
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
Definition: Index.cpp:23
faiss::IDSelector
Definition: AuxIndexStructures.h:52
faiss::IndexIVFPQCompact::Alloc_type_mmap
was mmapped
Definition: IndexIVFPQ.h:221
faiss::IndexIVFPQ::decode_multiple
void decode_multiple(size_t n, const long *keys, const uint8_t *xcodes, float *x) const
inverse of encode_multiple
Definition: IndexIVFPQ.cpp:169
faiss::IndexIVFPQ::train_residual_o
void train_residual_o(idx_t n, const float *x, float *residuals_2)
same as train_residual, also output 2nd level residuals
Definition: IndexIVFPQ.cpp:71
faiss::IndexIVFPQ::do_polysemous_training
bool do_polysemous_training
reorder PQ centroids after training?
Definition: IndexIVFPQ.h:34
faiss::IndexIVFPQ::scan_table_threshold
size_t scan_table_threshold
use table computation or on-the-fly?
Definition: IndexIVFPQ.h:38
faiss::IndexIVFPQR::train_residual
void train_residual(idx_t n, const float *x) override
trains the two product quantizers
Definition: IndexIVFPQ.cpp:1169
faiss::IndexIVFPQR::add_core
void add_core(idx_t n, const float *x, const long *xids, const long *precomputed_idx=nullptr)
same as add_with_ids, but optionally use the precomputed list ids
Definition: IndexIVFPQ.cpp:1193
faiss::ScopeDeleter
Definition: FaissException.h:39
faiss::IndexIVFPQCompact::limits
uint32_t * limits
size nlist + 1
Definition: IndexIVFPQ.h:226
faiss::ProductQuantizer::dsub
size_t dsub
dimensionality of each subvector
Definition: ProductQuantizer.h:32
faiss::ClusteringParameters::seed
int seed
seed for the random number generator
Definition: Clustering.h:36
faiss::IndexIVFPQ::precomputed_table
std::vector< float > precomputed_table
Definition: IndexIVFPQ.h:45
faiss::fvec_madd
void fvec_madd(size_t n, const float *a, float bf, const float *b, float *c)
Definition: utils.cpp:1887
faiss::IndexIVFPQ::polysemous_ht
int polysemous_ht
Hamming thresh for polysemous filtering.
Definition: IndexIVFPQ.h:40
faiss::IndexIVFPQR::reset
void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1160
faiss::IndexIVFPQ::add_with_ids
void add_with_ids(idx_t n, const float *x, const long *xids=nullptr) override
Definition: IndexIVFPQ.cpp:186
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:62
faiss::ProductQuantizer::compute_codes
void compute_codes(const float *x, uint8_t *codes, size_t n) const
same as compute_code for several vectors
Definition: ProductQuantizer.cpp:385
faiss::Index::d
int d
vector dimension
Definition: Index.h:64
faiss::IndexIVF::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:49
faiss::IndexIVFPQ::max_codes
size_t max_codes
max nb of codes to visit to do a query
Definition: IndexIVFPQ.h:39
faiss::CodeCmp
Definition: IndexIVFPQ.cpp:1091
faiss::IndexIVFPQR::refine_codes
std::vector< uint8_t > refine_codes
corresponding codes
Definition: IndexIVFPQ.h:170
faiss::ProductQuantizer::code_size
size_t code_size
byte per indexed vector
Definition: ProductQuantizer.h:34
faiss::IndexIVFPQ::train_residual
void train_residual(idx_t n, const float *x) override
trains the product quantizer
Definition: IndexIVFPQ.cpp:65
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:29
faiss::IndexIVFPQ::encode_multiple
void encode_multiple(size_t n, long *keys, const float *x, uint8_t *codes, bool compute_keys=false) const
Definition: IndexIVFPQ.cpp:150
faiss::IndexIVFPQCompact::train
void train(idx_t, const float *) override
Trains the quantizer and calls train_residual to train sub-quantizers.
Definition: IndexIVFPQ.cpp:1424
faiss::IndexIVF::own_fields
bool own_fields
whether object owns the quantizer
Definition: IndexIVF.h:57
faiss::ProductQuantizer::ksub
size_t ksub
number of centroids for each subquantizer
Definition: ProductQuantizer.h:35
faiss::HeapArray
Definition: Heap.h:350
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::ProductQuantizer::compute_code
void compute_code(const float *x, uint8_t *code) const
Quantize one vector with the product quantizer.
Definition: ProductQuantizer.cpp:311
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:65
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:66
faiss::IndexIVF::reset
void reset() override
removes all elements from the database.
Definition: IndexIVF.cpp:112
faiss::getmillisecs
double getmillisecs()
ms elapsed since some arbitrary epoch
Definition: utils.cpp:74
faiss::PolysemousTraining
optimizes the order of indices in a ProductQuantizer
Definition: PolysemousTraining.h:125
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::MultiIndexQuantizer
Definition: IndexPQ.h:135
faiss::fvec_norm_L2sqr
float fvec_norm_L2sqr(const float *x, size_t d)
Definition: utils.cpp:632
faiss::ProductQuantizer::cp
ClusteringParameters cp
parameters used during clustering
Definition: ProductQuantizer.h:49
faiss::IndexIVFPQ::by_residual
bool by_residual
Encode residual or plain vector?
Definition: IndexIVFPQ.h:30
faiss::HeapArray::get_ids
TI * get_ids(size_t key)
Correspponding identifiers.
Definition: Heap.h:363
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:72
faiss::IndexIVFPQCompact::Alloc_type_none
alloc from outside
Definition: IndexIVFPQ.h:219
faiss::IndexIVFPQ::search_and_reconstruct
void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *reconstructed)
Definition: IndexIVFPQ.cpp:1032
faiss::IndexIVFPQ::pq
ProductQuantizer pq
produces the codes
Definition: IndexIVFPQ.h:32
faiss::ProductQuantizer::M
size_t M
number of subquantizers
Definition: ProductQuantizer.h:28
faiss::IndexIVF::nlist
size_t nlist
number of possible key values
Definition: IndexIVF.h:46
faiss::Index
Definition: Index.h:60
faiss::IndexIVFPQ::reconstruct_n
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexIVFPQ.cpp:274
faiss::IndexIVFPQ::add_core_o
void add_core_o(idx_t n, const float *x, const long *xids, float *residuals_2, const long *precomputed_idx=nullptr)
Definition: IndexIVFPQ.cpp:192
faiss::fvec_madd_and_argmin
int fvec_madd_and_argmin(size_t n, const float *a, float bf, const float *b, float *c)
Definition: utils.cpp:1961
faiss::IndexIVFPQR::remove_ids
long remove_ids(const IDSelector &sel) override
Dataset manipulation functions.
Definition: IndexIVFPQ.cpp:1333
faiss::IndexIVFPQ::search_preassigned
void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs) const override
Definition: IndexIVFPQ.cpp:931
faiss::ProductQuantizer
Definition: ProductQuantizer.h:25
faiss::IndexIVFPQCompact::reset
void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1420
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:69
faiss::Index::compute_residual
void compute_residual(const float *x, float *residual, idx_t key) const
Definition: Index.cpp:56
faiss::ProductQuantizer::get_centroids
float * get_centroids(size_t m, size_t i)
return the centroids associated with subvector m
Definition: ProductQuantizer.h:55
faiss::Index::reconstruct
virtual void reconstruct(idx_t key, float *recons) const
Definition: Index.cpp:43
faiss::IndexIVFPQCompact::Alloc_type_new
was allocated with new
Definition: IndexIVFPQ.h:220
faiss::IndexIVF::maintain_direct_map
bool maintain_direct_map
map for direct access to the elements. Enables reconstruct().
Definition: IndexIVF.h:68
faiss::PolysemousTraining::optimize_pq_for_hamming
void optimize_pq_for_hamming(ProductQuantizer &pq, size_t n, const float *x) const
Definition: PolysemousTraining.cpp:933
faiss::ClusteringParameters::max_points_per_centroid
int max_points_per_centroid
to limit size of dataset
Definition: Clustering.h:34
faiss::ProductQuantizer::verbose
bool verbose
verbose during training?
Definition: ProductQuantizer.h:36
faiss::IndexIVFPQR::add_with_ids
void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: IndexIVFPQ.cpp:1189
faiss::IndexIVFPQR::search
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexIVFPQ.cpp:1212
faiss::IndexIVF::merge_from
virtual void merge_from(IndexIVF &other, idx_t add_id)
Definition: IndexIVF.cpp:245
faiss::IndexIVF::code_size
size_t code_size
code size per vector in bytes
Definition: IndexIVF.h:64
faiss::IndexIVFPQ::find_duplicates
size_t find_duplicates(idx_t *ids, size_t *lims) const
Definition: IndexIVFPQ.cpp:1104
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:43
faiss::IndexIVFPQR::k_factor
float k_factor
factor between k requested in search and the k requested from the IVFPQ
Definition: IndexIVFPQ.h:173
faiss::IndexIVFPQ::use_precomputed_table
int use_precomputed_table
if by_residual, build precompute tables
Definition: IndexIVFPQ.h:31