Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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