Updating stream caching (#614)

- Adding non-captured hipStream for use in setup

src/channel.cc

@@ -18,11 +18,11 @@ ncclResult_t initChannel(struct ncclComm* comm, int channelid) {
   channel->id = channelid;
 
   // Ring index to user rank table.
-  NCCLCHECK(ncclCudaCalloc(&channel->ring.devUserRanks, comm->nRanks));
+  NCCLCHECK(ncclCudaCalloc(&channel->ring.devUserRanks, comm->nRanks, comm->sideStream));
   NCCLCHECK(ncclCalloc(&channel->ring.userRanks, comm->nRanks));
 
   // Communication structures with peers.
-  NCCLCHECK(ncclCudaCalloc(&channel->devPeers, comm->nRanks+1)); // The extra one rank is for collnet root (i.e. network)
+  NCCLCHECK(ncclCudaCalloc(&channel->devPeers, comm->nRanks+1, comm->sideStream)); // The extra one rank is for collnet root (i.e. network)
   NCCLCHECK(ncclCalloc(&channel->peers, comm->nRanks+1));
   for (size_t i=0; i<comm->nRanks+1; ++i) {
     for (int b=0; b<NCCL_MAX_CONNS; b++) {

src/clique/CliqueManager.cc

@@ -190,7 +190,7 @@ ncclResult_t CliqueManager::Init(ncclUniqueId const* commId, int suffix)
       hipIpcMemHandle_t handle;
       // Allocate fine-grained device memory on rank 0 and get IPC handle for it
       // Re-usable barrier consists of (globalCount / globalSense) pair of integers
-      NCCLCHECKGOTO(ncclCudaCalloc(&m_fineGrainBarrierMem, NCCL_MAX_OPS * 2 * sizeof(int), true), res, dropback);
+      NCCLCHECKGOTO(ncclCudaCalloc(&m_fineGrainBarrierMem, NCCL_MAX_OPS * 2 * sizeof(int), nullptr, true), res, dropback);
       if (hipIpcGetMemHandle(&handle, m_fineGrainBarrierMem) != hipSuccess)
       {
         WARN("Unable to get IPC handle for barrier memory");
@@ -228,7 +228,7 @@ ncclResult_t CliqueManager::Init(ncclUniqueId const* commId, int suffix)
     // First rank prepares fine-grained memory shared across ranks used for the two barrier variables
     if (m_rank == 0)
     {
-      NCCLCHECKGOTO(ncclCudaCalloc(&m_staticGpuBarrierMem, NCCL_MAX_OPS * 2 * sizeof(int), true), res, dropback);
+      NCCLCHECKGOTO(ncclCudaCalloc(&m_staticGpuBarrierMem, NCCL_MAX_OPS * 2 * sizeof(int), nullptr, true), res, dropback);
       // Prepare all barriers
       for (int opIndex = 0; opIndex < NCCL_MAX_OPS; opIndex++)
       {

src/include/alloc.h

@@ -78,17 +78,7 @@ static_assert(sizeof(struct allocationTracker) == 64, "allocationTracker must be
 extern struct allocationTracker allocTracker[];
 
 template <typename T>
-static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
-
-  // Need async stream for P2P pre-connect + CUDA Graph
-  static bool streamCreated = false;
-  static hipStream_t stream;
-  if (rcclParamEnableHipGraph() && !streamCreated)
-  {
-    // Create stream only once to avoid performance penalty
-    CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
-    streamCreated = true;
-  }
+static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, hipStream_t sideStream = nullptr, bool isFineGrain = false) {
 
   if (isFineGrain)
     CUDACHECK(hipExtMallocWithFlags((void**)ptr, nelem*sizeof(T), hipDeviceMallocFinegrained));
@@ -96,10 +86,14 @@ static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr,
     CUDACHECK(hipMalloc(ptr, nelem*sizeof(T)));
 
   if (rcclParamEnableHipGraph()) {
+    hipStream_t stream = sideStream;
+    if (stream == nullptr)
+      CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
+
     CUDACHECK(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream));
     CUDACHECK(hipStreamSynchronize(stream));
-    // NOTE: Currently the re-used stream is not destroyed
-    //CUDACHECK(hipStreamDestroy(stream));
+    if (sideStream == nullptr)
+      CUDACHECK(hipStreamDestroy(stream));
   } else {
     CUDACHECK(hipMemset(*ptr, 0, nelem*sizeof(T)));
     CUDACHECK(hipStreamSynchronize(NULL));

src/include/comm.h

@@ -220,6 +220,7 @@ struct ncclComm {
   int p2pRecvCount;
 
   // [RCCL]
+  hipStream_t sideStream;
   //CliqueManager* cliqueManager;    // CliqueManager handles pointer collection / distribution for clique-based kernels
   //int rootPid;                     // Process ID of root
   // [/RCCL]

src/init.cc

@@ -303,6 +303,8 @@ static ncclResult_t commFree(ncclComm_t comm) {
     CUDACHECK(hipStreamDestroy(comm->groupStream));
   }
 
+  CUDACHECK(hipStreamDestroy(comm->sideStream));
+
   // Last rank frees shared resources between threads
   int isLast;
   NCCLCHECK(ncclCpuBarrierIn(comm, &isLast));
@@ -365,6 +367,9 @@ static ncclResult_t commAlloc(ncclComm_t* comret, int ndev, int rank, int virtua
   NCCLCHECK(getBusId(comm->cudaDev, &comm->busId));
   TRACE(NCCL_INIT,"comm %p rank %d nranks %d cudaDev %d busId %lx", comm, rank, ndev, comm->cudaDev, comm->busId);
 
+  // RCCL: create persistent stream for calloc
+  CUDACHECK(hipStreamCreateWithFlags(&comm->sideStream, hipStreamNonBlocking));
+
   comm->doneEvent = doneEvent;
   comm->intDoneEvent = intDoneEvent;
   comm->checkPointers = ncclParamCheckPointers() == 1 ? true : false;
@@ -385,7 +390,7 @@ static ncclResult_t commAlloc(ncclComm_t* comret, int ndev, int rank, int virtua
 
   comm->argsptrs[0] = &comm->devComm;
 #ifdef ENABLE_PROFILING
-  NCCLCHECK(ncclCudaCalloc(&comm->hostDevComm.devProf, MAXCHANNELS*PROFILE_NUM_LAUNCHES));
+  NCCLCHECK(ncclCudaCalloc(&comm->hostDevComm.devProf, MAXCHANNELS*PROFILE_NUM_LAUNCHES, comm->sideStream));
 #endif
 
 #ifdef ENABLE_COLLTRACE
@@ -445,7 +450,7 @@ static ncclResult_t commAlloc(ncclComm_t* comret, int ndev, int rank, int virtua
 
 static ncclResult_t devCommSetup(ncclComm_t comm) {
   ncclDevCommAndChannels *devCommAndChans;
-  NCCLCHECK(ncclCudaCalloc(&devCommAndChans, 1));
+  NCCLCHECK(ncclCudaCalloc(&devCommAndChans, 1, comm->sideStream));
   comm->devComm = &devCommAndChans->comm;
   comm->hostDevComm.channels = devCommAndChans->channels;
 

src/transport.cc

@@ -81,9 +81,6 @@ void dumpData(struct ncclConnect* data, int ndata) {
 }
 
 ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, int connIndex, int* highestTransportType/*=NULL*/) {
-  // Stream used during transport setup; need for P2P pre-connect + CUDA Graph
-  hipStream_t transportSetupStream;
-  CUDACHECK(hipStreamCreateWithFlags(&transportSetupStream, hipStreamNonBlocking));
   int highestType = TRANSPORT_P2P;  // track highest transport type
 
   struct ncclConnect data[2*MAXCHANNELS];
@@ -137,7 +134,7 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
         struct ncclConnector* conn = comm->channels[c].peers[sendPeer].send + connIndex;
         NCCLCHECK(conn->transportComm->connect(comm, sendData++, 1, comm->rank, conn));
         conn->connected = 1;
-        CUDACHECK(hipMemcpyAsync(comm->channels[c].devPeers[sendPeer].send+connIndex, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice, transportSetupStream));
+        CUDACHECK(hipMemcpyAsync(comm->channels[c].devPeers[sendPeer].send+connIndex, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice, comm->sideStream));
       }
     }
     TIME_STOP(3);
@@ -147,14 +144,13 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
         struct ncclConnector* conn = comm->channels[c].peers[recvPeer].recv + connIndex;
         NCCLCHECK(conn->transportComm->connect(comm, recvData++, 1, comm->rank, conn));
         conn->connected = 1;
-        CUDACHECK(hipMemcpyAsync(comm->channels[c].devPeers[recvPeer].recv+connIndex, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice, transportSetupStream));
+        CUDACHECK(hipMemcpyAsync(comm->channels[c].devPeers[recvPeer].recv+connIndex, conn, sizeof(struct ncclConnector), hipMemcpyHostToDevice, comm->sideStream));
       }
     }
     TIME_STOP(4);
     comm->connectRecv[recvPeer+comm->nRanks*connIndex] = comm->connectSend[sendPeer+comm->nRanks*connIndex] = 0;
   }
-  CUDACHECK(hipStreamSynchronize(transportSetupStream));
-  CUDACHECK(hipStreamDestroy(transportSetupStream));
+  CUDACHECK(hipStreamSynchronize(comm->sideStream));
   if (highestTransportType != NULL) *highestTransportType = highestType;
   TIME_PRINT("P2P Setup/Connect");
   return ncclSuccess;

src/transport/coll_net.cc

@@ -351,7 +351,7 @@ static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, char** gp
   *size = state->size;
 
   if (cuda && state->cudaBuff == NULL) {
-    NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, *size, cuda));
+    NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, *size, comm->sideStream, cuda));
   }
   if (!cuda && state->hostBuff == NULL) {
     NCCLCHECK(ncclCudaHostCalloc(&state->hostBuff, *size));

src/transport/net.cc

@@ -396,7 +396,7 @@ static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, int local
   if (size) *size = state->size;
 
   if (cuda && state->cudaBuff == NULL) {
-    NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, state->size, cuda));
+    NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, state->size, comm->sideStream, cuda));
     if (sameProcess == 0) {
       CUDACHECK(hipIpcGetMemHandle(&state->ipc, state->cudaBuff));
     }
@@ -575,7 +575,7 @@ static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, str
       if (!map->sameProcess) {
         ALIGN_SIZE(map->mems[NCCL_NET_MAP_DEVMEM].size, CUDA_IPC_MIN);
       }
-      NCCLCHECK(ncclCudaCalloc(&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].size, resources->useGdr));
+      NCCLCHECK(ncclCudaCalloc(&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].size, comm->sideStream, resources->useGdr));
       map->mems[NCCL_NET_MAP_DEVMEM].cpuPtr = map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr;
     }
     if (!map->sameProcess) {
@@ -689,7 +689,7 @@ static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, str
 
   if (map->mems[NCCL_NET_MAP_DEVMEM].size) {
     if (resources->shared == 0) {
-      NCCLCHECK(ncclCudaCalloc(&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].size, resources->useGdr));
+      NCCLCHECK(ncclCudaCalloc(&map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr, map->mems[NCCL_NET_MAP_DEVMEM].size, comm->sideStream, resources->useGdr));
       map->mems[NCCL_NET_MAP_DEVMEM].cpuPtr = map->mems[NCCL_NET_MAP_DEVMEM].gpuPtr;
     }
   }

src/transport/p2p.cc

@@ -115,7 +115,7 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
     // Check that legacy IPC support is available (WSL WAR)
     char *dummy;
     hipIpcMemHandle_t ipc;
-    NCCLCHECK(ncclCudaCalloc(&dummy, CUDA_IPC_MIN));
+    CUDACHECK(hipMalloc(&dummy, CUDA_IPC_MIN));
     if (hipIpcGetMemHandle(&ipc, dummy) != hipSuccess) {
       INFO(NCCL_INIT|NCCL_P2P,"Legacy IPC not supported");
       *ret = 0;
@@ -347,7 +347,7 @@ static ncclResult_t p2pProxySetup(struct ncclProxyConnection* connection, struct
   int size = *((int*)reqBuff);
   if (respSize != sizeof(struct ncclP2pBuff)) return ncclInternalError;
   struct ncclP2pBuff* p2pBuff = (struct ncclP2pBuff*)respBuff;
-  NCCLCHECK(ncclCudaCalloc((char**)&p2pBuff->directPtr, size, true));
+  NCCLCHECK(ncclCudaCalloc((char**)&p2pBuff->directPtr, size, comm->sideStream, true));
   connection->transportResources = p2pBuff->directPtr;
   hipError_t res = hipIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
   if (res != hipSuccess) {

tools/ib-test/ib_test.cpp

@@ -245,7 +245,7 @@ private:
     printf("GDR Read %s\n", use_gdr_read ? "enabled" : "disabled");
 
     if (use_gdr_read) {
-      NCCLCHECK(ncclCudaCalloc(&sendDevBuffer, sendBuffSize, 1));
+      NCCLCHECK(ncclCudaCalloc(&sendDevBuffer, sendBuffSize, nullptr, 1));
       printf("Allocated sendDevBuffer %p of %d bytes, sliceSteps %d\n",
                 sendDevBuffer, sendBuffSize, sliceSteps);
     }
@@ -397,7 +397,7 @@ private:
     printf("GDR Write %s\n", use_gdr_write ? "enabled" : "disabled");
 
     if (use_gdr_write) {
-      NCCLCHECK(ncclCudaCalloc(&recvDevBuffer, recvBuffSize, 1));
+      NCCLCHECK(ncclCudaCalloc(&recvDevBuffer, recvBuffSize, nullptr, 1));
       printf("Allocated recvDevBuffer %p of %d bytes, sliceSteps %d\n",
                 recvDevBuffer, recvBuffSize, sliceSteps);
     }