Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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