From 821d2e1f30c8e6ebf8eab19b62621deb0540d12e Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Fri, 11 Oct 2024 10:40:32 -0700
Subject: [PATCH] Allow zero byte sendrecv in alltoallv (#1349)

* Allow zero byte sendrecv in alltoallv

* Fix previous merge error
---
 src/collectives.cc | 4 ++--
 src/enqueue.cc     | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/collectives.cc b/src/collectives.cc
index 04f2683528..cb488d97b8 100644
--- a/src/collectives.cc
+++ b/src/collectives.cc
@@ -142,14 +142,14 @@ ncclResult_t ncclAllToAllv_impl(const void *sendbuff, const size_t sendcounts[],
   NCCLCHECK(ncclCommCount(comm, &nRanks));
   NCCLCHECK(ncclGroupStart());
   for (int r=0; r<nRanks; r++) {
-    if (sendcounts[r]) NCCLCHECK(ncclSend(
+    NCCLCHECK(ncclSend(
         ((char*)sendbuff) + sdispls[r]*ncclTypeSize(datatype),
         sendcounts[r],
         datatype,
         r,
         comm,
         stream));
-    if (recvcounts[r]) NCCLCHECK(ncclRecv(
+    NCCLCHECK(ncclRecv(
         ((char*)recvbuff) + rdispls[r]*ncclTypeSize(datatype),
         recvcounts[r],
         datatype,
diff --git a/src/enqueue.cc b/src/enqueue.cc
index 63f87af4fe..196653c94b 100644
--- a/src/enqueue.cc
+++ b/src/enqueue.cc
@@ -609,7 +609,7 @@ static ncclResult_t addP2pToPlan(
 
   struct ncclWorkElemP2p elem = {0};
   elem.proto = info.protocol;
-  elem.peer = addr == nullptr ? -1 : peer;
+  elem.peer = peer;
   elem.nWarps = NCCL_MAX_NTHREADS/comm->WarpSize;
   elem.reg = proxyOp.reg;
   elem.p2pType = isSendNotRecv ? ncclWorkP2pTypeSend : ncclWorkP2pTypeRecv;
@@ -628,7 +628,7 @@ static ncclResult_t addP2pToPlan(
   // Calculate the opCount after appendWorkElemP2p since it will always return
   // with channel->nWork equal to one plus the work index this p2p settled in.
   proxyOp.opCount = uint64_t(plan->channels[channelId].nWork)<<1 | 1;
-  if (addr != nullptr) NCCLCHECK(addProxyOpIfNeeded(comm, plan, &proxyOp));
+  NCCLCHECK(addProxyOpIfNeeded(comm, plan, &proxyOp));
   return ncclSuccess;
 }