Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFBase.cuh
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 "../GpuIndicesOptions.h"
13 #include "../utils/DeviceVector.cuh"
14 #include "../utils/DeviceTensor.cuh"
15 #include "../utils/MemorySpace.h"
16 #include <memory>
17 #include <thrust/device_vector.h>
18 #include <vector>
19
20 namespace faiss { namespace gpu {
21
22 class GpuResources;
23 struct FlatIndex;
24
25 /// Base inverted list functionality for IVFFlat and IVFPQ
26 class IVFBase {
27 public:
28 IVFBase(GpuResources* resources,
29 /// We do not own this reference
30 FlatIndex* quantizer,
31 int bytesPerVector,
32 IndicesOptions indicesOptions,
33 MemorySpace space);
34
35 virtual ~IVFBase();
36
37 /// Reserve GPU memory in our inverted lists for this number of vectors
38 void reserveMemory(size_t numVecs);
39
40 /// Clear out all inverted lists, but retain the coarse quantizer
41 /// and the product quantizer info
42 void reset();
43
44 /// Return the number of dimensions we are indexing
45 int getDim() const;
46
47 /// After adding vectors, one can call this to reclaim device memory
48 /// to exactly the amount needed. Returns space reclaimed in bytes
49 size_t reclaimMemory();
50
51 /// Returns the number of inverted lists
52 size_t getNumLists() const;
53
54 /// For debugging purposes, return the list length of a particular
55 /// list
56 int getListLength(int listId) const;
57
58 /// Return the list indices of a particular list back to the CPU
59 std::vector<long> getListIndices(int listId) const;
60
61 protected:
62 /// Reclaim memory consumed on the device for our inverted lists
63 /// `exact` means we trim exactly to the memory needed
64 size_t reclaimMemory_(bool exact);
65
66 /// Update all device-side list pointer and size information
67 void updateDeviceListInfo_(cudaStream_t stream);
68
69 /// For a set of list IDs, update device-side list pointer and size
70 /// information
71 void updateDeviceListInfo_(const std::vector<int>& listIds,
72 cudaStream_t stream);
73
74 /// Shared function to copy indices from CPU to GPU
75 void addIndicesFromCpu_(int listId,
76 const long* indices,
77 size_t numVecs);
78
79 protected:
80 /// Collection of GPU resources that we use
81 GpuResources* resources_;
82
83 /// Quantizer object
84 FlatIndex* quantizer_;
85
86 /// Expected dimensionality of the vectors
87 const int dim_;
88
89 /// Number of inverted lists we maintain
90 const int numLists_;
91
92 /// Number of bytes per vector in the list
93 const int bytesPerVector_;
94
95 /// How are user indices stored on the GPU?
96 const IndicesOptions indicesOptions_;
97
98 /// What memory space our inverted list storage is in
99 const MemorySpace space_;
100
101 /// Device representation of all inverted list data
102 /// id -> data
103 thrust::device_vector<void*> deviceListDataPointers_;
104
105 /// Device representation of all inverted list index pointers
106 /// id -> data
107 thrust::device_vector<void*> deviceListIndexPointers_;
108
109 /// Device representation of all inverted list lengths
110 /// id -> length
111 thrust::device_vector<int> deviceListLengths_;
112
113 /// Maximum list length seen
114 int maxListLength_;
115
116 /// Device memory for each separate list, as managed by the host.
117 /// Device memory as stored in DeviceVector is stored as unique_ptr
118 /// since deviceListSummary_ pointers must remain valid despite
119 /// resizing of deviceLists_
120 std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListData_;
121 std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListIndices_;
122
123 /// If we are storing indices on the CPU (indicesOptions_ is
124 /// INDICES_CPU), then this maintains a CPU-side map of what
125 /// (inverted list id, offset) maps to which user index
126 std::vector<std::vector<long>> listOffsetToUserIndex_;
127 };
128
129 } } // namespace
faiss::gpu::IVFBase::numLists_
const int numLists_
Number of inverted lists we maintain.
Definition: IVFBase.cuh:90
faiss::gpu::IVFBase::maxListLength_
int maxListLength_
Maximum list length seen.
Definition: IVFBase.cuh:114
faiss::gpu::IVFBase::IVFBase
IVFBase(GpuResources *resources, FlatIndex *quantizer, int bytesPerVector, IndicesOptions indicesOptions, MemorySpace space)
Definition: IVFBase.cu:24
faiss::gpu::IVFBase::listOffsetToUserIndex_
std::vector< std::vector< long > > listOffsetToUserIndex_
Definition: IVFBase.cuh:126
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:22
faiss::gpu::IVFBase::getDim
int getDim() const
Return the number of dimensions we are indexing.
Definition: IVFBase.cu:99
faiss::gpu::IVFBase::getListLength
int getListLength(int listId) const
Definition: IVFBase.cu:199
faiss::gpu::IVFBase::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: IVFBase.cu:44
faiss::gpu::IVFBase::reclaimMemory_
size_t reclaimMemory_(bool exact)
Definition: IVFBase.cu:110
faiss::gpu::IVFBase
Base inverted list functionality for IVFFlat and IVFPQ.
Definition: IVFBase.cuh:26
faiss::gpu::GpuResources
Definition: GpuResources.h:22
faiss::gpu::IVFBase::deviceListLengths_
thrust::device_vector< int > deviceListLengths_
Definition: IVFBase.cuh:111
faiss::gpu::IVFBase::deviceListIndexPointers_
thrust::device_vector< void * > deviceListIndexPointers_
Definition: IVFBase.cuh:107
faiss::gpu::IVFBase::quantizer_
FlatIndex * quantizer_
Quantizer object.
Definition: IVFBase.cuh:84
faiss::gpu::IVFBase::deviceListDataPointers_
thrust::device_vector< void * > deviceListDataPointers_
Definition: IVFBase.cuh:103
faiss::gpu::IVFBase::reset
void reset()
Definition: IVFBase.cu:74
faiss::gpu::IVFBase::resources_
GpuResources * resources_
Collection of GPU resources that we use.
Definition: IVFBase.cuh:81
faiss::gpu::IVFBase::bytesPerVector_
const int bytesPerVector_
Number of bytes per vector in the list.
Definition: IVFBase.cuh:93
faiss::gpu::IVFBase::updateDeviceListInfo_
void updateDeviceListInfo_(cudaStream_t stream)
Update all device-side list pointer and size information.
Definition: IVFBase.cu:137
faiss::gpu::IVFBase::getListIndices
std::vector< long > getListIndices(int listId) const
Return the list indices of a particular list back to the CPU.
Definition: IVFBase.cu:206
faiss::gpu::IVFBase::indicesOptions_
const IndicesOptions indicesOptions_
How are user indices stored on the GPU?
Definition: IVFBase.cuh:96
faiss::gpu::IVFBase::space_
const MemorySpace space_
What memory space our inverted list storage is in.
Definition: IVFBase.cuh:99
faiss::gpu::IVFBase::deviceListData_
std::vector< std::unique_ptr< DeviceVector< unsigned char > > > deviceListData_
Definition: IVFBase.cuh:120
faiss::gpu::IVFBase::dim_
const int dim_
Expected dimensionality of the vectors.
Definition: IVFBase.cuh:87
faiss::gpu::IVFBase::addIndicesFromCpu_
void addIndicesFromCpu_(int listId, const long *indices, size_t numVecs)
Shared function to copy indices from CPU to GPU.
Definition: IVFBase.cu:244
faiss::gpu::IVFBase::reclaimMemory
size_t reclaimMemory()
Definition: IVFBase.cu:104
faiss::gpu::IVFBase::getNumLists
size_t getNumLists() const
Returns the number of inverted lists.
Definition: IVFBase.cu:194