Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndex.h
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 "../Index.h"
12 #include "utils/MemorySpace.h"
13
14 namespace faiss { namespace gpu {
15
16 class GpuResources;
17
18 struct GpuIndexConfig {
19 inline GpuIndexConfig()
20 : device(0),
21 memorySpace(MemorySpace::Device) {
22 }
23
24 /// GPU device on which the index is resident
25 int device;
26
27 /// What memory space to use for primary storage.
28 /// On Pascal and above (CC 6+) architectures, allows GPUs to use
29 /// more memory than is available on the GPU.
30 MemorySpace memorySpace;
31 };
32
33 class GpuIndex : public faiss::Index {
34 public:
35 GpuIndex(GpuResources* resources,
36 int dims,
37 faiss::MetricType metric,
38 GpuIndexConfig config);
39
40 inline int getDevice() const {
41 return device_;
42 }
43
44 inline GpuResources* getResources() {
45 return resources_;
46 }
47
48 /// Set the minimum data size for searches (in MiB) for which we use
49 /// CPU -> GPU paging
50 void setMinPagingSize(size_t size);
51
52 /// Returns the current minimum data size for paged searches
53 size_t getMinPagingSize() const;
54
55 /// `x` can be resident on the CPU or any GPU; copies are performed
56 /// as needed
57 /// Handles paged adds if the add set is too large; calls addInternal_
58 void add(faiss::Index::idx_t, const float* x) override;
59
60 /// `x` and `ids` can be resident on the CPU or any GPU; copies are
61 /// performed as needed
62 /// Handles paged adds if the add set is too large; calls addInternal_
63 void add_with_ids(Index::idx_t n,
64 const float* x,
65 const Index::idx_t* ids) override;
66
67 /// `x`, `distances` and `labels` can be resident on the CPU or any
68 /// GPU; copies are performed as needed
69 void search(Index::idx_t n,
70 const float* x,
71 Index::idx_t k,
72 float* distances,
73 Index::idx_t* labels) const override;
74
75 protected:
76 /// Does addImpl_ require IDs? If so, and no IDs are provided, we will
77 /// generate them sequentially based on the order in which the IDs are added
78 virtual bool addImplRequiresIDs_() const = 0;
79
80 /// Overridden to actually perform the add
81 /// All data is guaranteed to be resident on our device
82 virtual void addImpl_(int n,
83 const float* x,
84 const Index::idx_t* ids) = 0;
85
86 /// Overridden to actually perform the search
87 /// All data is guaranteed to be resident on our device
88 virtual void searchImpl_(int n,
89 const float* x,
90 int k,
91 float* distances,
92 Index::idx_t* labels) const = 0;
93
94 private:
95 /// Handles paged adds if the add set is too large, passes to
96 /// addImpl_ to actually perform the add for the current page
97 void addPaged_(int n,
98 const float* x,
99 const Index::idx_t* ids);
100
101 /// Calls addImpl_ for a single page of GPU-resident data
102 void addPage_(int n,
103 const float* x,
104 const Index::idx_t* ids);
105
106 /// Calls searchImpl_ for a single page of GPU-resident data
107 void searchNonPaged_(int n,
108 const float* x,
109 int k,
110 float* outDistancesData,
111 Index::idx_t* outIndicesData) const;
112
113 /// Calls searchImpl_ for a single page of GPU-resident data,
114 /// handling paging of the data and copies from the CPU
115 void searchFromCpuPaged_(int n,
116 const float* x,
117 int k,
118 float* outDistancesData,
119 Index::idx_t* outIndicesData) const;
120
121 protected:
122 /// Manages streams, cuBLAS handles and scratch memory for devices
123 GpuResources* resources_;
124
125 /// The GPU device we are resident on
126 const int device_;
127
128 /// The memory space of our primary storage on the GPU
129 const MemorySpace memorySpace_;
130
131 /// Size above which we page copies from the CPU to GPU
132 size_t minPagedSize_;
133 };
134
135 } } // namespace
faiss::gpu::GpuIndex::getMinPagingSize
size_t getMinPagingSize() const
Returns the current minimum data size for paged searches.
Definition: GpuIndex.cu:74
faiss::gpu::GpuIndex::searchImpl_
virtual void searchImpl_(int n, const float *x, int k, float *distances, Index::idx_t *labels) const =0
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::GpuIndex::addImplRequiresIDs_
virtual bool addImplRequiresIDs_() const =0
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:25
faiss::gpu::GpuIndex
Definition: GpuIndex.h:33
faiss::gpu::GpuIndex::add_with_ids
void add_with_ids(Index::idx_t n, const float *x, const Index::idx_t *ids) override
Definition: GpuIndex.cu:85
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:30
faiss::gpu::GpuIndex::device_
const int device_
The GPU device we are resident on.
Definition: GpuIndex.h:126
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streams, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:123
faiss::gpu::GpuIndex::add
void add(faiss::Index::idx_t, const float *x) override
Definition: GpuIndex.cu:79
faiss::gpu::GpuIndex::memorySpace_
const MemorySpace memorySpace_
The memory space of our primary storage on the GPU.
Definition: GpuIndex.h:129
faiss::gpu::GpuIndexConfig
Definition: GpuIndex.h:18
faiss::Index
Definition: Index.h:61
faiss::gpu::GpuIndex::minPagedSize_
size_t minPagedSize_
Size above which we page copies from the CPU to GPU.
Definition: GpuIndex.h:132
faiss::gpu::GpuIndex::setMinPagingSize
void setMinPagingSize(size_t size)
Definition: GpuIndex.cu:69
faiss::gpu::GpuIndex::addImpl_
virtual void addImpl_(int n, const float *x, const Index::idx_t *ids)=0
faiss::gpu::GpuIndex::search
void search(Index::idx_t n, const float *x, Index::idx_t k, float *distances, Index::idx_t *labels) const override
Definition: GpuIndex.cu:176
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:44