Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFPQ.cuh
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 "IVFBase.cuh"
12 #include "../utils/Float16.cuh"
13
14 namespace faiss { namespace gpu {
15
16 /// Implementing class for IVFPQ on the GPU
17 class IVFPQ : public IVFBase {
18 public:
19 IVFPQ(GpuResources* resources,
20 /// We do not own this reference
21 FlatIndex* quantizer,
22 int numSubQuantizers,
23 int bitsPerSubQuantizer,
24 float* pqCentroidData,
25 IndicesOptions indicesOptions,
26 bool useFloat16LookupTables,
27 MemorySpace space);
28
29 /// Returns true if we support PQ in this size
30 static bool isSupportedPQCodeLength(int size);
31
32 /// For no precomputed codes, is this a supported sub-dimension
33 /// size?
34 /// FIXME: get MM implementation working again
35 static bool isSupportedNoPrecomputedSubDimSize(int dims);
36
37 ~IVFPQ() override;
38
39 /// Enable or disable pre-computed codes
40 void setPrecomputedCodes(bool enable);
41
42 /// Adds a set of codes and indices to a list; the data can be
43 /// resident on either the host or the device
44 void addCodeVectorsFromCpu(int listId,
45 const void* codes,
46 const long* indices,
47 size_t numVecs);
48
49 /// Calcuates the residual and quantizes the vectors, adding them to
50 /// this index
51 /// The input data must be on our current device.
52 /// Returns the number of vectors successfully added. Vectors may
53 /// not be able to be added because they contain NaNs.
54 int classifyAndAddVectors(Tensor<float, 2, true>& vecs,
55 Tensor<long, 1, true>& indices);
56
57 /// Find the approximate k nearest neigbors for `queries` against
58 /// our database
59 void query(Tensor<float, 2, true>& queries,
60 int nprobe,
61 int k,
62 Tensor<float, 2, true>& outDistances,
63 Tensor<long, 2, true>& outIndices);
64
65 /// Return the list codes of a particular list back to the CPU
66 std::vector<unsigned char> getListCodes(int listId) const;
67
68 /// Returns our set of sub-quantizers of the form
69 /// (sub q)(code id)(sub dim)
70 Tensor<float, 3, true> getPQCentroids();
71
72 private:
73 /// Sets the current product quantizer centroids; the data can be
74 /// resident on either the host or the device. It will be transposed
75 /// into our preferred data layout
76 /// Data must be a row-major, 3-d array of size
77 /// (numSubQuantizers, numSubQuantizerCodes, dim / numSubQuantizers)
78 void setPQCentroids_(float* data);
79
80 /// Calculate precomputed residual distance information
81 void precomputeCodes_();
82
83 /// Runs kernels for scanning inverted lists with precomputed codes
84 void runPQPrecomputedCodes_(Tensor<float, 2, true>& queries,
85 DeviceTensor<float, 2, true>& coarseDistances,
86 DeviceTensor<int, 2, true>& coarseIndices,
87 int k,
88 Tensor<float, 2, true>& outDistances,
89 Tensor<long, 2, true>& outIndices);
90
91 /// Runs kernels for scanning inverted lists without precomputed codes
92 void runPQNoPrecomputedCodes_(Tensor<float, 2, true>& queries,
93 DeviceTensor<float, 2, true>& coarseDistances,
94 DeviceTensor<int, 2, true>& coarseIndices,
95 int k,
96 Tensor<float, 2, true>& outDistances,
97 Tensor<long, 2, true>& outIndices);
98
99 private:
100 /// Number of sub-quantizers per vector
101 const int numSubQuantizers_;
102
103 /// Number of bits per sub-quantizer
104 const int bitsPerSubQuantizer_;
105
106 /// Number of per sub-quantizer codes (2^bits)
107 const int numSubQuantizerCodes_;
108
109 /// Number of dimensions per each sub-quantizer
110 const int dimPerSubQuantizer_;
111
112 /// Do we maintain precomputed terms and lookup tables in float16
113 /// form?
114 const bool useFloat16LookupTables_;
115
116 /// On the GPU, we prefer different PQ centroid data layouts for
117 /// different purposes.
118 ///
119 /// (sub q)(sub dim)(code id)
120 DeviceTensor<float, 3, true> pqCentroidsInnermostCode_;
121
122 /// (sub q)(code id)(sub dim)
123 DeviceTensor<float, 3, true> pqCentroidsMiddleCode_;
124
125 /// Are precomputed codes enabled? (additional factoring and
126 /// precomputation of the residual distance, to reduce query-time work)
127 bool precomputedCodes_;
128
129 /// Precomputed term 2 in float form
130 /// (centroid id)(sub q)(code id)
131 DeviceTensor<float, 3, true> precomputedCode_;
132
133 #ifdef FAISS_USE_FLOAT16
134 /// Precomputed term 2 in half form
135 DeviceTensor<half, 3, true> precomputedCodeHalf_;
136 #endif
137 };
138
139 } } // namespace
faiss::gpu::IVFPQ::addCodeVectorsFromCpu
void addCodeVectorsFromCpu(int listId, const void *codes, const long *indices, size_t numVecs)
Definition: IVFPQ.cu:345
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:21
faiss::gpu::IVFBase
Base inverted list functionality for IVFFlat and IVFPQ.
Definition: IVFBase.cuh:25
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::IVFPQ::isSupportedPQCodeLength
static bool isSupportedPQCodeLength(int size)
Returns true if we support PQ in this size.
Definition: IVFPQ.cu:70
faiss::gpu::IVFPQ::classifyAndAddVectors
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
Definition: IVFPQ.cu:118
faiss::gpu::IVFPQ::query
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
Definition: IVFPQ.cu:516
faiss::gpu::IVFPQ::getPQCentroids
Tensor< float, 3, true > getPQCentroids()
Definition: IVFPQ.cu:591
faiss::gpu::IVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: IVFPQ.cu:100
faiss::gpu::IVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Return the list codes of a particular list back to the CPU.
Definition: IVFPQ.cu:583
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::DeviceTensor< float, 2, true >
faiss::gpu::IVFPQ::IVFPQ
IVFPQ(GpuResources *resources, FlatIndex *quantizer, int numSubQuantizers, int bitsPerSubQuantizer, float *pqCentroidData, IndicesOptions indicesOptions, bool useFloat16LookupTables, MemorySpace space)
Definition: IVFPQ.cu:33
faiss::gpu::IVFPQ
Implementing class for IVFPQ on the GPU.
Definition: IVFPQ.cuh:17
faiss::gpu::IVFPQ::isSupportedNoPrecomputedSubDimSize
static bool isSupportedNoPrecomputedSubDimSize(int dims)
Definition: IVFPQ.cu:95