|
15 | 15 | #include "paddle/fluid/memory/allocation/allocator.h"
|
16 | 16 | #include <gflags/gflags.h>
|
17 | 17 | #include <map>
|
| 18 | +#include <string> |
18 | 19 | #include <unordered_map>
|
19 | 20 | #include <vector>
|
20 | 21 | #include "paddle/fluid/memory/allocation/aligned_allocator.h"
|
@@ -209,6 +210,7 @@ class AllocatorFacadePrivate {
|
209 | 210 | for (int dev_id = 0; dev_id < platform::GetCUDADeviceCount(); ++dev_id) {
|
210 | 211 | places.emplace_back(platform::CUDAPlace(dev_id));
|
211 | 212 | }
|
| 213 | + places.emplace_back(platform::CUDAPinnedPlace()); |
212 | 214 | #endif
|
213 | 215 | for (auto& p : places) {
|
214 | 216 | allocators_[p] = std::make_shared<LegacyAllocator>(p);
|
@@ -255,13 +257,17 @@ AllocatorFacade& AllocatorFacade::Instance() {
|
255 | 257 |
|
256 | 258 | std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
|
257 | 259 | const platform::Place& place, size_t size, Allocator::Attr attr) {
|
258 |
| - return std::shared_ptr<Allocation>( |
259 |
| - m_->allocators_.at(place)->Allocate(size, attr).release(), |
260 |
| - AllocationDeleter()); |
| 260 | + return std::shared_ptr<Allocation>(Alloc(place, size, attr).release(), |
| 261 | + AllocationDeleter()); |
261 | 262 | }
|
262 | 263 |
|
263 | 264 | AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size,
|
264 | 265 | Allocator::Attr attr) {
|
| 266 | + auto it = m_->allocators_.find(place); |
| 267 | + if (it == m_->allocators_.end()) { |
| 268 | + throw BadAlloc( |
| 269 | + string::Sprintf("No such allocator for the place, %s", place)); |
| 270 | + } |
265 | 271 | return m_->allocators_.at(place)->Allocate(size, attr);
|
266 | 272 | }
|
267 | 273 |
|
|
0 commit comments