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
30 bool getUseFloat16() const;
31
32 /// Returns the number of vectors we contain
33 int getSize() const;
34
35 int getDim() const;
36
37 /// Returns a reference to our vectors currently in use
38 Tensor<float, 2, true>& getVectorsFloat32Ref();
39
40 #ifdef FAISS_USE_FLOAT16
41 /// Returns a reference to our vectors currently in use (useFloat16 mode)
42 Tensor<half, 2, true>& getVectorsFloat16Ref();
43 #endif
44
45 /// Performs a copy of the vectors on the given device, converting
46 /// as needed from float16
47 DeviceTensor<float, 2, true> getVectorsFloat32Copy(cudaStream_t stream);
48
49 /// Returns only a subset of the vectors
50 DeviceTensor<float, 2, true> getVectorsFloat32Copy(int from,
51 int num,
52 cudaStream_t stream);
53
54 void query(Tensor<float, 2, true>& vecs,
55 int k,
56 Tensor<float, 2, true>& outDistances,
57 Tensor<int, 2, true>& outIndices,
58 bool exactDistance,
59 int tileSize = -1);
60
61 #ifdef FAISS_USE_FLOAT16
62 void query(Tensor<half, 2, true>& vecs,
63 int k,
64 Tensor<half, 2, true>& outDistances,
65 Tensor<int, 2, true>& outIndices,
66 bool exactDistance,
67 int tileSize = -1);
68 #endif
69
70 /// Add vectors to ourselves; the pointer passed can be on the host
71 /// or the device
72 void add(const float* data, int numVecs, cudaStream_t stream);
73
74 /// Free all storage
75 void reset();
76
77 private:
78 /// Collection of GPU resources that we use
79 GpuResources* resources_;
80
81 /// Dimensionality of our vectors
82 const int dim_;
83
84 /// Float16 data format
85 const bool useFloat16_;
86
87 /// L2 or inner product distance?
88 bool l2Distance_;
89
90 /// How many vectors we have
91 int num_;
92
93 /// The underlying expandable storage
94 DeviceVector<char> rawData_;
95
96 /// Vectors currently in rawData_
97 DeviceTensor<float, 2, true> vectors_;
98
99 #ifdef FAISS_USE_FLOAT16
100 /// Vectors currently in rawData_, float16 form
101 DeviceTensor<half, 2, true> vectorsHalf_;
102 #endif
103
104 /// Precomputed L2 norms
105 DeviceTensor<float, 1, true> norms_;
106
107 #ifdef FAISS_USE_FLOAT16
108 /// Precomputed L2 norms, float16 form
109 DeviceTensor<half, 1, true> normsHalf_;
110 #endif
111 };
112
113 } } // namespace
faiss::gpu::FlatIndex::getVectorsFloat32Copy
DeviceTensor< float, 2, true > getVectorsFloat32Copy(cudaStream_t stream)
Definition: FlatIndex.cu:73
faiss::gpu::FlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: FlatIndex.cu:40
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:177
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:61
faiss::gpu::FlatIndex::reset
void reset()
Free all storage.
Definition: FlatIndex.cu:235
faiss::gpu::DeviceTensor< float, 2, true >