UCP: SGL datatype implementation by michal-shalev · Pull Request #11344 · openucx/ucx

michal-shalev · 2026-04-15T11:51:29Z

What?

UCP SGL (scatter–gather list) datatype: iterators, packing, RMA put integration (proto / put offload), UCT v2 iface capability bits, and gtests.

Why?

Lets callers pass multi-buffer local and remote descriptors for one-sided put in a first-class way.

How?

dt_sgl + extensions to the datatype iterator and proto selection for SGL.
Put path hooks in put_offload / RMA send and UCT uct_iface_v2 fields for max SGL put segments / zcopy.
Tests: invalid-param coverage, optional zcopy functional tests gated on UCT_IFACE_FLAG_V2_PUT_SGL_ZCOPY.

…tialized in put/sgl/offload

brminich · 2026-04-17T09:10:25Z

+
+    iface_attr_v2.field_mask = UCT_IFACE_ATTR_FIELD_MAX_PUT_SGL_ZCOPY_COUNT;
+    status                   = uct_iface_query_v2(
+            ucp_worker_iface(ep->worker,


max_put_sgl_zcopy_count is queried on every send invocation but is static post-init. Consider caching it in the protocol private data during probe, the same way max_iov is handled in other protocols.

max_put_sgl_zcopy_count is now cached in ucp_proto_multi_lane_priv_t.

brminich · 2026-04-17T09:10:25Z

+        /* Multiple progress calls, translate memh + rkey per-chunk on stack */
+        uct_mem_h *uct_memhs;
+        uct_rkeys = ucs_alloca(elem_count * sizeof(uct_rkey_t));
+


ucs_alloca asserts if the allocation exceeds ~1KB, which fires at ~128 elements (8 bytes × 128). Consider a small fixed threshold and fall back to ucs_malloc for larger counts — similar to what the SIZE_MAX branch below already does.

is it safe to use alloca considering that request may end up in pending queue?

For ucs_alloca with fallback to malloc you can use ucs_alloc_on_stack (and free with ucs_free_on_stack)

brminich · 2026-04-17T09:10:25Z

+        ucp_datatype_iter_detect_mem_info(context, local->buffers[0],
+                                          local->lengths[0], dt_iter, param);
+        if (ENABLE_PARAMS_CHECK && (count > 1)) {
+            status = ucp_dt_sgl_memtype_check(context, local, count,


ucp_dt_sgl_memtype_check is only called under ENABLE_PARAMS_CHECK. In release builds, mixed HOST+GPU buffers pass through silently. Worth noting in the API docs at minimum.

This check matches the pre-existing IOV path (ucp_datatype_iov_iter_init).
IOV doesn't document this.
I added a documentation for it (see ucp_dt_local_sgl_t).

brminich · 2026-04-17T09:10:25Z

        break;
+    case UCP_DATATYPE_SGL:
+        if (!ucp_memh_is_buffer_in_range(memh, dt_iter->type.sgl.buffers[0],
+                                         dt_iter->type.sgl.lengths[0])) {


Only buffers[0] is range-checked against the user memh. If the intent is one memh per SGL element, this should loop over all entries.

ucp_datatype_iter_is_user_memh_valid signature only takes a single memh, so I had to check it from dt_iter.
Fixed + added tests.

brminich · 2026-04-17T09:12:33Z

    }

    UCP_REQUEST_CHECK_PARAM(param);
+    UCP_REQUEST_CHECK_PARAM_UNSUPPORTED_REMOTE(param);


can we unite it with UCP_REQUEST_CHECK_PARAM?

brminich · 2026-04-17T09:14:26Z

    return dst_iov_index;
 }

+ucs_status_t ucp_datatype_sgl_iter_init(ucp_context_h context,


seems the common naming is ucp_datatype_iter_*, lets rename accordingly
ucp_datatype_iov_iter_init is the only exception

Renamed to ucp_datatype_iter_sgl_init.

brminich · 2026-04-17T09:21:29Z

+            ucp_mem_h        *memhs;
+            const uint64_t   *remote_addrs;
+            ucp_rkey_h const *rkeys;
+            int              memhs_owned;


can we reuse ucp_memh_is_user_memh somehow instead of adding this field?

brminich · 2026-04-17T09:28:01Z


-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT)) ||
+        (init_params->select_param->dt_class == UCP_DATATYPE_SGL)) {


should we also restrict selecting all other protocols for SGL datatype?

brminich · 2026-04-17T09:34:26Z

+    elem_count = ucp_datatype_iter_next_sgl(dt_iter, max_sgl_count, next_iter);
+
+    if (max_sgl_count < SIZE_MAX) {
+        /* Multiple progress calls, translate memh + rkey per-chunk on stack */


why multiple?

Because for cuda_ipc it's a single progress call.
In any case, I simplified the logic here.

brminich · 2026-04-17T09:36:14Z

+        /* Multiple progress calls, translate memh + rkey per-chunk on stack */
+        uct_mem_h *uct_memhs;
+        uct_rkeys = ucs_alloca(elem_count * sizeof(uct_rkey_t));
+


is it safe to use alloca considering that request may end up in pending queue?

brminich · 2026-04-17T09:36:47Z

+                     uct_ep,
+                     &dt_iter->type.sgl.buffers[start_index],
+                     &dt_iter->type.sgl.lengths[start_index],
+                     NULL,


why no memh here?

cuda_ipc doesn't use memh.

Now that I think of this again, I decided to always fill the uct_memhs array, mirroring the regular zcopy path, and to make UCP/UCT layering cleaner.

guy-ealey-morag · 2026-04-20T07:03:44Z

+                    goto out_unlock;
+                }
+                remote = (const ucp_dt_remote_sgl_t *)param->remote;
+                rkey   = remote->rkeys[0];


This assumes count > 0, should we guard against it?

brminich · 2026-04-24T11:56:03Z

+
+#define UCP_REQUEST_CHECK_PARAM(_param) \
+    do { \
+        UCP_REQUEST_CHECK_PARAM_ALLOW_REMOTE(_param); \


seems naming is wrong, as checking remote is below in this macro.
also better to omit ALLOW in the name

Why wrong?
The _ALLOW_REMOTE suffix is meant to say "this variant allows remote/SGL params" (because it doesn't run the rejection check that UCP_REQUEST_CHECK_PARAM adds).
I renamed it to UCP_REQUEST_CHECK_PARAM_COMMON to make it clearer.

brminich · 2026-04-24T11:56:34Z

    void          *request = NULL;
    ucp_request_t *close_req;

+    UCP_REQUEST_CHECK_PARAM(param);


i'd not add it here as it is not relevant for ep_close

brminich · 2026-04-24T12:14:06Z

+        ucp_datatype_iter_detect_mem_info(context, local->buffers[0],
+                                          local->lengths[0], dt_iter, param);
+        if (ENABLE_PARAMS_CHECK && (count > 1)) {
+            status = ucp_dt_sgl_check_same_mem_info(context, local, count,


maybe pass local->buffers and local->lengths separately? Like

Suggested change

status = ucp_dt_sgl_check_same_mem_info(context, local, count,

status = ucp_dt_sgl_check_same_mem_info(context, local->buffers +1 , local->lengths +1 count - 1,

brminich · 2026-04-24T12:50:42Z

 {
    void *request;

+    UCP_REQUEST_CHECK_PARAM(param);


it is not needed here

brminich · 2026-04-24T12:50:48Z

 {
    void *request;

+    UCP_REQUEST_CHECK_PARAM(param);


it is not needed here

…E helpers

…tics

redbrick9 · 2026-05-25T05:59:27Z

@michal-shalev will ucx_perftest test case ucp_put_bw support sge datatype?

michal-shalev · 2026-05-25T11:30:29Z

will ucx_perftest test case ucp_put_bw support sge datatype?

Support for SGL datatype in ucx_perftest is planned, but it is outside the scope of this PR and will be addressed separately in the future.

redbrick9 · 2026-05-27T01:01:21Z

will ucx_perftest test case ucp_put_bw support sge datatype?

Support for SGL datatype in ucx_perftest is planned, but it is outside the scope of this PR and will be addressed separately in the future.

Got it. Thanks.

iyastreb · 2026-05-27T07:12:15Z

+            if (ucs_unlikely(status != UCS_OK)) {
+                ret = UCS_STATUS_PTR(status);
+                goto out_unlock;
+            }


newline after block missing

iyastreb · 2026-05-27T07:12:35Z

+                ret = UCS_STATUS_PTR(status);
+                goto out_unlock;
+            }
+            remote = (const ucp_dt_remote_sgl_t *)param->remote;


casting from void * is redundant

iyastreb · 2026-05-27T07:16:01Z

+
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_put_sgl_check_params(const void *buffer, size_t count,


Should this be a macro, so that we don't have extra if branch in ucp_put_nbx?
Another alternative is to call this function only when ENABLE_PARAMS_CHECK is set, so that in release mode it's never invoked:
ucp_put_nbx:

if (ENABLE_PARAMS_CHECK) { status = ucp_put_sgl_check_params(...) }

it is already no-op if ENABLE_PARAMS is not set (see branch below)

Changed to a macro to match existing RMA check style as a consistency change, not a performance fix.

iyastreb · 2026-05-27T07:35:06Z

+        uct_memhs[i] = (sgl_memhs != NULL) ?
+                       sgl_memhs[start_index + i]->uct[md_index] :
+                       UCT_MEM_HANDLE_NULL;
+        uct_rkeys[i] = ucp_rkey_get_tl_rkey(


Maybe we can validate rkey_index once, and then just assign in a loop?

if (rkey_index == UCP_NULL_RESOURCE) { for (i = 0; i < elem_count; i++) { uct_rkeys[i] = UCT_INVALID_RKEY; uct_memhs[i] = (sgl_memhs != NULL) ? sgl_memhs[start_index + i]->uct[md_index] : UCT_MEM_HANDLE_NULL; } } else { for (i = 0; i < elem_count; i++) { uct_rkeys[i] = dt_iter->type.sgl.rkeys[start_index + i]->tl_rkey[rkey_index].rkey.rkey; uct_memhs[i] = (sgl_memhs != NULL) ? sgl_memhs[start_index + i]->uct[md_index] : UCT_MEM_HANDLE_NULL; } }

Is this what you had in mind?

It would duplicate the memh lines, and since ucp_rkey_get_tl_rkey() is inline, rkey_index is the same every iteration, and elem_count is bounded by max_put_sgl_zcopy_count, so I'd expect the compiler to already optimize the current version.

brminich · 2026-05-29T13:21:47Z

+
+
+static UCS_F_ALWAYS_INLINE ucs_status_t
+ucp_put_sgl_check_params(const void *buffer, size_t count,


it is already no-op if ENABLE_PARAMS is not set (see branch below)

…engths

UCP: SGL datatype implementation

066ab45

michal-shalev requested review from brminich and yosefe April 15, 2026 11:51

michal-shalev self-assigned this Apr 15, 2026

michal-shalev added the Ready for Review label Apr 15, 2026

guy-ealey-morag reviewed Apr 15, 2026

View reviewed changes

Comment thread src/uct/api/v2/uct_v2.h

Comment thread test/gtest/ucp/ucp_test.cc Outdated

Comment thread src/ucp/proto/proto_common.c Outdated

Comment thread src/ucp/rma/put_offload.c Outdated

Comment thread src/uct/api/v2/uct_v2.h Outdated

guy-ealey-morag reviewed Apr 15, 2026

View reviewed changes

Comment thread src/ucp/dt/dt_sgl.c Outdated

guy-ealey-morag reviewed Apr 15, 2026

View reviewed changes

Comment thread src/ucp/dt/dt_sgl.h Outdated

guy-ealey-morag reviewed Apr 15, 2026

View reviewed changes

Comment thread test/gtest/ucp/test_ucp_rma.cc

Comment thread src/ucp/dt/datatype_iter.c Outdated

michal-shalev added 9 commits April 15, 2026 13:56

UCT: Remove 2x uct_ep_put_sgl_zcopy

c5afa20

UCP/UCT: Update copyrights

31339b6

UCP: Fix copyrights for new files

2a98569

UCP: Check return status of uct_iface_query_v2

731a918

UCT: Fix uct_iface_attr_v2_t layout and align v2 iface cap docs

3455df1

GTEST/UCP: Drop v2 lane caps query + add negative test put_zero_count

772f2b1

UCP/RMA: initialize next_iter before iface query to fix -Wmaybe-unini…

57ddf98

…tialized in put/sgl/offload

GTEST/UCP: Remove copyrights fix in files that didn't change

4c3114a

GTEST/UCP: Add mixed mem types test

24b9d43

michal-shalev force-pushed the ucp-sgl-put-impl branch from dc504a0 to 24b9d43 Compare April 15, 2026 20:56

michal-shalev requested a review from guy-ealey-morag April 15, 2026 21:01

michal-shalev force-pushed the ucp-sgl-put-impl branch from 0d9e37d to 4c64772 Compare April 16, 2026 10:40

UCP: Add UCP_REQUEST_CHECK_PARAM_UNSUPPORTED_REMOTE

7598a51

michal-shalev force-pushed the ucp-sgl-put-impl branch from 4c64772 to 7598a51 Compare April 16, 2026 11:41

michal-shalev added 2 commits April 16, 2026 07:09

UCP: Remove copyrights years fix for non-NVIDIA

6b762c2

UCP/DT: Remove ucp_datatype_iter_sgl_allocate_memh

d19b8fb

brminich reviewed Apr 17, 2026

View reviewed changes

michal-shalev added 3 commits April 17, 2026 07:30

UCP/PROTO: cache max_put_sgl_zcopy_count

6ce20ad

UCP/RMA/PUT: simplify put_sgl_offload_send_func

979021d

UCP/DT: Rename ucp_datatype_sgl_iter_init to ucp_datatype_iter_sgl_init

ef2a72e

UCP/PROTO: Fix Coverity overrun warning in proto select loop

6934aa2

michal-shalev force-pushed the ucp-sgl-put-impl branch from 906a1f9 to 6934aa2 Compare April 19, 2026 12:59

guy-ealey-morag reviewed Apr 20, 2026

View reviewed changes

brminich reviewed Apr 24, 2026

View reviewed changes

michal-shalev added 3 commits April 24, 2026 07:38

UCP/REQUEST: Split UCP_REQUEST_CHECK_PARAM into _COMMON and _NO_REMOT…

f01b9d6

…E helpers

UCP/CORE: Skip param validation for ops with no datatype/remote seman…

880cdc0

…tics

UCP/DT: Pass buffers/lengths directly to ucp_dt_sgl_check_same_mem_info

684a494

michal-shalev requested a review from brminich April 24, 2026 15:59

michal-shalev added 2 commits April 26, 2026 07:36

Merge branch 'master' into ucp-sgl-put-impl

d870e0a

GTEST/UCP: Skip SGL put tests when v2 put_sgl_zcopy cap is missing

3b1da4d

openucx deleted a comment from svc-nixl May 4, 2026

iyastreb reviewed May 20, 2026

View reviewed changes

Comment thread src/ucp/core/ucp_request.inl Outdated

Comment thread src/ucp/core/ucp_rkey.c Outdated

Comment thread test/gtest/ucp/test_ucp_rma.cc

Comment thread test/gtest/ucp/ucp_test.cc

michal-shalev added 2 commits May 25, 2026 06:41

Merge branch 'master' into ucp-sgl-put-impl

ef6950a

UCP/RKEY: Remove extra spaces

12eefc5

michal-shalev requested a review from iyastreb May 25, 2026 14:09

iyastreb reviewed May 27, 2026

View reviewed changes

brminich reviewed May 29, 2026

View reviewed changes

michal-shalev added 2 commits May 31, 2026 12:28

UCP/RMA: Convert SGL put param validation to macro and check remote l…

4916e62

…engths

UCP/GTEST: Add negative test for missing remote SGL lengths in put

a360442

michal-shalev requested review from brminich and iyastreb June 1, 2026 00:07

iyastreb approved these changes Jun 1, 2026

View reviewed changes

brminich approved these changes Jun 1, 2026

View reviewed changes

brminich merged commit d65ffe7 into openucx:master Jun 1, 2026
160 checks passed

michal-shalev deleted the ucp-sgl-put-impl branch June 1, 2026 11:18

michal-shalev removed the Ready for Review label Jun 1, 2026

guy-ealey-morag mentioned this pull request Jun 8, 2026

UCP/RMA: Fix build error in master branch #11533

Merged

	status = ucp_dt_sgl_check_same_mem_info(context, local, count,
	status = ucp_dt_sgl_check_same_mem_info(context, local->buffers +1 , local->lengths +1 count - 1,



		static UCS_F_ALWAYS_INLINE ucs_status_t
		ucp_put_sgl_check_params(const void *buffer, size_t count,

Conversation

michal-shalev commented Apr 15, 2026 • edited by yosefe Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What?

Why?

How?

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

michal-shalev commented Apr 15, 2026 •

edited by yosefe

Loading