CUDAUtilities/interface/HostAllocator.h

0001 #ifndef HeterogeneousCore_CUDAUtilities_HostAllocator_h
0002 #define HeterogeneousCore_CUDAUtilities_HostAllocator_h
0003
0004 #include <memory>
0005 #include <new>
0006 #include <cuda_runtime.h>
0007
0008 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0009
0010 namespace cms {
0011   namespace cuda {
0012
0013     class bad_alloc : public std::bad_alloc {
0014     public:
0015       bad_alloc(cudaError_t error) noexcept : error_(error) {}
0016
0017       const char* what() const noexcept override { return cudaGetErrorString(error_); }
0018
0019     private:
0020       cudaError_t error_;
0021     };
0022
0023     template <typename T, unsigned int FLAGS = cudaHostAllocDefault>
0024     class HostAllocator {
0025     public:
0026       using value_type = T;
0027
0028       template <typename U>
0029       struct rebind {
0030         using other = HostAllocator<U, FLAGS>;
0031       };
0032
0033       CMS_THREAD_SAFE T* allocate(std::size_t n) const __attribute__((warn_unused_result)) __attribute__((malloc))
0034       __attribute__((returns_nonnull)) {
0035         void* ptr = nullptr;
0036         cudaError_t status = cudaMallocHost(&ptr, n * sizeof(T), FLAGS);
0037         if (status != cudaSuccess) {
0038           throw bad_alloc(status);
0039         }
0040         if (ptr == nullptr) {
0041           throw std::bad_alloc();
0042         }
0043         return static_cast<T*>(ptr);
0044       }
0045
0046       void deallocate(T* p, std::size_t n) const {
0047         cudaError_t status = cudaFreeHost(p);
0048         if (status != cudaSuccess) {
0049           throw bad_alloc(status);
0050         }
0051       }
0052     };
0053
0054   }  // namespace cuda
0055 }  // namespace cms
0056
0057 #endif  // HeterogeneousCore_CUDAUtilities_HostAllocator_h