forked from cms-patatrack/pixeltrack-standalone
-
Notifications
You must be signed in to change notification settings - Fork 0
/
allocate_host.cc
36 lines (30 loc) · 1.05 KB
/
allocate_host.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
#include <limits>
#include "CUDACore/allocate_host.h"
#include "CUDACore/cudaCheck.h"
#include "getCachingHostAllocator.h"
namespace {
const size_t maxAllocationSize =
notcub::CachingDeviceAllocator::IntPow(cms::cuda::allocator::binGrowth, cms::cuda::allocator::maxBin);
}
namespace cms::cuda {
void *allocate_host(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::getCachingHostAllocator().HostAllocate(&ptr, nbytes, stream));
} else {
cudaCheck(cudaMallocHost(&ptr, nbytes));
}
return ptr;
}
void free_host(void *ptr) {
if constexpr (allocator::useCaching) {
cudaCheck(allocator::getCachingHostAllocator().HostFree(ptr));
} else {
cudaCheck(cudaFreeHost(ptr));
}
}
} // namespace cms::cuda