文字没有提及的代码内容,不需要太在意,当然也可以瞟两眼;
首先,总体而言函数 ncclAllReduce 的功能在于将携带了一个操作的info结构体,放入了队列中,待后面执行;
排队的函数调用是 ncclEnqueueCheck(&info),在 ncclAllReduce函数体中被调用。
其他几个类似机制的 api 是
ncclAllGather
ncclAllReduce
ncclBroadcast
ncclBcast
ncclReduce
ncclReduceScatter
ncclSend
ncclRecv
他们都在文件 nccl/src/collectives.cc 中定义;
那么,稍微深入一下函数? ncclResult_t ncclEnqueueCheck(struct ncclInfo* info)
它调用了?? NCCLCHECKGOTO(taskAppend(info->comm, info), ret, fail)
?? ??? ??? ??? ?而taskAppend() 又调用了两个函数:
?? ??? ??? ??? ?????? hostToDevRedOp() 将reduce的ncclSum操作,转换成dev的ncclDevSum操作,然后调用了
?? ??? ??? ??? ?????? ncclIntruQueueEnqueue(&tasks->collQueue, t); 将这个任务放入了comm的任务队列中。
那么需要看一下 ncclIntruQueueEnqueue 到底对t中的 t->op做了什么解析,t->op是这个函数的第二个参数的op成员;
ncclIntruQueueEnqueue() 仅仅是将 第二个参数t插入了一个链表info->comm->tasks中;
这个info是在ncclAllReduce()中定义的? struct ncclInfo info,其中info->comm 是ncclAllReduce 传递进来的第五个参数 ncclComm* comm。
综上所述,ncclAllReduce 仅仅是将一个 reduce 的任务插入到了 comm 的 tasks 链表中而已,并没有涉及到调用任何的 cuda 函数。
所以,启动相关的阿cuda kernel等,应该是在后面的 ncclGroupEnd() 中,通过解析 comm->tasks的数据元素来启动的。
接下来看一下 ncclGroupEnd() 的实现。