背景
2.28的nccl发布之后，也支持了最近很火的ibgda，nccl中称ibgda为GIN，即GPU-Initiated Networking，这节我们看下nccl是怎么做的。
用法
int main() {
[...]
memset(&reqs, 0, sizeof(ncclDevCommRequirements));
int nCTAs = 1;
reqs.railGinBarrierCount = nCTAs;
reqs.ginSignalCount = 1;
NCCLCHECK(ncclDevCommCreate(comm, &reqs, &devComm));
[...]
}
template <typename T>
__global__ void ginAlltoAllKernel(ncclDevComm devComm, ncclWindow_t win,
size_t inputOffset, size_t outputOffset, size_t count) {
int ginContext = 0;
ncclGinSignal_t signalIndex = 0;
ncclGin gin { devComm, ginContext };
uint64_t signalValue = gin.readSignal(signalIndex);
ncclGinBarrierSession<ncclCoopCta> bar { ncclCoopCta(), gin, ncclTeamWorld(devComm),
devComm.railGinBarrier, blockIdx.x };
bar.sync(ncclCoopCta(), cuda::memory_order_relaxed, ncclGinFenceLevel::Relaxed);
const int rank = devComm.rank, nRanks = devComm.nRanks;
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
const int nThreads = blockDim.x * gridDim.x;
const size_t size = count * sizeof(T);
for (int peer = tid; peer < nRanks; peer += nThreads) {
gin.put(ncclTeamWorld(devComm), peer, win, outputOffset + rank * size,
win, inputOffset + peer * size, size, ncclGin_SignalInc{signalIndex});
}
gin.waitSignal(ncclCoopCta(), signalIndex, signalValue + nRanks);
gin.flush(ncclCoopCta());
}
相对于对称内存，多了一步ncclDevCommCreate创建device端的comm，device端构造ncclGin结构体，执行put进行数据发送，通过waitSignal等待所有数据的接收，最后执行flush，保证当前rank执行的put对应的源数据buffer可以被重新使用，不保证数据在对端落盘，实现上用的pollcq。
host端初始化
创建context
ncclResult_t ncclGinConnectOnce(struct ncclComm* comm) {
struct ncclGinState* ginState = &comm->sharedRes->ginState;
NCCLCHECK(ginState->ncclGin->init(&ginState->ginInstance, comm->commHash, ncclDebugLog));
NCCLCHECK(ncclTopoGetLocalNets(comm->topo, comm->rank, localNets, &nLocalNets));
ginState->ginCommCount = std::min<int>(NCCL_GIN_MAX_CONTEXTS, ncclParamGinNcontexts());
ncclSpaceConstruct(&ginState->counterSpace);
ncclSpaceConstruct(&ginState->signalSpace);
}
ginState存了gin相关所有信息。ncclGin的init就是获取机器所有可用的网卡，然后ncclTopoGetLocalNets获取当前gpu本地的网卡，ginCommCount表示用多少个context，gin暂时每个context都是一个qp，多qp是通过多context实现的，后续都假设为1。然后初始化counter和signal的space。
ncclResult_t ncclGinConnectOnce(struct ncclComm* comm) {
...
for (int n = 0; n < ginState->ginCommCount; n++) {
void* listenComm;
ginState->ncclGin->listen(ginState->ginInstance, localNets[n%nLocalNets],
allHandles + NCCL_NET_HANDLE_MAXSIZE * comm->rank, &listenComm);
bootstrapAllGather(comm->bootstrap, allHandles, NCCL_NET_HANDLE_MAXSIZE);
ginState->ncclGin->connect(comm->ginContext, handles, comm->nRanks, comm->rank,
listenComm, ginState->ginComms + n);
ginState->ncclGin->createContext(
ginState->ginComms[n], ginState->signalSpaceSize, ginState->counterSpaceSize,
&ginState->ginCtx[n], &ginState->ginDevHandles[n]);
NCCLCHECKGOTO(ginState->ncclGin->closeListen(listenComm), ret, fail);
...
}
在创建context之前，首先要初始化ncclGinIbCollComm，即ginState->ginComms，就是gin的bootstrap网络，用于之后的元信息交换。
ncclGin->connect就是bootstrap网络的建链，建链用的就是ncclNetIb的方法，即之前介绍过的rdma建链，最后得到的bootstrap网络中组成改了一个环形网络，通过sendComm发消息给rank + 1，通过recvComm接收rank - 1的消息，使用这个环形网络可以执行allgather。进一步的，还建立了一个全连接的网络，即用于发送的fullSendComm和用于接收的fullRecvComm，用于执行all2all。
counter与signal
在deepep中，通过put_nbi接口将数据发送过去，然后通过atomic通知对端rank数据已经完成发送了。对于这种场景，gin抽象出来了counter和signal，signal就是deepep这个用法，用于通知对端rank，表示数据已经发送过去了。counter同样的作用，只不过是作用于当前rank，因此后续可以看到，在当前rank创建到其他peer的qp时，一个peer对应两个qp，称为qp_main和qp_companion，qp_main用于向其他rank发送数据和signal，qp_companion为self-loop qp，用于通知当前rank，然后我们先看下ncclGinGdakiCreateContext中关于signal和counter的操作。因为signal和counter流程一致，所以只看下counter。
ncclResult_t ncclGinGdakiCreateContext(void *collComm, int nSignals, int nCounters,
void **outGinCtx, ncclNetDeviceHandle_v11_t **outDevHandle) {
const int num_counters = nCounters;
GdakiGlobalGPUBufferTable<uint64_t> *counters_table =
new GdakiGlobalGPUBufferTable<uint64_t>(num_counters, nranks);
NCCLCHECKGOTO(counters_table->register_mr(gdaki_ctx->ib_pd, true), status, out);
NCCLCHECKGOTO(counters_table->exchange_info(cComm), status, out);
}
num_counters由环境变量指定，默认是65536，GdakiGlobalGPUBufferTable维护了一块gpu内存，支持内存的注册和rkey的全局allgather，如下所示。
template <typename T>
class GdakiGlobalGPUBufferTable {
CUmemGenericAllocationHandle cumemhandle;
GdakiHostGPUMemHandle<__be32> rkeys_hd_mhandle;
T *gpu_ptr;
struct ibv_mr *mr;
ncclResult_t allocate(unsigned int num_elements, unsigned int num_ranks) {
NCCLCHECK(ncclCuMemAlloc((void **)&this->gpu_ptr, &this->cumemhandle, CU_MEM_HANDLE_TYPE_NONE,
num_elements * sizeof(T)));
NCCLCHECK(this->rkeys_hd_mhandle.allocate(num_ranks));
}
GdakiGlobalGPUBufferTable()
: gpu_ptr(nullptr), mr(nullptr), cumemhandle(nullptr), num_elements(0), next_unused_idx(0){};
GdakiGlobalGPUBufferTable(unsigned int num_elements, unsigned int num_ranks) {
this->allocate(num_elements, num_ranks);
};
num_elements就是num_counters，表示counters_table的大小，首先分配counters_table到gpu_ptr。rkeys_hd_mhandle用于保存所有rank gpu_ptr的rkey，类型为GdakiHostGPUMemHandle，维护了一个hostbuf和devbuf，用于host和device的同步，可以理解为一个长度为nranks的数组，然后通过register_mr注册gpu_ptr。
static ncclResult_t gdakiRegMr(struct ibv_mr **mr, struct ibv_pd *pd, void *addr, size_t length,
int access, bool force_strict_ordering = false) {
if (!force_strict_ordering && gdakiRelaxedOrderingEnabled())
access |= IBV_ACCESS_RELAXED_ORDERING;
NOWARN(status = gdakiRegMrDmaBuf(mr, pd, addr, length, access), NCCL_NET);
if (status == ncclSuccess) return ncclSuccess;
NCCLCHECK(wrap_ibv_reg_mr_iova2(mr, pd, addr, length, 0, access));
return ncclSuccess;
}
gdakiRegMr负责注册gpu_ptr，首先尝试dmabuf，如果不成功则尝试peermem，不只是counters_table会通过这个函数注册内存，上节提到的window也会用gdakiRegMr注册，可以看到这里iova传的是0，因此，后续执行数据收发时填的地址是个相对于addr的offset。
因为counter和signal会使用atomic，因此需要一块buffer存储返回值，就是sink_buffer，但是对于返回值不会使用，因此大小为sizeof(uint64_t)。
ncclResult_t ncclGinGdakiCreateContext(...) {
ncclCuMemAlloc((void **)&sink_buffer, &sink_buffer_mhandle, CU_MEM_HANDLE_TYPE_NONE,
sizeof(uint64_t));
gdakiRegMr(&sink_buffer_mr, gdaki_ctx->ib_pd, sink_buffer, sizeof(uint64_t),
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE |
IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_ATOMIC);
}
qp和cq的创建
ncclResult_t ncclGinGdakiCreateContext(...) {
const int ncontexts = 1;
const int nqps_per_rank = ncontexts;
const int nqps_for_comm = nqps_per_rank * nranks; // Number of QPs for communication
const int ncompanion_qps = nqps_for_comm * 2; // Number of companion QPs for communication
// Double because we connect to self.
const int nqps =
nqps_per_rank * (nranks + 1); // +1 for the local rank.
// The last group is the responder of the local rank.
}
nic_handler表示分配uar的方式，比如用BF还是NONCACHE，默认为DOCA_GPUNETIO_VERBS_NIC_HANDLER_AUTO。
mreg_type表示注册umem的方式，用dmabuf还是peermem，默认为DOCA_GPUNETIO_VERBS_MEM_REG_TYPE_DEFAULT。
doca_error_t doca_gpu_verbs_create_qp_group_hl(struct doca_gpu_verbs_qp_init_attr_hl *qp_init_attr,
struct doca_gpu_verbs_qp_group_hl **qpg) {
if (qp_init_attr->nic_handler == DOCA_GPUNETIO_VERBS_NIC_HANDLER_GPU_SM_BF) {
return DOCA_ERROR_INVALID_VALUE;
}
struct doca_gpu_verbs_qp_group_hl *qpg_ =
(struct doca_gpu_verbs_qp_group_hl *)calloc(1, sizeof(struct doca_gpu_verbs_qp_group_hl));
}
值得注意的是现在gin不支持使用BF。然后分配qpg_，qpg_包含两个qp，即qp_main和qp_companion。
分配uar
static doca_error_t create_uar(struct ibv_context *ibctx,
enum doca_gpu_dev_verbs_nic_handler nic_handler,
struct doca_verbs_uar **external_uar, bool bf_supported) {
if (nic_handler != DOCA_GPUNETIO_VERBS_NIC_HANDLER_GPU_SM_BF) {
status = doca_verbs_uar_create(ibctx, DOCA_VERBS_UAR_ALLOCATION_TYPE_NONCACHE_DEDICATED,
external_uar);
if (status != DOCA_SUCCESS) {
doca_verbs_uar_create(ibctx, DOCA_VERBS_UAR_ALLOCATION_TYPE_NONCACHE, external_uar);
return DOCA_SUCCESS;
} else
return DOCA_SUCCESS;
}
}
首先尝试分配NONCACHE_DEDICATED的uar，如果分配失败，则分配普通的NONCACHE的uar。doca_verbs_uar_create就是通过mlx5dv_devx_alloc_uar分配uar，然后将page_id和reg_addr记录到doca_verbs_uar，返回给qpg_->qp_main.external_uar。
创建cq
然后开始创建cq。
static doca_error_t create_cq(...) {
status = doca_verbs_cq_attr_create(&verbs_cq_attr);
status = doca_gpu_mem_alloc(gpu_dev, external_umem_size, priv_get_page_size(),
DOCA_GPU_MEM_TYPE_GPU, (void **)gpu_umem_dev_ptr, NULL);
cq_ring_haddr = (doca_gpunetio_ib_mlx5_cqe64 *)(calloc(external_umem_size, sizeof(uint8_t)));
mlx5_init_cqes(cq_ring_haddr, ncqes);
status_cuda = cudaMemcpy((*gpu_umem_dev_ptr), (void *)(cq_ring_haddr), external_umem_size,
cudaMemcpyDefault);
}
verbs_cq_attr记录了cq的各种属性，比如cq的uar，buf大小等，用于最后创建cq。
接着通过doca_gpu_mem_alloc分配位于gpu的cq buff，即gpu_umem_dev_ptr，doca_gpu_mem_alloc就是cudaMalloc。
然后开始初始化gpu_umem_dev_ptr，通过calloc分配位于cpu内存的cq_ring_haddr，即cq buff，初始化cpu的cq buff的op_own为invalid，然后cudaMemcpy到gpu_umem_dev_ptr。
static doca_error_t create_gpu_umem(...) {
if (mreg_type == DOCA_GPUNETIO_VERBS_MEM_REG_TYPE_DEFAULT) {
status = doca_gpu_dmabuf_fd(gpu_dev, umem_ptr, umem_sz, &dmabuf_fd);
if (status != DOCA_SUCCESS) {
dmabuf_fd = DOCA_VERBS_DMABUF_INVALID_FD;
}
status = doca_verbs_umem_create(ibctx, umem_ptr, umem_sz, IBV_ACCESS_LOCAL_WRITE, dmabuf_fd,
0, umem);
if (status != DOCA_SUCCESS) {
if (dmabuf_fd > 0) {
status = doca_verbs_umem_create(ibctx, umem_ptr, umem_sz, IBV_ACCESS_LOCAL_WRITE,
DOCA_VERBS_DMABUF_INVALID_FD, 0, umem);
if (status != DOCA_SUCCESS) {
goto destroy_resources;
}
} else {
goto destroy_resources;
}
}
}
然后执行create_gpu_umem对gpu_umem_dev_ptr注册umem，支持dmabuf和peermem，如果是default，则首先尝试dmabuf，不成功回退到peermem。
doca_gpu_dmabuf_fd会通过cuMemGetHandleForAddressRange获取这段显存的fd，doca_verbs_umem_create通过执行mlx5dv_devx_umem_reg_ex进行umem的注册，如果输入的dmabuf_fd有效，那么将通过dmabuf的方式注册，否则peer_mem。
到现在就完成了cq buf的分配，初始化和umem的注册，然后同样的流程分配了dbr显存，初始化以及umem的注册。
有了cq buff，dbr，然后通过mlx5dv_devx_obj_create创建cq，保存到qpg_->qp_main.cq_sq，和nvshmem的不同点是gin的cq不用collapsed。
创建qp
static doca_error_t create_qp(...) {
status = doca_verbs_qp_init_attr_create(&verbs_qp_init_attr);
status = doca_gpu_mem_alloc(gpu_dev, external_umem_size, priv_get_page_size(),
DOCA_GPU_MEM_TYPE_GPU, gpu_umem_dev_ptr, NULL);
status =
create_gpu_umem(gpu_dev, ibpd, mreg_type, external_umem_size, *gpu_umem_dev_ptr, gpu_umem);
// dbr
...
status = doca_verbs_qp_create(ibctx, verbs_qp_init_attr, &new_qp);
}
然后是qp的创建，和cq差不多，先通过doca_gpu_mem_alloc分配qp buf，然后注册umem，同理对dbr。将相关信息set到qp_init_attr，如qp_buf，dbr，uar，cq_sq等，然后执行doca_verbs_qp_create创建qp，保存到qpg_->qp_main.qp，和nvshmem不同的在于现在gin不使用rq。
到这里就完成了qp的创建，结构体为doca_verbs_qp，因此还需要转换成device格式，这个操作就是通过doca_gpu_verbs_export_qp做的。
doca_error_t doca_gpu_verbs_export_qp(struct doca_gpu *gpu_dev, struct doca_verbs_qp *qp,
enum doca_gpu_dev_verbs_nic_handler nic_handler,
void *gpu_qp_umem_dev_ptr, struct doca_verbs_cq *cq_sq,
struct doca_gpu_verbs_qp **qp_out) {
*qp_out = (struct doca_gpu_verbs_qp *)calloc(1, sizeof(struct doca_gpu_verbs_qp));
(*qp_out)->qp_cpu = (doca_gpu_dev_verbs_qp *)calloc(1, sizeof(struct doca_gpu_dev_verbs_qp));
qp_cpu_ = (*qp_out)->qp_cpu;
doca_verbs_qp_get_wq(qp, ...);
uint32_t *dbrec = reinterpret_cast<uint32_t *>(doca_verbs_qp_get_dbr_addr(qp));
qp_cpu_->sq_wqe_pi = 0;
qp_cpu_->sq_rsvd_index = 0;
qp_cpu_->sq_ready_index = 0;
qp_cpu_->sq_lock = 0;
qp_cpu_->sq_dbrec = (__be32 *)(dbrec + DOCA_GPUNETIO_IB_MLX5_SND_DBR);
status = doca_gpu_verbs_export_uar(sq_db, (uint64_t **)&(qp_cpu_->sq_db));
qp_cpu_->nic_handler = DOCA_GPUNETIO_VERBS_NIC_HANDLER_GPU_SM_DB;
doca_verbs_cq_get_wq(cq_sq, ...);
doca_verbs_cq_get_dbr_addr(cq_sq, &uar_db_reg, (uint32_t **)&(cq_dbrec), &arm_dbr);
qp_cpu_->cq_sq.dbrec = (__be32 *)cq_dbrec;
}
分配doca_gpu_verbs_qp qp_out，qp_out有doca_gpu_dev_verbs_qp类型的qp_cpu和qp_gpu，接下来要做的事情就是通过qp初始化qp_out->qp_cpu，之后qp_cpu会拷贝到device上的qp_gpu提供给kernel使用。
通过doca_verbs_qp_get_wq获取qp相关信息，赋值到qp_cpu_，doca_gpu_verbs_export_uar会通过cudaHostRegister注册uar，然后通过cudaHostGetDevicePointer获取uar的device地址，之后kernel会通过这个device地址写uar。同理对于cq。
到这里就完成qp_main的创建了，然后开始创建qp_companion，流程完全一致，创建cq到qpg_->qp_companion.cq_sq，创建qp到qpg_->qp_companion.qp，还是通过doca_gpu_verbs_export_qp将qp_companion export到qpg_->qp_companion.qp_gverbs。在一个qpg_中，
qp_main.cq_sq
建链
ncclResult_t ncclGinGdakiCreateContext(......)
for (int rank_idx = 0; rank_idx < nranks; rank_idx++) {
int qp_idx = rank_idx + ctx_idx * nranks;
gdakiFillExchInfo(&local_exch_info[rank_idx], gdaki_ctx, gdaki_ctx->gqps[qp_idx]);
}
cComm->allToAll(cComm, local_exch_info, remote_exch_info, sizeof(struct gdaki_exch_info)),
for (int rank_idx = 0; rank_idx < nranks; rank_idx++) {
int qp_idx = rank_idx + ctx_idx * nranks;
if (rank_idx == rank)
gdakiFillExchInfo(&remote_exch_info[rank_idx], gdaki_ctx,
gdaki_ctx->gqps[nqps_for_comm + ctx_idx]);
NCCLCHECKGOTO(gdakiConnectQp(gdaki_ctx, gdaki_ctx->gqps[qp_idx], &remote_exch_info[rank_idx]),
status, out);
}
}
对qp_main进行建链，对gqps的每个qp执行gdakiFillExchInfo，将lid，qpn等信息记录到local_exch_info，然后执行一个all2all获取其他rank对应的qp信息，然后开始执行建链，遍历gqps，如果rank_idx为当前rank，那么将最后一个qp，即gqps[nqps_for_comm]的信息赋值到remote_exch_info[rank_idx]，最后通过gdakiConnectQp进行建链，即modify_qp，这里不再赘述。
ncclResult_t ncclGinGdakiCreateContext(...) {
for (int qp_idx = 0; qp_idx < nqps_per_rank; qp_idx++) {
int peer_qp_idx = nqps_for_comm + qp_idx;
struct gdaki_exch_info exch_info;
gdakiFillExchInfo(&exch_info, gdaki_ctx, gdaki_ctx->gqps[qp_idx * nqps_per_rank + rank]);
NCCLCHECKGOTO(gdakiConnectQp(gdaki_ctx, gdaki_ctx->gqps[peer_qp_idx], &exch_info), status, out)
}
}
相反的，还需将gqps[nqps_for_comm]和gqps[rank]进行建链。
ncclResult_t ncclGinGdakiCreateContext(...) {
for (int qp_idx = 0; qp_idx < nqps_for_comm; qp_idx++) {
int peer_qp_idx = nqps_for_comm + qp_idx;
struct gdaki_exch_info exch_info;
gdakiFillExchInfo(&exch_info, gdaki_ctx, gdaki_ctx->companion_gqps[peer_qp_idx]);
NCCLCHECKGOTO(gdakiConnectQp(gdaki_ctx, gdaki_ctx->companion_gqps[qp_idx], &exch_info), status,
out);
gdakiFillExchInfo(&exch_info, gdaki_ctx, gdaki_ctx->companion_gqps[qp_idx]);
NCCLCHECKGOTO(gdakiConnectQp(gdaki_ctx, gdaki_ctx->companion_gqps[peer_qp_idx], &exch_info),
status, out);
}
}
然后开始对qp_companion进行建链，qp_companion都是self-loop的，companion_gqps[i]和companion_gqps[i + nqps_for_comm ]进行建链。
DevComm的创建
建链完成之后，需要将位于host内存的qp，counters_table等信息拷贝到device。
ncclResult_t ncclGinGdakiCreateContext(..., void **outGinCtx, ncclNetDeviceHandle_v11_t **outDevHandle) {
struct ncclGinGdakiGPUContext *gin_gdaki_gpu_ctx =
&gin_gdaki_gpu_ctx_hd_mhandle->host_buf[ctx_idx];
tmp_qp = (struct doca_gpu_dev_verbs_qp *)calloc(nranks, sizeof(struct doca_gpu_dev_verbs_qp));
tmp_qp_companion = (struct doca_gpu_dev_verbs_qp *)calloc(nranks, sizeof(struct doca_gpu_dev_verbs_qp));
for (int qp_idx = 0; qp_idx < nranks; qp_idx++) {
struct doca_gpu_dev_verbs_qp *qp_cpu =
gdaki_ctx->gqps[(ctx_idx * nranks) + qp_idx]->qp_gverbs->qp_cpu;
memcpy(&tmp_qp[qp_idx], qp_cpu, sizeof(struct doca_gpu_dev_verbs_qp));
qp_cpu = gdaki_ctx->companion_gqps[(ctx_idx * nranks) + qp_idx]->qp_gverbs->qp_cpu;
memcpy(&tmp_qp_companion[qp_idx], qp_cpu, sizeof(struct doca_gpu_dev_verbs_qp));
}
doca_gpu_mem_alloc(gdaki_ctx->gdev, sizeof(struct doca_gpu_dev_verbs_qp) * nranks,
host_page_size, DOCA_GPU_MEM_TYPE_GPU,
(void **)&gin_gdaki_gpu_ctx->gdqp, nullptr);
ncclCudaMemcpy<struct doca_gpu_dev_verbs_qp>(gin_gdaki_gpu_ctx->gdqp, tmp_qp, nranks);
doca_gpu_mem_alloc(gdaki_ctx->gdev, sizeof(struct doca_gpu_dev_verbs_qp) * nranks,
host_page_size, DOCA_GPU_MEM_TYPE_GPU,
(void **)&gin_gdaki_gpu_ctx->companion_gdqp, nullptr);
ncclCudaMemcpy<struct doca_gpu_dev_verbs_qp>(
gin_gdaki_gpu_ctx->companion_gdqp, tmp_qp_companion, nranks);
gin_gdaki_gpu_ctx->counters_table.buffer = counters_table->gpu_ptr;
gin_gdaki_gpu_ctx->counters_table.rkeys = counters_table->get_rkeys_d();
devHandle->netDeviceType = NCCL_NET_DEVICE_GIN_GDAKI;
devHandle->netDeviceVersion = NCCL_GIN_GDAKI_VERSION;
devHandle->handle = (void *)gin_gdaki_gpu_ctx_hd_mhandle->gpu_buf;
*outDevHandle = devHandle;
}
ncclGinGdakiCreateContext返回的是outDevHandle和outGinCtx，非proxy场景下我们只关注outDevHandle。gin_gdaki_gpu_ctx_hd_mhandle同样是个GdakiHostGPUMemHandle，维护了hostbuf和devbuf，
对于qp，将qp_main和qp_companion拷贝到连续的cpu空间tmp_qp和tmp_qp_companion，再拷贝到device上gin_gdaki_gpu_ctx的gdqp和companion_gdqp。
对于counters_table，将counters_tabl的gpu_ptr和rkeys赋值给gin_gdaki_gpu_ctx。
对称内存的注册
上节中介绍了对称内存，因为gin只能操作对称内存，那么还需对之前执行过ncclCommWindowRegister的内存进行rdma的注册。
ncclResult_t ncclDevrCommCreateInternal(...) {
if (ginActivated) {
NCCLCHECKGOTO(ncclGinConnectOnce(comm), ret, fail);
for (struct ncclDevrMemory* mem = devr->memHead; mem != nullptr; mem = mem->next) {
NCCLCHECKGOTO(symMemoryRegisterGin(comm, mem), ret, fail);
}
}
static ncclResult_t symMemoryRegisterGin(struct ncclComm* comm, struct ncclDevrMemory* mem) {
NCCLCHECK(ncclGinRegister(comm, mem->primaryAddr, mem->size, mem->ginHostWins, mem->ginDevWins));
}
这里就是遍历所有对称内存对应的mem，执行rdma注册，得到的handle存到mem的ginHostWins和ginDevWins，handle就是注册得到的mr相关的信息。
ncclResult_t ncclGinGdakiRegMrSym(void *collComm, void *data, size_t size, int type, void **mhandle,
void **ginHandle) {
struct ncclGinIbCollComm *cComm = (struct ncclGinIbCollComm *)collComm;
struct gdaki_context *gdaki_ctx = (struct gdaki_context *)cComm->ginCtx;
GdakiHostGPUMemHandle<struct ncclGinGdakiMemHandle> *gdaki_mhandle_hd_mhandle =
new GdakiHostGPUMemHandle<struct ncclGinGdakiMemHandle>(1);
GdakiHostGPUMemHandle<__be32> *rkeys_hd_mhandle =
new GdakiHostGPUMemHandle<__be32>(cComm->nranks);
NCCLCHECK(gdakiRegMr(&mr, gdaki_ctx->ib_pd, data, size,
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ |
IBV_ACCESS_REMOTE_ATOMIC));
rkey = htobe32(mr->rkey);
NCCLCHECK(cComm->allGather(cComm, &rkey, rkeys_hd_mhandle->host_buf, sizeof(__be32)));
NCCLCHECK(rkeys_hd_mhandle->copy_h_to_d());
gdaki_mhandle_hd_mhandle->host_buf->rkeys = rkeys_hd_mhandle->gpu_buf;
gdaki_mhandle_hd_mhandle->host_buf->lkey = htobe32(mr->lkey);
NCCLCHECK(gdaki_mhandle_hd_mhandle->copy_h_to_d());
*ginHandle = (void *)gdaki_mhandle_hd_mhandle->gpu_buf;
}
通过gdakiRegMr对data进行注册，将rkey进行allgather到rkeys_hd_mhandle->host_buf，然后将lkey和rkey填充到ginHandle的gpu_buf，然后返回ginHandle。
如上一章节所述，window记录的是userPtr和userSize，mem记录的是userPtr对应的内存块的起始位置memPtr，rdma注册的是memPtr，因此这里更新了一下window的ginOffset4K，即userPtr - memPtr，然后记录ginWins，即lkey和rkey，最后再拷贝回device。
ncclResult_t ncclDevrCommCreateInternal(...) {
if (ginActivated) {
for (int i=0; i < devr->winSortedCount; i++) {
struct ncclDevrWindow* win = devr->winSorted[i].win;
struct ncclWindow_vidmem* winHost;
NCCLCHECKGOTO(ncclShadowPoolToHost(&devr->shadows, win->vidmem, &winHost), ret, fail_stream);
winHost->ginOffset4K = (win->bigOffset - win->memory->bigOffset)>>12;
for (int i=0; i < NCCL_GIN_MAX_CONTEXTS; i++) {
winHost->ginWins[i] = win->memory->ginDevWins[i];
}
CUDACHECKGOTO(cudaMemcpyAsync(win->vidmem, winHost, sizeof(struct ncclWindow_vidmem), cudaMemcpyHostToDevice, stream), ret, fail_stream);
}
}
}
分配counter和signal
在创建context的时候分配了一片counter和signal的内存，这里要根据用户传入实际的counter数量进行分配
ncclResult_t ncclDevrCommCreateInternal(...) {
outDevComm->ginContextCount = nGinContexts;
outDevComm->ginSignalCount = ginSignalTotal;
outDevComm->ginCounterCount = ginCounterTotal;
NCCLCHECKGOTO(ncclGinAllocSignalsCounters(comm,
ginSignalTotal, &outDevComm->ginSignalBase,
ginCounterTotal, &outDevComm->ginCounterBase
), ret, fail_stream_mem_win);
}
ncclResult_t ncclGinAllocSignalsCounters(struct ncclComm* comm, int nSignals, uint32_t* outSignal0,
int nCounters, uint32_t* outCounter0) {
ncclSpaceAlloc(&ginState->counterSpace, ginState->counterSpaceSize, nCounters, 1, &start);
*outCounter0 = (uint32_t)start;
}
对counter的管理也是通过Space，这里通过ncclSpaceAlloc分配出nCounters个counter，首地址为outCounter0。
最后再赋值给outDevComm，就可以在device端用到gin的所有信息了。
ncclResult_t ncclDevrCommCreateInternal(...) {
for (int ctx=0; ctx < nGinContexts; ctx++) {
outDevComm->ginTypes[ctx] = (int)comm->sharedRes->ginState.ginDevHandles[ctx]->netDeviceType;
outDevComm->ginHandles[ctx] = comm->sharedRes->ginState.ginDevHandles[ctx]->handle;
}
}
device执行
ncclGin
ncclGin为device端的结构体，gin相关操作如put之类的都是通过ncclGin的方法执行。
using ncclGin = ncclGin_BackendMask<NCCL_GIN_BACKEND_MASK_ALL>;
template<unsigned backendMask>
struct ncclGin_BackendMask {
ncclDevComm const& comm;
uint32_t nContexts:8, contextId:8, _ginBackend:8;
}
backendMask为支持哪些ncclNetDeviceType，默认是支持GIN_PROXY和GIN_GDAKI，我们只关注GDAKI就好。
template<unsigned beMask>
NCCL_DEVICE_INLINE ncclGin_BackendMask<beMask>::ncclGin_BackendMask(ncclDevComm const& comm, int contextIndex):
comm(comm) {
this->nContexts = comm.ginContextCount;
this->contextId = comm.ginContextCount == 3
? uint32_t(contextIndex)%3 // 3 is only non power of 2
: contextIndex & (comm.ginContextCount-1); // powers of 2
this->_ginBackend = comm.ginTypes[this->contextId];
this->_ginHandle = comm.ginHandles[this->contextId];
this->_signalShadows = comm.ginSignalShadows + this->contextId*comm.ginSignalCount;
}
构造函数就是记录下devComm，保存对应的context到_ginHandle。
put
首先看下put的定义
template<
// Action to take on peer when put completes. If a signalling action is used
// then that signal will be visible only after the payload of this put as well as
// the payloads of preceding puts on this netContext to the same peer are settled.
typename RemoteAction = ncclGin_None, // one of ncclGin_{None|SignalInc|SignalAdd|SignalSet}
// Action to take locally when source has been consumed.
typename LocalAction = ncclGin_None, // one of ncclGin_{None|CounterInc}
// Set of threads participating in this put. Must be a subset of Coop.
typename Coop = ncclCoopThread,
// Optional smem descriptor space to use. Either ncclGin_{None|DescriptorSmem}
typename DescriptorSmem = ncclGin_None
>
NCCL_DEVICE_INLINE void put(
ncclTeam, int peer,
ncclWindow_t dstWnd, size_t dstOffset,
ncclWindow_t srcWnd, size_t srcOffset, size_t bytes,
RemoteAction remoteAction = ncclGin_None{},
LocalAction localAction = ncclGin_None{},
Coop coop = ncclCoopThread{},
DescriptorSmem descriptor = ncclGin_None{},
cuda::thread_scope alreadyReleased = cuda::thread_scope_thread,
cuda::thread_scope expected_scope = cuda::thread_scope_device
) const;
当put的数据在peer可见之后，会对peer执行RemoteAction的操作，就是signal，默认是ncclGin_None，表示啥也不做，ncclGin_SignalAdd 会对指定的signal加上value，ncclGin_SignalInc 会对指定signal加一。当put的数据在当前rank被消费后，会对当前rank执行LocalAction的操作，默认是啥也不做，ncclGin_CounterInc 会对指定的counter加一。
struct ncclGin_None {};
struct ncclGin_SignalAdd { ncclGinSignal_t signal; uint64_t value; };
struct ncclGin_SignalInc { ncclGinSignal_t signal; };
struct ncclGin_CounterInc { ncclGinCounter_t counter; };
然后看下put的实现。
struct ncclGinApi_Put<NCCL_NET_DEVICE_GIN_GDAKI> {
template <typename Coop>
NCCL_DEVICE_INLINE static void call(...) {
coop.sync();
if (coop.thread_rank() == 0) {
ncclGinGdakiGPUContext* gdaki = (struct ncclGinGdakiGPUContext*)ctx.handle;
doca_gpu_dev_verbs_qp* qp = loadConst(&gdaki->gdqp) + peer;
doca_gpu_dev_verbs_qp* companion_qp;
ncclGinGdakiMemHandle* dstMh = (ncclGinGdakiMemHandle*)dstWin;
ncclGinGdakiMemHandle* srcMh = (ncclGinGdakiMemHandle*)srcWin;
raddr.addr = dstOff;
raddr.key = loadConst(loadConst(&dstMh->rkeys) + peer);
laddr.addr = srcOff, laddr.key = loadConst(&srcMh->lkey);
...
}
}
}
coop表示协作组，put默认用thread，即单线程。gdaki就是初始化过程中存下来的ncclGinGdakiGPUContext，存储了qp，couters_table等信息。
dstWin和srcWin就是对称内存注册时的lkey和rkey，前边有说过，通信时用的地址都是offset，因此设置laddr和raddr为srcOff和dstOff，win中获取lkey和rkey。
template <typename Coop>
NCCL_DEVICE_INLINE static void call(...) {
coop.sync();
if (coop.thread_rank() == 0) {
ncclGinGdakiGPUContext* gdaki = (struct ncclGinGdakiGPUContext*)ctx.handle;
doca_gpu_dev_verbs_qp* qp = loadConst(&gdaki->gdqp) + peer;
doca_gpu_dev_verbs_qp* companion_qp;
ncclGinGdakiMemHandle* dstMh = (ncclGinGdakiMemHandle*)dstWin;
ncclGinGdakiMemHandle* srcMh = (ncclGinGdakiMemHandle*)srcWin;
raddr.addr = dstOff;
raddr.key = loadConst(loadConst(&dstMh->rkeys) + peer);
laddr.addr = srcOff, laddr.key = loadConst(&srcMh->lkey);
...
}
}
template <typename Coop>
NCCL_DEVICE_INLINE static void call(...) {
if (coop.thread_rank() == 0) {
doca_gpu_dev_verbs_addr sig_raddr, sig_laddr;
if (hasSignal) {
if (signalOp == ncclGinSignalInc) signalOpArg = 1;
sig_raddr.addr = sizeof(uint64_t) * signalId;
sig_raddr.key = loadConst(loadConst(&gdaki->signals_table.rkeys) + peer);
sig_laddr.addr = 0;
sig_laddr.key = loadConst(&gdaki->sink_buffer_lkey);
}
doca_gpu_dev_verbs_addr counter_raddr, counter_laddr;
if (hasCounter) {
companion_qp = loadConst(&gdaki->companion_gdqp) + peer;
counter_raddr.addr = sizeof(uint64_t) * counterId;
counter_raddr.key = loadConst(loadConst(&gdaki->counters_table.rkeys) + ctx.rank);
counter_laddr.addr = 0;
counter_laddr.key = loadConst(&gdaki->sink_buffer_lkey);
}
doca_gpu_dev_verbs_put_signal_counter<DOCA_GPUNETIO_VERBS_SIGNAL_OP_ADD>(
qp, raddr, laddr, bytes, sig_raddr, sig_laddr, signalOpArg, companion_qp, counter_raddr,
counter_laddr, 1);
}
}
然后开始获取counter和signal相关信息。对于signal，通过signals_table获取peer的rkey，通过用户指定的signalId获取sig_raddr，同理对于counter，返回值均为sink_buffer。唯一不同的是counter使用的是companion_qp，即self-loop qp。
最后通过doca_gpu_dev_verbs_put_signal_counter完成数据的put。对wq，cq，db的同步过程和nvshmem几乎完全一样。
__device__ static __forceinline__ void doca_gpu_dev_verbs_put_signal_counter(...) {
base_wqe_idx = doca_gpu_dev_verbs_reserve_wq_slots<resource_sharing_mode>(qp, num_chunks + 1);
for (uint64_t i = 0; i < num_chunks; i++) {
wqe_idx = base_wqe_idx + i;
size_ = remaining_size > DOCA_GPUNETIO_VERBS_MAX_TRANSFER_SIZE
? DOCA_GPUNETIO_VERBS_MAX_TRANSFER_SIZE
: remaining_size;
wqe_ptr = doca_gpu_dev_verbs_get_wqe_ptr(qp, wqe_idx);
[[likely]] if (size_ > 0) {
doca_gpu_dev_verbs_wqe_prepare_write(...);
} else {
doca_gpu_dev_verbs_wqe_prepare_nop(qp, wqe_ptr, wqe_idx,
DOCA_GPUNETIO_IB_MLX5_WQE_CTRL_CQ_UPDATE);
}
remaining_size -= size_;
}
++wqe_idx;
wqe_ptr = doca_gpu_dev_verbs_get_wqe_ptr(qp, wqe_idx);
doca_gpu_dev_verbs_wqe_prepare_atomic(...);
doca_gpu_dev_verbs_mark_wqes_ready<resource_sharing_mode>(qp, base_wqe_idx, wqe_idx);
}
非常眼熟，doca_gpu_dev_verbs_reserve_wq_slots就是在qp中通过atomic预留num_chunks + 1个wqe，num_chunks为write wqe，1个为atomic wqe。
写wqe完成之后，开始执行doca_gpu_dev_verbs_mark_wqes_ready，就是nvshmem的ready_idx，表示这段连续的wqe填充完成。
__device__ static __forceinline__ void doca_gpu_dev_verbs_mark_wqes_ready(
struct doca_gpu_dev_verbs_qp *qp, uint64_t from_wqe_idx, uint64_t to_wqe_idx) {
if (resource_sharing_mode == DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_EXCLUSIVE)
qp->sq_ready_index = to_wqe_idx + 1;
else if (resource_sharing_mode == DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_CTA) {
doca_gpu_dev_verbs_fence_release<DOCA_GPUNETIO_VERBS_SYNC_SCOPE_CTA>();
cuda::atomic_ref<uint64_t, cuda::thread_scope_block> ready_index_aref(qp->sq_ready_index);
while (ready_index_aref.load(cuda::memory_order_relaxed) != from_wqe_idx) continue;
doca_gpu_dev_verbs_fence_acquire<DOCA_GPUNETIO_VERBS_SYNC_SCOPE_CTA>();
ready_index_aref.store(to_wqe_idx + 1, cuda::memory_order_relaxed);
} else if (resource_sharing_mode == DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_GPU) {
doca_gpu_dev_verbs_fence_release<DOCA_GPUNETIO_VERBS_SYNC_SCOPE_GPU>();
cuda::atomic_ref<uint64_t, cuda::thread_scope_device> ready_index_aref(qp->sq_ready_index);
while (ready_index_aref.load(cuda::memory_order_relaxed) != from_wqe_idx) continue;
doca_gpu_dev_verbs_fence_acquire<DOCA_GPUNETIO_VERBS_SYNC_SCOPE_GPU>();
ready_index_aref.store(to_wqe_idx + 1, cuda::memory_order_relaxed);
}
}
对于单线程的qp，直接对sq_ready_index赋值即可，对于多线程共享的，需要先等待sq_ready_index等于from_wqe_idx，然后才可以写入to_wqe_idx。
__device__ static __forceinline__ void doca_gpu_dev_verbs_put_signal_counter(...) {
uint64_t companion_base_wqe_idx =
doca_gpu_dev_verbs_reserve_wq_slots<resource_sharing_mode>(companion_qp, 2);
uint64_t companion_wqe_idx = companion_base_wqe_idx;
wqe_ptr = doca_gpu_dev_verbs_get_wqe_ptr(companion_qp, companion_wqe_idx);
doca_gpu_dev_verbs_wqe_prepare_wait(companion_qp, wqe_ptr, companion_wqe_idx,
DOCA_GPUNETIO_IB_MLX5_WQE_CTRL_CQ_UPDATE, wqe_idx,
qp->cq_sq.cq_num);
++companion_wqe_idx;
wqe_ptr = doca_gpu_dev_verbs_get_wqe_ptr(companion_qp, companion_wqe_idx);
doca_gpu_dev_verbs_wqe_prepare_atomic(
companion_qp, wqe_ptr, companion_wqe_idx, DOCA_GPUNETIO_IB_MLX5_OPCODE_ATOMIC_FA,
DOCA_GPUNETIO_IB_MLX5_WQE_CTRL_CQ_UPDATE, counter_raddr.addr, counter_raddr.key,
counter_laddr.addr, counter_laddr.key, sizeof(uint64_t), counter_val, 0);
doca_gpu_dev_verbs_mark_wqes_ready<resource_sharing_mode>(companion_qp, companion_base_wqe_idx,
companion_wqe_idx);
doca_gpu_dev_verbs_qp *qps[num_qps] = {qp, companion_qp};
uint64_t prod_indices[num_qps] = {wqe_idx + 1, companion_wqe_idx + 1};
doca_gpu_dev_verbs_submit_multi_qps<num_qps, resource_sharing_mode,
DOCA_GPUNETIO_VERBS_SYNC_SCOPE_GPU, nic_handler>(
qps, prod_indices);
}
然后开始写counter，这里用了rdma wait，max_index为qp_main的wqe_idx，因此当qp_main的atomic完成之后，qp_companion的wait完成，qp_companion的atomic开始执行，之后写入当前rank的counter。最后写db触发执行，这个过程和nvshmem一致，不再赘述。
waitSignal
waitSignal就是while轮询直到signal等于expect。
flush
flush就是pollcq，通过pollcq保证所有操作执行完成，最后更新cons_index + 1到cqe_ci。
__device__ static __forceinline__ int doca_gpu_dev_verbs_poll_cq_at(
struct doca_gpu_dev_verbs_cq *cq, uint64_t cons_index) {
int status = doca_priv_gpu_dev_verbs_poll_cq_at<resource_sharing_mode, qp_type>(cq, cons_index);
if (status == 0) {
doca_gpu_dev_verbs_fence_acquire<DOCA_GPUNETIO_VERBS_SYNC_SCOPE_SYS>();
doca_gpu_dev_verbs_atomic_max<uint64_t, resource_sharing_mode>(&cq->cqe_ci, cons_index + 1);
}
return status;
}
pollcq的过程就是等待cons_index位置的cqe的op_own。
__device__ static __forceinline__ int doca_priv_gpu_dev_verbs_poll_cq_at(
struct doca_gpu_dev_verbs_cq *cq, uint64_t cons_index) {
struct doca_gpunetio_ib_mlx5_cqe64 *cqe =
(struct doca_gpunetio_ib_mlx5_cqe64 *)__ldg((uintptr_t *)&cq->cqe_daddr);
const uint32_t cqe_num = __ldg(&cq->cqe_num);
uint32_t idx = cons_index & (cqe_num - 1);
struct doca_gpunetio_ib_mlx5_cqe64 *cqe64 = &cqe[idx];
do {
cqe_ci = doca_gpu_dev_verbs_load_relaxed<resource_sharing_mode>(&cq->cqe_ci);
[[unlikely]] if (cons_index < cqe_ci)
return 0;
opown = doca_gpu_dev_verbs_load_relaxed_sys_global((uint8_t *)&cqe64->op_own);
} while ((cons_index >= cqe_ci + cqe_num) ||
((cqe_ci <= cons_index) &&
((opown & DOCA_GPUNETIO_IB_MLX5_CQE_OWNER_MASK) ^ !!(cons_index & cqe_num))));
...
}
总结
2.28之后，nccl也拥有了nvshmem的所有特色，后续应该会有gin实现的allreduce等接口，另外感觉也可以对外暴露device的ring，tree等结构，方便用户通过device api利用nccl的这些能力。
扫一扫
285
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
