NCCL源码1：官网案例详解，单进程单设备使用/调用案例

int main(int argc, char* argv[])  

{  

  // 定义一个整数变量size，代表缓冲区大小为32MB  
  int size = 32*1024*1024;  

  // 定义MPI相关的变量，包括当前进程的排名（myRank）、总进程数（nRanks）和本地排名（localRank）  
  int myRank, nRanks, localRank = 0;  

  

  //////// 1、初试化和启动MPI通信////////////////
  MPICHECK(MPI_Init(&argc, &argv));  

  // 获取当前进程的排名  
  MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank));  
  // 获取总进程数  

  MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks));  

  
  ///////// 2、计算主机名的哈希值，并MPI_allgather通信使得每个rank（进程）都获取其它rank的哈希值。////////////////
  // 基于主机名哈希计算localRank，用于选择GPU  
  uint64_t hostHashs[nRanks];  
  char hostname[1024];  
  getHostName(hostname, 1024); // 获取主机名  
  hostHashs[myRank] = getHostHash(hostname); // 计算主机名的哈希值  

  // 使用MPI_Allgather收集所有进程的哈希值  
  MPICHECK(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, hostHashs, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD));  

  
  ///////////////3、根据获取的哈希值，计算得到（哈希值相同的rank在同一主机上）该rank所在的主机参与通信的rank总数。（哈希值就是主机名，其实可以用主机名来获取主机上参与通信的总rank数，只是主机命名五花八门，哈希值更容易比较）//////////////////////
  // 计算localRank，即具有相同主机哈希值的进程数（不包括当前进程）  
  for (int p=0; p<nRanks; p++) {  

     if (p == myRank) break; // 如果是当前进程，则跳出循环  

     if (hostHashs[p] == hostHashs[myRank]) localRank++; // 如果哈希值相同，则localRank加1  

  }  

  

  // NCCL相关的变量，包括唯一ID（id）和通信器（comm）  

  ncclUniqueId id;  

  ncclComm_t comm;  

  // 定义CUDA相关的变量，包括发送和接收缓冲区（sendbuff, recvbuff）以及CUDA流（s）  

  float *sendbuff, *recvbuff;  

  cudaStream_t s;  

  
  ///////////////////4、rank0上获取NCCL的唯一ID，并MPI_Bcast广播给其它rank。（这个唯一的ID是用来标识通信组，因此所有通信组中的rank有相同的ID）////////////
  // 在rank 0上获取NCCL的唯一ID，并使用MPI_Bcast广播给所有其他进程  
  if (myRank == 0) ncclGetUniqueId(&id);  
  MPICHECK(MPI_Bcast((void *)&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD));  

  
  ////////////////5、基于localrank绑定GPU，并分配发送接收缓冲区，创建CUDA流。////////////
  // 基于localRank选择GPU，并分配设备缓冲区  
  // CUDACHECK是一个宏，用于检查CUDA函数的返回值  

  CUDACHECK(cudaSetDevice(localRank)); // 设置CUDA设备  

  CUDACHECK(cudaMalloc(&sendbuff, size * sizeof(float))); // 分配发送缓冲区  

  CUDACHECK(cudaMalloc(&recvbuff, size * sizeof(float))); // 分配接收缓冲区  

  CUDACHECK(cudaStreamCreate(&s)); // 创建一个CUDA流  

  

  /////////////6. 初始化NCCL通信器///////////////////////  

  NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank));  

  

  //////////7、使用NCCL进行AllReduce操作 //////////////////////

  // 此操作将sendbuff中的值在所有进程中求和，并将结果放在recvbuff中  

  NCCLCHECK(ncclAllReduce((const void*)sendbuff, (void*)recvbuff, size / sizeof(float), ncclFloat, ncclSum, comm, s));  

  // 注意：size / sizeof(float) 是因为ncclAllReduce需要元素数量，而不是字节数  

  

  // 同步CUDA流以确保NCCL操作完成  

  CUDACHECK(cudaStreamSynchronize(s));  

  

  //////////////8、释放设备缓冲区///////////////  

  CUDACHECK(cudaFree(sendbuff));  

  CUDACHECK(cudaFree(recvbuff));  

  

  ////////////// 9、销毁NCCL通信器/////////////////  

  ncclCommDestroy(comm);  

  

  ///////////////10、 终止MPI环境 ///////////////////// 

  MPICHECK(MPI_Finalize());


  printf("[MPI Rank %d] Success \n", myRank);
  return 0;
}