代码解读基于版本2.11.4；后续会基于最新版本重构，敬请期待！

NCCL里数据通信是通过CUDA kernel执行的, 其中每一个channel会绑定一个block（也就是一个SM）. kernel的代码在src/collectives/device 目录的各个.h中，

包含了all_reduce，all_gather等各种通信原语, 我们就以all_reduce入手, 着重讲讲其原理, 其他原语比较类似, 有区别的地方会讲到. 另外all_reduce原语还分别有ring及tree算法, 每种算法实现还有simple/ll(low latency)/ll_128三种协议

我们接下来以all_reduce原语的ring算法中的simple协议的实现入手, 逐行解析, 后面会对其中不一样的地方进行补充

Simple协议Ring算法的all_reduce原语

我们首先需要了解ring all_reduce的实现原理( ‣), 假设当前有4张卡, 有4个channel进行数据传输, 从上帝视角看, 其传输过程如下图:

ring all-reduce示意图

ring all-reduce示意图

这时候我们从时序的角度, 看下每张卡做的事情:

image.png

也就是说在任意时刻, 任意一张卡都是在做send/recv, 而且其要发送的rank及接收的rank是确定且固定的.

接下来我们看下代码

代码片段0

template<typename T, typename RedOp, typename Proto>
  __device__ __forceinline__ void runRing(ncclWorkElem *args) {
    const int tid = threadIdx.x;
    const int nthreads = args->nThreads;
    const int bid = args->coll.bid;
    const int nChannels = args->coll.nChannels;
    ncclRing *ring = &ncclShmem.channel.ring;
    int ringIx = ring->index;
    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLREDUCE_CHUNKSTEPS : 1));
    const int nranks = ncclShmem.comm.nRanks;
    const ssize_t loopSize = nChannels*nranks*chunkSize;
    const ssize_t size = args->coll.count;

这里分别解释下含义: