Enable collnet in RCCL (#333)

* Enable CollNet and use different number of channels

* topo_expl: enable collnet

[ROCm/rccl commit: 1d6244b18d]

projects/rccl/src/collectives/device/all_reduce.h

@@ -258,7 +258,7 @@ class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_COLLNET, NCCL_PROTO_SIMPLE, FUNC
       }
     }
 
-    if (blockIdx.x >= nChannels) { // second half of the channels do broadcast
+    if (blockIdx.x >= nChannels && blockIdx.x < 2*nChannels) { // second half of the channels do broadcast
       ncclPrimitives<UNROLL, 1, 1, T, 1, 1, 0, FUNC>
         prims(tid, nthreads, &tree->up, tree->down, NULL, stepSize, channel, comm, ncclShmem->ptrs, 0);
       for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
@@ -455,7 +455,7 @@ class ncclFunction<ncclFuncAllReduce, NCCL_ALGO_COLLNET, NCCL_PROTO_LL, FUNC, T,
       }
     }
 
-    if (blockIdx.x >= nChannels) { // second half of the channels do broadcast
+    if (blockIdx.x >= nChannels && blockIdx.x < 2*nChannels) { // second half of the channels do broadcast
       ncclLLPrimitives<T, FUNC, 1, 1> LLprims(tid, nthreads, &tree->up, tree->down, stepLines, channel, comm);
       for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
         // Down

projects/rccl/src/enqueue.cc

@@ -308,7 +308,7 @@ static ncclResult_t getAlgoInfo(struct ncclInfo* info) {
   TRACE(NCCL_COLL, "%ld Bytes -> Algo %d proto %d time %f", info->nBytes, info->algorithm, info->protocol, minTime);
 
   int nc = (info->nChannels > 0) ? info->nChannels :
-           (info->algorithm == NCCL_ALGO_COLLNET) ? comm->nChannels/2 : comm->nChannels; // CollNet uses one channel for up and one channel for down
+           (info->algorithm == NCCL_ALGO_COLLNET) ? comm->collNetnChannels/2 : comm->nChannels; // CollNet uses one channel for up and one channel for down
   int nt = comm->maxThreads[info->algorithm][info->protocol];
   int threadThreshold = comm->threadThresholds[info->algorithm][info->protocol];
   while (info->nBytes < nc*nt*threadThreshold) {
@@ -511,7 +511,7 @@ ncclResult_t ncclSaveKernel(struct ncclInfo* info) {
     proxyArgs.channel = channel;
     // Adjust pattern for CollNet based on channel index
     if (nSubChannels == 2) {
-      info->pattern = (channelId < info->comm->nChannels/nSubChannels) ? ncclPatternCollTreeUp : ncclPatternCollTreeDown;
+      info->pattern = (channelId < info->comm->collNetnChannels/nSubChannels) ? ncclPatternCollTreeUp : ncclPatternCollTreeDown;
     }
 
     if (proxyArgs.nsteps) NCCLCHECK(ncclProxySaveColl(&proxyArgs, info->pattern, info->root, info->comm->nRanks));

projects/rccl/src/graph/connect.cc

@@ -174,7 +174,7 @@ ncclResult_t ncclTopoConnectCollNet(struct ncclComm* comm, struct ncclTopoGraph*
   int depth = nranks/comm->nNodes;
   int sendIndex = collNetGraph->pattern == NCCL_TOPO_PATTERN_TREE ? 0 : 1;  // send GPU index depends on topo pattern
   int sendEndIndex = (sendIndex+comm->localRanks-1)%comm->localRanks;
-  for (int c=0; c<comm->nChannels/2; c++) {
+  for (int c=0; c<comm->collNetnChannels/2; c++) {
     struct ncclChannel* channel = comm->channels+c;
     // Set root of collTree to id nranks
     if (rank == collNetGraph->intra[sendIndex+c*comm->localRanks]) { // is master
@@ -188,8 +188,8 @@ ncclResult_t ncclTopoConnectCollNet(struct ncclComm* comm, struct ncclTopoGraph*
   }
   int recvIndex = 0;  // recv GPU index is always 0
   int recvEndIndex = (recvIndex+comm->localRanks-1)%comm->localRanks;
-  for (int c=0; c<comm->nChannels/2; c++) {
-    struct ncclChannel* channel = comm->channels+comm->nChannels/2+c;
+  for (int c=0; c<comm->collNetnChannels/2; c++) {
+    struct ncclChannel* channel = comm->channels+comm->collNetnChannels/2+c;
     // Set root of collTree to id nranks
     if (rank == collNetGraph->intra[recvIndex+c*comm->localRanks]) { // is master
       channel->collTree.up = nranks;
@@ -198,7 +198,7 @@ ncclResult_t ncclTopoConnectCollNet(struct ncclComm* comm, struct ncclTopoGraph*
       channel->collTree.down[0] = -1;
     }
     channel->collTree.depth = depth;
-    INFO(NCCL_GRAPH, "CollNet Channel %d rank %d up %d down %d", comm->nChannels/2+c, rank, channel->collTree.up, channel->collTree.down[0]);
+    INFO(NCCL_GRAPH, "CollNet Channel %d rank %d up %d down %d", comm->collNetnChannels/2+c, rank, channel->collTree.up, channel->collTree.down[0]);
   }
   return ncclSuccess;
 }

projects/rccl/src/graph/paths.cc

@@ -572,6 +572,6 @@ ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm) {
     for (int b=1, mb=(comm->p2pnChannels>>1); b<comm->p2pnChannels; b<<=1, mb>>=1) if (c & b) mirror |= mb;
     comm->p2pChannels[c] = mirror;
   }
-  INFO(NCCL_INIT, "%d coll channels, %d p2p channels, %d p2p channels per peer", comm->nChannels, comm->p2pnChannels, comm->p2pnChannelsPerPeer);
+  INFO(NCCL_INIT, "%d coll channels, %d collnet channels, %d p2p channels, %d p2p channels per peer", comm->nChannels, comm->collNetnChannels, comm->p2pnChannels, comm->p2pnChannelsPerPeer);
   return ncclSuccess;
 }

projects/rccl/src/graph/search.cc

@@ -1112,6 +1112,8 @@ float speedArray[] = { 42.0, 30.0, 24.0, 21.0, 18.0, 15.0, 12.0, 10.0, 9.0, 7.0,
 #endif
 #define NSPEEDS (sizeof(speedArray)/sizeof(float))
 
+RCCL_PARAM(ModelMatchingDisable, "MODEL_MATCHING_DISABLE", 0);
+
 ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph) {
   int ngpus = system->nodes[GPU].count;
   int nnets = system->nodes[NET].count;
@@ -1143,13 +1145,18 @@ ncclResult_t ncclTopoCompute(ncclTopoSystem* system, struct ncclTopoGraph* graph
     if (graph->nChannels) {
       system->type |= RCCL_TOPO_4P2H_ROME;
     }
-  } else {
+  } else if (!rcclParamModelMatchingDisable()) {
     // try to match 8P6L
     NCCLCHECK(parseChordalRing(system, graph));
     if (graph->nChannels) return ncclSuccess;
     // try to match Rome 4P2H
     NCCLCHECK(parseRome4P2H(system, graph));
   }
+  if (graph->collNet && graph->nChannels) {
+    graph->nChannels = 1;
+    memcpy(graph->intra+graph->nChannels*ngpus, graph->intra, ngpus*sizeof(int));
+    memcpy(graph->inter+graph->nChannels*2, graph->inter, 2*sizeof(int));
+  }
   if (graph->nChannels) return ncclSuccess;
 
   if (ngpus == 1) if (graph->pattern != NCCL_TOPO_PATTERN_RING) graph->pattern = NCCL_TOPO_PATTERN_TREE;

projects/rccl/src/include/comm.h

@@ -101,6 +101,8 @@ struct ncclComm {
   int p2pnChannels;
   int p2pnChannelsPerPeer;
   int p2pChannels[MAXCHANNELS];
+  //Channels for collnet
+  int collNetnChannels;
 
   // Buffer sizes
   int buffSizes[NCCL_NUM_PROTOCOLS];

projects/rccl/src/init.cc

@@ -717,7 +717,7 @@ static ncclResult_t checkCollNetSetup(struct ncclComm* comm, int rank, int collN
   if (collNetSetupFail) {
     if (rank == 0) WARN("Cannot initialize CollNet, using %s instead", ncclNetName());
     // Free collNet resources
-    for (int r=0; r<comm->nChannels; r++) {
+    for (int r=0; r<comm->collNetnChannels; r++) {
       struct ncclChannel* channel = comm->channels+r;
       struct ncclPeer* peer = channel->peers+nranks;
       if (peer->send.transportResources && peer->send.transportComm) NCCLCHECK(peer->send.transportComm->free(peer->send.transportResources));
@@ -1035,6 +1035,8 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   if (comm->nNodes > 1 &&
       ncclParamCollNetEnable() == 1 &&
       collNetSupport() && collNetGraph.nChannels) {
+    // Force 2 channels for CollNet
+    comm->collNetnChannels = collNetGraph.nChannels = 2;
     NCCLCHECK(ncclTopoConnectCollNet(comm, &collNetGraph, rank));
   }
 
@@ -1092,7 +1094,7 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   if (comm->nNodes > 1 &&
       ncclParamCollNetEnable() == 1 &&
       collNetSupport() && collNetGraph.nChannels) {
-    int logicChannels = comm->nChannels/2;
+    int logicChannels = comm->collNetnChannels/2;
     int collNetSetupFail = 0;
     const int recvIndex = 0;  // recv GPU index is always 0
     const int sendIndex = collNetGraph.pattern == NCCL_TOPO_PATTERN_TREE ? 0 : 1;  // send GPU index depends on topo pattern

projects/rccl/src/transport/coll_net.cc

@@ -9,12 +9,14 @@
 #include "coll_net.h"
 #include "graph.h"
 #include <assert.h>
+#include <hsa/hsa_ext_amd.h>
 
 struct collNetRecvConnectInfo {
   collNetHandle_t collNetHandle;
 };
 
 struct collNetSendConnectInfo {
+  collNetHandle_t collNetHandle;
   void* collNetComm;
   void* mhandles[NCCL_NUM_PROTOCOLS];
   struct reqSlot* reqFifo;
@@ -39,6 +41,7 @@ struct collNetSendResources {
   uint64_t llLastCleaning;
   struct reqSlot* reqFifo;
   int collNetRank;
+  uint32_t* curr_hdp_reg;  // Curr GPU in ring (for rdma transport use only)
 };
 
 struct collNetRecvResources {
@@ -55,6 +58,7 @@ struct collNetRecvResources {
   uint64_t llLastCleaning;
   struct reqSlot* reqFifo;
   int collNetRank;
+  uint32_t* curr_hdp_reg;  // Curr GPU in ring (for rdma transport use only)
 };
 
 /* Determine if we can communicate with the peer */
@@ -78,7 +82,28 @@ ncclResult_t collNetSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) recvSize += send->comm->buffSizes[p];
 
   if (resources->useGdr) {
-    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize));
+    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize, resources->useGdr));
+    //CUDACHECK(hipDeviceGetAttribute((int*)&resources->curr_hdp_reg, hipDeviceAttributeHdpMemFlushCntl, myInfo->cudaDev));
+    struct data_struct {hsa_agent_t agent; int counter;} out;
+    out.counter = 0;
+    out.agent.handle = myInfo->cudaDev;
+    hsa_iterate_agents([](hsa_agent_t agent, void* data) {
+      int devId = ((struct data_struct *)data)->agent.handle;
+      hsa_device_type_t type;
+      hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &type);
+      if(type != HSA_DEVICE_TYPE_GPU)
+        return HSA_STATUS_SUCCESS;
+      if(((struct data_struct *)data)->counter!=devId) {
+        ((struct data_struct *)data)->counter++;
+        return HSA_STATUS_SUCCESS;
+      }
+      ((struct data_struct *)data)->agent = agent;
+      return HSA_STATUS_SUCCESS;
+    }, (void*)&out);
+    hsa_amd_hdp_flush_t hdpinfo;
+    hsa_status_t err = hsa_agent_get_info(out.agent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_HDP_FLUSH, &hdpinfo);
+    resources->curr_hdp_reg = hdpinfo.HDP_MEM_FLUSH_CNTL;
+    send->conn.curr_hdp_reg = resources->curr_hdp_reg;
   }
   NCCLCHECK(ncclCudaHostCalloc((char**)&resources->recvMem, recvSize));
   NCCLCHECK(ncclIbMalloc((void**)&(resources->llData), send->comm->buffSizes[NCCL_PROTO_LL]/2));
@@ -103,7 +128,7 @@ ncclResult_t collNetRecvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) recvSize += recv->comm->buffSizes[p];
 
   if (resources->useGdr) {
-    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize));
+    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize, resources->useGdr));
   }
   NCCLCHECK(ncclCudaHostCalloc((char**)&resources->recvMem, recvSize));
 
@@ -291,6 +316,11 @@ ncclResult_t collNetSendProxy(struct ncclProxyArgs* args) {
           size = nFifoLines*2*sizeof(uint32_t);
         }
         if (ready) {
+          // flush HDP if not done
+          if (resources->curr_hdp_reg && args->hdp_flushed < LOAD(recvTail)) {
+            args->hdp_flushed = LOAD(recvTail);
+            STORE(resources->curr_hdp_reg, 1);
+          }
           // Data is ready, try to send.
           int count = size/ncclTypeSize(args->dtype);
           NCCLCHECK(collNetIallreduce(resources->collNetSendComm, (void*) buff, (void*)(reqFifo[buffSlot].recvBuff), count, args->dtype, args->redOp, sendMhandle, recvMhandle, args->requests+buffSlot));