Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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