Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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