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_(DynMem_ &&other) : num_elem(other.num_elem)
  36     {
  37         ptr_h = other.ptr_h;
  38         other.ptr_h = nullptr;
  39 #ifdef CUFFT
  40         ptr_d = other.ptr_d;
  41         other.ptr_d = nullptr;
  42 #endif
  43     }
  44     ~DynMem_()
  45     {
  46 #ifdef CUFFT
  47         CudaSafeCall(cudaFreeHost(ptr_h));
  48 #else
  49         delete[] ptr_h;
  50 #endif
  51     }
  52     T *hostMem() { return ptr_h; }
  53 #ifdef CUFFT
  54     T *deviceMem() { return ptr_d; }
  55 #endif
  56     void operator=(DynMem_ &rhs) {
  57         assert(num_elem == rhs.num_elem);
  58         memcpy(ptr_h, rhs.ptr_h, num_elem * sizeof(T));
  59     }
  60     void operator=(DynMem_ &&rhs)
  61     {
  62         assert(num_elem == rhs.num_elem);
  63         ptr_h = rhs.ptr_h;
  64         rhs.ptr_h = nullptr;
  65 #ifdef CUFFT
  66         ptr_d = rhs.ptr_d;
  67         rhs.ptr_d = nullptr;
  68 #endif
  69     }
  70     T operator[](uint i) const { return ptr_h[i]; }
  71 };
  72
  73 typedef DynMem_<float> DynMem;
  74
  75
  76 class MatDynMem : public DynMem, public cv::Mat {
  77   public:
  78     MatDynMem(cv::Size size, int type)
  79         : DynMem(size.area() * CV_MAT_CN(type)), cv::Mat(size, type, hostMem())
  80     {
  81         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  82     }
  83     MatDynMem(int height, int width, int type)
  84         : DynMem(width * height * CV_MAT_CN(type)), cv::Mat(height, width, type, hostMem())
  85     {
  86         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  87     }
  88     MatDynMem(int ndims, const int *sizes, int type)
  89         : DynMem(volume(ndims, sizes) * CV_MAT_CN(type)), cv::Mat(ndims, sizes, type, hostMem())
  90     {
  91         assert((type & CV_MAT_DEPTH_MASK) == CV_32F);
  92     }
  93     MatDynMem(std::vector<int> size, int type)
  94         : DynMem(std::accumulate(size.begin(), size.end(), 1, std::multiplies<int>()))
  95         , cv::Mat(size.size(), size.data(), type, hostMem()) {}
  96     MatDynMem(MatDynMem &&other) = default;
  97     MatDynMem(const cv::Mat &other)
  98         : DynMem(other.total()) , cv::Mat(other) {}
  99
 100     void operator=(const cv::MatExpr &expr) {
 101         static_cast<cv::Mat>(*this) = expr;
 102     }
 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 class Mat3d : public MatDynMem
 117 {
 118 public:
 119     Mat3d(uint dim0, cv::Size size) : MatDynMem({{int(dim0), size.height, size.width}}, CV_32F) {}
 120
 121     cv::Mat plane(uint idx) {
 122         assert(dims == 3);
 123         assert(int(idx) < size[0]);
 124         return cv::Mat(size[1], size[2], cv::Mat::type(), ptr(idx));
 125     }
 126     const cv::Mat plane(uint idx) const {
 127         assert(dims == 3);
 128         assert(int(idx) < size[0]);
 129         return cv::Mat(size[1], size[2], cv::Mat::type(), const_cast<uchar*>(ptr(idx)));
 130     }
 131
 132 };
 133
 134 class MatFeats : public Mat3d
 135 {
 136 public:
 137     MatFeats(uint num_features, cv::Size size) : Mat3d(num_features, size) {}
 138 };
 139 class MatScales : public Mat3d
 140 {
 141 public:
 142     MatScales(uint num_scales, cv::Size size) : Mat3d(num_scales, size) {}
 143 };
 144
 145 class MatScaleFeats : public MatDynMem
 146 {
 147 public:
 148     MatScaleFeats(uint num_scales, uint num_features, cv::Size size)
 149         : MatDynMem({{int(num_scales), int(num_features), size.height, size.width}}, CV_32F) {}
 150
 151     cv::Mat plane(uint scale, uint feature) {
 152         assert(dims == 4);
 153         assert(int(scale) < size[0]);
 154         assert(int(feature) < size[1]);
 155         return cv::Mat(size[2], size[3], cv::Mat::type(), ptr(scale, feature));
 156     }
 157     cv::Mat scale(uint scale) {
 158         assert(dims == 4);
 159         assert(int(scale) < size[0]);
 160         return cv::Mat(3, std::vector<int>({size[1], size[2], size[3]}).data(), cv::Mat::type(), ptr(scale));
 161     }
 162 };
 163
 164 #endif // DYNMEM_HPP