NCCL源码解析: 共享内存连接的建立

前言

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
};

本文重点关注 shmTransport。

接口如下：

struct ncclTransport shmTransport = {
  "SHM",
  shmCanConnect,
  { shmSendSetup, shmSendConnect, shmSendFree, NULL, NULL, NULL, NULL, NULL },
  { shmRecvSetup, shmRecvConnect, shmRecvFree, NULL, NULL, NULL, NULL, NULL }
};

发送建立流程为 p2pCanConnect() -> shmSendSetup()
接收建立流程为 p2pCanConnect() -> shmRecvSetup()

共享内存相对比较简单，就是在两个 GPU 设备不能进行 nvlink 或者通过switch 进行 P2P 的时候，在系统内存中分配一段空间，两张卡都通过操作这一段共享内存进行数据通信。

详解

1. shmCanConnect()

检查能不能使用共享内存：

1. 网络快的话，用网络；
2. 两张卡在一个主机上，要不然看不到同一块内存；
3. 两张卡都能通过访问 /dev/shm 文件共享内存；

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(transport->canConnect(&ret, comm->topo, graph, myInfo, peerInfo));
		connector->transportComm = transportComm;
		NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
	}


// 第一步，检查能不能用共享能存
NCCLCHECK(transport->canConnect(&ret, comm->topo, graph, myInfo, peerInfo));
static ncclResult_t shmCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1, struct ncclPeerInfo* info2) 
{
	// 环境变量不能禁止 NCCL_SHM_DISABLE
	if (ncclParamShmDisable() == 1) 
		return ncclSuccess;
	// 检查用网卡是不是速度更快
	// 比较 GPU1 到网卡的 bw 与 GPU1 到 GPU2 的 bw
	// 比较 GPU2 到网卡的 bw 与 GPU1 到 GPU2 的 bw
	// 如果两个都快，那么用网络
	int useNet = 0;
  	NCCLCHECK(ncclTopoCheckNet(topo, info1->busId, info2->busId, &useNet));
  	if (useNet) 
		return ncclSuccess;

	// 还要保证两张卡再同一个主机上
	if (info1->hostHash != info2->hostHash) 
		return ncclSuccess;

	// 确保两张卡的环境能访问同一块内存 /dev/shm
	// info->shmDev = statbuf.st_dev;
	// shmDev 保存的是设备文件的主次设备号，根据这个信息可以决定容器环境中是否可以使用共享内存
	if (info1->shmDev != info2->shmDev) 
		return ncclSuccess;

	// 以上都没问题，就可以用共享内存了
	*ret = 1;
}

2. shmSendSetup

发送方向设置：

1. 在 /dev/shm 中创建文件，就是在内存上分配内存；
2. 通过 mmap 映射文件，就是映射内存，返回的首地址 hptr，线程可以通过 hptr 读写此段内存；
3. 通过 CUDA API 建立映射 hptr ，返回地址 dptr，设备可以通过 dptr 访问该地址；
4. 记录信息保存在struct shmSendResources* resources->hostHandle 中，资源信息 resources

// 发送设置
// connect : 连接信息缓冲区首地址
// connector : 发送或者接收软件抽象
NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
static ncclResult_t shmSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* send, int channelId, int connIndex) 
{
	// 发送缓冲区申请内存
	struct shmSendResources* resources;
	NCCLCHECK(ncclCalloc(&resources, 1));
	send->transportResources = resources;
	// 连接缓冲区，通过强制转换，填充不同类型的连接信息
	struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;

	int shmSize = sizeof(struct ncclSendMem);
	info->shmSize = resources->shmSize = shmSize;
	// /dev/shm 可以认为是内存设备的实体文件，对此目录的操作会落到内存的读写上
	// 在该目录创建的文件保存在内存中，大小是由限制，最大内存的一半
	// resources->hostMem 保存内存首地址
	// resources->devHostMem 保存设备要访问共享内存时使用的首地址
	// resources->hostHandle 保存共享内存所有信息
	NCCLCHECK(ncclShmOpen(shmPath, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1, &resources->hostHandle));
	ncclResult_t ncclShmOpen(char* shmPath, size_t shmSize, void** shmPtr, void** devShmPtr, int refcount, ncclShmHandle_t* handle)
	{
		// 多申请 4 个字节，要保存 refcount 信息， 引用计数
		const size_t refSize = sizeof(int); /* extra sizeof(int) bytes for reference count */
		const size_t realShmSize = shmSize + refSize;
		// 打开文件，申请内存
		SYSCHECKGOTO(fd = open(shmPath, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR), ret, fail);
		(ftruncate(fd, realShmSize) != 0)
		// 内存申请成功之后，使用 mmap 进行映射，程序操作映射返回的首地址 hptr，就是操作申请的这段空间的内存
		hptr = (char*)mmap(NULL, realShmSize, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
		// 最后 4个字节记录引用计数，新创建的为 1
		*(int*)(hptr + shmSize) = refcount;
		if (devShmPtr) {
			// cudaHostRegister() 把可分页内存标记为锁页内存, 内存不可换出
			// cudaHostRegisterMapped 表示把锁页内存地址映射到设备地址空间
			// 这样，这块存储会有两个地址：一个是从cudaHostAlloc() 或 malloc() 返回的在主机内存地址空间上
			// 另一个在设备存储器上，可以通过cudaHostGetDevicePointer() 取得, 核函数可以通过这个地址访问这段空间
			CUDACHECKGOTO(cudaHostRegister((void*)hptr, realShmSize, cudaHostRegisterMapped), ret, fail);
			// GPU 内部使用 dptr 访问这段空间
			CUDACHECKGOTO(cudaHostGetDevicePointer(&dptr, (void*)hptr, 0), ret, fail);
		}
		// 所有信息保存在 tmphandle 中
		shmHandleInit(fd, shmPath, shmSize, realShmSize, hptr, dptr, create, tmphandle);
		static void shmHandleInit(int fd, char* shmPath, size_t shmSize, size_t realShmSize, char* hptr, void* dptr, bool create, struct shmHandleInternal* handle) 
		{
			handle->fd = fd;
			handle->shmPtr = hptr;
			handle->devShmPtr = dptr;
			handle->shmSize = shmSize;
			handle->realShmSize = realShmSize;
			handle->refcount = (hptr != NULL) ? (int*)(hptr + shmSize) : NULL;
			if (create) {
				int slen = strlen(shmPath);
				handle->shmPath = (char*)malloc(slen + 1);
				memcpy(handle->shmPath, shmPath, slen + 1);
				if (hptr) memset(hptr, 0, shmSize);
			} else {
				handle->shmPath = NULL;
			}
		}
		*shmPtr = hptr;
		if (devShmPtr) 
			*devShmPtr = dptr;
  		*handle = (ncclShmHandle_t)tmphandle;
	}
}

3. shmRecvSetup

接收方向设置：，与发送一样：

1. 在 /dev/shm 中创建文件，就是在内存上分配内存；
2. 通过 mmap 映射文件，就是映射内存，返回的首地址 hptr，线程可以通过 hptr 读写此段内存；
3. 通过 CUDA API 建立映射 hptr ，返回地址 dptr，设备可以通过 dptr 访问该地址；
4. 记录信息保存在 struct shmRecvResources* resources->hostHandle 中，资源信息 resources

// 接收设置
NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
static ncclResult_t shmRecvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* recv, int channelId, int connIndex) {
	// 为 resources 申请内存空间
	struct shmRecvResources* resources;
	NCCLCHECK(ncclCalloc(&resources, 1));
	recv->transportResources = resources;
	// 连接信息，准备按照特定格式填充 connectInfo
	struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;

	char shmPath[PATH_MAX];
	shmPath[0] = '\0';
	int shmSize = sizeof(struct ncclRecvMem);
	// shmLocality = ncclParamShmLocality();
	// 全局变量 NCCL_SHM_LOCALITY 定义
	// #define SHM_RECV_SIDE 2
	if (shmLocality == SHM_RECV_SIDE) {
		for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) 
			shmSize += comm->buffSizes[p];
	}
	info->shmSize = resources->shmSize = shmSize;
	// 通过 /dev/shm 创建共享内存
	// 保存在 resources->hostMem 中
	// 设备侧操作内存的地址保存在 resources->devHostMem 中
	// resources->hostHandle 保存所有信息
	NCCLCHECK(ncclShmOpen(shmPath, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1, &resources->hostHandle));
	TRACE(NCCL_SHM,"Opened shmName %s shmSize %d", shmPath, info->shmSize);
	memcpy(info->shmName, shmPath+sizeof("/dev/shm/nccl-")-1, sizeof(info->shmName));

	return ncclSuccess;
}