src/dynmem.hpp

   1 #ifndef DYNMEM_HPP
   2 #define DYNMEM_HPP
   3
   4 #include <cstdlib>
   5 #include <opencv2/opencv.hpp>
   6 #include <cassert>
   7 #include <numeric>
   8
   9 #if defined(CUFFT) || defined(CUFFTW)
  10 #include "cuda_runtime.h"
  11 #ifdef CUFFT
  12 #include "cuda_error_check.hpp"
  13 #endif
  14 #endif
  15
  16 template <typename T> class DynMem_ {
  17   private:
  18     T *ptr_h = nullptr;
  19 #ifdef CUFFT
  20     T *ptr_d = nullptr;
  21 #endif
  22   public:
  23     typedef T value_type;
  24     const size_t num_elem;
  25
  26     DynMem_(size_t num_elem) : num_elem(num_elem)
  27     {
  28 #ifdef CUFFT
  29         CudaSafeCall(cudaHostAlloc(reinterpret_cast<void **>(&ptr_h), num_elem * sizeof(T), cudaHostAllocMapped));
  30         CudaSafeCall(cudaHostGetDevicePointer(reinterpret_cast<void **>(&ptr_d), reinterpret_cast<void *>(ptr_h), 0));
  31 #else
  32         ptr_h = new T[num_elem];
  33 #endif
  34     }
  35     DynMem_(const DynMem_ &other) : DynMem_(other.num_elem)
  36     {
  37         memcpy(ptr_h, other.ptr_h, num_elem * sizeof(T));
  38     }
  39     DynMem_(DynMem_ &&other) : num_elem(other.num_elem)
  40     {
  41         ptr_h = other.ptr_h;
  42         other.ptr_h = nullptr;
  43 #ifdef CUFFT
  44         ptr_d = other.ptr_d;
  45         other.ptr_d = nullptr;
  46 #endif
  47     }
  48     ~DynMem_()
  49     {
  50         release();
  51     }
  52     T *hostMem() { return ptr_h; }
  53     const T *hostMem() const { return ptr_h; }
  54 #ifdef CUFFT
  55     T *deviceMem() { return ptr_d; }
  56 #endif
  57     void operator=(DynMem_ &rhs) {
  58         assert(num_elem == rhs.num_elem);
  59         memcpy(ptr_h, rhs.ptr_h, num_elem * sizeof(T));
  60     }
  61     void operator=(DynMem_ &&rhs)
  62     {
  63         assert(num_elem == rhs.num_elem);
  64         release();
  65         ptr_h = rhs.ptr_h;
  66         rhs.ptr_h = nullptr;
  67 #ifdef CUFFT
  68         ptr_d = rhs.ptr_d;
  69         rhs.ptr_d = nullptr;
  70 #endif
  71     }
  72     T operator[](uint i) const { return ptr_h[i]; }
  73 private:
  74     void release()
  75     {
  76 #ifdef CUFFT
  77         CudaSafeCall(cudaFreeHost(ptr_h));
  78 #else
  79         delete[] ptr_h;
  80 #endif
  81     }
  82 };
  83
  84 typedef DynMem_<float> DynMem;
  85
  86
  87 class MatDynMem : public DynMem, public cv::Mat {
  88   public:
  89     MatDynMem(cv::Size size, int type)
  90         : DynMem(size.area() * CV_MAT_CN(type)), cv::Mat(size, type, hostMem())
  91     {
  92         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  93     }
  94     MatDynMem(int height, int width, int type)
  95         : DynMem(width * height * CV_MAT_CN(type)), cv::Mat(height, width, type, hostMem())
  96     {
  97         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  98     }
  99     MatDynMem(int ndims, const int *sizes, int type)
 100         : DynMem(volume(ndims, sizes) * CV_MAT_CN(type)), cv::Mat(ndims, sizes, type, hostMem())
 101     {
 102         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
 103     }
 104     MatDynMem(std::vector<int> size, int type)
 105         : DynMem(std::accumulate(size.begin(), size.end(), 1, std::multiplies<int>()))
 106         , cv::Mat(size.size(), size.data(), type, hostMem()) {}
 107     MatDynMem(MatDynMem &&other) = default;
 108     MatDynMem(const cv::Mat &other)
 109         : DynMem(other.total()) , cv::Mat(other) {}
 110
 111     void operator=(const cv::MatExpr &expr) {
 112         static_cast<cv::Mat>(*this) = expr;
 113     }
 114
 115   private:
 116     static int volume(int ndims, const int *sizes)
 117     {
 118         int vol = 1;
 119         for (int i = 0; i < ndims; i++)
 120             vol *= sizes[i];
 121         return vol;
 122     }
 123
 124     using cv::Mat::create;
 125 };
 126
 127 class Mat3d : public MatDynMem
 128 {
 129 public:
 130     Mat3d(uint dim0, cv::Size size) : MatDynMem({{int(dim0), size.height, size.width}}, CV_32F) {}
 131
 132     cv::Mat plane(uint idx) {
 133         assert(dims == 3);
 134         assert(int(idx) < size[0]);
 135         return cv::Mat(size[1], size[2], cv::Mat::type(), ptr(idx));
 136     }
 137     const cv::Mat plane(uint idx) const {
 138         assert(dims == 3);
 139         assert(int(idx) < size[0]);
 140         return cv::Mat(size[1], size[2], cv::Mat::type(), const_cast<uchar*>(ptr(idx)));
 141     }
 142
 143 };
 144
 145 class MatFeats : public Mat3d
 146 {
 147 public:
 148     MatFeats(uint num_features, cv::Size size) : Mat3d(num_features, size) {}
 149 };
 150 class MatScales : public Mat3d
 151 {
 152 public:
 153     MatScales(uint num_scales, cv::Size size) : Mat3d(num_scales, size) {}
 154 };
 155
 156 class MatScaleFeats : public MatDynMem
 157 {
 158 public:
 159     MatScaleFeats(uint num_scales, uint num_features, cv::Size size)
 160         : MatDynMem({{int(num_scales), int(num_features), size.height, size.width}}, CV_32F) {}
 161
 162     cv::Mat plane(uint scale, uint feature) {
 163         assert(dims == 4);
 164         assert(int(scale) < size[0]);
 165         assert(int(feature) < size[1]);
 166         return cv::Mat(size[2], size[3], cv::Mat::type(), ptr(scale, feature));
 167     }
 168     cv::Mat scale(uint scale) {
 169         assert(dims == 4);
 170         assert(int(scale) < size[0]);
 171         return cv::Mat(3, std::vector<int>({size[1], size[2], size[3]}).data(), cv::Mat::type(), ptr(scale));
 172     }
 173 };
 174
 175 #endif // DYNMEM_HPP