volume_gpu.h

#pragma once

#include "volume.h"
#include "volume_allocator_gpu.h"
#include <thrust/device_ptr.h>

template<typename, int> class ndarray_ref;

namespace iu {

template<typename PixelType, class Allocator>
class VolumeGpu : public Volume
{
public:
  typedef PixelType pixel_type;

  VolumeGpu() :
    Volume(),
    data_(0), pitch_(0), ext_data_pointer_(false)
  {
  }

  virtual ~VolumeGpu()
  {
    if(!ext_data_pointer_)
    {
      // do not delete externally handeled data pointers.
      Allocator::free(data_);
      data_ = 0;
    }
    pitch_ = 0;
  }

  VolumeGpu(unsigned int _width, unsigned int _height, unsigned int _depth) :
    Volume(_width, _height, _depth), data_(0), pitch_(0),
    ext_data_pointer_(false)
  {
    data_ = Allocator::alloc(this->size(), &pitch_);
  }

  VolumeGpu(const iu::Size<3>& size) :
    Volume(size), data_(0), pitch_(0),
    ext_data_pointer_(false)
  {
    data_ = Allocator::alloc(this->size(), &pitch_);
  }

  VolumeGpu(PixelType* _data, unsigned int _width, unsigned int _height, unsigned int _depth,
            size_t _pitch, bool ext_data_pointer = false) :
    Volume(_width, _height, _depth), data_(0), pitch_(0),
    ext_data_pointer_(ext_data_pointer)
  {
    if(ext_data_pointer_)
    {
      // This uses the external data pointer as internal data pointer.
      data_ = _data;
      pitch_ = _pitch;
    }
    else
    {
      // allocates an internal data pointer and copies the external data onto it.
      if(_data == 0)
        return;

      data_ = Allocator::alloc(this->size(), &pitch_);
      Allocator::copy(_data, _pitch, data_, pitch_, this->size());
    }
  }

  size_t bytes() const
  {
    return depth()*height()*pitch_;
  }

  size_t pitch() const
  {
    return pitch_;
  }

  size_t slice_pitch() const
  {
    return height()*pitch_;
  }

  size_t stride() const
  {
    return pitch_/sizeof(PixelType);
  }

  size_t slice_stride() const
  {
    return height()*pitch_/sizeof(PixelType);
  }

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

  virtual bool onDevice() const
  {
    return true;
  }


  PixelType* data(int ox = 0, int oy = 0, int oz = 0)
  {
    return &data_[oz*slice_stride() + oy*stride() + ox];
  }

  const PixelType* data(int ox = 0, int oy = 0, int oz = 0) const
  {
    return reinterpret_cast<const PixelType*>(
          &data_[oz*slice_stride() + oy*stride() + ox]);
  }

  ImageGpu<PixelType, iuprivate::ImageAllocatorGpu<PixelType> > getSlice(int oz)
  {
    return ImageGpu<PixelType, iuprivate::ImageAllocatorGpu<PixelType> >(
          &data_[oz*slice_stride()], width(), height(), pitch_, true);
  }

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

  thrust::device_ptr<PixelType> end(void)
  {
      return thrust::device_ptr<PixelType>(data()+slice_stride()*depth());
  }

  struct KernelData
  {
      PixelType* data_;
      int width_;
      int height_;
      int depth_;
      int stride_;

      __device__ PixelType& operator()(int x, int y, int z)
      {
          return data_[z*height_*stride_ + y*stride_ + x];
      }

      __host__ KernelData(const VolumeGpu<PixelType, Allocator> &vol)
          : data_(const_cast<PixelType*>(vol.data())), width_(vol.width()), height_(vol.height()),
            depth_(vol.depth()), stride_(vol.stride())
      { }
  };

  ndarray_ref<PixelType,3> ref() const;

  VolumeGpu(const ndarray_ref<PixelType,3> &x);

protected:

private:
  PixelType* data_;
  size_t pitch_;
  bool ext_data_pointer_;

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

} // namespace iu