[P2P] Have connIdx for both send and recv (#1524)

[ROCm/rccl commit: 387c973b5d]

projects/rccl/src/device/common.h

@@ -63,7 +63,8 @@
       collTrace->p2p.nSendChannels = p2pWork->nSendChannels; \
       collTrace->p2p.nRecvChannels = p2pWork->nRecvChannels; \
       collTrace->p2p.channelBase = p2pWork->channelBase; \
-      collTrace->p2p.connIndex = p2pWork->connIndex; \
+      collTrace->p2p.sendConnIndex = p2pWork->sendConnIndex; \
+      collTrace->p2p.recvConnIndex = p2pWork->recvConnIndex; \
       collTrace->p2p.sendProtoLL = p2pWork->sendProtoLL; \
       collTrace->p2p.recvProtoLL = p2pWork->recvProtoLL; \
       collTrace->p2pOpCount[0] = p2pWork->sendOpCount; \

projects/rccl/src/device/sendrecv.h

@@ -42,7 +42,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
 
     Primitives<T, RedOp, FanAsymmetric<0, 1>, 0, Proto, 1>
       prims(tid, tn, nullptr, &work->sendRank, work->sendAddr, nullptr,
-            /*redOpArg(ignored)=*/0, group, work->connIndex, work->connIndex, nullptr,
+            /*redOpArg(ignored)=*/0, group, work->sendConnIndex, work->sendConnIndex, nullptr,
             /*userBufferMode=*/work->sendRegistered, ncclShmem.comm.p2pChunkSize);
 
 #if defined(ENABLE_NPKIT)
@@ -100,7 +100,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
 
     Primitives<T, RedOp, FanAsymmetric<1, 0>, 0, Proto, 1>
       prims(tid, tn, &work->recvRank, nullptr, nullptr, work->recvAddr,
-            /*redOpArg(ignored)=*/0, group, work->connIndex, work->connIndex, nullptr,
+            /*redOpArg(ignored)=*/0, group, work->recvConnIndex, work->recvConnIndex, nullptr,
             /*userBufferMode=*/work->recvRegistered, ncclShmem.comm.p2pChunkSize);
 
 #if defined(ENABLE_NPKIT)

projects/rccl/src/enqueue.cc

@@ -871,7 +871,7 @@ static ncclResult_t addP2pToPlan(
     int recvRank, void* recvAddr, ssize_t recvBytes,
     uint64_t sendOpCount, uint64_t recvOpCount
   ) {
-  int connIndex = 1;
+  int connIndex[2] = {1, 1};
   bool selfSend = (sendRank == comm->rank);
   // recv: dir=0, send: dir=1
   void* addrs[2] = {recvAddr, sendAddr};
@@ -881,8 +881,12 @@ static ncclResult_t addP2pToPlan(
   bool proxySameProcess[2] = {true, true};
   uint8_t base = ncclP2pChannelBaseForRound(comm, p2pRound);
 
-  if (comm->p2pNet && (sendBytes > rcclParamP2pNetThreshold() || recvBytes > rcclParamP2pNetThreshold()))
-    connIndex = NCCL_CONN_IDX_P2P_NET;
+  if (comm->p2pNet) {
+    for (int dir = 0; dir <= 1; dir++) {
+      if (bytes[dir] > rcclParamP2pNetThreshold())
+        connIndex[dir] = NCCL_CONN_IDX_P2P_NET;
+    }
+  }
   
   if (!selfSend) {
     for (int part=0; part < nChannelsMax; part++) {
@@ -890,8 +894,8 @@ static ncclResult_t addP2pToPlan(
       struct ncclChannelPeer** channelPeers = comm->channels[channelId].peers;
       for (int dir=0; dir <= 1; dir++) {
         int peerRank = dir ? sendRank : recvRank;
-        struct ncclConnector* conn = dir ? &channelPeers[peerRank]->send[connIndex]
-                                         : &channelPeers[peerRank]->recv[connIndex];
+        struct ncclConnector* conn = dir ? &channelPeers[peerRank]->send[connIndex[dir]]
+                                         : &channelPeers[peerRank]->recv[connIndex[dir]];
         protoLL[dir] &= conn->conn.buffs[NCCL_PROTO_LL] != nullptr && !IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx12");
         network[dir] |= conn->transportComm == (dir ? &netTransport.send : &netTransport.recv);
         proxySameProcess[dir] &= conn->proxyConn.sameProcess;
@@ -970,6 +974,8 @@ static ncclResult_t addP2pToPlan(
   work->sendRank = sendRank;
   work->sendAddr = sendAddr;
   work->sendBytes = sendBytes==-1 ? 0 : sendBytes;
+  work->sendConnIndex = connIndex[1];
+  work->sendOpCount = sendOpCount;
   work->nRecvChannels = nChannels[0];
   work->recvProtoLL = protoLL[0];
   work->recvRegistered = registered[0];
@@ -977,8 +983,7 @@ static ncclResult_t addP2pToPlan(
   work->recvRank = recvRank;
   work->recvAddr = recvAddr;
   work->recvBytes = recvBytes==-1 ? 0 : recvBytes;
-  work->connIndex = connIndex;
-  work->sendOpCount = sendOpCount;
+  work->recvConnIndex = connIndex[0];
   work->recvOpCount = recvOpCount;
 
   struct ncclProxyOp proxyOps[2] = {};
@@ -994,7 +999,7 @@ static ncclResult_t addP2pToPlan(
     op->pattern = dir ? ncclPatternSend : ncclPatternRecv;
     op->chunkSize = chunkSize[dir];
     op->reg = registered[dir];
-    op->connIndex = connIndex;
+    op->connIndex = connIndex[dir];
     // The following are modified per channel part in addWorkToChannels():
     // op->buffer, op->nbytes, op->nsteps = ...;
   }

projects/rccl/src/include/device.h

@@ -247,7 +247,7 @@ struct alignas(16) ncclDevWorkP2p {
   uint8_t sendProtoLL:1, recvProtoLL:1;
   uint8_t sendRegistered:1, recvRegistered:1;
 
-  uint8_t connIndex:2;
+  uint8_t sendConnIndex:2, recvConnIndex:2;
 };
 
 // Compute the subset of the data transfer corresponding to the given part index.
@@ -441,7 +441,8 @@ struct ncclCollTrace {
       uint8_t nSendChannels;
       uint8_t nRecvChannels;
       uint8_t channelBase;
-      uint8_t connIndex;
+      uint8_t sendConnIndex:2;
+      uint8_t recvConnIndex:2;
       uint8_t sendProtoLL:1;
       uint8_t recvProtoLL:1;
     } p2p;

projects/rccl/src/init.cc

@@ -260,9 +260,9 @@ void *ncclCommThreadMain(void *arg) {
           if (type == ncclCollTraceCollElemType) {
             sprintf(line+offset, " CE %s nw %d bi %d nc %d root %d busId %lx nRanks %d", funcNames[fIdx], td->coll.nWarps, td->coll.bid, td->coll.nChannels, td->coll.root, comm->busId, comm->nRanks);
           } else if (type == ncclCollTraceP2pElemType) {
-            sprintf(line+offset, " Send %d -> Recv %d connIdx/nc/cb/sLL/rLL %d/%d/%d/%d/%d busId %lx nRanks %d",
-              td->p2p.sendRank, td->p2p.recvRank, td->p2p.connIndex, td->p2p.nP2pChannels, td->p2p.channelBase,
-              td->p2p.sendProtoLL, td->p2p.recvProtoLL, comm->busId, comm->nRanks);
+            sprintf(line+offset, " Send %d -> %d/%d connIdx/LL %d/%d -> Recv %d nc %d cb %d busId %lx nRanks %d",
+              td->p2p.sendRank, td->p2p.sendConnIndex, td->p2p.sendProtoLL, td->p2p.recvConnIndex, td->p2p.recvProtoLL, td->p2p.recvRank, td->p2p.nP2pChannels, td->p2p.channelBase,
+              comm->busId, comm->nRanks);
           } else {
             switch (type&0xf) {
               case ncclCollTraceKernelLaunchType:
@@ -275,9 +275,9 @@ void *ncclCommThreadMain(void *arg) {
                 if ((type&0xf0) == ncclCollTraceCollElemType)
                   sprintf(line+offset, " nw %d bi %d nc %d root %d busId %lx nRanks %d", td->coll.nWarps, td->coll.bid, td->coll.nChannels, td->coll.root, comm->busId, comm->nRanks);
                 else if ((type&0xf0) == ncclCollTraceP2pElemType)
-                  sprintf(line+offset, " Send %d -> Recv %d connIdx/nc/cb/sLL/rLL %d/%d/%d/%d/%d busId %lx nRanks %d",
-                    td->p2p.sendRank, td->p2p.recvRank, td->p2p.connIndex, td->p2p.nP2pChannels, td->p2p.channelBase,
-                    td->p2p.sendProtoLL, td->p2p.recvProtoLL, comm->busId, comm->nRanks);
+                  sprintf(line+offset, " Send %d -> %d/%d ConnIdx/LL %d/%d -> Recv %d nc %d cb %d busId %lx nRanks %d",
+                    td->p2p.sendRank, td->p2p.sendConnIndex, td->p2p.sendProtoLL, td->p2p.recvConnIndex, td->p2p.recvProtoLL, td->p2p.recvRank, td->p2p.nP2pChannels, td->p2p.channelBase,
+                    comm->busId, comm->nRanks);
                 break;
               case ncclCollTraceKernelEndType:
                 sprintf(line+offset, " KE busId %lx nRanks %d", comm->busId, comm->nRanks);