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