File indexing completed on 2024-04-06 12:15:45
0001 #include <limits>
0002
0003 #include "FWCore/Utilities/interface/Likely.h"
0004 #include "HeterogeneousCore/CUDAUtilities/interface/ScopedSetDevice.h"
0005 #include "HeterogeneousCore/CUDAUtilities/interface/allocate_device.h"
0006 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0007
0008 #include "getCachingDeviceAllocator.h"
0009
0010 namespace {
0011 const size_t maxAllocationSize =
0012 notcub::CachingDeviceAllocator::IntPow(cms::cuda::allocator::binGrowth, cms::cuda::allocator::maxBin);
0013 }
0014
0015 namespace cms::cuda {
0016 void *allocate_device(int dev, size_t nbytes, cudaStream_t stream) {
0017 void *ptr = nullptr;
0018 if constexpr (allocator::useCaching) {
0019 if (UNLIKELY(nbytes > maxAllocationSize)) {
0020 throw std::runtime_error("Tried to allocate " + std::to_string(nbytes) +
0021 " bytes, but the allocator maximum is " + std::to_string(maxAllocationSize));
0022 }
0023 cudaCheck(allocator::getCachingDeviceAllocator().DeviceAllocate(dev, &ptr, nbytes, stream));
0024 } else {
0025 ScopedSetDevice setDeviceForThisScope(dev);
0026 cudaCheck(cudaMalloc(&ptr, nbytes));
0027 }
0028 return ptr;
0029 }
0030
0031 void free_device(int device, void *ptr) {
0032 if constexpr (allocator::useCaching) {
0033 cudaCheck(allocator::getCachingDeviceAllocator().DeviceFree(device, ptr));
0034 } else {
0035 ScopedSetDevice setDeviceForThisScope(device);
0036 cudaCheck(cudaFree(ptr));
0037 }
0038 }
0039
0040 }
