Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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