Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
FlatIndex.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 "../utils/DeviceTensor.cuh"
12 #include "../utils/DeviceVector.cuh"
13 #include "../utils/Float16.cuh"
14 #include "../utils/MemorySpace.h"
15
16 namespace faiss { namespace gpu {
17
18 class GpuResources;
19
20 /// Holder of GPU resources for a particular flat index
21 class FlatIndex {
22 public:
23 FlatIndex(GpuResources* res,
24 int dim,
25 bool l2Distance,
26 bool useFloat16,
27 bool useFloat16Accumulator,
28 bool storeTransposed,
29 MemorySpace space);
30
31 bool getUseFloat16() const;
32
33 /// Returns the number of vectors we contain
34 int getSize() const;
35
36 int getDim() const;
37
38 /// Reserve storage that can contain at least this many vectors
39 void reserve(size_t numVecs, cudaStream_t stream);
40
41 /// Returns a reference to our vectors currently in use
42 Tensor<float, 2, true>& getVectorsFloat32Ref();
43
44 #ifdef FAISS_USE_FLOAT16
45 /// Returns a reference to our vectors currently in use (useFloat16 mode)
46 Tensor<half, 2, true>& getVectorsFloat16Ref();
47 #endif
48
49 /// Performs a copy of the vectors on the given device, converting
50 /// as needed from float16
51 DeviceTensor<float, 2, true> getVectorsFloat32Copy(cudaStream_t stream);
52
53 /// Returns only a subset of the vectors
54 DeviceTensor<float, 2, true> getVectorsFloat32Copy(int from,
55 int num,
56 cudaStream_t stream);
57
58 void query(Tensor<float, 2, true>& vecs,
59 int k,
60 Tensor<float, 2, true>& outDistances,
61 Tensor<int, 2, true>& outIndices,
62 bool exactDistance);
63
64 #ifdef FAISS_USE_FLOAT16
65 void query(Tensor<half, 2, true>& vecs,
66 int k,
67 Tensor<half, 2, true>& outDistances,
68 Tensor<int, 2, true>& outIndices,
69 bool exactDistance);
70 #endif
71
72 /// Add vectors to ourselves; the pointer passed can be on the host
73 /// or the device
74 void add(const float* data, int numVecs, cudaStream_t stream);
75
76 /// Free all storage
77 void reset();
78
79 private:
80 /// Collection of GPU resources that we use
81 GpuResources* resources_;
82
83 /// Dimensionality of our vectors
84 const int dim_;
85
86 /// Float16 data format
87 const bool useFloat16_;
88
89 /// For supporting hardware, whether or not we use Hgemm
90 const bool useFloat16Accumulator_;
91
92 /// Store vectors in transposed layout for speed; makes addition to
93 /// the index slower
94 const bool storeTransposed_;
95
96 /// L2 or inner product distance?
97 bool l2Distance_;
98
99 /// Memory space for our allocations
100 MemorySpace space_;
101
102 /// How many vectors we have
103 int num_;
104
105 /// The underlying expandable storage
106 DeviceVector<char> rawData_;
107
108 /// Vectors currently in rawData_
109 DeviceTensor<float, 2, true> vectors_;
110 DeviceTensor<float, 2, true> vectorsTransposed_;
111
112 #ifdef FAISS_USE_FLOAT16
113 /// Vectors currently in rawData_, float16 form
114 DeviceTensor<half, 2, true> vectorsHalf_;
115 DeviceTensor<half, 2, true> vectorsHalfTransposed_;
116 #endif
117
118 /// Precomputed L2 norms
119 DeviceTensor<float, 1, true> norms_;
120
121 #ifdef FAISS_USE_FLOAT16
122 /// Precomputed L2 norms, float16 form
123 DeviceTensor<half, 1, true> normsHalf_;
124 #endif
125 };
126
127 } } // namespace
faiss::gpu::FlatIndex::getVectorsFloat32Copy
DeviceTensor< float, 2, true > getVectorsFloat32Copy(cudaStream_t stream)
Definition: FlatIndex.cu:89
faiss::gpu::FlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: FlatIndex.cu:45
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:21
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::DeviceVector< char >
faiss::gpu::FlatIndex::reserve
void reserve(size_t numVecs, cudaStream_t stream)
Reserve storage that can contain at least this many vectors.
Definition: FlatIndex.cu:66
faiss::gpu::FlatIndex::add
void add(const float *data, int numVecs, cudaStream_t stream)
Definition: FlatIndex.cu:196
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::FlatIndex::getVectorsFloat32Ref
Tensor< float, 2, true > & getVectorsFloat32Ref()
Returns a reference to our vectors currently in use.
Definition: FlatIndex.cu:77
faiss::gpu::FlatIndex::reset
void reset()
Free all storage.
Definition: FlatIndex.cu:270
faiss::gpu::DeviceTensor< float, 2, true >