niftkCUDAManager.h

Go to the documentation of this file.

/*=============================================================================

  NifTK: A software platform for medical image computing.

  Copyright (c) University College London (UCL). All rights reserved.

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.

  See LICENSE.txt in the top level directory for details.

=============================================================================*/

#ifndef niftkCUDAManager_h
#define niftkCUDAManager_h

#include "niftkCUDAExports.h"
#include <CUDAImage/niftkCUDAImage.h>
#include <CUDAImage/niftkLightweightCUDAImage.h>
#include <QThread>
#include <QMutex>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <vector>
#include <list>
#include <set>
#include <string>
#include <boost/lockfree/queue.hpp>

namespace niftk
{

// FIXME: not yet implemented
struct ScopedCUDADevice
{
  ScopedCUDADevice(int dev);
  ~ScopedCUDADevice();
};


struct ReadAccessor
{
  const void*     m_DevicePointer;
  std::size_t     m_SizeInBytes;
  unsigned int    m_BytePitch;
  unsigned int    m_PixelWidth;
  unsigned int    m_PixelHeight;            // obviously the unit is lines of pixels
  int             m_FIXME_pixeltype;        // still havent thought about this one...

  unsigned int    m_Id;
  cudaEvent_t     m_ReadyEvent;
};


struct WriteAccessor
{
  void*           m_DevicePointer;
  std::size_t     m_SizeInBytes;
  unsigned int    m_BytePitch;
  unsigned int    m_PixelWidth;
  unsigned int    m_PixelHeight;            // obviously the unit is lines of pixels
  int             m_FIXME_pixeltype;        // still havent thought about this one...

  unsigned int    m_Id;
  cudaEvent_t     m_ReadyEvent;
};


// forward-decl
namespace impldetail
{
struct ModuleCleanup;
struct StreamCallbackReleasePOD;
}


class NIFTKCUDA_EXPORT CUDAManager : public QThread
{
  friend class LightweightCUDAImage;
  friend struct impldetail::ModuleCleanup;

public:
  static CUDAManager* GetInstance();


  // FIXME: not yet implemented
  ScopedCUDADevice ActivateDevice(int dev);

  cudaStream_t GetStream(const std::string& name);

  ReadAccessor RequestReadAccess(const LightweightCUDAImage& lwci);

  WriteAccessor RequestOutputImage(unsigned int width, unsigned int height, int FIXME_pixeltype);

  // when done with queueing commands to fill output image, call this.
  // it will give you a LightweightCUDAImage that can be stuffed in CUDAImage,
  // which in turn can go to a DataNode.
  LightweightCUDAImage Finalise(WriteAccessor& writeAccessor, cudaStream_t stream);

  LightweightCUDAImage FinaliseAndAutorelease(
      WriteAccessor& writeAccessor,
      ReadAccessor& readAccessor,
      cudaStream_t stream);

  void Autorelease(ReadAccessor& readAccessor, cudaStream_t stream);


  void Autorelease(WriteAccessor& writeAccessor, cudaStream_t stream);


protected:
  CUDAManager();
  virtual ~CUDAManager();


  void AllRefsDropped(LightweightCUDAImage& lwci);


private:
  CUDAManager(const CUDAManager& copyme);
  CUDAManager& operator=(const CUDAManager& assignme);


  std::size_t TierToSize(unsigned int tier) const;

  unsigned int SizeToTier(std::size_t size) const;


  static void CUDART_CB AutoReleaseStreamCallback(cudaStream_t stream, cudaError_t status, void* userData);


  void ReleaseReadAccess(unsigned int id);

  void ProcessAutoreleaseQueue();


  static CUDAManager*           s_Instance;
  // there's only one instance of our class (singleton), so a single mutex is ok too.
  static QMutex                 s_Lock;

  unsigned int                  m_LastIssuedId;

  // vector is a size tier, followed by linked list for that tier.
  std::vector<std::list<LightweightCUDAImage> >     m_AvailableImagePool;

  // images currently in use via WriteAccessor, i.e. work is being queued.
  std::map<unsigned int, LightweightCUDAImage>    m_InFlightOutputImages;

  // images that can be requested with RequestReadAccess.
  std::map<unsigned int, LightweightCUDAImage>    m_ValidImages;

  std::map<std::string, cudaStream_t>     m_Streams;

  // the auto-release callback cannot acquire s_Lock because that will deadlock within the cuda driver.
  boost::lockfree::queue<impldetail::StreamCallbackReleasePOD*>     m_AutoreleaseQueue;
};

} // end namespace
#endif