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