Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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