SyncedMemory.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
 
namespace MyCaffe.common
{
    public class SyncedMemory<T> : IDisposable 
    {
        Log m_log;
        CudaDnn<T> m_cuda;
        int m_nDeviceID = -1;
        long m_lCapacity = 0;
        long m_lCount = 0;
        long m_hGpuData = 0;
        T[] m_rgCpuData = null;
        bool m_bOwnData = true;
        bool m_bHalfSize = false;
        object m_tag = null;
 
        public SyncedMemory(CudaDnn<T> cuda, Log log, long lCapacity = 0, object tag = null, bool bUseHalfSize = false)
        {
            if (bUseHalfSize && typeof(T) != typeof(float))
            {
                bUseHalfSize = false;
 
                if (log != null)
                    log.WriteLine("Half size disabled for non 'float' basetypes!");
            }
 
            m_bHalfSize = bUseHalfSize;
            m_cuda = cuda;
            m_log = log;
            m_tag = tag;
 
            if (lCapacity > 0)
            {
                m_nDeviceID = m_cuda.GetDeviceID();
                m_hGpuData = m_cuda.AllocMemory(lCapacity);
                m_lCapacity = lCapacity;
                m_lCount = lCapacity;
            }
        }
 
        private void free()
        {
            if (m_hGpuData != 0)
            {
                check_device();
                if (m_bOwnData)
                    m_cuda.FreeMemory(m_hGpuData);
                else
                    m_cuda.FreeMemoryPointer(m_hGpuData);
            }
        }
 
        public void Dispose()
        {
            free();
            m_hGpuData = 0;
            m_nDeviceID = -1;
            m_lCapacity = 0;
            m_lCount = 0;
            m_rgCpuData = null;
        }
 
        public void Allocate(long lCount, bool bUseHalfSize = false)
        {
            if (bUseHalfSize && typeof(T) != typeof(float))
            {
                bUseHalfSize = false;
 
                if (m_log != null)
                    m_log.WriteLine("Half size disabled for non 'float' basetypes!");
            }
 
            free();
            m_nDeviceID = m_cuda.GetDeviceID();
            m_bHalfSize = bUseHalfSize;
            m_hGpuData = m_cuda.AllocMemory(lCount, m_bHalfSize);
            m_lCapacity = lCount;
            m_lCount = 0;
            m_bOwnData = true;
            return;
        }
 
        public void Allocate(T[] rg, bool bUseHalfSize = false)
        {
            if (bUseHalfSize && typeof(T) != typeof(float))
            {
                bUseHalfSize = false;
 
                if (m_log != null)
                    m_log.WriteLine("Half size disabled for non 'float' basetypes!");
            }
 
            free();
            m_nDeviceID = m_cuda.GetDeviceID();
            m_bHalfSize = bUseHalfSize;
            m_hGpuData = m_cuda.AllocMemory(rg, 0, m_bHalfSize);
            m_lCapacity = rg.Length;
            m_lCount = rg.Length;
            m_bOwnData = true;
            check_device();
 
            return;
        }
 
        public void Zero()
        {
            if (m_lCount > 0)
            {
                check_device();
                m_cuda.set((int)m_lCount, m_hGpuData, 0.0);
            }
        }
 
        public void ZeroAll()
        {
            if (m_lCapacity > 0)
            {
                check_device();
                m_cuda.set((int)m_lCapacity, m_hGpuData, 0.0);
            }
        }
 
        public void Set(double dfVal)
        {
            if (m_lCount > 0)
            {
                check_device();
                m_cuda.set((int)m_lCount, m_hGpuData, dfVal);
            }
        }
 
        public void SetAt(int nIdx, T fVal)
        {
            check_device();
            m_cuda.set((int)m_lCount, m_hGpuData, fVal, nIdx);
        }
 
        public T GetAt(int nIdx)
        {
            check_device();
            T[] rg = m_cuda.get((int)m_lCount, m_hGpuData, nIdx);
            return rg[0];
        }
 
        public long Copy(SyncedMemory<T> src, long hDstHostBuffer = 0)
        {
            if (src == null)
            {
                m_lCount = 0;
                return hDstHostBuffer;
            }
 
            if (m_lCapacity < src.m_lCount)
                Allocate(src.m_lCount);
 
            m_lCount = src.m_lCount;
 
            if (m_lCount > 0)
            {
                if (m_cuda.KernelHandle == src.m_cuda.KernelHandle)
                {
                    check_device();
                    m_cuda.copy((int)m_lCount, src.m_hGpuData, m_hGpuData);
                }
                else
                {
                    if (hDstHostBuffer == 0)
                    {
                        hDstHostBuffer = m_cuda.AllocHostBuffer(m_lCount);
                    }
                    else
                    {
                        long lCount = m_cuda.GetHostBufferCapacity(hDstHostBuffer);
                        if (lCount < m_lCount)
                        {
                            m_cuda.FreeHostBuffer(hDstHostBuffer);
                            hDstHostBuffer = m_cuda.AllocHostBuffer(m_lCount);
                        }
                    }
 
                    src.m_cuda.KernelCopy((int)m_lCount, src.m_hGpuData, 0, m_cuda.KernelHandle, m_hGpuData, 0, hDstHostBuffer, m_cuda.KernelHandle);
                }
            }
 
            return hDstHostBuffer;
        }
 
        public SyncedMemory<T> Clone()
        {
            SyncedMemory<T> dst = new SyncedMemory<T>(m_cuda, m_log, m_lCapacity, null, m_bHalfSize);
 
            if (m_lCount > 0)
                dst.Copy(this);
 
            return dst;
        }
 
        public bool HalfSize
        {
            get { return m_bHalfSize; }
        }
 
        public object Tag
        {
            get { return m_tag; }
            set { m_tag = value; }
        }
 
        public int DeviceID
        {
            get { return m_nDeviceID; }
        }
 
        public long Capacity
        {
            get { return m_lCapacity; }
        }
 
        public long Count
        {
            get { return m_lCount; }
            set { m_lCount= value; }
        }
 
        public long gpu_data
        {
            get { return m_hGpuData; }
        }
 
        public void set_gpu_data(long hData, long lCount, long lOffset)
        {
            free();
            m_hGpuData = m_cuda.CreateMemoryPointer(hData, lOffset, lCount);
            m_lCapacity = lCount;
            m_lCount = lCount;
            m_bOwnData = false;
            check_device();
        }
 
        public long mutable_gpu_data
        {
            get
            {
                check_device();
                return m_hGpuData;
            }
//          set { m_hGpuData = value; }
        }
 
        public T[] cpu_data
        {
            get { return m_rgCpuData; }
        }
 
        public void SetData(T[] rgData, int nCount, bool bSetCount = true)
        {
            if (nCount == -1)
                nCount = rgData.Length;
 
            if (nCount > m_lCapacity || m_hGpuData == 0)
            {
                bSetCount = true;
                Allocate(nCount);
            }
 
            m_cuda.SetMemory(m_hGpuData, rgData, 0, nCount);
 
            if (bSetCount)
                m_lCount = nCount;
 
            check_device();
        }
 
        public T[] mutable_cpu_data
        {
            get { return update_cpu_data(); }
            set
            {
                check_device();
                if (value.Length > m_lCapacity || m_hGpuData == 0)
                {
                    Allocate(value);
                }
                else
                {
                    m_cuda.SetMemory(m_hGpuData, value);
                    m_lCount = value.Length;
                    check_device();
                }
            }
        }
 
        public T[] update_cpu_data(long lCount = -1)
        {
            if (lCount >= 0)
            {
                if (lCount > m_lCapacity)
                    throw new ArgumentOutOfRangeException();
 
                m_lCount = lCount;
            }
 
            check_device();
 
            if (m_lCount == 0)
                m_rgCpuData = new List<T>().ToArray();
            else
                m_rgCpuData = m_cuda.GetMemory(m_hGpuData, m_lCount);
 
            return m_rgCpuData;
        }
 
        public void set_cpu_data_locally(T[] rg)
        {
            m_rgCpuData = rg;
        }
 
        public void async_gpu_push(long hStream, T[] rg)
        {
            check_device();
            if (m_hGpuData == 0)
            {
                m_hGpuData = m_cuda.AllocMemory(rg, hStream);
                m_lCapacity = rg.Length;
            }
            else
            {
                m_cuda.SetMemory(m_hGpuData, rg, hStream);
            }
 
            m_lCount = rg.Length;
            check_device();
        }
 
        private void check_device()
        {
#if DEBUG
            if (m_lCount > 0)
            {
                int nDeviceId = m_cuda.GetDeviceID();
                m_log.CHECK_EQ(nDeviceId, m_nDeviceID, "The current device DOES'NT match the device for which the memory was allocated!");
            }
#endif
        }
    }
}