Parallel.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using MyCaffe.basecode;
using MyCaffe.param;
using MyCaffe.solvers;
 
namespace MyCaffe.common
{
    public class Params<T>
    {
        protected long m_lCount;  
 
        protected long m_lExtra;  
 
        protected long m_hData;
        protected long m_hDiff;
        protected int m_nDeviceID;
 
        public Params(Solver<T> root_solver)
        {
            m_lCount = total_size(root_solver.net.learnable_parameters);
            m_lExtra = 1000;
            m_hData = 0;
            m_hDiff = 0;
 
            m_lCount += m_lExtra;
        }
 
        public long count
        {
            get { return m_lCount; }
        }
 
        public long data
        {
            get { return m_hData; }
        }
 
        public long diff
        {
            get { return m_hDiff; }
        }
 
        private long total_size(BlobCollection<T> rgParam)
        {
            long nSize = 0;
 
            for (int i = 0; i < rgParam.Count; i++)
            {
                nSize += (long)rgParam[i].count();               
            }
 
            // Size should have at least one byte, otherwise malloc fails
            //  if net has no learnable parameters.
            if (nSize == 0)
                nSize++;
 
            return nSize;
        }
    }
 
    public class GPUParams<T> : Params<T>, IDisposable
    {
        protected CudaDnn<T> m_cuda;
        protected Log m_log;
        protected long m_hStream;
 
        public enum Op
        {
            copy,
            replace_gpu,
            replace_gpu_diff
        }
 
        public GPUParams(CudaDnn<T> cuda, Log log, Solver<T> root_solver, int nDeviceID)
            : base(root_solver)
        {
            m_cuda = cuda;
            m_log = log;
 
            m_nDeviceID = m_cuda.GetDeviceID();
 
            if (nDeviceID != m_nDeviceID)
                m_cuda.SetDeviceID(nDeviceID);
 
            // Allocate device buffers
            m_hData = m_cuda.AllocMemory(m_lCount);
 
            // Copy blob values
            BlobCollection<T> net = root_solver.net.learnable_parameters;
            apply_buffers(net, m_hData, m_lCount, Op.copy);
 
            m_hDiff = m_cuda.AllocMemory(m_lCount);
            m_cuda.set((int)m_lCount, m_hDiff, 0);
 
            m_hStream = m_cuda.CreateStream();
 
            if (m_nDeviceID != nDeviceID)
                m_cuda.SetDeviceID(m_nDeviceID);
        }
 
        public void Dispose()
        {
            if (m_hData != 0)
            {
                m_cuda.FreeMemory(m_hData);
                m_hData = 0;
            }
 
            if (m_hDiff != 0)
            {
                m_cuda.FreeMemory(m_hDiff);
                m_hDiff = 0;
            }
 
            if (m_hStream != 0)
            {
                m_cuda.FreeStream(m_hStream);
                m_hStream = 0;
            }
        }
 
        public void SynchronizeStream()
        {
            m_cuda.SynchronizeStream(m_hStream);
        }
 
        public void Configure(Solver<T> solver)
        {
            BlobCollection<T> net = solver.net.learnable_parameters;
            apply_buffers(net, m_hData, m_lCount, Op.replace_gpu);
            apply_buffers(net, m_hDiff, m_lCount, Op.replace_gpu_diff);
        }
 
        public void apply_buffers(BlobCollection<T> rgBlobs, long hBuffer, long lTotalSize, Op op)
        {
            long lOffset = 0;
 
            for (int i = 0; i < rgBlobs.Count; i++)
            {
                int nCount = rgBlobs[i].count();
 
                switch (op)
                {
                    // Init buffer to current values of blobs
                    case Op.copy:
                        m_cuda.copy(nCount, rgBlobs[i].data.gpu_data, hBuffer, 0, (int)lOffset);
                        break;
 
                    case Op.replace_gpu:
                        rgBlobs[i].data.set_gpu_data(hBuffer, nCount, lOffset);
                        break;
 
                    case Op.replace_gpu_diff:
                        if (rgBlobs[i].DiffExists)
                            rgBlobs[i].diff.set_gpu_data(hBuffer, nCount, lOffset);
                        break;
                }
 
                lOffset += nCount;
            }
 
            // total_size is at least one byte
            // We allocate extra items past the items used as a pad.
            m_log.CHECK_EQ(lTotalSize - m_lExtra, (lOffset == 0) ? 1 : lOffset, "The total memory doesn't match.");
        }
    }
 
    public class NCCL<T> : GPUParams<T>, IDisposable 
    {
        long m_hNccl;
        Solver<T> m_solver;
        ManualResetEvent m_evtGradientsReady = new ManualResetEvent(false);
        List<ManualResetEvent> m_rgGradientReady = new List<ManualResetEvent>();
 
        public NCCL(CudaDnn<T> cuda, Log log, Solver<T> root_solver, int nDeviceID, long hNccl, List<ManualResetEvent> rgGradientReadyEvents)
            : base(cuda, log, root_solver, nDeviceID)
        {
            m_rgGradientReady = rgGradientReadyEvents;
            if (rgGradientReadyEvents != null && rgGradientReadyEvents.Count > 0)
                m_evtGradientsReady = rgGradientReadyEvents[root_solver.solver_rank];
 
            m_solver = root_solver;
            m_hNccl = hNccl;
            Configure(root_solver);
 
            root_solver.OnGradientsReady += Solver_OnGradientsReady;
        }
 
        public new void Dispose()
        {
            base.Dispose();
 
            if (m_hNccl != 0)
            {
                m_cuda.FreeNCCL(m_hNccl);
                m_hNccl = 0;
            }
        }
 
        public void Broadcast()
        {
            m_cuda.NcclBroadcast(m_hNccl, m_hStream, m_hData, (int)m_lCount);
            m_cuda.SynchronizeStream(m_hStream);
        }
 
        private void Solver_OnGradientsReady(object sender, GradientsReadyArgs e)
        {
            try
            {
                m_cuda.SynchronizeStream();
                m_evtGradientsReady.Set();
 
                while (!WaitHandle.WaitAll(m_rgGradientReady.ToArray(), 250))
                {
                    if (m_solver.CancelEvent.WaitOne(0))
                        return;
                }
 
                double dfScale = 1.0 / m_solver.solver_count;
                m_cuda.NcclAllReduce(m_hNccl, m_hStream, m_hDiff, (int)m_lCount, NCCL_REDUCTION_OP.SUM, dfScale);
                m_cuda.SynchronizeStream(m_hStream);
            }
            finally
            {
                m_evtGradientsReady.Reset();
            }
        }
 
        public void Run(List<int> rgGpus, int nIterationOverride = -1)
        {
            List<long> rghNccl = new List<long>();
            Guid guid = Guid.NewGuid();
 
            m_rgGradientReady = new List<ManualResetEvent>();
 
            for (int i = 0; i < rgGpus.Count; i++)
            {
                long hNccl = m_cuda.CreateNCCL(rgGpus[i], m_solver.solver_count, i, guid);
                rghNccl.Add(hNccl);
                m_rgGradientReady.Add(new ManualResetEvent(false));
            }
 
            m_cuda.NcclInitializeSingleProcess(rghNccl.ToArray());
            m_hNccl = rghNccl[0];
            m_evtGradientsReady = m_rgGradientReady[0];
 
            List<WaitHandle> rgWaitAllInit = new List<WaitHandle>();
            List<Worker<T>> rgWorkers = new List<common.Worker<T>>();
            ManualResetEvent evtAllCreated = new ManualResetEvent(false);
 
            for (int i = 1; i < rghNccl.Count; i++)
            {
                Worker<T> worker = new Worker<T>();
 
                SolverInfo<T> info = new common.SolverInfo<T>(m_solver, m_cuda.KernelHandle, rghNccl[i], i, nIterationOverride, m_cuda.Path, m_rgGradientReady, evtAllCreated);
                worker.StartInternalThread(null, null, rgGpus[i], info);
 
                List<WaitHandle> rgWait = new List<WaitHandle>();
                rgWait.AddRange(m_solver.CancelEvent.Handles);
                rgWait.Add(info.ErrorEvent);
                rgWait.Add(info.StartedEvent);
 
                int nWait = WaitHandle.WaitAny(rgWait.ToArray());
                if (nWait < rgWait.Count - 2)
                    return;
 
                if (nWait == rgWait.Count - 2)
                {
                    if (info.Error != null)
                        throw info.Error;
                    else
                        throw new Exception("Error starting the solver.");
                }
 
                rgWaitAllInit.Add(info.InitializedEvent);
                rgWorkers.Add(worker);
            }
 
            // Wait for all worksers to initialize
            while (!WaitHandle.WaitAll(rgWaitAllInit.ToArray(), 250))
            {
                if (m_solver.CancelEvent.WaitOne(0))
                    return;
            }
 
            m_cuda.SynchronizeDevice();
            evtAllCreated.Set();
 
            // Run first solver on current thread.
            Broadcast();
 
            m_solver.Solve(nIterationOverride);
 
            // Wait for shutdown
            for (int i = 0; i < rgWorkers.Count; i++)
            {
                rgWorkers[i].StopInternalThread();
            }
        }
    }
 
    public class Worker<T> : InternalThread<T>
    {
        CudaDnn<T> m_cuda;
 
        public Worker()
        {
            this.DoWork += Worker_DoWork;
        }
 
        private void Worker_DoWork(object sender, ActionStateArgs<T> e)
        {
            SolverInfo<T> info = e.Arg as SolverInfo<T>;
            NCCL<T> nccl = null;
 
            m_cuda = new common.CudaDnn<T>(e.DeviceID, DEVINIT.CUBLAS | DEVINIT.CURAND, null, info.CudaPath);
 
            try
            {
                Solver<T> rank0 = info.Rank0;
                Log log = new Log("Worker solver for DeviceID = " + e.DeviceID.ToString());
 
                //-----------------------------------------
                //  Transfer the NCCL handle from the 
                //  main kernel that created it to the
                //  one used by the CudaDnn on this thread.
                //
                //  After the copy, this thread will 'own'
                //  the nccl and be responsible for its 
                //  destruction.
                //-----------------------------------------
                long hNccl = m_cuda.KernelCopyNccl(info.KernelHandle, info.NcclHandle);
 
                // Create solver and install callbacks
                SolverParameter param = rank0.parameter.Clone();
                param.device_id = e.DeviceID;
                param.type = rank0.parameter.type;
                Solver<T> solver = Solver<T>.Create(m_cuda, log, param, rank0.CancelEvent, null, null, rank0.Database, null, rank0.solver_count, info.SolverRank);
                info.StartedEvent.Set();
                log.CHECK_EQ((int)solver.type, (int)rank0.type, "The solver types should be the same.");
 
                //-----------------------------------------
                //  Turn off logging for all other 
                //  operations on the worker thread.
                //-----------------------------------------
                log.Enable = false;
 
                nccl = new NCCL<T>(m_cuda, log, solver, e.DeviceID, hNccl, info.GradientReadyEvents);
 
                info.InitializedEvent.Set();
                m_cuda.SynchronizeDevice();
 
                List<WaitHandle> rgWait = new List<WaitHandle>();
                rgWait.AddRange(rank0.CancelEvent.Handles);
                rgWait.Add(info.AllCreatedEvent);
 
                int nWait = WaitHandle.WaitAny(rgWait.ToArray());
                if (nWait < rgWait.Count - 1)
                    return;
 
                nccl.Broadcast();
 
                int nIterations = param.max_iter - solver.iter;
                if (info.IterationOverride > 0)
                    nIterations = info.IterationOverride;
 
                solver.Step(nIterations);
                solver.Dispose();
            }
            catch (Exception excpt)
            {
                info.Error = excpt;
                info.ErrorEvent.Set();
            }
            finally
            {
                if (nccl != null)
                    nccl.Dispose();
 
                m_cuda.Dispose();
                m_cuda = null;
            }
        }
    }
 
    public class SolverInfo<T>
    {
        string m_strCudaPath;
        Solver<T> m_rank0;
        long m_hSrcKernel;
        long m_hSrcNccl;
        int m_nSolverRank;
        int m_nIterationOverride;
        ManualResetEvent m_evtInitialized = new ManualResetEvent(false);
        ManualResetEvent m_evtStarted = new ManualResetEvent(false);
        ManualResetEvent m_evtAllCreated = new ManualResetEvent(false);
        AutoResetEvent m_evtError = new AutoResetEvent(false);
        List<ManualResetEvent> m_rgGradientReadyEvents = null;
        Exception m_error = null;
 
        public SolverInfo(Solver<T> rank0, long hSrcKernel, long hSrcNccl, int nSolverRank, int nIterationOverride, string strCudaPath, List<ManualResetEvent> rgGradientReadyEvents, ManualResetEvent evtAllCreated)
        {
            m_strCudaPath = strCudaPath;
            m_rank0 = rank0;
            m_hSrcKernel = hSrcKernel;
            m_hSrcNccl = hSrcNccl;
            m_nSolverRank = nSolverRank;
            m_nIterationOverride = nIterationOverride;
            m_rgGradientReadyEvents = rgGradientReadyEvents;
            m_evtAllCreated = evtAllCreated;
        }
 
        public Solver<T> Rank0
        {
            get { return m_rank0; }
        }
 
        public string CudaPath
        {
            get { return m_strCudaPath; }
        }
 
        public int IterationOverride
        {
            get { return m_nIterationOverride; }
        }
 
        public long KernelHandle
        {
            get { return m_hSrcKernel; }
        }
 
        public long NcclHandle
        {
            get { return m_hSrcNccl; }
        }
 
        public int SolverRank
        {
            get { return m_nSolverRank; }
        }
 
        public ManualResetEvent InitializedEvent
        {
            get { return m_evtInitialized; }
        }
 
        public ManualResetEvent StartedEvent
        {
            get { return m_evtStarted; }
        }
 
        public ManualResetEvent AllCreatedEvent
        {
            get { return m_evtAllCreated; }
        }
 
        public List<ManualResetEvent> GradientReadyEvents
        {
            get { return m_rgGradientReadyEvents; }
        }
 
        public Exception Error
        {
            get { return m_error; }
            set { m_error = value; }
        }
 
        public AutoResetEvent ErrorEvent
        {
            get { return m_evtError; }
        }
    }
}