CUDAUtilities/src/allocate_device.cc

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 }  // namespace cms::cuda