Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
StandardGpuResources.cpp
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 #include "StandardGpuResources.h"
12 #include "../FaissAssert.h"
13
14 namespace faiss { namespace gpu {
15
16 namespace {
17
18 constexpr int kNumStreams = 2;
19
20 /// Use 18% of GPU memory for temporary space by default
21 constexpr float kDefaultTempMemFraction = 0.18f;
22
23 /// Default pinned memory allocation size
24 constexpr size_t kDefaultPinnedMemoryAllocation = (size_t) 256 * 1024 * 1024;
25
26 }
27
28 StandardGpuResources::StandardGpuResources() :
29 pinnedMemAlloc_(nullptr),
30 pinnedMemAllocSize_(0),
31 tempMemFraction_(kDefaultTempMemFraction),
32 tempMemSize_(0),
33 useFraction_(true),
34 pinnedMemSize_(kDefaultPinnedMemoryAllocation) {
35 }
36
37 StandardGpuResources::~StandardGpuResources() {
38 for (auto& entry : defaultStreams_) {
39 DeviceScope scope(entry.first);
40
41 auto it = userDefaultStreams_.find(entry.first);
42 if (it == userDefaultStreams_.end()) {
43 // The user did not specify this stream, thus we are the ones
44 // who have created it
45 CUDA_VERIFY(cudaStreamDestroy(entry.second));
46 }
47 }
48
49 for (auto& entry : alternateStreams_) {
50 DeviceScope scope(entry.first);
51
52 for (auto stream : entry.second) {
53 CUDA_VERIFY(cudaStreamDestroy(stream));
54 }
55 }
56
57 for (auto& entry : asyncCopyStreams_) {
58 DeviceScope scope(entry.first);
59
60 CUDA_VERIFY(cudaStreamDestroy(entry.second));
61 }
62
63 for (auto& entry : blasHandles_) {
64 DeviceScope scope(entry.first);
65
66 auto blasStatus = cublasDestroy(entry.second);
67 FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
68 }
69
70 if (pinnedMemAlloc_) {
71 CUDA_VERIFY(cudaFreeHost(pinnedMemAlloc_));
72 }
73 }
74
75 void
76 StandardGpuResources::noTempMemory() {
77 setTempMemory(0);
78 }
79
80 void
81 StandardGpuResources::setTempMemory(size_t size) {
82 useFraction_ = false;
83 tempMemSize_ = size;
84 }
85
86 void
87 StandardGpuResources::setTempMemoryFraction(float fraction) {
88 FAISS_ASSERT(fraction >= 0.0f && fraction <= 0.5f);
89 useFraction_ = true;
90 tempMemFraction_ = fraction;
91 }
92
93 void
94 StandardGpuResources::setPinnedMemory(size_t size) {
95 // Should not call this after devices have been initialized
96 FAISS_ASSERT(defaultStreams_.size() == 0);
97 FAISS_ASSERT(!pinnedMemAlloc_);
98
99 pinnedMemSize_ = size;
100 }
101
102 void
103 StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
104 auto it = defaultStreams_.find(device);
105 if (it != defaultStreams_.end()) {
106 // Replace this stream with the user stream
107 CUDA_VERIFY(cudaStreamDestroy(it->second));
108 it->second = stream;
109 }
110
111 userDefaultStreams_[device] = stream;
112 }
113
114 void
115 StandardGpuResources::setDefaultNullStreamAllDevices() {
116 for (int dev = 0; dev < getNumDevices(); ++dev) {
117 setDefaultStream(dev, nullptr);
118 }
119 }
120
121 void
122 StandardGpuResources::initializeForDevice(int device) {
123 // Use default streams as a marker for whether or not a certain
124 // device has been initialized
125 if (defaultStreams_.count(device) != 0) {
126 return;
127 }
128
129 // If this is the first device that we're initializing, create our
130 // pinned memory allocation
131 if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
132 CUDA_VERIFY(cudaHostAlloc(&pinnedMemAlloc_,
133 pinnedMemSize_,
134 cudaHostAllocDefault));
135 pinnedMemAllocSize_ = pinnedMemSize_;
136 }
137
138 FAISS_ASSERT(device < getNumDevices());
139 DeviceScope scope(device);
140
141 // Make sure that device properties for all devices are cached
142 auto& prop = getDeviceProperties(device);
143
144 // Also check to make sure we meet our minimum compute capability (3.0)
145 FAISS_ASSERT_FMT(prop.major >= 3,
146 "Device id %d with CC %d.%d not supported, "
147 "need 3.0+ compute capability",
148 device, prop.major, prop.minor);
149
150 // Create streams
151 cudaStream_t defaultStream = 0;
152 auto it = userDefaultStreams_.find(device);
153 if (it != userDefaultStreams_.end()) {
154 // We already have a stream provided by the user
155 defaultStream = it->second;
156 } else {
157 CUDA_VERIFY(cudaStreamCreateWithFlags(&defaultStream,
158 cudaStreamNonBlocking));
159 }
160
161 defaultStreams_[device] = defaultStream;
162
163 cudaStream_t asyncCopyStream = 0;
164 CUDA_VERIFY(cudaStreamCreateWithFlags(&asyncCopyStream,
165 cudaStreamNonBlocking));
166
167 asyncCopyStreams_[device] = asyncCopyStream;
168
169 std::vector<cudaStream_t> deviceStreams;
170 for (int j = 0; j < kNumStreams; ++j) {
171 cudaStream_t stream = 0;
172 CUDA_VERIFY(cudaStreamCreateWithFlags(&stream,
173 cudaStreamNonBlocking));
174
175 deviceStreams.push_back(stream);
176 }
177
178 alternateStreams_[device] = std::move(deviceStreams);
179
180 // Create cuBLAS handle
181 cublasHandle_t blasHandle = 0;
182 auto blasStatus = cublasCreate(&blasHandle);
183 FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
184 blasHandles_[device] = blasHandle;
185
186 size_t toAlloc = 0;
187 if (useFraction_) {
188 size_t devFree = 0;
189 size_t devTotal = 0;
190
191 CUDA_VERIFY(cudaMemGetInfo(&devFree, &devTotal));
192
193 toAlloc = (size_t) (tempMemFraction_ * devTotal);
194 } else {
195 toAlloc = tempMemSize_;
196 }
197
198 FAISS_ASSERT(memory_.count(device) == 0);
199 memory_.emplace(device,
200 std::unique_ptr<StackDeviceMemory>(
201 new StackDeviceMemory(device, toAlloc)));
202 }
203
204 cublasHandle_t
205 StandardGpuResources::getBlasHandle(int device) {
206 initializeForDevice(device);
207 return blasHandles_[device];
208 }
209
210 cudaStream_t
211 StandardGpuResources::getDefaultStream(int device) {
212 initializeForDevice(device);
213 return defaultStreams_[device];
214 }
215
216 std::vector<cudaStream_t>
217 StandardGpuResources::getAlternateStreams(int device) {
218 initializeForDevice(device);
219 return alternateStreams_[device];
220 }
221
222 DeviceMemory& StandardGpuResources::getMemoryManager(int device) {
223 initializeForDevice(device);
224 return *memory_[device];
225 }
226
227 std::pair<void*, size_t>
228 StandardGpuResources::getPinnedMemory() {
229 return std::make_pair(pinnedMemAlloc_, pinnedMemAllocSize_);
230 }
231
232 cudaStream_t
233 StandardGpuResources::getAsyncCopyStream(int device) {
234 initializeForDevice(device);
235 return asyncCopyStreams_[device];
236 }
237
238 } } // namespace
faiss::gpu::StandardGpuResources::setDefaultStream
void setDefaultStream(int device, cudaStream_t stream)
Called to change the stream for work ordering.
Definition: StandardGpuResources.cpp:103
faiss::gpu::StandardGpuResources::getBlasHandle
cublasHandle_t getBlasHandle(int device) override
Returns the cuBLAS handle that we use for the given device.
Definition: StandardGpuResources.cpp:205
faiss::gpu::StandardGpuResources::setTempMemoryFraction
void setTempMemoryFraction(float fraction)
Definition: StandardGpuResources.cpp:87
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:76
faiss::gpu::StandardGpuResources::initializeForDevice
void initializeForDevice(int device) override
Internal system calls.
Definition: StandardGpuResources.cpp:122
faiss::gpu::StandardGpuResources::getAsyncCopyStream
cudaStream_t getAsyncCopyStream(int device) override
Returns the stream on which we perform async CPU &lt;-&gt; GPU copies.
Definition: StandardGpuResources.cpp:233
faiss::gpu::StandardGpuResources::getMemoryManager
DeviceMemory & getMemoryManager(int device) override
Returns the temporary memory manager for the given device.
Definition: StandardGpuResources.cpp:222
faiss::gpu::StandardGpuResources::setTempMemory
void setTempMemory(size_t size)
Definition: StandardGpuResources.cpp:81
faiss::gpu::StandardGpuResources::setPinnedMemory
void setPinnedMemory(size_t size)
Definition: StandardGpuResources.cpp:94
faiss::gpu::StandardGpuResources::getDefaultStream
cudaStream_t getDefaultStream(int device) override
Definition: StandardGpuResources.cpp:211
faiss::gpu::StandardGpuResources::setDefaultNullStreamAllDevices
void setDefaultNullStreamAllDevices()
Definition: StandardGpuResources.cpp:115
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:65
faiss::gpu::StackDeviceMemory
Definition: StackDeviceMemory.h:22
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:46
faiss::gpu::StandardGpuResources::getPinnedMemory
std::pair< void *, size_t > getPinnedMemory() override
Returns the available CPU pinned memory buffer.
Definition: StandardGpuResources.cpp:228
faiss::gpu::StandardGpuResources::getAlternateStreams
std::vector< cudaStream_t > getAlternateStreams(int device) override
Returns the set of alternative streams that we use for the given device.
Definition: StandardGpuResources.cpp:217