Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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