src/dynmem.hpp

   1 #ifndef DYNMEM_HPP
   2 #define DYNMEM_HPP
   3
   4 #include <cstdlib>
   5 #include <opencv2/opencv.hpp>
   6 #include <cassert>
   7
   8 #if defined(CUFFT) || defined(CUFFTW)
   9 #include "cuda_runtime.h"
  10 #ifdef CUFFT
  11 #include "cuda/cuda_error_check.cuh"
  12 #endif
  13 #endif
  14
  15 template <typename T> class DynMem_ {
  16   private:
  17     T *ptr_h = nullptr;
  18 #ifdef CUFFT
  19     T *ptr_d = nullptr;
  20 #endif
  21     size_t num_elem;
  22   public:
  23     typedef T type;
  24     DynMem_(size_t num_elem) : num_elem(num_elem)
  25     {
  26 #ifdef CUFFT
  27         CudaSafeCall(cudaHostAlloc(reinterpret_cast<void **>(&ptr_h), num_elem * sizeof(T), cudaHostAllocMapped));
  28         CudaSafeCall(cudaHostGetDevicePointer(reinterpret_cast<void **>(&ptr_d), reinterpret_cast<void *>(ptr_h), 0));
  29 #else
  30         ptr_h = new T[num_elem];
  31 #endif
  32     }
  33     DynMem_(DynMem_&& other) {
  34         ptr_h = other.ptr_h;
  35         other.ptr_h = nullptr;
  36 #ifdef CUFFT
  37         ptr_d = other.ptr_d;
  38         other.ptr_d = nullptr;
  39 #endif
  40     }
  41     ~DynMem_()
  42     {
  43 #ifdef CUFFT
  44         CudaSafeCall(cudaFreeHost(ptr_h));
  45 #else
  46         delete[] ptr_h;
  47 #endif
  48     }
  49     T *hostMem() { return ptr_h; }
  50 #ifdef CUFFT
  51     T *deviceMem() { return ptr_d; }
  52 #endif
  53     void operator=(DynMem_ &rhs) {
  54         memcpy(ptr_h, rhs.ptr_h, num_elem * sizeof(T));
  55     }
  56     void operator=(DynMem_ &&rhs)
  57     {
  58         ptr_h = rhs.ptr_h;
  59         rhs.ptr_h = nullptr;
  60 #ifdef CUFFT
  61         ptr_d = rhs.ptr_d;
  62         rhs.ptr_d = nullptr;
  63 #endif
  64     }
  65 };
  66
  67 typedef DynMem_<float> DynMem;
  68
  69
  70 class MatDynMem : public DynMem, public cv::Mat {
  71   public:
  72     MatDynMem(cv::Size size, int type)
  73         : DynMem(size.area() * CV_MAT_CN(type)), cv::Mat(size, type, hostMem())
  74     {
  75         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  76     }
  77     MatDynMem(int height, int width, int type)
  78         : DynMem(width * height * sizeof(DynMem::type) * CV_MAT_CN(type)), cv::Mat(height, width, type, hostMem())
  79     {
  80         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  81     }
  82     MatDynMem(int ndims, const int *sizes, int type)
  83         : DynMem(volume(ndims, sizes) * CV_MAT_CN(type)), cv::Mat(ndims, sizes, type, hostMem())
  84     {
  85         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  86     }
  87     void operator=(const cv::MatExpr &expr) {
  88         static_cast<cv::Mat>(*this) = expr;
  89     }
  90
  91   private:
  92     static int volume(int ndims, const int *sizes)
  93     {
  94         int vol = 1;
  95         for (int i = 0; i < ndims; i++)
  96             vol *= sizes[i];
  97         return vol;
  98     }
  99 };
 100
 101 #endif // DYNMEM_HPP