HostBuffer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
 
namespace MyCaffe.common
{
    public class HostBuffer<T> : IDisposable
    {
        CudaDnn<T> m_cuda;
        long m_hBuffer;
        long m_lCapacity;
 
        public HostBuffer(CudaDnn<T> cuda)
        {
            m_cuda = cuda;
            m_hBuffer = 0;
            m_lCapacity = 0;
        }
 
        public void Dispose()
        {
            Free();
        }
 
        public long Handle
        {
            get { return m_hBuffer; }
        }
 
        public long Capacity
        {
            get { return m_lCapacity; }
        }
 
        public void Free()
        {
            if (m_hBuffer != 0)
            {
                m_cuda.FreeHostBuffer(m_hBuffer);
                m_hBuffer = 0;
                m_lCapacity = 0;
            }
        }
 
        public void CopyFrom(Blob<T> b, bool bFromDiff = false)
        {
            CopyFromGpu(b.count(), (bFromDiff) ? b.gpu_diff : b.gpu_data);
        }
 
        public void CopyTo(Blob<T> b, bool bToDiff = false)
        {
            CopyToGpu(b.count(), (bToDiff) ? b.mutable_gpu_diff : b.mutable_gpu_data);
        }
 
        public void CopyFromGpu(int nCount, long hGpu)
        {
            if (nCount > m_lCapacity)
            {
                Free();
                m_hBuffer = m_cuda.AllocHostBuffer(nCount);
                m_lCapacity = nCount;
            }
 
            m_cuda.CopyDeviceToHost(nCount, hGpu, m_hBuffer);
        }
 
        public void CopyToGpu(int nCount, long hGpu)
        {
            m_cuda.CopyHostToDevice(nCount, m_hBuffer, hGpu);
        }
 
        public T[] GetHostData()
        {
            return m_cuda.GetHostMemory(m_hBuffer);
        }
 
        public double[] GetHostDataAsDouble()
        {
            return m_cuda.GetHostMemoryDouble(m_hBuffer);
        }
 
        public float[] GetHostDataAsFloat()
        {
            return m_cuda.GetHostMemoryFloat(m_hBuffer);
        }
 
        public void SetHostData(T[] rgSrc)
        {
            if (m_lCapacity < rgSrc.Length)
            {
                Free();
                m_hBuffer = m_cuda.AllocHostBuffer(rgSrc.Length);
            }
 
            m_cuda.SetHostMemory(m_hBuffer, rgSrc);
        }
    }
}