lineardevicememory.h

#pragma once

#include <cuda_runtime_api.h>
#include <thrust/device_ptr.h>
#include <type_traits>

#include "../iucutil.h"
#include "linearmemory.h"

template<typename, int> class ndarray_ref;

namespace iu {

template<typename PixelType, unsigned int Ndim>
class LinearDeviceMemory: public LinearMemory<Ndim>
{
public:
  typedef PixelType pixel_type;
  static const unsigned int ndim = Ndim;

  LinearDeviceMemory() :
      LinearMemory<Ndim>(), data_(0), ext_data_pointer_(false)
  {
  }

  virtual ~LinearDeviceMemory()
  {
    if ((!ext_data_pointer_) && (data_ != NULL))
    {
      IU_CUDA_SAFE_CALL(cudaFree(data_));
      data_ = 0;
    }
  }

  LinearDeviceMemory(const Size<Ndim>& size) :
      LinearMemory<Ndim>(size), data_(0), ext_data_pointer_(false)
  {
    IU_CUDA_SAFE_CALL(
        cudaMalloc((void** )&data_, this->numel() * sizeof(PixelType)));
    if (data_ == 0)
      throw std::bad_alloc();
  }

  LinearDeviceMemory(const unsigned int& numel) :
      LinearMemory<Ndim>(numel), data_(0), ext_data_pointer_(false)
  {
    IU_CUDA_SAFE_CALL(
        cudaMalloc((void** )&data_, this->numel() * sizeof(PixelType)));
    if (data_ == 0)
      throw std::bad_alloc();
  }

  LinearDeviceMemory(PixelType* device_data, const Size<Ndim>& size,
                     bool ext_data_pointer = false) :
      LinearMemory<Ndim>(size), data_(0), ext_data_pointer_(ext_data_pointer)
  {
    if (device_data == 0)
      throw IuException("input data not valid", __FILE__, __FUNCTION__,
                        __LINE__);
    if (ext_data_pointer_)
    {
      // This uses the external data pointer as internal data pointer.
      data_ = device_data;
    }
    else
    {
      // allocates an internal data pointer and copies the external data onto it.
      IU_CUDA_SAFE_CALL(
          cudaMalloc((void** )&data_, this->numel() * sizeof(PixelType)));
      if (data_ == 0)
        throw std::bad_alloc();
      IU_CUDA_SAFE_CALL(
          cudaMemcpy(data_, device_data, this->numel() * sizeof(PixelType),
                     cudaMemcpyHostToDevice));
    }
  }

  LinearDeviceMemory(PixelType* device_data, const unsigned int& numel,
                     bool ext_data_pointer = false) :
      LinearMemory<Ndim>(numel), data_(0), ext_data_pointer_(ext_data_pointer)
  {
    if (device_data == 0)
      throw IuException("input data not valid", __FILE__, __FUNCTION__,
                        __LINE__);
    if (ext_data_pointer_)
    {
      // This uses the external data pointer as internal data pointer.
      data_ = device_data;
    }
    else
    {
      // allocates an internal data pointer and copies the external data onto it.
      IU_CUDA_SAFE_CALL(
          cudaMalloc((void** )&data_, this->numel() * sizeof(PixelType)));
      if (data_ == 0)
        throw std::bad_alloc();
      IU_CUDA_SAFE_CALL(
          cudaMemcpy(data_, device_data, this->numel() * sizeof(PixelType),
                     cudaMemcpyHostToDevice));
    }
  }

  PixelType* data(unsigned int offset = 0)
  {
    if (offset >= this->numel())
    {
      std::stringstream msg;
      msg << "Offset (" << offset << ") out of range (" << this->numel() << ").";
      throw IuException(msg.str(), __FILE__, __FUNCTION__, __LINE__);
    }
    return &(data_[offset]);
  }

  const PixelType* data(unsigned int offset = 0) const
  {
    if (offset >= this->numel())
    {
      std::stringstream msg;
      msg << "Offset (" << offset << ") out of range (" << this->numel() << ").";
      throw IuException(msg.str(), __FILE__, __FUNCTION__, __LINE__);
    }
    return reinterpret_cast<const PixelType*>(&(data_[offset]));
  }

  thrust::device_ptr<PixelType> begin(void)
  {
    return thrust::device_ptr<PixelType>(data());
  }

  thrust::device_ptr<PixelType> end(void)
  {
    return thrust::device_ptr<PixelType>(data() + this->numel());
  }

  virtual size_t bytes() const
  {
    return this->numel() * sizeof(PixelType);
  }

  virtual unsigned int bitDepth() const
  {
    return 8 * sizeof(PixelType);
  }

  virtual bool onDevice() const
  {
    return true;
  }

  struct KernelData
  {
    PixelType* data_;

    int numel_;

    int* size_;

    int* stride_;

    __host__ KernelData(const LinearDeviceMemory<PixelType, Ndim> &mem) :
        data_(const_cast<PixelType*>(mem.data())), numel_(mem.numel())
    {
      IU_CUDA_SAFE_CALL(cudaMalloc((void** )&size_, Ndim * sizeof(unsigned int)));
      IU_CUDA_SAFE_CALL(
          cudaMemcpy(size_, mem.size().ptr(), Ndim * sizeof(unsigned int),
                     cudaMemcpyHostToDevice));
      IU_CUDA_SAFE_CALL(cudaMalloc((void** )&stride_, Ndim * sizeof(unsigned int)));
      IU_CUDA_SAFE_CALL(
          cudaMemcpy(stride_, mem.stride().ptr(), Ndim * sizeof(unsigned int),
                     cudaMemcpyHostToDevice));
    }

    __host__ ~KernelData()
    {
      IU_CUDA_SAFE_CALL(cudaFree(size_));
      size_ = 0;
      IU_CUDA_SAFE_CALL(cudaFree(stride_));
      stride_ = 0;
    }

    template<typename ResultType = void>
    __device__ typename std::enable_if<(Ndim == 2), ResultType>::type getPosition(
        const unsigned int& linear_idx, unsigned int& idx0, unsigned int& idx1)
    {
      idx1 = linear_idx / stride_[1];
      idx0 = linear_idx % stride_[1];
    }

    template<typename ResultType = void>
    __device__ typename std::enable_if<(Ndim == 3), ResultType>::type getPosition(
        const unsigned int& linear_idx, unsigned int& idx0, unsigned int& idx1,
        unsigned int& idx2)
    {
      idx2 = linear_idx / stride_[2];
      idx1 = (linear_idx % stride_[2]) / stride_[1];
      idx0 = (linear_idx % stride_[2]) % stride_[1];
    }

    template<typename ResultType = void>
    __device__ typename std::enable_if<(Ndim == 4), ResultType>::type getPosition(
        const unsigned int& linear_idx, unsigned int& idx0, unsigned int& idx1,
        unsigned int& idx2, unsigned int& idx3)
    {
      idx3 = linear_idx / stride_[3];
      idx2 = (linear_idx % stride_[3]) / stride_[2];
      idx1 = ((linear_idx % stride_[3]) % stride_[2]) / stride_[1];
      idx0 = ((linear_idx % stride_[3]) % stride_[2]) % stride_[1];
    }

    template<typename ResultType = void>
    __device__ typename std::enable_if<(Ndim == 5), ResultType>::type getPosition(
        const unsigned int& linear_idx, unsigned int& idx0, unsigned int& idx1,
        unsigned int& idx2, unsigned int& idx3, unsigned int& idx4)
    {
      idx4 = linear_idx / stride_[4];
      idx3 = (linear_idx % stride_[4]) / stride_[3];
      idx2= ((linear_idx % stride_[4]) % stride_[3]) / stride_[2];
      idx1 = (((linear_idx % stride_[4]) % stride_[3]) % stride_[2]) / stride_[1];
      idx0 = (((linear_idx % stride_[4]) % stride_[3]) % stride_[2]) % stride_[1];
    }

    template<typename ResultType = unsigned int>
    __device__ typename std::enable_if<(Ndim > 1), ResultType>::type getLinearIndex(
        const unsigned int& idx0, const unsigned int& idx1)
    {
      unsigned int linear_idx = idx0;
      linear_idx += stride_[1] * idx1;
      return linear_idx;
    }

    template<typename ResultType = unsigned int>
    __device__ typename std::enable_if<(Ndim > 2), ResultType>::type getLinearIndex(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2)
    {
      unsigned int linear_idx = idx0;
      linear_idx += stride_[1] * idx1;
      linear_idx += stride_[2] * idx2;
      return linear_idx;
    }

    template<typename ResultType = unsigned int>
    __device__ typename std::enable_if<(Ndim > 3), ResultType>::type getLinearIndex(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2, const unsigned int& idx3)
    {
      unsigned int linear_idx = idx0;
      linear_idx += stride_[1] * idx1;
      linear_idx += stride_[2] * idx2;
      linear_idx += stride_[3] * idx3;
      return linear_idx;
    }

    template<typename ResultType = unsigned int>
    __device__ typename std::enable_if<(Ndim > 4), ResultType>::type getLinearIndex(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2, const unsigned int& idx3, const unsigned int& idx4)
    {
      unsigned int linear_idx = idx0;
      linear_idx += stride_[1] * idx1;
      linear_idx += stride_[2] * idx2;
      linear_idx += stride_[3] * idx3;
      linear_idx += stride_[4] * idx4;
      return linear_idx;
    }

    __device__ PixelType& operator()(const unsigned int& idx)
    {
      return data_[idx];
    }

    template<typename ResultType = PixelType>
    __device__ typename std::enable_if<(Ndim > 1), ResultType&>::type operator()(
        const unsigned int& idx0, const unsigned int& idx1)
    {
      return data_[getLinearIndex(idx0, idx1)];
    }

    template<typename ResultType = PixelType>
    __device__ typename std::enable_if<(Ndim > 2), ResultType&>::type operator()(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2)
    {
      return data_[getLinearIndex(idx0, idx1, idx2)];
    }

    template<typename ResultType = PixelType>
    __device__ typename std::enable_if<(Ndim > 3), ResultType&>::type operator()(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2, const unsigned int& idx3)
    {
      return data_[getLinearIndex(idx0, idx1, idx2, idx3)];
    }

    template<typename ResultType = PixelType>
    __device__ typename std::enable_if<(Ndim > 4), ResultType&>::type operator()(
        const unsigned int& idx0, const unsigned int& idx1, const unsigned int& idx2, const unsigned int& idx3, const unsigned int& idx4)
    {
      return data_[getLinearIndex(idx0, idx1, idx2, idx3, idx4)];
    }
  };

  ndarray_ref<PixelType, Ndim> ref() const;

  LinearDeviceMemory(const ndarray_ref<PixelType, Ndim> &x);

private:
  PixelType* data_;
  bool ext_data_pointer_;

private:
  LinearDeviceMemory(const LinearDeviceMemory&);
  LinearDeviceMemory& operator=(const LinearDeviceMemory&);
};

template<> LinearDeviceMemory<float, 1>::LinearDeviceMemory(const ndarray_ref<float, 1> &x);

}  // namespace iu