UCT/CUDA_IPC: Implement uct_ep_put_sgl_zcopy

src/uct/cuda/cuda_ipc/cuda_ipc_ep.c

@@ -99,6 +99,60 @@ static UCS_F_ALWAYS_INLINE ucs_status_t uct_cuda_ipc_ctx_rsc_get(
     return UCS_OK;
 }
 
+static UCS_F_ALWAYS_INLINE ucs_status_t
+uct_cuda_ipc_get_stream_and_event(uct_cuda_ipc_iface_t *iface,
+                                  unsigned stream_id,
+                                  uct_cuda_queue_desc_t **q_desc_p,
+                                  CUstream **stream_p,
+                                  uct_cuda_ipc_event_desc_t **event_p)
+{
+    uct_cuda_ipc_ctx_rsc_t *ctx_rsc;
+    ucs_status_t status;
+
+    status = uct_cuda_ipc_ctx_rsc_get(iface, &ctx_rsc);
+    if (ucs_unlikely(status != UCS_OK)) {
+        return status;
+    }
+
+    *q_desc_p = &ctx_rsc->queue_desc[stream_id % iface->config.max_streams];
+    *stream_p = &(*q_desc_p)->stream;
+
+    status = uct_cuda_base_init_stream(*stream_p);
+    if (ucs_unlikely(status != UCS_OK)) {
+        return status;
+    }
+
+    *event_p = uct_cuda_base_event_desc_mpool_get(&ctx_rsc->super.event_mp);
+    if (ucs_unlikely(*event_p == NULL)) {
+        ucs_error("Failed to allocate cuda_ipc event object");
+        return UCS_ERR_NO_MEMORY;
+    }
+
+    return UCS_OK;
+}
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+uct_cuda_ipc_event_record_and_enqueue(uct_cuda_ipc_iface_t *iface,
+                                      uct_cuda_queue_desc_t *q_desc,
+                                      uct_cuda_ipc_event_desc_t *event,
+                                      CUstream stream)
+{
+    ucs_status_t status;
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(cuEventRecord(event->super.event,
+                                                    stream));
+    if (ucs_unlikely(status != UCS_OK)) {
+        return status;
+    }
+
+    if (ucs_queue_is_empty(&q_desc->event_queue)) {
+        ucs_queue_push(&iface->super.active_queue, &q_desc->queue);
+    }
+
+    ucs_queue_push(&q_desc->event_queue, &event->super.queue);
+    return UCS_OK;
+}
+
 static UCS_F_ALWAYS_INLINE ucs_status_t
 uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
                                   const uct_iov_t *iov, uct_rkey_t rkey,
@@ -112,11 +166,9 @@ uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
     void *mapped_rem_addr;
     void *mapped_addr;
     uct_cuda_ipc_event_desc_t *cuda_ipc_event;
-    uct_cuda_ipc_ctx_rsc_t *ctx_rsc;
     uct_cuda_queue_desc_t *q_desc;
     ucs_status_t status;
     CUdeviceptr dst, src;
-    CUcontext UCS_V_UNUSED cuda_context;
     CUstream *stream;
 
     if (ucs_unlikely(0 == iov[0].length)) {
@@ -137,34 +189,13 @@ uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
         goto out;
     }
 
-    status = uct_cuda_ipc_ctx_rsc_get(iface, &ctx_rsc);
+    status = uct_cuda_ipc_get_stream_and_event(iface, key->stream_id,
+                                               &q_desc, &stream,
+                                               &cuda_ipc_event);
     if (ucs_unlikely(status != UCS_OK)) {
         goto out;
     }
 
-    /* round-robin */
-    q_desc = &ctx_rsc->queue_desc[key->stream_id % iface->config.max_streams];
-    stream = &q_desc->stream;
-    status = uct_cuda_base_init_stream(stream);
-    if (ucs_unlikely(status != UCS_OK)) {
-        goto out;
-    }
-
-    if (ucs_unlikely(stream == NULL)) {
-        ucs_error("stream=%d for dev_num=%d not available", key->stream_id,
-                  cuda_device);
-        status = UCS_ERR_IO_ERROR;
-        goto out;
-    }
-
-    cuda_ipc_event = uct_cuda_base_event_desc_mpool_get(
-            &ctx_rsc->super.event_mp);
-    if (ucs_unlikely(cuda_ipc_event == NULL)) {
-        ucs_error("Failed to allocate cuda_ipc event object");
-        status = UCS_ERR_NO_MEMORY;
-        goto out;
-    }
-
     dst = (CUdeviceptr)
         ((direction == UCT_CUDA_IPC_PUT) ? mapped_rem_addr : iov[0].buffer);
     src = (CUdeviceptr)
@@ -177,24 +208,22 @@ uct_cuda_ipc_post_cuda_async_copy(uct_ep_h tl_ep, uint64_t remote_addr,
         goto out;
     }
 
-    status = UCT_CUDADRV_FUNC_LOG_ERR(cuEventRecord(cuda_ipc_event->super.event,
-                                                    *stream));
-    if (UCS_OK != status) {
+    status = uct_cuda_ipc_event_record_and_enqueue(iface, q_desc,
+                                                   cuda_ipc_event, *stream);
+    if (ucs_unlikely(status != UCS_OK)) {
         ucs_mpool_put(cuda_ipc_event);
         goto out;
     }
 
-    if (ucs_queue_is_empty(&q_desc->event_queue)) {
-        ucs_queue_push(&iface->super.active_queue, &q_desc->queue);
-    }
-
-    ucs_queue_push(&q_desc->event_queue, &cuda_ipc_event->super.queue);
     cuda_ipc_event->super.comp  = comp;
     cuda_ipc_event->mapped_addr = mapped_addr;
     cuda_ipc_event->d_bptr      = (uintptr_t)key->super.super.d_bptr;
     cuda_ipc_event->pid         = key->super.super.pid;
     cuda_ipc_event->pid_ns      = key->super.pid_ns;
     cuda_ipc_event->cuda_device = cuda_device;
+#if CUDA_VERSION >= 13000
+    cuda_ipc_event->sgl_mapping = NULL;
+#endif
     ucs_trace("cuMemcpyDtoDAsync issued :%p dst:%p, src:%p  len:%ld",
              cuda_ipc_event, (void *) dst, (void *) src, iov[0].length);
     status = UCS_INPROGRESS;
@@ -246,6 +275,140 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_ep_put_zcopy,
     return status;
 }
 
+#if CUDA_VERSION >= 13000
+UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_ep_put_sgl_zcopy,
+                 (tl_ep, buffers, lengths, memhs, remote_addrs, rkeys, counts,
+                  strides, count, comp),
+                 uct_ep_h tl_ep, void * const *buffers,
+                 const size_t *lengths, uct_mem_h const *memhs UCS_V_UNUSED,
+                 const uint64_t *remote_addrs, uct_rkey_t const *rkeys,
+                 const size_t *counts, const size_t *strides,
+                 size_t count, uct_completion_t *comp)
+{
+    uct_cuda_ipc_iface_t *iface = ucs_derived_of(tl_ep->iface,
+                                                 uct_cuda_ipc_iface_t);
+    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_queue_desc_t *q_desc;
+    CUmemcpyAttributes attr;
+    void *mapped_rem_addr, *mapped_addr;
+    CUdeviceptr *cuda_dsts;
+    size_t attrs_idx;
+    size_t i, total_length;
+    CUdevice cuda_device;
+    CUstream *stream;
+    ucs_status_t status;
+    int is_ctx_pushed;
+
+    /* TODO: add strided elements support */
+    if (ucs_unlikely((counts != NULL) || (strides != NULL))) {
+        ucs_error("cuda_ipc put_sgl_zcopy does not support strided elements");
+        return UCS_ERR_UNSUPPORTED;
+    }
+
+    total_length = 0;
+    for (i = 0; i < count; i++) {
+        total_length += lengths[i];
+    }
+
+    if (ucs_unlikely(total_length == 0)) {
+        ucs_trace_data("Zero length put_sgl_zcopy: skip it");
+        return UCS_OK;
+    }
+
+    status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[0],
+                                             &cuda_device, &is_ctx_pushed);
+    if (ucs_unlikely(status != UCS_OK)) {
+        return status;
+    }
+
+    mapping = ucs_malloc(sizeof(*mapping) +
+                         count * (sizeof(*mapping->entries) +
+                                  sizeof(CUdeviceptr)),
+                         "cuda_ipc_sgl_mapping");
+    if (ucs_unlikely(mapping == NULL)) {
+        ucs_error("Failed to allocate cuda_ipc sgl mapping");
+        status = UCS_ERR_NO_MEMORY;
+        goto out_ctx;
+    }
+
+    mapping->count   = 0;
+    mapping->entries = (uct_cuda_ipc_sgl_entry_t *)(mapping + 1);
+    cuda_dsts        = (CUdeviceptr *)(mapping->entries + count);
+
+    memset(&attr, 0, sizeof(attr));
+    attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
+    attrs_idx           = 0;
+
+    for (i = 0; i < count; i++) {
+        key = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[i];
+
+        status = uct_cuda_ipc_get_remote_address(&key->super, remote_addrs[i],
+                                                 cuda_device, &mapped_rem_addr,
+                                                 &mapped_addr);
+        if (ucs_unlikely(status != UCS_OK)) {
+            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_addr;
+        cuda_dsts[i]                    = (CUdeviceptr)mapped_rem_addr;
+        mapping->count++;
+    }
+
+    key    = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[0];
+    status = uct_cuda_ipc_get_stream_and_event(iface, key->stream_id,
+                                               &q_desc, &stream,
+                                               &cuda_ipc_event);
+    if (ucs_unlikely(status != UCS_OK)) {
+        goto out_unmap;
+    }
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuMemcpyBatchAsync(cuda_dsts, (CUdeviceptr *)buffers,
+                               (size_t *)lengths, count, &attr, &attrs_idx, 1,
+                               *stream));
+    if (ucs_unlikely(status != UCS_OK)) {
+        ucs_mpool_put(cuda_ipc_event);
+        goto out_unmap;
+    }
+
+    status = uct_cuda_ipc_event_record_and_enqueue(iface, q_desc,
+                                                   cuda_ipc_event, *stream);
+    if (ucs_unlikely(status != UCS_OK)) {
+        ucs_mpool_put(cuda_ipc_event);
+        goto out_unmap;
+    }
+
+    cuda_ipc_event->super.comp  = comp;
+    cuda_ipc_event->sgl_mapping = mapping;
+    cuda_ipc_event->mapped_addr = NULL;
+    cuda_ipc_event->d_bptr      = 0;
+    cuda_ipc_event->pid         = 0;
+    cuda_ipc_event->cuda_device = cuda_device;
+
+    ucs_trace("cuMemcpyBatchAsync issued: count=%zu total_length=%zu",
+              count, total_length);
+
+    UCT_TL_EP_STAT_OP(ucs_derived_of(tl_ep, uct_base_ep_t), PUT, ZCOPY,
+                      total_length);
+
+    status = UCS_INPROGRESS;
+    goto out_ctx;
+
+out_unmap:
+    uct_cuda_ipc_sgl_mapping_destroy(mapping, cuda_device,
+                                     iface->config.enable_cache);
+
+out_ctx:
+    uct_cuda_ipc_check_and_pop_ctx(is_ctx_pushed);
+    return status;
+}
+#endif /* CUDA_VERSION >= 13000 */
+
 ucs_status_t uct_cuda_ipc_ep_get_device_ep(uct_ep_h tl_ep,
                                            uct_device_ep_h *device_ep_p)
 {

src/uct/cuda/cuda_ipc/cuda_ipc_ep.h

@@ -33,6 +33,17 @@ ucs_status_t uct_cuda_ipc_ep_put_zcopy(uct_ep_h tl_ep,
                                        uint64_t remote_addr, uct_rkey_t rkey,
                                        uct_completion_t *comp);
 
+ucs_status_t uct_cuda_ipc_ep_put_sgl_zcopy(uct_ep_h tl_ep,
+                                           void * const *buffers,
+                                           const size_t *lengths,
+                                           uct_mem_h const *memhs,
+                                           const uint64_t *remote_addrs,
+                                           uct_rkey_t const *rkeys,
+                                           const size_t *counts,
+                                           const size_t *strides,
+                                           size_t count,
+                                           uct_completion_t *comp);
+
 int uct_cuda_ipc_ep_is_connected(const uct_ep_h tl_ep,
                                  const uct_ep_is_connected_params_t *params);
 

src/uct/cuda/cuda_ipc/cuda_ipc_iface.c

@@ -334,6 +334,15 @@ static void uct_cuda_ipc_complete_event(uct_iface_h tl_iface,
     uct_cuda_ipc_event_desc_t *cuda_ipc_event = ucs_derived_of(cuda_event,
                                                                uct_cuda_ipc_event_desc_t);
 
+#if CUDA_VERSION >= 13000
+    if (cuda_ipc_event->sgl_mapping != NULL) {
+        uct_cuda_ipc_sgl_mapping_destroy(cuda_ipc_event->sgl_mapping,
+                                         cuda_ipc_event->cuda_device,
+                                         iface->config.enable_cache);
+        return;
+    }
+#endif
+
     uct_cuda_ipc_unmap_memhandle(cuda_ipc_event->pid, cuda_ipc_event->pid_ns,
                                  cuda_ipc_event->d_bptr,
                                  cuda_ipc_event->mapped_addr,
@@ -423,16 +432,44 @@ uct_cuda_ipc_estimate_perf(uct_iface_h tl_iface, uct_perf_attr_t *perf_attr)
     return UCS_OK;
 }
 
+static ucs_status_t
+uct_cuda_ipc_iface_query_v2(uct_iface_h iface, uct_iface_attr_v2_t *iface_attr)
+{
+    ucs_status_t status;
+
+    status = uct_iface_base_query_v2(iface, iface_attr);
+    if (status != UCS_OK) {
+        return status;
+    }
+
+#if CUDA_VERSION >= 13000
+    if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_CAP_FLAGS) {
+        iface_attr->cap.flags |= UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY;
+    }
+
+    if (iface_attr->field_mask & UCT_IFACE_ATTR_FIELD_MAX_PUT_SGL_ZCOPY_COUNT) {
+        iface_attr->max_put_sgl_zcopy_count = SIZE_MAX;
+    }
+#endif
+
+    return UCS_OK;
+}
+
 static uct_iface_internal_ops_t uct_cuda_ipc_iface_internal_ops = {
-    .iface_query_v2         = uct_iface_base_query_v2,
+    .iface_query_v2         = uct_cuda_ipc_iface_query_v2,
     .iface_estimate_perf    = uct_cuda_ipc_estimate_perf,
     .iface_vfs_refresh      = (uct_iface_vfs_refresh_func_t)ucs_empty_function,
     .ep_query               = (uct_ep_query_func_t)ucs_empty_function_return_unsupported,
     .ep_invalidate          = (uct_ep_invalidate_func_t)ucs_empty_function_return_unsupported,
     .ep_connect_to_ep_v2    = (uct_ep_connect_to_ep_v2_func_t)ucs_empty_function_return_unsupported,
     .iface_is_reachable_v2  = uct_cuda_ipc_iface_is_reachable_v2,
     .ep_is_connected        = uct_cuda_ipc_ep_is_connected,
-    .ep_get_device_ep       = uct_cuda_ipc_ep_get_device_ep
+    .ep_get_device_ep       = uct_cuda_ipc_ep_get_device_ep,
+#if CUDA_VERSION >= 13000
+    .ep_put_sgl_zcopy       = uct_cuda_ipc_ep_put_sgl_zcopy
+#else
+    .ep_put_sgl_zcopy       = (uct_ep_put_sgl_zcopy_func_t)ucs_empty_function_return_unsupported
+#endif
 };
 
 static uct_cuda_ctx_rsc_t * uct_cuda_ipc_ctx_rsc_create(uct_iface_h tl_iface)