Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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