NCCL源码图解之Primitives类的初始化

date
Jul 12, 2025
slug
nccl_primitives_initialize
status
Published
tags
NCCL
源码图解
summary
Primitives类是NCCL中最基本的通信行为, 常用的集合通信原语都是这些行为的组合
type
Post
上节我们讲到all-reduce通信kernel在实际执行过程中会被拆解成多个send/recv步骤, 也就是下面这张图:
notion image
这里的Send, recvReduceSend等操作在NCCL里面就是基本的数据传输Op(这里列的只有一部分, 后面会讲到所有的), 这些op都是primitives这个类的成员函数, 而primitives这个泛型类的实例里保存着实际通信过程使用到的资源, 通信的行为; 所以我们今天主要讲解这个类.
这个类的代码位于src/collectives/device/ 目录下, 一共有三个, 分别是simple, ll(low lantency), ll128三种协议的实现, 这里我们以simple协议入手(其他两种大同小异), 看下prims_simple.h 这个文件

类成员变量

我们先看下代码, 参数及变量的含义已经标注了:
上面的标注做一个简单的了解, 后面遇到也会结合上下文讲解, 更容易理解
我们了解了大概含义就可以看看all-reduce中是怎么用的了:
这里有两个信息: 因为是ring算法, 所以Fan设置为1; 启用了Direct
接下来看下初始化

初始化

代码片段1:

  1. 首先看下函数声明里的变量, 这里的recvPeers 和sendPeers 表示当前rank的接收rank及发送rank的buffer的地址; 而inputBuf 和outputBuf 则表示用户给的输入输出, 是在当前rank上的地址.
  1. 这里的group 和connIndex不用管, all-reduce中=0, 暂时用不到
  1. 这里计算了stepSize, 也就是一个step要处理的数据个数

代码片段2

  1. 首先看下ThreadPerSync , 这里是把线程进行了分类, 然后赋予了不同的角色, 其实没有起到同步的作用 每类8个线程
  1. 这里的g表示当前线程在哪个类别
  1. index表示线程在类别中的index
下面的代码我们看下下面的图:
notion image
可以看到这里是为几个线程赋予了角色, 这里假设nthreads=512, 所以:
  • 第0号(g=0, index=0)线程的角色就是WaitRecv
  • 第1号(g=0, index=1)线程的角色就是input
  • 第8号(g=1, index=0)线程的角色就是WaitSend
  • 第0号(g=1, index=1)线程的角色就是output
  • 第1f0号(g=ng-2, index=0)线程的角色就是PostRecv
  • 第1f8号(g=ng-1, index=0)线程的角色就是PostSend
接下来的两步就是分别为Recv/send角色的线程获取到对应的peer rank

代码片段3

这两个函数做的事情类似, 我们就以loadRecvConn 进行讲解:

loadRecvConn

首先我们看下面的图, 方便理解不同角色干的事情:
notion image
我们看函数参数: peer, 这里是从共享内存中取出了peer的conn, 也就是recv rank的device侧的context.
然后看这里有个if, 所以只有RoleWaitRecv 和RolePostRecv 两个线程才会进入, 其他线程直接返回
然后就是从conn中取出了step,并对step向上取整, 也就是说我们一次通信任务完成SlicePerChunk*StepPerSlice数据量
接着看下RolePostRecv , 首先取出远程userbuf的指针head, 并将step赋予

RoleWaitRecv

最后看下RoleWaitRecv, 先取出了conn, 然后获取了tail指针, 也就是当前rank的buffer的尾指针, 而且把该指针赋给了uint64_t volatile *connStepPtr , 也就是说明这个指针会被另一个地方修改, 那么是哪里呢? 看下图:
notion image
就是recv peer, 该rank会发送数据并修改tail指针, 因为我们的buffer是一个先入先出(FIFO)的队列, 所以需要有一个head指针表示当前消费到了什么地方, 需要一个tail指针表示当前生产到了什么地方, 所以我们就知道了RoleWaitRecv 这个线程主要的功能就是当一个数据step发送完后, 修改send peer的tail指针.
接下来就是Direct操作, 这段逻辑只会在RoleWaitRecv中, 可以看到就是在做判断, 然后写flag. 其中read和write模式可以看
最后就是取出buffer的指针

loadSendConn

这个函数和上面的loadRecvConn 基本一样, 我们就不进行讲解了, 看上面的图明白是哪个角色在处理就可以了

setDataPtrs

该函数是初始化的最后一步, 我们先看下代码:
先贴图方便理解:
notion image
首先是RoleInput 把输入数据地址inputBuf赋值给类成员变量userBuff, 同时RoleOutput 获取了输出数据地址
接下来是根据角色设置了几个变量, DirectWrite和DirectRead逻辑区别不大(通信过程的write和read模式的异同见 ), 这里我们主要看下DirectWrite方式, 也就是recvProvider 和sendAcceptor

Direct

可以看到这里只对Direct生效, 非Direct模式会直接退出该函数

recvProvider

先是通过 volatile load获取了recvConn的ptrExchange, 如果是非direct模式, 我们在NCCL初始化的时候就把这个指针指向了send peer的buffer了, 现在direct模式, 我们就需要把数据直接写到用户提供的输出地址, 所以要把这个地址进行修改.
这个指针如果不是null, 代表了send peer还没有处理上一次的数据,就通过while循环等待, 直到send peer把这个值设置为null; 然后把output与ptrExchange指针的地址(也就是slot)进行异或(先验知识: 数a与数b异或两次, 还等于a)

sendAcceptor

先是通过 volatile load获取了sendConn的ptrExchange, 然后也是个while循环, 等待recv peer写入这个指针, 不为NULL说明等到了, 那么获取内容并退出循环, 然后再与ptrExchange指针的地址进行异或, 那么就得到了output的地址了. 最后把ptrExchange 的值设置为NULL
所以sendAcceptor需要和recvProvider配合使用, 功能是接收端告诉发送端, 如果是Direct模式, 应该要把数据写到哪
最后我们用一张图来看下Direct做的事情:
notion image
这里Primitives就初始化完成了, 数据传输所需要的指针交换也完成了, 接下来就可以直接发送数据了

© 木白 2024 - 2025