Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFBase.cuh
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 "../GpuIndicesOptions.h"
15 #include "../utils/DeviceVector.cuh"
16 #include "../utils/DeviceTensor.cuh"
17 #include <memory>
18 #include <thrust/device_vector.h>
19 #include <vector>
20
21 namespace faiss { namespace gpu {
22
23 class GpuResources;
24 struct FlatIndex;
25
26 /// Base inverted list functionality for IVFFlat and IVFPQ
27 class IVFBase {
28 public:
29 IVFBase(GpuResources* resources,
30 /// We do not own this reference
31 FlatIndex* quantizer,
32 int bytesPerVector,
33 IndicesOptions indicesOptions);
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 /// Device representation of all inverted list data
99 /// id -> data
100 thrust::device_vector<void*> deviceListDataPointers_;
101
102 /// Device representation of all inverted list index pointers
103 /// id -> data
104 thrust::device_vector<void*> deviceListIndexPointers_;
105
106 /// Device representation of all inverted list lengths
107 /// id -> length
108 thrust::device_vector<int> deviceListLengths_;
109
110 /// Maximum list length seen
111 int maxListLength_;
112
113 /// Device memory for each separate list, as managed by the host.
114 /// Device memory as stored in DeviceVector is stored as unique_ptr
115 /// since deviceListSummary_ pointers must remain valid despite
116 /// resizing of deviceLists_
117 std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListData_;
118 std::vector<std::unique_ptr<DeviceVector<unsigned char>>> deviceListIndices_;
119
120 /// If we are storing indices on the CPU (indicesOptions_ is
121 /// INDICES_CPU), then this maintains a CPU-side map of what
122 /// (inverted list id, offset) maps to which user index
123 std::vector<std::vector<long>> listOffsetToUserIndex_;
124 };
125
126 } } // 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:111
faiss::gpu::IVFBase::listOffsetToUserIndex_
std::vector< std::vector< long > > listOffsetToUserIndex_
Definition: IVFBase.cuh:123
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:23
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:27
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::IVFBase::deviceListLengths_
thrust::device_vector< int > deviceListLengths_
Definition: IVFBase.cuh:108
faiss::gpu::IVFBase::deviceListIndexPointers_
thrust::device_vector< void * > deviceListIndexPointers_
Definition: IVFBase.cuh:104
faiss::gpu::IVFBase::IVFBase
IVFBase(GpuResources *resources, FlatIndex *quantizer, int bytesPerVector, IndicesOptions indicesOptions)
Definition: IVFBase.cu:26
faiss::gpu::IVFBase::quantizer_
FlatIndex * quantizer_
Quantizer object.
Definition: IVFBase.cuh:84
faiss::gpu::IVFBase::deviceListDataPointers_
thrust::device_vector< void * > deviceListDataPointers_
Definition: IVFBase.cuh:100
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::deviceListData_
std::vector< std::unique_ptr< DeviceVector< unsigned char > > > deviceListData_
Definition: IVFBase.cuh:117
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