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 #include <faiss/gpu/GpuResources.h>
9 #include <faiss/gpu/utils/DeviceUtils.h>
10 #include <sstream>
11
12 namespace faiss {
13 namespace gpu {
14
allocTypeToString(AllocType t)15 std::string allocTypeToString(AllocType t) {
16 switch (t) {
17 case AllocType::Other:
18 return "Other";
19 case AllocType::FlatData:
20 return "FlatData";
21 case AllocType::IVFLists:
22 return "IVFLists";
23 case AllocType::Quantizer:
24 return "Quantizer";
25 case AllocType::QuantizerPrecomputedCodes:
26 return "QuantizerPrecomputedCodes";
27 case AllocType::TemporaryMemoryBuffer:
28 return "TemporaryMemoryBuffer";
29 case AllocType::TemporaryMemoryOverflow:
30 return "TemporaryMemoryOverflow";
31 default:
32 return "Unknown";
33 }
34 }
35
memorySpaceToString(MemorySpace s)36 std::string memorySpaceToString(MemorySpace s) {
37 switch (s) {
38 case MemorySpace::Temporary:
39 return "Temporary";
40 case MemorySpace::Device:
41 return "Device";
42 case MemorySpace::Unified:
43 return "Unified";
44 default:
45 return "Unknown";
46 }
47 }
48
toString() const49 std::string AllocInfo::toString() const {
50 std::stringstream ss;
51 ss << "type " << allocTypeToString(type) << " dev " << device << " space "
52 << memorySpaceToString(space) << " stream " << (void*)stream;
53
54 return ss.str();
55 }
56
toString() const57 std::string AllocRequest::toString() const {
58 std::stringstream ss;
59 ss << AllocInfo::toString() << " size " << size << " bytes";
60
61 return ss.str();
62 }
63
makeDevAlloc(AllocType at,cudaStream_t st)64 AllocInfo makeDevAlloc(AllocType at, cudaStream_t st) {
65 return AllocInfo(at, getCurrentDevice(), MemorySpace::Device, st);
66 }
67
makeTempAlloc(AllocType at,cudaStream_t st)68 AllocInfo makeTempAlloc(AllocType at, cudaStream_t st) {
69 return AllocInfo(at, getCurrentDevice(), MemorySpace::Temporary, st);
70 }
71
makeSpaceAlloc(AllocType at,MemorySpace sp,cudaStream_t st)72 AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st) {
73 return AllocInfo(at, getCurrentDevice(), sp, st);
74 }
75
76 //
77 // GpuMemoryReservation
78 //
79
GpuMemoryReservation()80 GpuMemoryReservation::GpuMemoryReservation()
81 : res(nullptr), device(0), stream(nullptr), data(nullptr), size(0) {}
82
GpuMemoryReservation(GpuResources * r,int dev,cudaStream_t str,void * p,size_t sz)83 GpuMemoryReservation::GpuMemoryReservation(
84 GpuResources* r,
85 int dev,
86 cudaStream_t str,
87 void* p,
88 size_t sz)
89 : res(r), device(dev), stream(str), data(p), size(sz) {}
90
GpuMemoryReservation(GpuMemoryReservation && m)91 GpuMemoryReservation::GpuMemoryReservation(GpuMemoryReservation&& m) noexcept {
92 res = m.res;
93 m.res = nullptr;
94 device = m.device;
95 m.device = 0;
96 stream = m.stream;
97 m.stream = nullptr;
98 data = m.data;
99 m.data = nullptr;
100 size = m.size;
101 m.size = 0;
102 }
103
operator =(GpuMemoryReservation && m)104 GpuMemoryReservation& GpuMemoryReservation::operator=(
105 GpuMemoryReservation&& m) {
106 // Can't be both a valid allocation and the same allocation
107 FAISS_ASSERT(
108 !(res && res == m.res && device == m.device && data == m.data));
109
110 release();
111 res = m.res;
112 m.res = nullptr;
113 device = m.device;
114 m.device = 0;
115 stream = m.stream;
116 m.stream = nullptr;
117 data = m.data;
118 m.data = nullptr;
119 size = m.size;
120 m.size = 0;
121
122 return *this;
123 }
124
release()125 void GpuMemoryReservation::release() {
126 if (res) {
127 res->deallocMemory(device, data);
128 res = nullptr;
129 device = 0;
130 stream = nullptr;
131 data = nullptr;
132 size = 0;
133 }
134 }
135
~GpuMemoryReservation()136 GpuMemoryReservation::~GpuMemoryReservation() {
137 if (res) {
138 res->deallocMemory(device, data);
139 }
140 }
141
142 //
143 // GpuResources
144 //
145
~GpuResources()146 GpuResources::~GpuResources() {}
147
getBlasHandleCurrentDevice()148 cublasHandle_t GpuResources::getBlasHandleCurrentDevice() {
149 return getBlasHandle(getCurrentDevice());
150 }
151
getDefaultStreamCurrentDevice()152 cudaStream_t GpuResources::getDefaultStreamCurrentDevice() {
153 return getDefaultStream(getCurrentDevice());
154 }
155
getAlternateStreamsCurrentDevice()156 std::vector<cudaStream_t> GpuResources::getAlternateStreamsCurrentDevice() {
157 return getAlternateStreams(getCurrentDevice());
158 }
159
getAsyncCopyStreamCurrentDevice()160 cudaStream_t GpuResources::getAsyncCopyStreamCurrentDevice() {
161 return getAsyncCopyStream(getCurrentDevice());
162 }
163
syncDefaultStream(int device)164 void GpuResources::syncDefaultStream(int device) {
165 CUDA_VERIFY(cudaStreamSynchronize(getDefaultStream(device)));
166 }
167
syncDefaultStreamCurrentDevice()168 void GpuResources::syncDefaultStreamCurrentDevice() {
169 syncDefaultStream(getCurrentDevice());
170 }
171
allocMemoryHandle(const AllocRequest & req)172 GpuMemoryReservation GpuResources::allocMemoryHandle(const AllocRequest& req) {
173 return GpuMemoryReservation(
174 this, req.device, req.stream, allocMemory(req), req.size);
175 }
176
getTempMemoryAvailableCurrentDevice() const177 size_t GpuResources::getTempMemoryAvailableCurrentDevice() const {
178 return getTempMemoryAvailable(getCurrentDevice());
179 }
180
181 //
182 // GpuResourcesProvider
183 //
184
~GpuResourcesProvider()185 GpuResourcesProvider::~GpuResourcesProvider() {}
186
187 } // namespace gpu
188 } // namespace faiss
189