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     const T *deviceMem() const { return ptr_d; }
  57 #endif
  58     void operator=(DynMem_ &rhs) {
  59         assert(num_elem == rhs.num_elem);
  60         memcpy(ptr_h, rhs.ptr_h, num_elem * sizeof(T));
  61     }
  62     void operator=(DynMem_ &&rhs)
  63     {
  64         assert(num_elem == rhs.num_elem);
  65         release();
  66         ptr_h = rhs.ptr_h;
  67         rhs.ptr_h = nullptr;
  68 #ifdef CUFFT
  69         ptr_d = rhs.ptr_d;
  70         rhs.ptr_d = nullptr;
  71 #endif
  72     }
  73     T operator[](uint i) const { return ptr_h[i]; }
  74 private:
  75     void release()
  76     {
  77 #ifdef CUFFT
  78         CudaSafeCall(cudaFreeHost(ptr_h));
  79 #else
  80         delete[] ptr_h;
  81 #endif
  82     }
  83 };
  84
  85 typedef DynMem_<float> DynMem;
  86
  87
  88 class MatDynMem : public DynMem, public cv::Mat {
  89   public:
  90     MatDynMem(cv::Size size, int type)
  91         : DynMem(size.area() * CV_MAT_CN(type)), cv::Mat(size, type, hostMem())
  92     {
  93         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  94     }
  95     MatDynMem(int height, int width, int type)
  96         : DynMem(width * height * CV_MAT_CN(type)), cv::Mat(height, width, type, hostMem())
  97     {
  98         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  99     }
 100     MatDynMem(int ndims, const int *sizes, int type)
 101         : DynMem(volume(ndims, sizes) * CV_MAT_CN(type)), cv::Mat(ndims, sizes, type, hostMem())
 102     {
 103         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
 104     }
 105     MatDynMem(std::vector<int> size, int type)
 106         : DynMem(std::accumulate(size.begin(), size.end(), 1, std::multiplies<int>()))
 107         , cv::Mat(size.size(), size.data(), type, hostMem()) {}
 108     MatDynMem(MatDynMem &&other) = default;
 109     MatDynMem(const cv::Mat &other)
 110         : DynMem(other.total()) , cv::Mat(other) {}
 111
 112     void operator=(const cv::MatExpr &expr) {
 113         static_cast<cv::Mat>(*this) = expr;
 114     }
 115
 116   private:
 117     static int volume(int ndims, const int *sizes)
 118     {
 119         int vol = 1;
 120         for (int i = 0; i < ndims; i++)
 121             vol *= sizes[i];
 122         return vol;
 123     }
 124
 125     using cv::Mat::create;
 126 };
 127
 128 class Mat3d : public MatDynMem
 129 {
 130 public:
 131     Mat3d(uint dim0, cv::Size size) : MatDynMem({{int(dim0), size.height, size.width}}, CV_32F) {}
 132
 133     cv::Mat plane(uint idx) {
 134         assert(dims == 3);
 135         assert(int(idx) < size[0]);
 136         return cv::Mat(size[1], size[2], cv::Mat::type(), ptr(idx));
 137     }
 138     const cv::Mat plane(uint idx) const {
 139         assert(dims == 3);
 140         assert(int(idx) < size[0]);
 141         return cv::Mat(size[1], size[2], cv::Mat::type(), const_cast<uchar*>(ptr(idx)));
 142     }
 143
 144 };
 145
 146 class MatFeats : public Mat3d
 147 {
 148 public:
 149     MatFeats(uint num_features, cv::Size size) : Mat3d(num_features, size) {}
 150 };
 151 class MatScales : public Mat3d
 152 {
 153 public:
 154     MatScales(uint num_scales, cv::Size size) : Mat3d(num_scales, size) {}
 155 };
 156
 157 class MatScaleFeats : public MatDynMem
 158 {
 159 public:
 160     MatScaleFeats(uint num_scales, uint num_features, cv::Size size)
 161         : MatDynMem({{int(num_scales), int(num_features), size.height, size.width}}, CV_32F) {}
 162
 163     cv::Mat plane(uint scale, uint feature) {
 164         assert(dims == 4);
 165         assert(int(scale) < size[0]);
 166         assert(int(feature) < size[1]);
 167         return cv::Mat(size[2], size[3], cv::Mat::type(), ptr(scale, feature));
 168     }
 169     cv::Mat scale(uint scale) {
 170         assert(dims == 4);
 171         assert(int(scale) < size[0]);
 172         return cv::Mat(3, std::vector<int>({size[1], size[2], size[3]}).data(), cv::Mat::type(), ptr(scale));
 173     }
 174 };
 175
 176 #endif // DYNMEM_HPP