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