Merge pull request #2 from parasailteam/connection_setup

Connection setup

src/collectives/device/all_to_all.h

@@ -18,18 +18,24 @@ class ncclFunction<ncclFuncAllToAll, ALGO, PROTO, FUNC, T, UNROLL> {
       const int nChannels = args->coll.nChannels;
       struct ncclDevComm* comm = args->comm;
       struct ncclChannel* channel = comm->channels+blockIdx.x;
-      struct ncclRing* ring = &channel->ring;
+      struct scklGraph* sGraph = &channel->sGraph;
       const int stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / (sizeof(T)*NCCL_STEPS);
       const int chunkSize = stepSize * ALLTOALL_CHUNKSTEPS;
       const int nranks = comm->nRanks;
       const ssize_t loopSize = nChannels*(ssize_t)chunkSize;
       const ssize_t size = args->coll.count;
-
-
+      return;
       // Compute pointers
       const T * __restrict__ thisInput = (const T*)args->sendbuff;
       T * __restrict__ thisOutput = (T*)args->recvbuff;
-      printf("This is an empty function! %d %d %d %d\n", (int) size, tid, bid, (int) sizeof(T));
+      ncclPrimitives<UNROLL, ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS, T, 3, 3, 1, FUNC>
+        prims(tid, nthreads, sGraph->recv, sGraph->send, thisOutput, stepSize, channel, comm, ncclShmem->ptrs, 0);
+      if (tid == 0 && bid == 0){
+        printf("connected to %d %d %d\n", sGraph->send[0], sGraph->send[1], sGraph->send[2]);
+      }
+      int testSize = min(chunkSize, (int)size/nChannels/nranks);
+      prims.directSend(thisInput, 0, testSize);
+      prims.directRecv(thisOutput, 0, testSize);
       return;
     }
 };

src/graph/topo.cc

@@ -607,6 +607,73 @@ ncclResult_t ncclTopoGetLocalNet(struct ncclTopoSystem* system, int rank, int64_
   return ncclSuccess;
 }
 
+ncclResult_t scklGetTopoFromXMLAndSetChannels(struct ncclComm* comm) {
+  char* str = getenv("SCKL_XML_FILE");
+  if (str){
+    INFO(NCCL_ENV, "SCKL_XML_FILE set by environment to %s", str);
+    struct ncclXml* xml;
+    NCCLCHECK(ncclCalloc(&xml, 1));
+    NCCLCHECK(scklTopoGetXmlGraphFromFile(str, xml));
+    int rank = comm->rank;
+
+    for (int c=0; c<comm->nChannels; c++){
+      comm->channels[c].sGraph.nRecvPeers = 0;
+      comm->channels[c].sGraph.nSendPeers = 0;
+    }
+    
+    struct ncclXmlNode* topNode;
+    NCCLCHECK(xmlFindTag(xml, "system", &topNode));
+    for (int s=0; s<topNode->nSubs; s++) {
+      struct ncclXmlNode* node = topNode->subs[s];
+      if (strcmp(node->name, "gpu") == 0){
+        int id;
+        NCCLCHECK(xmlGetAttrInt(node, "id", &id));
+        if (id == rank){
+          for (int p=0; p<node->nSubs; p++) {
+            struct ncclXmlNode* typeOfComm = node->subs[p];
+            if (strcmp(typeOfComm->name, "conn") == 0){
+              const char* type;
+              NCCLCHECK(xmlGetAttrStr(typeOfComm, "type", &type));
+
+              bool isRecv = false;
+              bool isSend = false;
+              if (strcmp(type, "recv") == 0){
+                isRecv = true;
+              } else if (strcmp(type, "send") == 0){
+                isSend = true;
+              }
+              for (int p=0; p<typeOfComm->nSubs; p++) {
+                struct ncclXmlNode* peer = typeOfComm->subs[p];
+                int peerId;
+                NCCLCHECK(xmlGetAttrInt(peer, "id", &peerId));
+                // SCKL generates the same scklGraph for all channels for now. This will change in the future
+                for (int c=0; c<comm->nChannels; c++){
+                  if (isRecv) {
+                    if (comm->channels[c].sGraph.nRecvPeers < SCKL_MAX_NUM_CONN){
+                      comm->channels[c].sGraph.recv[comm->channels[c].sGraph.nRecvPeers++] = peerId;
+                    } else {
+                      WARN("Too many recv connections for device %d channel %d -- connection to %d is ignored. This may cause deadlock in initialization.", rank, c, peerId);
+                    }                    
+                  } else if (isSend){
+                    if (comm->channels[c].sGraph.nSendPeers < SCKL_MAX_NUM_CONN){
+                      comm->channels[c].sGraph.send[comm->channels[c].sGraph.nSendPeers++] = peerId;
+                    } else {
+                      WARN("Too many recv connections for device %d channel %d -- connection to %d is ignored.  This may cause deadlock in initialization.", rank, c, peerId);
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+    free(xml);
+  }
+  return ncclSuccess;
+}
+
+
 /****************************/
 /* External query functions */
 /****************************/

src/graph/xml.cc

@@ -805,3 +805,46 @@ ncclResult_t ncclTopoGetXmlGraphFromFile(const char* xmlGraphFile, struct ncclXm
   fclose(file);
   return ncclSuccess;
 }
+
+ncclResult_t scklTopoXmlPeerLoad(FILE* file, struct ncclXml* xml, struct ncclXmlNode* head) {
+  int id;
+  NCCLCHECK(xmlGetAttrInt(head, "id", &id));
+  struct xmlHandler handlers[] = { };
+  NCCLCHECK(xmlLoadSub(file, xml, head, handlers, 1));
+  return ncclSuccess;
+}
+
+ncclResult_t scklTopoXmlConnLoad(FILE* file, struct ncclXml* xmlGraph, struct ncclXmlNode* head) {
+  const char* type;
+  NCCLCHECK(xmlGetAttrStr(head, "type", &type));
+  struct xmlHandler handlers[] = { { "peer", scklTopoXmlPeerLoad } };
+  NCCLCHECK(xmlLoadSub(file, xmlGraph, head, handlers, 1));
+  return ncclSuccess;
+}
+
+ncclResult_t scklTopoXmlGraphLoad(FILE* file, struct ncclXml* xmlGraph, struct ncclXmlNode* head) {
+  int id;
+  NCCLCHECK(xmlGetAttrInt(head, "id", &id));
+  struct xmlHandler handlers[] = { { "conn", scklTopoXmlConnLoad } };
+  NCCLCHECK(xmlLoadSub(file, xmlGraph, head, handlers, 1));
+  return ncclSuccess;
+}
+
+ncclResult_t scklTopoXmlSystemLoad(FILE* file, struct ncclXml* xmlGraph, struct ncclXmlNode* head) {
+  struct xmlHandler handlers[] = { { "gpu", scklTopoXmlGraphLoad } };
+  NCCLCHECK(xmlLoadSub(file, xmlGraph, head, handlers, 1));
+  return ncclSuccess;
+}
+
+ncclResult_t scklTopoGetXmlGraphFromFile(const char* xmlGraphFile, struct ncclXml* xml) {
+  FILE* file = fopen(xmlGraphFile, "r");
+  if (file == NULL) {
+    WARN("Could not open XML SCKL graph file %s : %s", xmlGraphFile, strerror(errno));
+    return ncclSystemError;
+  }
+  struct xmlHandler handlers[] = { { "system", scklTopoXmlSystemLoad } };
+  xml->maxIndex = 0;
+  NCCLCHECK(xmlLoadSub(file, xml, NULL, handlers, 1));
+  fclose(file);
+  return ncclSuccess;
+}

src/graph/xml.h

@@ -1,3 +1,4 @@
+
 /*************************************************************************
  * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
  *
@@ -42,6 +43,7 @@ ncclResult_t ncclTopoGetXmlFromFile(const char* xmlTopoFile, struct ncclXml* xml
 ncclResult_t ncclTopoDumpXmlToFile(const char* xmlTopoFile, struct ncclXml* xml);
 #define NCCL_GRAPH_XML_VERSION 1
 ncclResult_t ncclTopoGetXmlGraphFromFile(const char* xmlGraphFile, struct ncclXml* xml);
+ncclResult_t scklTopoGetXmlGraphFromFile(const char* xmlGraphFile, struct ncclXml* xml);
 
 /* Auto-detect functions */
 ncclResult_t ncclTopoFillGpu(struct ncclXml* xml, const char* busId, struct ncclXmlNode** gpuNode);

src/include/devcomm.h

@@ -116,6 +116,14 @@ struct ncclRing {
   int* devUserRanks;
 };
 
+#define SCKL_MAX_NUM_CONN 16
+
+struct scklGraph {
+  int nRecvPeers;
+  int nSendPeers;
+  int recv[SCKL_MAX_NUM_CONN];
+  int send[SCKL_MAX_NUM_CONN];
+};
 
 #define NCCL_MAX_TREE_ARITY 3
 struct ncclTree {
@@ -177,6 +185,7 @@ struct ncclChannel {
       struct ncclRing ring;
       struct ncclTree tree;
       struct ncclTree collTree;
+      struct scklGraph sGraph;
 
       int id;
 

src/include/graph.h

@@ -35,6 +35,9 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 // Set CPU affinity
 ncclResult_t ncclTopoSetAffinity(struct ncclTopoSystem* system, int rank);
 
+// SCKL setup peers
+ncclResult_t scklGetTopoFromXMLAndSetChannels(struct ncclComm* comm);
+
 #define NCCL_TOPO_CPU_ARCH_X86 1
 #define NCCL_TOPO_CPU_ARCH_POWER 2
 #define NCCL_TOPO_CPU_ARCH_ARM 3

src/init.cc

@@ -823,6 +823,18 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   NCCLCHECKGOTO(ncclTransportP2pSetup(comm, &treeGraph), ret, affinity_restore);
   INFO(NCCL_INIT, "Connected all trees");
 
+  // NetSharedBuffers needs to be set for this to work across nodes.
+  NCCLCHECK(scklGetTopoFromXMLAndSetChannels(comm));
+  // Connect SCKL graph
+  for (int c=0; c<comm->nChannels; c++) {
+    struct ncclChannel* channel = comm->channels+c;
+    if (comm->nRanks == 1) continue;
+    NCCLCHECKGOTO(ncclTransportP2pConnect(comm, channel, channel->sGraph.nRecvPeers, channel->sGraph.recv, channel->sGraph.nSendPeers, channel->sGraph.send), ret, affinity_restore);
+  }
+  // It appears that graph is not really needed for P2pSetup. The only place that actually uses it is in ncclTopoGetNetDev which has a bypass for when it is set to NULL.
+  NCCLCHECKGOTO(ncclTransportP2pSetup(comm, NULL), ret, affinity_restore);
+  INFO(NCCL_INIT, "Connected SCKL graph");  
+
   // Check if we can setup CollNet
   if (comm->nNodes > 1 &&
       ncclParamCollNetEnable() == 1 &&

src/transport.cc

@@ -67,7 +67,6 @@ ncclResult_t ncclTransportP2pSetup(struct ncclComm* comm, struct ncclTopoGraph*
     int sendPeer = (comm->rank + i) % comm->nRanks;
     uint32_t recvMask = comm->connectRecv[recvPeer];
     uint32_t sendMask = comm->connectSend[sendPeer];
-
     struct ncclConnect* recvData = data;
     int sendChannels = 0, recvChannels = 0;
     for (int c=0; c<MAXCHANNELS; c++) {