From 9fb6cedc7a72d9d0867df03004b1200831eaf597 Mon Sep 17 00:00:00 2001
From: Thomas Vegas <tvegas@nvidia.com>
Date: Fri, 24 May 2024 11:07:03 +0300
Subject: [PATCH] Don't queue zero-length request into pending

---
 src/ucx_uct_rd_plugin.c | 33 +++++++++++++++++----------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/ucx_uct_rd_plugin.c b/src/ucx_uct_rd_plugin.c
index 68c8081d..c58e0e70 100644
--- a/src/ucx_uct_rd_plugin.c
+++ b/src/ucx_uct_rd_plugin.c
@@ -93,19 +93,25 @@ static void nccl_uct_rd_send_ats(nccl_uct_rd_req_t *req) {
 static void nccl_uct_rd_pending_add(nccl_uct_rd_comm_t *comm,
                                     nccl_uct_mem_t *src, nccl_uct_mem_t *dst) {
   nccl_uct_rd_req_t *req = dst->req;
-  unsigned i             = comm->pending.last++;
   nccl_uct_get_param_t *param;
 
-  param = &comm->pending.param[i & NCCL_UCT_PENDING_MASK];
-
   assert(src->size <= dst->size);
-  assert((comm->pending.first & NCCL_UCT_PENDING_MASK) !=
-         (comm->pending.last & NCCL_UCT_PENDING_MASK));
   assert(req->rts_count < NCCL_UCX_UCT_MAX_RECVS);
 
   req->sizes[dst->index]            = src->size;
   req->remote_req[req->rts_count++] = src->req; /* src->req is a cookie */
 
+  if (src->size == 0) {
+    req->completion.count--;
+    return;
+  }
+
+  param = &comm->pending.param[comm->pending.last & NCCL_UCT_PENDING_MASK];
+  comm->pending.last++;
+
+  assert((comm->pending.first & NCCL_UCT_PENDING_MASK) !=
+         (comm->pending.last & NCCL_UCT_PENDING_MASK));
+
   param->iov.buffer = dst->data;
   param->iov.length = src->size;
   param->iov.memh   = dst->u.uct_memh->memh;
@@ -116,24 +122,19 @@ static void nccl_uct_rd_pending_add(nccl_uct_rd_comm_t *comm,
   param->req        = req;
 }
 
-static inline void nccl_uct_rd_pending_drain(nccl_uct_rd_comm_t *comm) {
+static void nccl_uct_rd_pending_drain(nccl_uct_rd_comm_t *comm) {
   ucs_status_t status;
   nccl_uct_get_param_t *param;
 
   for (; comm->pending.first != comm->pending.last; comm->pending.first++) {
     param = &comm->pending.param[comm->pending.first & NCCL_UCT_PENDING_MASK];
 
-    if (param->iov.length == 0) {
+    status = uct_ep_get_zcopy(comm->base.uct_ep->ep, &param->iov, 1, param->rva,
+                              param->rkey, &param->req->completion);
+    if (status == UCS_OK) {
       param->req->completion.count--;
-    } else {
-      status = uct_ep_get_zcopy(comm->base.uct_ep->ep, &param->iov, 1,
-                                param->rva, param->rkey,
-                                &param->req->completion);
-      if (status == UCS_OK) {
-        param->req->completion.count--;
-      } else if (status != UCS_INPROGRESS) {
-        break;
-      }
+    } else if (status != UCS_INPROGRESS) {
+      break;
     }
 
     if (param->req->completion.count == 1) {