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