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