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     T operator[](uint i) const { return ptr_h[i]; }
  66 };
  67
  68 typedef DynMem_<float> DynMem;
  69
  70
  71 class MatDynMem : public DynMem, public cv::Mat {
  72   public:
  73     MatDynMem(cv::Size size, int type)
  74         : DynMem(size.area() * CV_MAT_CN(type)), cv::Mat(size, type, hostMem())
  75     {
  76         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  77     }
  78     MatDynMem(int height, int width, int type)
  79         : DynMem(width * height * CV_MAT_CN(type)), cv::Mat(height, width, type, hostMem())
  80     {
  81         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  82     }
  83     MatDynMem(int ndims, const int *sizes, int type)
  84         : DynMem(volume(ndims, sizes) * CV_MAT_CN(type)), cv::Mat(ndims, sizes, type, hostMem())
  85     {
  86         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  87     }
  88     MatDynMem(std::array<int, 3> size, int type)
  89         : DynMem(size[0] * size[1] * size[2]), cv::Mat(3, (int*)&size, type, hostMem())
  90     {}
  91     MatDynMem(MatDynMem &&other) = default;
  92     MatDynMem(const cv::Mat &other)
  93         : DynMem(other.total()) , cv::Mat(other) {}
  94
  95     void operator=(const cv::MatExpr &expr) {
  96         static_cast<cv::Mat>(*this) = expr;
  97     }
  98
  99     cv::Mat plane(uint i) {
 100         assert(dims == 3);
 101         assert(int(i) < size[0]);
 102         return cv::Mat(dims - 1, size + 1, cv::Mat::type(), ptr(i));
 103     }
 104   private:
 105     static int volume(int ndims, const int *sizes)
 106     {
 107         int vol = 1;
 108         for (int i = 0; i < ndims; i++)
 109             vol *= sizes[i];
 110         return vol;
 111     }
 112
 113     using cv::Mat::create;
 114 };
 115
 116 #endif // DYNMEM_HPP