Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndex.h
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 "../Index.h"
13 #include "utils/MemorySpace.h"
14
15 namespace faiss { namespace gpu {
16
17 class GpuResources;
18
19 struct GpuIndexConfig {
20 inline GpuIndexConfig()
21 : device(0),
22 memorySpace(MemorySpace::Device) {
23 }
24
25 /// GPU device on which the index is resident
26 int device;
27
28 /// What memory space to use for primary storage.
29 /// On Pascal and above (CC 6+) architectures, allows GPUs to use
30 /// more memory than is available on the GPU.
31 MemorySpace memorySpace;
32 };
33
34 class GpuIndex : public faiss::Index {
35 public:
36 GpuIndex(GpuResources* resources,
37 int dims,
38 faiss::MetricType metric,
39 GpuIndexConfig config);
40
41 int getDevice() const {
42 return device_;
43 }
44
45 GpuResources* getResources() {
46 return resources_;
47 }
48
49 /// `x` can be resident on the CPU or any GPU; copies are performed
50 /// as needed
51 /// Handles paged adds if the add set is too large; calls addInternal_
52 void add(faiss::Index::idx_t, const float* x) override;
53
54 /// `x` and `ids` can be resident on the CPU or any GPU; copies are
55 /// performed as needed
56 /// Handles paged adds if the add set is too large; calls addInternal_
57 void add_with_ids(Index::idx_t n, const float* x, const Index::idx_t* ids)
58 override;
59
60 /// `x`, `distances` and `labels` can be resident on the CPU or any
61 /// GPU; copies are performed as needed
62 void search(
63 faiss::Index::idx_t n,
64 const float* x,
65 faiss::Index::idx_t k,
66 float* distances,
67 faiss::Index::idx_t* labels) const override;
68
69 protected:
70 /// Handles paged adds if the add set is too large, passes to
71 /// addImpl_ to actually perform the add for the current page
72 void addInternal_(Index::idx_t n,
73 const float* x,
74 const Index::idx_t* ids);
75
76 /// Overridden to actually perform the add
77 virtual void addImpl_(Index::idx_t n,
78 const float* x,
79 const Index::idx_t* ids) = 0;
80
81 /// Overridden to actually perform the search
82 virtual void searchImpl_(faiss::Index::idx_t n,
83 const float* x,
84 faiss::Index::idx_t k,
85 float* distances,
86 faiss::Index::idx_t* labels) const = 0;
87
88 protected:
89 /// Manages streans, cuBLAS handles and scratch memory for devices
90 GpuResources* resources_;
91
92 /// The GPU device we are resident on
93 const int device_;
94
95 /// The memory space of our primary storage on the GPU
96 const MemorySpace memorySpace_;
97 };
98
99 } } // namespace
faiss::gpu::GpuIndex::search
void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Definition: GpuIndex.cu:107
faiss::gpu::GpuIndex::searchImpl_
virtual void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const =0
Overridden to actually perform the search.
faiss::gpu::GpuIndex::addInternal_
void addInternal_(Index::idx_t n, const float *x, const Index::idx_t *ids)
Definition: GpuIndex.cu:72
faiss::gpu::GpuResources
Definition: GpuResources.h:22
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::GpuIndex
Definition: GpuIndex.h:34
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:65
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:31
faiss::gpu::GpuIndex::device_
const int device_
The GPU device we are resident on.
Definition: GpuIndex.h:93
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streans, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:90
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndex::addImpl_
virtual void addImpl_(Index::idx_t n, const float *x, const Index::idx_t *ids)=0
Overridden to actually perform the add.
faiss::gpu::GpuIndex::add
void add(faiss::Index::idx_t, const float *x) override
Definition: GpuIndex.cu:60
faiss::gpu::GpuIndex::memorySpace_
const MemorySpace memorySpace_
The memory space of our primary storage on the GPU.
Definition: GpuIndex.h:96
faiss::gpu::GpuIndexConfig
Definition: GpuIndex.h:19
faiss::Index
Definition: Index.h:62
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:45