UCT/CUDA_IPC: Skip zero-length elements in SGL zcopy

src/uct/cuda/cuda_ipc/cuda_ipc_ep.c

@@ -288,14 +288,16 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     uct_cuda_ipc_unpacked_rkey_t *key;
     uct_cuda_ipc_event_desc_t *cuda_ipc_event;
     uct_cuda_ipc_sgl_mapping_t *mapping;
+    uct_cuda_ipc_sgl_entry_t *entry;
     uct_cuda_queue_desc_t *q_desc;
     CUmemcpyAttributes attr;
     void *mapped_addr;
     const void *mapped_base_addr;
-    CUdeviceptr *mapped_addrs;
+    CUdeviceptr *mapped_addrs, *local_addrs;
     CUdeviceptr *dsts, *srcs;
+    size_t *copy_lengths;
     size_t attrs_idx;
-    size_t i, total_length;
+    size_t i, first, total_length;
     CUdevice cuda_device;
     CUstream *stream;
     ucs_status_t status;
@@ -308,7 +310,12 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     }
 
     total_length = 0;
+    first        = count;
     for (i = 0; i < count; i++) {
+        if ((first == count) && (lengths[i] != 0)) {
+            first = i;
+        }
+
         total_length += lengths[i];
     }
 
@@ -318,15 +325,15 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
         return UCS_OK;
     }
 
-    status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[0],
+    status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[first],
                                              &cuda_device, &is_ctx_pushed);
     if (ucs_unlikely(status != UCS_OK)) {
         return status;
     }
 
     mapping = ucs_malloc(sizeof(*mapping) +
                          count * (sizeof(*mapping->entries) +
-                                  sizeof(CUdeviceptr)),
+                                  (2 * sizeof(CUdeviceptr)) + sizeof(size_t)),
                          "cuda_ipc_sgl_mapping");
     if (ucs_unlikely(mapping == NULL)) {
         ucs_error("Failed to allocate cuda_ipc sgl mapping");
@@ -337,12 +344,19 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     mapping->count   = 0;
     mapping->entries = (uct_cuda_ipc_sgl_entry_t *)(mapping + 1);
     mapped_addrs     = (CUdeviceptr *)(mapping->entries + count);
+    local_addrs      = mapped_addrs + count;
+    copy_lengths     = (size_t *)(local_addrs + count);
 
     memset(&attr, 0, sizeof(attr));
     attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
     attrs_idx           = 0;
 
     for (i = 0; i < count; i++) {
+        /* cuMemcpyBatchAsync does not accept copies of zero length */
+        if (lengths[i] == 0) {
+            continue;
+        }
+
         key = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[i];
 
         status = uct_cuda_ipc_get_remote_address(&key->super, remote_addrs[i],
@@ -352,23 +366,26 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
             goto out_unmap;
         }
 
-        mapping->entries[i].pid         = key->super.super.pid;
-        mapping->entries[i].pid_ns      = key->super.pid_ns;
-        mapping->entries[i].d_bptr      = (uintptr_t)key->super.super.d_bptr;
-        mapping->entries[i].mapped_addr = mapped_base_addr;
-        mapped_addrs[i]                 = (CUdeviceptr)mapped_addr;
+        entry                        = &mapping->entries[mapping->count];
+        entry->pid                   = key->super.super.pid;
+        entry->pid_ns                = key->super.pid_ns;
+        entry->d_bptr                = (uintptr_t)key->super.super.d_bptr;
+        entry->mapped_addr           = mapped_base_addr;
+        mapped_addrs[mapping->count] = (CUdeviceptr)mapped_addr;
+        local_addrs[mapping->count]  = (CUdeviceptr)buffers[i];
+        copy_lengths[mapping->count] = lengths[i];
         mapping->count++;
     }
 
     if (direction == UCT_CUDA_IPC_PUT) {
         dsts = mapped_addrs;
-        srcs = (CUdeviceptr*)buffers;
+        srcs = local_addrs;
     } else {
-        dsts = (CUdeviceptr*)buffers;
+        dsts = local_addrs;
         srcs = mapped_addrs;
     }
 
-    key    = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[0];
+    key    = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[first];
     status = uct_cuda_ipc_get_stream_and_event(iface, key->stream_id,
                                                &q_desc, &stream,
                                                &cuda_ipc_event);
@@ -377,7 +394,7 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     }
 
     status = UCT_CUDADRV_FUNC_LOG_ERR(
-            cuMemcpyBatchAsync(dsts, srcs, (size_t *)lengths, count, &attr,
+            cuMemcpyBatchAsync(dsts, srcs, copy_lengths, mapping->count, &attr,
                                &attrs_idx, 1, *stream));
     if (ucs_unlikely(status != UCS_OK)) {
         ucs_mpool_put(cuda_ipc_event);

test/gtest/uct/v2/test_uct_sgl.cc

@@ -130,6 +130,34 @@ class test_uct_sgl : public uct_test {
         test_sgl(std::vector<size_t>(count, elem_size(UCS_KBYTE)), &comp);
     }
 
+    void test_sgl_zero_length(const std::vector<size_t> &empty_indices)
+    {
+        static const size_t count = 3;
+        sgl_completion comp       = {};
+
+        connect_or_skip();
+
+        if (max_count() < count) {
+            UCS_TEST_SKIP_R("not enough sgl elements");
+        }
+
+        if (elem_size(0) > 0) {
+            UCS_TEST_SKIP_R("zero length is not supported");
+        }
+
+        comp.uct.func   = completion_cb;
+        comp.uct.count  = 1;
+        comp.uct.status = UCS_OK;
+
+        sgl_arrays sgl(*m_sender, *m_receiver, first_supported_mem_type(), op(),
+                       std::vector<size_t>(count, elem_size(UCS_KBYTE)));
+        for (size_t index : empty_indices) {
+            sgl.lengths[index] = 0;
+        }
+
+        test_sgl(sgl, &comp);
+    }
+
 private:
     enum sgl_op_t {
         SGL_OP_PUT,
@@ -261,19 +289,19 @@ class test_uct_sgl : public uct_test {
     void check_sgl(const sgl_arrays &sgl)
     {
         for (size_t i = 0; i < sgl.lengths.size(); ++i) {
+            /* The target of a zero-length element keeps its initial data */
+            uint64_t seed = (sgl.lengths[i] == 0) ? SEED2 : (SEED1 + i);
+
             if (op() == SGL_OP_PUT) {
-                sgl.remote[i]->pattern_check(SEED1 + i);
+                sgl.remote[i]->pattern_check(seed);
             } else {
-                sgl.local[i]->pattern_check(SEED1 + i);
+                sgl.local[i]->pattern_check(seed);
             }
         }
     }
 
-    void test_sgl(const std::vector<size_t> &sizes,
-                  sgl_completion *comp = NULL)
+    void test_sgl(const sgl_arrays &sgl, sgl_completion *comp = NULL)
     {
-        sgl_arrays sgl(*m_sender, *m_receiver, first_supported_mem_type(),
-                       op(), sizes);
         ucs_status_t status = sgl_op(sgl, (comp == NULL) ? NULL : &comp->uct);
 
         ASSERT_UCS_OK_OR_INPROGRESS(status);
@@ -285,9 +313,24 @@ class test_uct_sgl : public uct_test {
         }
 
         m_sender->flush();
+
+        if ((comp != NULL) && (status == UCS_OK)) {
+            /* Completion is not called for an operation completed in place */
+            EXPECT_EQ(0u, comp->done);
+        }
+
         check_sgl(sgl);
     }
 
+    void test_sgl(const std::vector<size_t> &sizes,
+                  sgl_completion *comp = NULL)
+    {
+        sgl_arrays sgl(*m_sender, *m_receiver, first_supported_mem_type(),
+                       op(), sizes);
+
+        test_sgl(sgl, comp);
+    }
+
     static void completion_cb(uct_completion_t *self)
     {
         ucs_container_of(self, sgl_completion, uct)->done = 1;
@@ -327,4 +370,12 @@ UCS_TEST_P(test_uct_sgl, with_callback) {
     test_sgl_with_callback();
 }
 
+UCS_TEST_P(test_uct_sgl, zero_length_element) {
+    test_sgl_zero_length({1});
+}
+
+UCS_TEST_P(test_uct_sgl, zero_length_all_elements) {
+    test_sgl_zero_length({0, 1, 2});
+}
+
 UCT_INSTANTIATE_TEST_CASE(test_uct_sgl)