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, num_nonempty, total_length;
     CUdevice cuda_device;
     CUstream *stream;
     ucs_status_t status;
@@ -308,8 +310,10 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     }
 
     total_length = 0;
+    num_nonempty = 0;
     for (i = 0; i < count; i++) {
         total_length += lengths[i];
+        num_nonempty += (lengths[i] != 0);
     }
 
     if (ucs_unlikely(total_length == 0)) {
@@ -318,15 +322,26 @@ 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],
+    /* The context and the stream are taken from the first non-empty element,
+     * which exists because total_length != 0 */
+    first = 0;
+    while (lengths[first] == 0) {
+        first++;
+    }
+
+    status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[first],
                                              &cuda_device, &is_ctx_pushed);
     if (ucs_unlikely(status != UCS_OK)) {
         return status;
     }
 
+    /* Only the elements with data go to the batch, so their local addresses
+     * and lengths are stored in the mapping next to the remote addresses */
     mapping = ucs_malloc(sizeof(*mapping) +
-                         count * (sizeof(*mapping->entries) +
-                                  sizeof(CUdeviceptr)),
+                         num_nonempty * (sizeof(*mapping->entries) +
+                                         sizeof(*mapped_addrs) +
+                                         sizeof(*local_addrs) +
+                                         sizeof(*copy_lengths)),
                          "cuda_ipc_sgl_mapping");
     if (ucs_unlikely(mapping == NULL)) {
         ucs_error("Failed to allocate cuda_ipc sgl mapping");
@@ -336,13 +351,20 @@ 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);
+    mapped_addrs     = (CUdeviceptr *)(mapping->entries + num_nonempty);
+    local_addrs      = mapped_addrs + num_nonempty;
+    copy_lengths     = (size_t *)(local_addrs + num_nonempty);
 
     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 (ucs_unlikely(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 +374,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 +402,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);
@@ -399,7 +424,7 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
     cuda_ipc_event->cuda_device = cuda_device;
 
     ucs_trace("cuMemcpyBatchAsync issued: count=%zu total_length=%zu",
-              count, total_length);
+              mapping->count, total_length);
 
     *total_length_p = total_length;
     status          = UCS_INPROGRESS;

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 (min_zcopy() > 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,
@@ -213,17 +241,21 @@ class test_uct_sgl : public uct_test {
                sgl_iface_attr().max_get_sgl_zcopy_count;
     }
 
+    size_t min_zcopy()
+    {
+        return (op() == SGL_OP_PUT) ?
+               m_sender->iface_attr().cap.put.min_zcopy :
+               m_sender->iface_attr().cap.get.min_zcopy;
+    }
+
     size_t elem_size(size_t size)
     {
         const uct_iface_attr &attr = m_sender->iface_attr();
-        size_t min_zcopy           = (op() == SGL_OP_PUT) ?
-                                     attr.cap.put.min_zcopy :
-                                     attr.cap.get.min_zcopy;
         size_t max_zcopy           = (op() == SGL_OP_PUT) ?
                                      attr.cap.put.max_zcopy :
                                      attr.cap.get.max_zcopy;
 
-        return ucs_min(ucs_max(size, min_zcopy), max_zcopy);
+        return ucs_min(ucs_max(size, min_zcopy()), max_zcopy);
     }
 
     ucs_memory_type_t first_supported_mem_type()
@@ -261,19 +293,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 +317,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 +374,16 @@ 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_first_last_elements) {
+    test_sgl_zero_length({0, 2});
+}
+
+UCS_TEST_P(test_uct_sgl, zero_length_all_elements) {
+    test_sgl_zero_length({0, 1, 2});
+}
+
 UCT_INSTANTIATE_TEST_CASE(test_uct_sgl)