Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFPQ.h
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the CC-by-NC 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
11 #pragma once
12
13 #include "GpuIndexIVF.h"
14 #include <vector>
15
16 namespace faiss { struct IndexIVFPQ; }
17
18 namespace faiss { namespace gpu {
19
20 class GpuIndexFlat;
21 class IVFPQ;
22
23 struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
24 inline GpuIndexIVFPQConfig()
25 : useFloat16LookupTables(false),
26 usePrecomputedTables(false) {
27 }
28
29 /// Whether or not float16 residual distance tables are used in the
30 /// list scanning kernels. When subQuantizers * 2^bitsPerCode >
31 /// 16384, this is required.
32 bool useFloat16LookupTables;
33
34 /// Whether or not we enable the precomputed table option for
35 /// search, which can substantially increase the memory requirement.
36 bool usePrecomputedTables;
37 };
38
39 /// IVFPQ index for the GPU
40 class GpuIndexIVFPQ : public GpuIndexIVF {
41 public:
42 /// Construct from a pre-existing faiss::IndexIVFPQ instance, copying
43 /// data over to the given GPU, if the input index is trained.
44 GpuIndexIVFPQ(GpuResources* resources,
45 const faiss::IndexIVFPQ* index,
46 GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
47
48 /// Construct an empty index
49 GpuIndexIVFPQ(GpuResources* resources,
50 int dims,
51 int nlist,
52 int subQuantizers,
53 int bitsPerCode,
54 faiss::MetricType metric,
55 GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
56
57 ~GpuIndexIVFPQ() override;
58
59 /// Reserve space on the GPU for the inverted lists for `num`
60 /// vectors, assumed equally distributed among
61
62 /// Initialize ourselves from the given CPU index; will overwrite
63 /// all data in ourselves
64 void copyFrom(const faiss::IndexIVFPQ* index);
65
66 /// Copy ourselves to the given CPU index; will overwrite all data
67 /// in the index instance
68 void copyTo(faiss::IndexIVFPQ* index) const;
69
70 /// Reserve GPU memory in our inverted lists for this number of vectors
71 void reserveMemory(size_t numVecs);
72
73 /// Enable or disable pre-computed codes
74 void setPrecomputedCodes(bool enable);
75
76 /// Are pre-computed codes enabled?
77 bool getPrecomputedCodes() const;
78
79 /// Return the number of sub-quantizers we are using
80 int getNumSubQuantizers() const;
81
82 /// Return the number of bits per PQ code
83 int getBitsPerCode() const;
84
85 /// Return the number of centroids per PQ code (2^bits per code)
86 int getCentroidsPerSubQuantizer() const;
87
88 /// After adding vectors, one can call this to reclaim device memory
89 /// to exactly the amount needed. Returns space reclaimed in bytes
90 size_t reclaimMemory();
91
92 /// Clears out all inverted lists, but retains the coarse and
93 /// product centroid information
94 void reset() override;
95
96 void train(Index::idx_t n, const float* x) override;
97
98 /// For debugging purposes, return the list length of a particular
99 /// list
100 int getListLength(int listId) const;
101
102 /// For debugging purposes, return the list codes of a particular
103 /// list
104 std::vector<unsigned char> getListCodes(int listId) const;
105
106 /// For debugging purposes, return the list indices of a particular
107 /// list
108 std::vector<long> getListIndices(int listId) const;
109
110 protected:
111 /// Called from GpuIndex for add/add_with_ids
112 void addImpl_(
113 faiss::Index::idx_t n,
114 const float* x,
115 const faiss::Index::idx_t* ids) override;
116
117 /// Called from GpuIndex for search
118 void searchImpl_(
119 faiss::Index::idx_t n,
120 const float* x,
121 faiss::Index::idx_t k,
122 float* distances,
123 faiss::Index::idx_t* labels) const override;
124
125 private:
126 void verifySettings_() const;
127
128 void trainResidualQuantizer_(Index::idx_t n, const float* x);
129
130 private:
131 GpuIndexIVFPQConfig ivfpqConfig_;
132
133 /// Number of sub-quantizers per encoded vector
134 int subQuantizers_;
135
136 /// Bits per sub-quantizer code
137 int bitsPerCode_;
138
139 /// Desired inverted list memory reservation
140 size_t reserveMemoryVecs_;
141
142 /// The product quantizer instance that we own; contains the
143 /// inverted lists
144 IVFPQ* index_;
145 };
146
147 } } // namespace
faiss::gpu::GpuIndexIVFPQ::getListIndices
std::vector< long > getListIndices(int listId) const
Definition: GpuIndexIVFPQ.cu:408
faiss::gpu::GpuIndexIVFPQ::searchImpl_
void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Called from GpuIndex for search.
Definition: GpuIndexIVFPQ.cu:349
faiss::gpu::GpuIndexIVFPQ::reset
void reset() override
Definition: GpuIndexIVFPQ.cu:245
faiss::gpu::GpuIndexIVFPQConfig::usePrecomputedTables
bool usePrecomputedTables
Definition: GpuIndexIVFPQ.h:36
faiss::gpu::GpuIndexIVFPQ::GpuIndexIVFPQ
GpuIndexIVFPQ(GpuResources *resources, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config=GpuIndexIVFPQConfig())
Definition: GpuIndexIVFPQ.cu:25
faiss::gpu::GpuIndexIVFPQ::getListLength
int getListLength(int listId) const
Definition: GpuIndexIVFPQ.cu:394
faiss::gpu::GpuIndexIVFConfig
Definition: GpuIndexIVF.h:25
faiss::gpu::GpuResources
Definition: GpuResources.h:23
faiss::gpu::GpuIndexIVFPQ::getBitsPerCode
int getBitsPerCode() const
Return the number of bits per PQ code.
Definition: GpuIndexIVFPQ.cu:225
faiss::gpu::GpuIndexIVFPQConfig::useFloat16LookupTables
bool useFloat16LookupTables
Definition: GpuIndexIVFPQ.h:32
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:37
faiss::gpu::GpuIndexIVFPQ::train
void train(Index::idx_t n, const float *x) override
Definition: GpuIndexIVFPQ.cu:301
faiss::gpu::GpuIndexIVFPQ::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFPQ.cu:235
faiss::gpu::GpuIndexIVFPQConfig
Definition: GpuIndexIVFPQ.h:23
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:29
faiss::gpu::GpuIndexIVFPQ::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFPQ.cu:195
faiss::gpu::GpuIndexIVFPQ::copyFrom
void copyFrom(const faiss::IndexIVFPQ *index)
Definition: GpuIndexIVFPQ.cu:80
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexIVFPQ::addImpl_
void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids) override
Called from GpuIndex for add/add_with_ids.
Definition: GpuIndexIVFPQ.cu:320
faiss::gpu::GpuIndexIVFPQ::copyTo
void copyTo(faiss::IndexIVFPQ *index) const
Definition: GpuIndexIVFPQ.cu:145
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:40
faiss::gpu::GpuIndexIVFPQ::getNumSubQuantizers
int getNumSubQuantizers() const
Return the number of sub-quantizers we are using.
Definition: GpuIndexIVFPQ.cu:220
faiss::gpu::GpuIndexIVFPQ::getCentroidsPerSubQuantizer
int getCentroidsPerSubQuantizer() const
Return the number of centroids per PQ code (2^bits per code)
Definition: GpuIndexIVFPQ.cu:230
faiss::gpu::GpuIndexIVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: GpuIndexIVFPQ.cu:204
faiss::gpu::GpuIndexIVFPQ::getPrecomputedCodes
bool getPrecomputedCodes() const
Are pre-computed codes enabled?
Definition: GpuIndexIVFPQ.cu:215
faiss::gpu::IVFPQ
Implementing class for IVFPQ on the GPU.
Definition: IVFPQ.cuh:19
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:43
faiss::gpu::GpuIndexIVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Definition: GpuIndexIVFPQ.cu:400