NCCL源码解析: P2P 连接的建立

前言

NCCL 源码解析总目录

我尽量在每个函数之前介绍每个函数的作用，建议先不要投入到函数内部实现，先把函数作用搞清楚，有了整体框架，再回归到细节。

习惯： 我的笔记习惯：为了便于快速理解，函数调用关系通过缩进表示，也可能是函数展开，根据情况而定。

如下

// 调用 proxyConnInit
NCCLCHECK(proxyConnInit(peer, connectionPool, proxyState, (ncclProxyInitReq*) op->reqBuff, (ncclProxyInitResp*) op->respBuff, &op->connection));
// 对函数 proxyConnInit 进行展开，可方便看参数
static ncclResult_t proxyConnInit(struct ncclProxyLocalPeer* peer, struct ncclProxyConnectionPool* connectionPool, struct ncclProxyState* proxyState, ncclProxyInitReq* req, ncclProxyInitResp* resp, struct 

如有问题，请留言指正。

图后面再补；
有些遗漏之处，还没涉及，后面补；
闲话后面再补。

概括

recvpeer 表示本卡作为接收端的对端
sendpeer 表示本卡作为发送端的对端

对于每个 channel ，卡与卡之间要建立通信，先通过调用 selectTransport<0>() 建立接收通道，0 表示与 recvpeer 建立通信，再通过selectTransport<1>() 建立发送通道，1表示与 sendpeer 建立通信。
建立通道时会遍历 NTRANSPORTS 4种情况：P2P、共享内存、网络、collNet(collective Network, 还没看，不了解)

struct ncclTransport* ncclTransports[NTRANSPORTS] = {
  &p2pTransport,
  &shmTransport,
  &netTransport,
  &collNetTransport
};

本文重点关注 P2P。

接口如下：

struct ncclTransport p2pTransport = {
  "P2P",
  p2pCanConnect,
  { p2pSendSetup, p2pSendConnect, p2pSendFree, NULL, p2pSendProxySetup, NULL, p2pSendProxyFree, NULL },
  { p2pRecvSetup, p2pRecvConnect, p2pRecvFree, NULL, p2pRecvProxySetup, NULL, p2pRecvProxyFree, NULL }
};

发送建立流程为 p2pCanConnect() -> p2pSendSetup() -> p2pSendProxySetup()
接收建立流程为 p2pCanConnect() -> p2pRecvSetup() -> p2pRecvProxySetup()

先检查两个卡支不支持 P2P，主要检查两项：设备支不支持、路径支不支持，路径类型要小于 PATH_PXB，即不通过主桥的路径；
然后 p2pSendSetup() 填充一下 p2pConnectInfo, 向 proxy 线程请求 ncclProxyMsgSetup;
proxy 线程调用 p2pSendProxySetup(), 在本卡内申请显存，返回首地址以及相应的句柄devIpc, 其他进程或者线程可以通过这个句柄，获得此显存的操作地址。（我猜是让对端 GPU 卡也来操作这段内存，从而完成通信，还没看到那一步，完事来更新）
接收与发送机制一样。

详解

ncclTransportP2pSetup()

此P2P 非彼 P2P。ncclTransportP2pSetup 的 P2P 是广义上的两个设备之间的通信设置，包含 P2P、网络以及共享内存等。
建立两卡通信的入口函数。
因为要与 proxy 双线程操作，但是又是同步的，所以下文把两个线程的操作线性展开了，请注意。

ncclTransportP2pSetup(comm, &ringGraph, 0)
ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType/*=NULL*/)
{
	// 信息保存在 data[i] 中， i 为 rank
	// data[i] 大小为 2 * 64 个 connect, 先存放 recv, 再存放 send
	recvData[i] = data[i];
	// recvData[]  所有 recvChannels 的 ncclConnect 缓冲区的首地址
	// 首先 <0> 表示处理的是接收，处理与前一个 rank 的连接
	selectTransport<0>(comm, graph, recvData[i]+recvChannels++, c, recvPeer, connIndex, &type)
	static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType)
	{
		struct ncclPeerInfo* myInfo = comm->peerInfo+comm->rank;
  		struct ncclPeerInfo* peerInfo = comm->peerInfo+peer;
		struct ncclConnector* connector = (type == 1) ? comm->channels[channelId].peers[peer]->send + connIndex :
                                                  comm->channels[channelId].peers[peer]->recv + connIndex;
		NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex))
		{
			NCCLCHECK(ncclCalloc(&resources, 1));
  			recv->transportResources = resources;
			struct p2pConnectInfo* info = (struct p2pConnectInfo*)connectInfo;
			// 如果使用nvlink, 且两个GPU 计算能力一样，(gpu1->gpu.cudaCompCap == 80), 那么 useRead = 1
			// 如果通过参数 P2P_READ_ENABLE 设置该值，  P2P 使用 read 而不是 write
			info->read = useRead; 
			for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) if (!(info->read && p == NCCL_PROTO_SIMPLE)) 
				recvSize += comm->buffSizes[p];
			// 如果同一个进程内的，且 DirectDisable 没有设置，P2P_USE_CUDA_MEMCPY 参数没有设置，并且ncclCuMemEnable 为假
			// 那么
			{
			resources->type = P2P_DIRECT;
      		recv->conn.flags |= info->read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
			}
			// 建立到 proxy 的连接, 连接信息在 recv->proxyConn
			NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_P2P, 0, tpProxyRank, &recv->proxyConn));
			// 请求 proxy 执行 ncclProxyMsgSetup
			// recvSize += comm->buffSizes[p];
			// 发送数据 4字节 recvSize  10485760 = 4096 + NCCL_NUM_PROTOCOLS 3 类型的缓冲区大小
			// 接收数据缓冲区 info->p2pBuff
			// 要接收的大小 sizeof(struct ncclP2pBuff)
			// info->p2pBuff 保存 buf 信息
  			NCCLCHECK(ncclProxyCallBlocking(comm, &recv->proxyConn, ncclProxyMsgSetup, &recvSize, sizeof(int), &info->p2pBuff, sizeof(struct ncclP2pBuff)));
                // 下面为 proxy 线程
				// proxy 线程接收数据进行处理
				{
					// op->connection : 设备与 proxy 连接的控制对象
					// proxyState : rank 的 ncclProxyState
					// op->reqBuff : proxy 本地的接收缓冲区首地址， 按照  op->reqSize 大小申请
					// op->reqSize ： 客户端发送的发送数据的大小
					// op->respBuff: proxy 本地的发送缓冲区的首地址，按照 op->respSize 大小申请
					// p2pRecvProxySetup ： 设备申请内存，首地址信息存入 respBuff
					NCCLCHECK(op->connection->tcomm->proxySetup(op->connection, proxyState, op->reqBuff, op->reqSize, op->respBuff, op->respSize, &done));
					static ncclResult_t p2pRecvProxySetup(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) 
					{
						// 获取设备侧告知的 recvSize 的值 10485760
						int size = *((int*)reqBuff);
						struct ncclP2pBuff* p2pBuff = (struct ncclP2pBuff*)respBuff;
						NCCLCHECK(ncclP2pAllocateShareableBuffer(size, &p2pBuff->ipcDesc, &p2pBuff->directPtr));
						ncclResult_t ncclP2pAllocateShareableBuffer(size_t size, ncclIpcDesc *ipcDesc, void **ptr) 
						{
							// 在设备侧申请内存，地址保存在 ptr
							NCCLCHECK(ncclCudaCalloc((char **)ptr, size));
							// cudaIpcGetMemHandle ： 获取现有设备内存分配的进程间内存句柄
							// 获取指向使用cudaMalloc创建的现有设备内存分配的基址的指针，并将其导出以供另一个进程使用
							// __host__ cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr )
							// 获取现有设备内存分配的进程间内存句柄。
							// 参数:
							// handle - 指向用户分配的 cudaIpcMemHandle 以返回句柄的指针。
							// devPtr - 指向先前分配的设备内存的基指针

    						cudaError_t res = cudaIpcGetMemHandle(&ipcDesc->devIpc, *ptr);
							// cudaIpcOpenMemHandle : 打开从另一个进程导出的进程间内存句柄并返回可用于本地进程的设备指针
							// __host__ cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags )
​
							// 打开从另一个进程导出的进程间内存句柄并返回可用于本地进程的设备指针。
							// 参数
							// devPtr - 返回设备指针
							// handle - cudaIpcMemHandle 打开
							// flags - 此操作的标志。必须指定为cudaIpcMemLazyEnablePeerAccess
						}
						p2pBuff->size = size;
						connection->transportResources = p2pBuff->directPtr;
					}
				}
            // 下面不是 proxy 线程
			// 设备收到 proxy 返回的信息： 设备内部申请的缓冲区首地址，以及地址句柄 ipcDesc->devIpc
			// comm->peerInfo AllGather1 时保存的所有 rank 的信息: rank cudaDev hostHash pidHash busId
			
			// info->rank = myInfo->rank
			// p2pBuff : info->p2pBuff
			// devMem : (void**)&resources->recvDevMem 设备接收资源的接收缓冲区内存指针地址
			// ipcPtr : &resources->recvMemIpc 设备接收资源的接收 内存Ipc 指针地址
			NCCLCHECK(p2pMap(comm, myInfo, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&resources->recvDevMem, &resources->recvMemIpc));
			static ncclResult_t p2pMap(struct ncclComm *comm, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclP2pBuff* p2pBuff, void** devMem, void** ipcPtr)
			{
				// 如果 ncclCuMemEnable 为假，且两个 GPU 设备在同一进程中
				// 那么
				{
					// 如果本设备与对端设备不是同一设备
					if (peerInfo->cudaDev != myInfo->cudaDev) {
						// 如果可以从设备直接访问 peerDevice，则可以通过调用 cudaDeviceEnablePeerAccess() 来启用访问
						cudaError_t err = cudaDeviceEnablePeerAccess(peerInfo->cudaDev, 0);
					}
					// 把 proxy 从设备申请的内存首地址赋值给 *devMem，即 resources->recvDevMem
					// resources->recvDevMem = p2pBuff->directPtr;
					*devMem = p2pBuff->directPtr;
					// 同一个设备不用 ipc
    				*ipcPtr = NULL;
				}
				else
				{
					if ((myInfo->pidHash == peerInfo->pidHash) && (peerInfo->cudaDev == myInfo->cudaDev)) {
						// 同一个进程，同一个设备
						// Same PID and GPU
						*devMem = p2pBuff->directPtr;
						*ipcPtr = NULL;
					} else {
						// 不同进程或者不同设备
						// Different PID or different GPU
						NCCLCHECK(ncclP2pImportShareableBuffer(comm, comm->topParentRanks[peerInfo->rank], p2pBuff->size, &p2pBuff->ipcDesc, devMem));
						ncclResult_t ncclP2pImportShareableBuffer(struct ncclComm *comm, int tpPeer, size_t size, ncclIpcDesc *ipcDesc, void **devMemPtr) 
						{
							// cudaIpcOpenMemHandle : 打开从另一个进程导出的进程间内存句柄并返回可用于本地进程的设备指针
							// __host__ cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags )
​
							// 打开从另一个进程导出的进程间内存句柄并返回可用于本地进程的设备指针。
							// 参数
							// devPtr - 返回设备指针
							// handle - cudaIpcMemHandle 打开
							// flags - 此操作的标志。必须指定为cudaIpcMemLazyEnablePeerAccess
							// 通过 ipcDesc->devIpc 获取设备内存首地址 devMemPtr
							CUDACHECK(cudaIpcOpenMemHandle(devMemPtr, ipcDesc->devIpc, cudaIpcMemLazyEnablePeerAccess));
						}
						// devMem 已经赋值为设备内存首地址
						*ipcPtr = *devMem;
					}
				}
			}
		}
	}


	// 发送
	// 信息保存在 data[i] 中， i 为 rank
	// data[i] 大小为 2 * 64 个 connect, 先存放 recv, 再存放 send
	// sendData[]  所有 recvChannels 的发送 ncclConnect 缓冲区的首地址
	sendData[i] = recvData[i] + recvChannels;
	// 调用发送，处理与后一个 rank 的连接
	NCCLCHECKGOTO(selectTransport<1>(comm, graph, sendData[i]+sendChannels++, c, sendPeer, connIndex, &type), ret, fail);
	static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect, int channelId, int peer, int connIndex, int* transportType) 
	{
		NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
		ncclResult_t p2pSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* send, int channelId, int connIndex)
		{
			NCCLCHECK(ncclCalloc(&resources, 1));
  			send->transportResources = resources;
			info->read = useRead;
			if (graph && connIndex == 1) 
				info->read = 0;
  			const char* useReadStr = info->read ? "/read" : "";
			// For P2P Read the SIMPLE buffer is tagged on the end of the ncclSendMem structure
			if (info->read) 
				// 只有读的时候，使用缓冲区 NCCL_PROTO_SIMPLE
				sendSize += comm->buffSizes[NCCL_PROTO_SIMPLE];
			info->rank = myInfo->rank;
			resources->type = P2P_DIRECT;
      		send->conn.flags |= info->read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
			// 与接收一样的操作
			// 设备收到 proxy 返回的信息保存在 p2pBuff中： 设备内部申请的缓冲区首地址，以及地址句柄 ipcDesc->devIpc
			NCCLCHECK(ncclProxyCallBlocking(comm, &send->proxyConn, ncclProxyMsgSetup, &sendSize, sizeof(int), &info->p2pBuff, sizeof(struct ncclP2pBuff)));
			// p2pMap : 根据接收到的信息做一个发送缓冲区的首地址解析，得到 sendDevMem 或者 sendMemIpc
			// info->rank = myInfo->rank
			// p2pBuff : info->p2pBuff
			// devMem : (void**)&resources->recvDevMem 设备接收资源的接收缓冲区内存指针地址
			// ipcPtr : &resources->recvMemIpc 设备接收资源的接收 内存Ipc 指针地址
    		NCCLCHECK(p2pMap(comm, myInfo, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&resources->sendDevMem, &resources->sendMemIpc));
		}
	}

	if (sendPeer == recvPeer) {
		if (recvChannels+sendChannels) {
		NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, recvPeer, bootstrapTag, data[i], sizeof(struct ncclConnect)*(recvChannels+sendChannels)), ret, fail);
		NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, recvPeer, bootstrapTag, data[i], sizeof(struct ncclConnect)*(recvChannels+sendChannels)), ret, fail);
		sendData[i] = data[i];
		recvData[i] = data[i]+sendChannels;
		}
	} else {
		// 如果 sendPeer recvPeer 不是同一个
		// 假设 0 -> 1 -> 2, 当前 rank 为 1
		// sendPeer = 2， recvPeer = 0
		// sendPeer ： 我作为发送的对端 rank
		// recvPeer ： 我作为接收的对端 rank

		if (recvChannels) 
			// 向前一个 rank 发送 recvChannels 个接收连接信息
			NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, recvPeer, bootstrapTag, recvData[i], sizeof(struct ncclConnect)*recvChannels), ret, fail);

		if (sendChannels) 
			// 向后一个 rank 发送 recvChannels 个发送连接信息
			NCCLCHECKGOTO(bootstrapSend(comm->bootstrap, sendPeer, bootstrapTag, sendData[i], sizeof(struct ncclConnect)*sendChannels), ret, fail);

		if (sendChannels) 
			// 接收后一个 rank 的接收连接信息到 sendData
			NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, sendPeer, bootstrapTag, sendData[i], sizeof(struct ncclConnect)*sendChannels), ret, fail);

		if (recvChannels) 
			// 接收前一个 rank 的接收发送信息
			NCCLCHECKGOTO(bootstrapRecv(comm->bootstrap, recvPeer, bootstrapTag, recvData[i], sizeof(struct ncclConnect)*recvChannels), ret, fail);
	}
}