Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndex.h
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 "../Index.h"
15
16 namespace faiss { namespace gpu {
17
18 class GpuResources;
19
20 class GpuIndex : public faiss::Index {
21 public:
22 GpuIndex(GpuResources* resources,
23 int device,
24 int dims,
25 faiss::MetricType metric);
26
27 inline int getDevice() const {
28 return device_;
29 }
30
31 GpuResources* getResources() {
32 return resources_;
33 }
34
35 /// `x` can be resident on the CPU or any GPU; copies are performed
36 /// as needed
37 /// Handles paged adds if the add set is too large; calls addInternal_
38 virtual void add(faiss::Index::idx_t, const float* x);
39
40 /// `x` and `ids` can be resident on the CPU or any GPU; copies are
41 /// performed as needed
42 /// Handles paged adds if the add set is too large; calls addInternal_
43 virtual void add_with_ids(Index::idx_t n,
44 const float* x,
45 const Index::idx_t* ids);
46
47 /// `x`, `distances` and `labels` can be resident on the CPU or any
48 /// GPU; copies are performed as needed
49 virtual void search(faiss::Index::idx_t n,
50 const float* x,
51 faiss::Index::idx_t k,
52 float* distances,
53 faiss::Index::idx_t* labels) const;
54
55
56 protected:
57 /// Handles paged adds if the add set is too large, passes to
58 /// addImpl_ to actually perform the add for the current page
59 void addInternal_(Index::idx_t n,
60 const float* x,
61 const Index::idx_t* ids);
62
63 /// Overridden to actually perform the add
64 virtual void addImpl_(Index::idx_t n,
65 const float* x,
66 const Index::idx_t* ids) = 0;
67
68 /// Overridden to actually perform the search
69 virtual void searchImpl_(faiss::Index::idx_t n,
70 const float* x,
71 faiss::Index::idx_t k,
72 float* distances,
73 faiss::Index::idx_t* labels) const = 0;
74
75 protected:
76 /// Manages streans, cuBLAS handles and scratch memory for devices
77 GpuResources* resources_;
78
79 /// The GPU device we are resident on
80 int device_;
81 };
82
83 } } // namespace
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:49
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::GpuIndex
Definition: GpuIndex.h:20
faiss::gpu::GpuIndex::device_
int device_
The GPU device we are resident on.
Definition: GpuIndex.h:80
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streans, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:77
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
virtual void add(faiss::Index::idx_t, const float *x)
Definition: GpuIndex.cu:37
faiss::gpu::GpuIndex::add_with_ids
virtual void add_with_ids(Index::idx_t n, const float *x, const Index::idx_t *ids)
Definition: GpuIndex.cu:42
faiss::gpu::GpuIndex::search
virtual void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const
Definition: GpuIndex.cu:80
faiss::Index
Definition: Index.h:61
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44