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