Make network isend/irecv non blocking

1 files changed, 13 insertions, 7 deletions

diff --git a/src/transport/net.cu b/src/transport/net.cu
index 8a7e3b8..165187c 100644
--- a/src/transport/net.cu
+++ b/src/transport/net.cu
@@ -416,17 +416,21 @@ ncclResult_t netSendProxy(struct ncclProxyArgs* args) {
             while (f1[0] != flag || f2[0] != flag);
           }
           NCCLCHECK(ncclNetIsend(resources->netSendComm, lines, size, ptrType, requests+slot));
-          sizesFifo[slot] = size;
-          tail++;
-          idle = 0;
+          if (requests[slot] != NULL) {
+            sizesFifo[slot] = size;
+            tail++;
+            idle = 0;
+          }
         }
       }
     } else while (tail < *prevTail) {
         // Send through network
         int slot = tail%args->substeps;
         NCCLCHECK(ncclNetIsend(resources->netSendComm, localBuff+slot*sliceSize, sizesFifo[slot], ptrType, requests+slot));
-        tail++;
-        idle = 0;
+        if (requests[slot] != NULL) {
+          tail++;
+          idle = 0;
+        }
       }
     if (head < tail) {
       int done;
@@ -502,8 +506,10 @@ ncclResult_t netRecvProxy(struct ncclProxyArgs* args) {
     if ((tail < head + args->substeps) && (tail < *nextHead + args->substeps) && (tail < end)) {
       int slot = tail%args->substeps;
       NCCLCHECK(ncclNetIrecv(resources->netRecvComm, localBuff+slot*sliceSize, sliceSize, ptrType, requests+slot));
-      tail++;
-      idle = 0;
+      if (requests[slot] != NULL) {
+        tail++;
+        idle = 0;
+      }
     }
     if (tail > head) {
       int done;