forked from cms-patatrack/pixeltrack-standalone
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathallocate_device.cc
39 lines (33 loc) · 1.23 KB
/
allocate_device.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
#include <limits>
#include "CUDACore/ScopedSetDevice.h"
#include "CUDACore/allocate_device.h"
#include "CUDACore/cudaCheck.h"
#include "getCachingDeviceAllocator.h"
namespace {
const size_t maxAllocationSize =
notcub::CachingDeviceAllocator::IntPow(cms::cuda::allocator::binGrowth, cms::cuda::allocator::maxBin);
}
namespace cms::cuda {
void *allocate_device(int dev, size_t nbytes, cudaStream_t stream) {
void *ptr = nullptr;
if constexpr (allocator::useCaching) {
if (nbytes > maxAllocationSize) {
throw std::runtime_error("Tried to allocate " + std::to_string(nbytes) +
" bytes, but the allocator maximum is " + std::to_string(maxAllocationSize));
}
cudaCheck(allocator::getCachingDeviceAllocator().DeviceAllocate(dev, &ptr, nbytes, stream));
} else {
ScopedSetDevice setDeviceForThisScope(dev);
cudaCheck(cudaMalloc(&ptr, nbytes));
}
return ptr;
}
void free_device(int device, void *ptr) {
if constexpr (allocator::useCaching) {
cudaCheck(allocator::getCachingDeviceAllocator().DeviceFree(device, ptr));
} else {
ScopedSetDevice setDeviceForThisScope(device);
cudaCheck(cudaFree(ptr));
}
}
} // namespace cms::cuda