CUDAUtilities/src/cachingAllocatorCommon.h

0001 #ifndef HeterogeneousCore_CUDACore_src_cachingAllocatorCommon
0002 #define HeterogeneousCore_CUDACore_src_cachingAllocatorCommon
0003
0004 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0005 #include "HeterogeneousCore/CUDAUtilities/interface/deviceCount.h"
0006
0007 #include <algorithm>
0008 #include <limits>
0009
0010 namespace cms::cuda::allocator {
0011   // Growth factor (bin_growth in cub::CachingDeviceAllocator
0012   constexpr unsigned int binGrowth = 2;
0013   // Smallest bin, corresponds to binGrowth^minBin bytes (min_bin in cub::CacingDeviceAllocator
0014   constexpr unsigned int minBin = 8;
0015   // Largest bin, corresponds to binGrowth^maxBin bytes (max_bin in cub::CachingDeviceAllocator). Note that unlike in cub, allocations larger than binGrowth^maxBin are set to fail.
0016   constexpr unsigned int maxBin = 30;
0017   // Total storage for the allocator. 0 means no limit.
0018   constexpr size_t maxCachedBytes = 0;
0019   // Fraction of total device memory taken for the allocator. In case there are multiple devices with different amounts of memory, the smallest of them is taken. If maxCachedBytes is non-zero, the smallest of them is taken.
0020   constexpr double maxCachedFraction = 0.8;
0021   constexpr bool debug = false;
0022
0023   inline size_t minCachedBytes() {
0024     size_t ret = std::numeric_limits<size_t>::max();
0025     int currentDevice;
0026     cudaCheck(cudaGetDevice(&currentDevice));
0027     const int numberOfDevices = deviceCount();
0028     for (int i = 0; i < numberOfDevices; ++i) {
0029       size_t freeMemory, totalMemory;
0030       cudaCheck(cudaSetDevice(i));
0031       cudaCheck(cudaMemGetInfo(&freeMemory, &totalMemory));
0032       ret = std::min(ret, static_cast<size_t>(maxCachedFraction * freeMemory));
0033     }
0034     cudaCheck(cudaSetDevice(currentDevice));
0035     if (maxCachedBytes > 0) {
0036       ret = std::min(ret, maxCachedBytes);
0037     }
0038     return ret;
0039   }
0040 }  // namespace cms::cuda::allocator
0041
0042 #endif