Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFFlat.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 "GpuIndexIVF.h"
13
14 namespace faiss { struct IndexIVFFlat; }
15
16 namespace faiss { namespace gpu {
17
18 class IVFFlat;
19 class GpuIndexFlat;
20
21 struct GpuIndexIVFFlatConfig : public GpuIndexIVFConfig {
22 inline GpuIndexIVFFlatConfig()
23 : useFloat16IVFStorage(false) {
24 }
25
26 /// Whether or not IVFFlat inverted list storage is in float16;
27 /// supported on all architectures
28 bool useFloat16IVFStorage;
29 };
30
31 /// Wrapper around the GPU implementation that looks like
32 /// faiss::IndexIVFFlat
33 class GpuIndexIVFFlat : public GpuIndexIVF {
34 public:
35 /// Construct from a pre-existing faiss::IndexIVFFlat instance, copying
36 /// data over to the given GPU, if the input index is trained.
37 GpuIndexIVFFlat(GpuResources* resources,
38 const faiss::IndexIVFFlat* index,
39 GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
40
41 /// Constructs a new instance with an empty flat quantizer; the user
42 /// provides the number of lists desired.
43 GpuIndexIVFFlat(GpuResources* resources,
44 int dims,
45 int nlist,
46 faiss::MetricType metric,
47 GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig());
48
49 ~GpuIndexIVFFlat() override;
50
51 /// Reserve GPU memory in our inverted lists for this number of vectors
52 void reserveMemory(size_t numVecs);
53
54 /// Initialize ourselves from the given CPU index; will overwrite
55 /// all data in ourselves
56 void copyFrom(const faiss::IndexIVFFlat* index);
57
58 /// Copy ourselves to the given CPU index; will overwrite all data
59 /// in the index instance
60 void copyTo(faiss::IndexIVFFlat* index) const;
61
62 /// After adding vectors, one can call this to reclaim device memory
63 /// to exactly the amount needed. Returns space reclaimed in bytes
64 size_t reclaimMemory();
65
66 void reset() override;
67
68 void train(Index::idx_t n, const float* x) override;
69
70 protected:
71 /// Called from GpuIndex for add/add_with_ids
72 void addImpl_(
73 faiss::Index::idx_t n,
74 const float* x,
75 const faiss::Index::idx_t* ids) override;
76
77 /// Called from GpuIndex for search
78 void searchImpl_(
79 faiss::Index::idx_t n,
80 const float* x,
81 faiss::Index::idx_t k,
82 float* distances,
83 faiss::Index::idx_t* labels) const override;
84
85 private:
86 GpuIndexIVFFlatConfig ivfFlatConfig_;
87
88 /// Desired inverted list memory reservation
89 size_t reserveMemoryVecs_;
90
91 /// Instance that we own; contains the inverted list
92 IVFFlat* index_;
93 };
94
95 } } // namespace
faiss::gpu::IVFFlat
Definition: IVFFlat.cuh:16
faiss::gpu::GpuIndexIVFFlat::GpuIndexIVFFlat
GpuIndexIVFFlat(GpuResources *resources, const faiss::IndexIVFFlat *index, GpuIndexIVFFlatConfig config=GpuIndexIVFFlatConfig())
Definition: GpuIndexIVFFlat.cu:24
faiss::gpu::GpuIndexIVFFlat::searchImpl_
void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Called from GpuIndex for search.
Definition: GpuIndexIVFFlat.cu:233
faiss::gpu::GpuIndexIVFFlat::copyFrom
void copyFrom(const faiss::IndexIVFFlat *index)
Definition: GpuIndexIVFFlat.cu:78
faiss::gpu::GpuIndexIVFConfig
Definition: GpuIndexIVF.h:24
faiss::gpu::GpuIndexIVFFlat::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFFlat.cu:150
faiss::gpu::GpuIndexIVFFlatConfig::useFloat16IVFStorage
bool useFloat16IVFStorage
Definition: GpuIndexIVFFlat.h:28
faiss::gpu::GpuResources
Definition: GpuResources.h:22
faiss::IndexIVFFlat
Definition: IndexIVFFlat.h:25
faiss::gpu::GpuIndexIVFFlat::addImpl_
void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids) override
Called from GpuIndex for add/add_with_ids.
Definition: GpuIndexIVFFlat.cu:203
faiss::gpu::GpuIndexIVFFlat::train
void train(Index::idx_t n, const float *x) override
Definition: GpuIndexIVFFlat.cu:173
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:36
faiss::gpu::GpuIndexIVFFlat::copyTo
void copyTo(faiss::IndexIVFFlat *index) const
Definition: GpuIndexIVFFlat.cu:122
faiss::gpu::GpuIndexIVFFlat::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFFlat.cu:70
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndexIVFFlatConfig
Definition: GpuIndexIVFFlat.h:21
faiss::gpu::GpuIndexIVFFlat
Definition: GpuIndexIVFFlat.h:33
faiss::gpu::GpuIndexIVFFlat::reset
void reset() override
removes all elements from the database.
Definition: GpuIndexIVFFlat.cu:161
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:45