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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritk3DCuda3DArrayImportImageContainer.txx
	Version:		$Revision: 3136 $
	Date:			$Date: 2012-04-19 14:42:59 +0200 (Do, 19 Apr 2012) $

	Copyright (c) 2010-2011
	Pattern Recognition Lab, 
    University Erlangen-Nuremberg, Germany
    All rights reserved.

	This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
	See RITKCopyright.txt for more information.

	http://www5.informatik.uni-erlangen.de/research/software/ritk

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

#include <cuda_runtime.h>

#include "ritkCuda3DArrayImportImageContainer.h"
#include "ritkCuda3DPitchedMemoryImportImageContainer.h"

#include "ritkDebugManager.h"

namespace ritk
{

//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>
::Cuda3DArrayImportImageContainer()
{
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>
::~Cuda3DArrayImportImageContainer()
{
	this->DeallocateManagedMemory();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
PrintSelf(std::ostream& os, itk::Indent indent) const
{
	os << indent << "Array @" << m_CudaArray.front().first << std::endl;
	Superclass::PrintSelf(os,indent);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
const cudaChannelFormatDesc&
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
GetChannelDesc() const
{
	return m_CudaArray.front().second;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
cudaChannelFormatDesc&
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
GetChannelDesc()
{
	return m_CudaArray.front().second;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
typename Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::Element*
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
AllocateElements(ElementIdentifier size) const
{
	// Sanity checks. This ensures that we can operate on the internal 2-D size parameters
	if ( Superclass::m_ContainerSize[0] == 0 || Superclass::m_ContainerSize[1] == 0 || Superclass::m_ContainerSize[2] == 0 )
	{
		LOG_DEB("Invalid width/height/depth detected: did you forget to set them via SetContainerSize()?");
		itkExceptionMacro(<<"Invalid width/height/depth detected: did you forget to set them via SetContainerSize()?");
	}
	if ( Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1]*Superclass::m_ContainerSize[2] != size )
	{
		LOG_DEB("Mismatch between widht/height/depth and requested size detected: " << Superclass::m_ContainerSize[0] << "*" << Superclass::m_ContainerSize[1] << "*" << Superclass::m_ContainerSize[2] << "=" << (Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1]*Superclass::m_ContainerSize[2]) << " vs. " << size);
		itkExceptionMacro(<<"Mismatch between widht/height/depth and requested size detected: " << Superclass::m_ContainerSize[0] << "*" << Superclass::m_ContainerSize[1] << "*" << Superclass::m_ContainerSize[2] << "=" << (Superclass::m_ContainerSize[0]*Superclass::m_ContainerSize[1]*Superclass::m_ContainerSize[2]) << " vs. " << size);
	}

	// Execute in the global CUDA context. Note that we can't use the size argument directly. Instead, we use the internal 3-D size!
	CudaBlockedConstMethodCallMacro(AllocateCudaMemory);

	// Call the superclass' implementation. This also returns the correct host memory pointer
	return Superclass::AllocateElements(size);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
DeallocateManagedMemory()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(DeallocateCudaMemory);

	// Call the superclass' implementation to deallocate host memory
	Superclass::DeallocateManagedMemory();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
SynchronizeDevice()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(CopyHostToDevice);

	// Synchronize MTimes
	Superclass::SynchronizeDevice();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
SynchronizeHost()
{
	// Execute in the global CUDA context
	CudaBlockedMethodCallMacro(CopyDeviceToHost);

	// Synchronize MTimes
	Superclass::SynchronizeHost();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
DeepCopy(const ImportImageContainerType *Container)
{
	// Cast and pass
	ImportImageContainerType *NonConstContainer = const_cast<ImportImageContainerType*>(Container);
	if ( !NonConstContainer )
	{
		LOG_DEB("Could not remove const qualifier!");
		itkExceptionMacro(<<"Could not remove const qualifier!");
	}
	DeepCopy(NonConstContainer);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
DeepCopy(ImportImageContainerType *Container)
{
	// Call the superclass implementation first. This will ensure correct memory allocation and host memory copy.
	Superclass::DeepCopy(Container);

	// Check if we have a cuda container at all. If no, the GenericImportImageContainer should call a Synchronize() triggered by Superclass::DeepCopy() on us
	Superclass *CudaContainer = dynamic_cast<Superclass*>(Container);
	if ( !CudaContainer )
		return;

	// Check for MTimes: we do not need to copy device memory that is out-of-date. Eventually, this has to be resolved by a synch!
	if ( CudaContainer->GetDeviceMTime() < CudaContainer->GetHostMTime() )
	{
		GenericImportImageContainerType::m_DeviceMTime = CudaContainer->GetDeviceTimeStamp();
		return;
	}

	// Device copy
	m_TmpCudaContainer = CudaContainer;
	CudaBlockedMethodCallMacro(DeepCopyCudaMemory);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
CopyToMappedResource(void *d_ptr, size_t size) const
{
	itkExceptionMacro(<<"Copying to mapped resources not supported by an array container!");
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
cudaArray*
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryArray()
{
	GenericImportImageContainerType::DeviceModified();
	return m_CudaArray.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
const cudaArray*
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
GetCudaMemoryArray() const
{
	return m_CudaArray.front().first;
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
AllocateCudaMemory() const
{
	// Create the channel descriptor for the TElement argument.
	cudaChannelFormatDesc Desc = cudaCreateChannelDesc<TElement>();

	// Allocate. Note that a correct size is checked in AllocateElements()
	cudaExtent VolumeSize = make_cudaExtent(Superclass::m_ContainerSize[0],Superclass::m_ContainerSize[1],Superclass::m_ContainerSize[2]);
	cudaArray *Array;
    ritkCudaSafeCall(cudaMalloc3DArray((cudaArray**)&Array, &Desc, VolumeSize));

	// Save the allocated memory at the end of our internal list such that it won't be deleted through the itk (de)allocate scheme
	std::pair<cudaArray*, cudaChannelFormatDesc> tmp;
	tmp.first = Array;
	tmp.second = Desc;
	m_CudaArray.push_back(tmp);
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
DeallocateCudaMemory()
{
	// Sanity check
	if ( m_CudaArray.size() == 0 )
	{
		return;
	}

	// Delete the first element in our internal list
	cudaArray *Array = m_CudaArray.front().first;
	ritkCudaSafeCall(cudaFreeArray(Array));
	m_CudaArray.pop_front();
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
CopyDeviceToHost()
{
	// Copy data from device to host
	cudaExtent Extent = make_cudaExtent(Superclass::m_ContainerSize[0],Superclass::m_ContainerSize[1],Superclass::m_ContainerSize[2]);
    cudaMemcpy3DParms CopyParams = {0};
	CopyParams.dstPtr = make_cudaPitchedPtr(GenericImportImageContainerType::GetBufferPointer(), Extent.width*sizeof(Element), Extent.width, Extent.height);
    CopyParams.srcArray = this->GetCudaMemoryArray();
    CopyParams.extent = Extent;
    CopyParams.kind  = cudaMemcpyDeviceToHost;
    ritkCudaSafeCall(cudaMemcpy3D(&CopyParams));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
CopyHostToDevice()
{
	// Copy data from host to device
	cudaExtent Extent = make_cudaExtent(Superclass::m_ContainerSize[0],Superclass::m_ContainerSize[1],Superclass::m_ContainerSize[2]);
    cudaMemcpy3DParms CopyParams = {0};
	CopyParams.srcPtr = make_cudaPitchedPtr(GenericImportImageContainerType::GetBufferPointer(), Extent.width*sizeof(Element), Extent.width, Extent.height);
    CopyParams.dstArray = this->GetCudaMemoryArray();
    CopyParams.extent = Extent;
    CopyParams.kind  = cudaMemcpyHostToDevice;
    ritkCudaSafeCall(cudaMemcpy3DAsync(&CopyParams));
}


//----------------------------------------------------------------------------
template<typename TElementIdentifier, typename TElement>
void
Cuda3DArrayImportImageContainer<TElementIdentifier, TElement>::
DeepCopyCudaMemory()
{
	// 3DArray -> 3DArray
	const Cuda3DArrayImportImageContainer *ArrayContainer = dynamic_cast<const Cuda3DArrayImportImageContainer*>(m_TmpCudaContainer);
	if ( ArrayContainer )
	{
		cudaExtent VolumeSize = make_cudaExtent(Superclass::m_ContainerSize[0],Superclass::m_ContainerSize[1],Superclass::m_ContainerSize[2]);
		cudaMemcpy3DParms CopyParams = {0};
		CopyParams.srcArray = const_cast<cudaArray*>(ArrayContainer->GetCudaMemoryArray()); // ugly but copy cudaMemcpy3DParms doesn't know const...
		CopyParams.dstArray = this->GetCudaMemoryArray();
		CopyParams.extent = VolumeSize;
		CopyParams.kind  = cudaMemcpyDeviceToDevice;
		ritkCudaSafeCall(cudaMemcpy3DAsync(&CopyParams));

		// Update the timestamps. Take care if host and device are already synchronized!
		GenericImportImageContainerType::DeviceModified();
		if ( m_TmpCudaContainer->GetDeviceTimeStamp() == m_TmpCudaContainer->GetHostTimeStamp() )
			GenericImportImageContainerType::m_HostMTime = GenericImportImageContainerType::m_DeviceMTime;

		// Clean up
		m_TmpCudaContainer = NULL;
		return;
	}

	// 3DPitched -> 3DArray
	typedef Cuda3DPitchedMemoryImportImageContainer<ElementIdentifier, Element> Cuda3DPitchedMemoryImportImageContainerType;
	const Cuda3DPitchedMemoryImportImageContainerType *PitchedMemoryContainer = dynamic_cast<const Cuda3DPitchedMemoryImportImageContainerType*>(m_TmpCudaContainer);
	if ( PitchedMemoryContainer )
	{
		cudaMemcpy3DParms CopyParams = {0};
		CopyParams.srcPtr = PitchedMemoryContainer->GetCudaPitchedMemoryPointer();
		CopyParams.dstArray = this->GetCudaMemoryArray();
		CopyParams.extent = make_cudaExtent(Superclass::m_ContainerSize[0],Superclass::m_ContainerSize[1],Superclass::m_ContainerSize[2]);
		CopyParams.kind  = cudaMemcpyDeviceToDevice;
		ritkCudaSafeCall(cudaMemcpy3DAsync(&CopyParams));

		// Update the timestamps. Take care if host and device are already synchronized!
		GenericImportImageContainerType::DeviceModified();
		if ( m_TmpCudaContainer->GetDeviceTimeStamp() == m_TmpCudaContainer->GetHostTimeStamp() )
			GenericImportImageContainerType::m_HostMTime = GenericImportImageContainerType::m_DeviceMTime;

		// Clean up
		m_TmpCudaContainer = NULL;
		return;
	}

	// We can't handle this container
	LOG_DEB("Invalid container detected: " << m_TmpCudaContainer);
	itkExceptionMacro(<<"Invalid container detected: " << m_TmpCudaContainer);
}


} // ritk
