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