diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index 6629640c002..7f96bd0a8f2 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -371,7 +371,7 @@ ucp_ep_peer_mem_get(ucp_context_h context, ucp_ep_h ep, uint64_t address, data->size = size; ucp_ep_rkey_unpack_internal(ep, rkey_buf, 0, UCS_BIT(rkey_ptr_md_index), 0, - UCS_SYS_DEVICE_ID_UNKNOWN, &data->rkey); + UCS_SYS_DEVICE_ID_UNKNOWN, 0, &data->rkey); rkey_index = ucs_bitmap2idx(data->rkey->md_map, rkey_ptr_md_index); status = uct_rkey_ptr(data->rkey->tl_rkey[rkey_index].cmpt, &data->rkey->tl_rkey[rkey_index].rkey, address, diff --git a/src/ucp/core/ucp_rkey.c b/src/ucp/core/ucp_rkey.c index e3dea958b66..54ee070d004 100644 --- a/src/ucp/core/ucp_rkey.c +++ b/src/ucp/core/ucp_rkey.c @@ -914,10 +914,11 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rkey_proto_resolve, UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, (ep, buffer, length, unpack_md_map, skip_md_map, sys_dev, - rkey_p), + uct_flags, rkey_p), ucp_ep_h ep, const void *buffer, size_t length, ucp_md_map_t unpack_md_map, ucp_md_map_t skip_md_map, - ucs_sys_device_t sys_dev, ucp_rkey_h *rkey_p) + ucs_sys_device_t sys_dev, uint64_t uct_flags, + ucp_rkey_h *rkey_p) { ucp_worker_h worker = ep->worker; const ucp_ep_config_t *ep_config = ucp_ep_config(ep); @@ -974,8 +975,10 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, rkey->ep = ep; #endif - unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE; + unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE | + UCT_RKEY_UNPACK_FIELD_FLAGS; unpack_params.sys_device = sys_dev; + unpack_params.flags = uct_flags; /* Go over remote MD indices and unpack rkey of each UCT MD */ rkey_index = 0; /* Index of the rkey in the array */ @@ -1056,18 +1059,34 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, return status; } -ucs_status_t ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer, - ucp_rkey_h *rkey_p) +static ucs_status_t +ucp_ep_rkey_unpack_flags(ucp_ep_h ep, const void *rkey_buffer, + uint64_t uct_flags, ucp_rkey_h *rkey_p) { ucs_status_t status; UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(ep->worker); - status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, rkey_p); + status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, uct_flags, + rkey_p); UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(ep->worker); return status; } +ucs_status_t +ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer, ucp_rkey_h *rkey_p) +{ + return ucp_ep_rkey_unpack_flags(ep, rkey_buffer, + UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING, + rkey_p); +} + +ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer, + ucp_rkey_h *rkey_p) +{ + return ucp_ep_rkey_unpack_flags(ep, rkey_buffer, 0, rkey_p); +} + void ucp_rkey_dump_packed(const void *buffer, size_t length, ucs_string_buffer_t *strb) { diff --git a/src/ucp/core/ucp_rkey.h b/src/ucp/core/ucp_rkey.h index 4b1402a8028..5c54b5c944f 100644 --- a/src/ucp/core/ucp_rkey.h +++ b/src/ucp/core/ucp_rkey.h @@ -231,11 +231,15 @@ ucp_memh_exported_unpack(ucp_context_h context, const void *export_mkey_buffer, int ucp_memh_buffer_is_dummy(const void *exported_memh_buffer); -ucs_status_t ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer, - size_t length, - ucp_md_map_t unpack_md_map, - ucp_md_map_t skip_md_map, - ucs_sys_device_t sys_device, +ucs_status_t +ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer, size_t length, + ucp_md_map_t unpack_md_map, + ucp_md_map_t skip_md_map, + ucs_sys_device_t sys_device, uint64_t uct_flags, + ucp_rkey_h *rkey_p); + +/* Unpack a remote key passing zero flags to uct_rkey_unpack */ +ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer, ucp_rkey_h *rkey_p); diff --git a/src/ucp/core/ucp_rkey.inl b/src/ucp/core/ucp_rkey.inl index 0552a193460..7b1d8c68ada 100644 --- a/src/ucp/core/ucp_rkey.inl +++ b/src/ucp/core/ucp_rkey.inl @@ -67,13 +67,14 @@ ucp_rkey_get_tl_rkey(ucp_rkey_h rkey, ucp_md_index_t rkey_index) static UCS_F_ALWAYS_INLINE ucs_status_t ucp_ep_rkey_unpack_reachable(ucp_ep_h ep, const void *buffer, size_t length, - ucp_rkey_h *rkey_p) + uint64_t uct_flags, ucp_rkey_h *rkey_p) { ucp_ep_config_t *config = &ucs_array_elem(&ep->worker->ep_config, ep->cfg_index); return ucp_ep_rkey_unpack_internal(ep, buffer, length, config->key.reachable_md_map, 0, - UCS_SYS_DEVICE_ID_UNKNOWN, rkey_p); + UCS_SYS_DEVICE_ID_UNKNOWN, uct_flags, + rkey_p); } static UCS_F_ALWAYS_INLINE int diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index ea5485e1ea2..f790386ce51 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -857,7 +857,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_proto_rndv_send_reply, */ status = ucp_ep_rkey_unpack_internal( ep, rkey_buffer, rkey_length, ep_config->key.reachable_md_map, - ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, &rkey); + ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, 0, &rkey); if (status != UCS_OK) { goto err; } diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 31042738a40..80d47e02d4c 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -952,7 +952,8 @@ static ucs_status_t ucp_rndv_req_send_rma_get(ucp_request_t *rndv_req, ucp_request_set_super(rndv_req, rreq); - status = ucp_ep_rkey_unpack(ep, rkey_buf, &rndv_req->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf, + &rndv_req->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -1124,7 +1125,7 @@ ucp_rndv_send_frag_update_get_rkey(ucp_worker_h worker, ucp_request_t *freq, &rkey_buffer, &rkey_size); ucs_assert_always(status == UCS_OK); - status = ucp_ep_rkey_unpack(mem_type_ep, rkey_buffer, rkey_p); + status = ucp_ep_rkey_unpack_wo_flags(mem_type_ep, rkey_buffer, rkey_p); ucs_assert_always(status == UCS_OK); ucp_rkey_buffer_release(rkey_buffer); } @@ -1303,8 +1304,8 @@ ucp_rndv_recv_start_get_pipeline(ucp_worker_h worker, ucp_request_t *rndv_req, * Step 3: Send ATS for RNDV request */ - status = ucp_ep_rkey_unpack(rndv_req->send.ep, rkey_buffer, - &rndv_req->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(rndv_req->send.ep, rkey_buffer, + &rndv_req->send.rndv.rkey); if (ucs_unlikely(status != UCS_OK)) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(rndv_req->send.ep), ucs_status_string(status)); @@ -1500,7 +1501,7 @@ static void ucp_rndv_do_rkey_ptr(ucp_request_t *rndv_req, ucp_request_t *rreq, ucp_trace_req(rndv_req, "start rkey_ptr rndv rreq %p", rreq); - status = ucp_ep_rkey_unpack(ep, rkey_buf, &rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf, &rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -2123,7 +2124,8 @@ static ucs_status_t ucp_rndv_send_start_put_pipeline(ucp_request_t *sreq, ucp_trace_req(sreq, "using put rndv pipeline protocol: va %p size %zu", (void*)rndv_rtr_hdr->address, rndv_rtr_hdr->size); - status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1, &sreq->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1, + &sreq->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -2325,8 +2327,8 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler, * and we have to use PUT_ZCOPY RNDV scheme instead */ } - status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1, - &sreq->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1, + &sreq->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); diff --git a/src/uct/api/v2/uct_v2.h b/src/uct/api/v2/uct_v2.h index a7115711a7a..b1a6db5f9d2 100644 --- a/src/uct/api/v2/uct_v2.h +++ b/src/uct/api/v2/uct_v2.h @@ -764,10 +764,25 @@ typedef struct uct_rkey_compare_params { * @brief Rkey unpack parameters field mask. */ typedef enum { - UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0) /**< sys_device field */ + UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0), /**< sys_device field */ + UCT_RKEY_UNPACK_FIELD_FLAGS = UCS_BIT(1) /**< flags field */ } uct_rkey_unpack_field_mask_t; +/** + * @ingroup UCT_MD + * @brief Rkey unpack flags. + */ +typedef enum { + /** + * When the remote key is released (see @ref uct_rkey_release), free the + * local mapping of the remote memory which was created for this remote key + * and is not used by any other remote key or operation. + */ + UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING = UCS_BIT(0) +} uct_rkey_unpack_flags_t; + + /** * @ingroup UCT_MD * @brief Parameters for unpacking remote key using @ref uct_rkey_unpack_v2. @@ -785,6 +800,12 @@ typedef struct uct_rkey_unpack_params { * (default behavior). */ ucs_sys_device_t sys_device; + + /** + * Flags to unpack the rkey with, using bits from + * @ref uct_rkey_unpack_flags_t. + */ + uint64_t flags; } uct_rkey_unpack_params_t; diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index bafd8048132..69a25dc0e83 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -647,14 +647,23 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, /* use write lock because cache maybe modified */ pthread_rwlock_wrlock(&cache->lock); pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr); - ucs_assert(pgt_region != NULL); - region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); + if (pgt_region == NULL) { + ucs_assert(mapped_addr == NULL); + pthread_rwlock_unlock(&cache->lock); + return; + } - ucs_assert(region->refcount >= 1); - region->refcount--; + region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); + if (mapped_addr != NULL) { + ucs_assert(region->refcount >= 1); + region->refcount--; + } if (!region->refcount && !cache_enabled) { - ucs_assert(region->mapped_addr == mapped_addr); + if (mapped_addr != NULL) { + ucs_assert(region->mapped_addr == mapped_addr); + } + uct_cuda_ipc_cache_region_destroy(cache, region); } diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index ba7e1e649aa..f34f3ed62ef 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -39,6 +39,7 @@ typedef struct { const void *mapped_addr; CUdevice cu_dev; + int free_unused_mapping; } uct_cuda_ipc_rkey_handle_t; typedef struct { @@ -474,7 +475,8 @@ uct_cuda_ipc_mkey_pack(uct_md_h md, uct_mem_h tl_memh, void *address, memh->dev_num)); } -static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle() +static uct_cuda_ipc_rkey_handle_t * +uct_cuda_ipc_create_rkey_handle(uint64_t flags) { uct_cuda_ipc_rkey_handle_t *rkey_handle; @@ -484,8 +486,10 @@ static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle() return NULL; } - rkey_handle->mapped_addr = NULL; - rkey_handle->cu_dev = CU_DEVICE_INVALID; + rkey_handle->mapped_addr = NULL; + rkey_handle->cu_dev = CU_DEVICE_INVALID; + rkey_handle->free_unused_mapping = !!( + flags & UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING); return rkey_handle; } @@ -562,6 +566,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, uct_cuda_ipc_unpacked_rkey_t *unpacked; uct_cuda_ipc_rkey_handle_t *rkey_handle; ucs_sys_device_t sys_dev; + uint64_t flags; CUdevice cuda_device; CUdevice avail_cuda_device; ucs_status_t status; @@ -569,6 +574,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, sys_dev = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, sys_device, SYS_DEVICE, UCS_SYS_DEVICE_ID_UNKNOWN); + flags = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, flags, FLAGS, 0); unpacked = ucs_malloc(sizeof(*unpacked), "uct_cuda_ipc_unpacked_rkey_t"); if (NULL == unpacked) { @@ -580,7 +586,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, unpacked->super.super = *packed; unpacked->super.super.pid &= ~UCT_CUDA_IPC_RKEY_FLAG_PID_NS; - rkey_handle = uct_cuda_ipc_create_rkey_handle(); + rkey_handle = uct_cuda_ipc_create_rkey_handle(flags); if (rkey_handle == NULL) { status = UCS_ERR_NO_MEMORY; goto err_free_key; @@ -633,8 +639,12 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, uct_cuda_ipc_extended_rkey_t *extended_rkey; uct_cuda_ipc_rkey_t *rkey; - if ((uct_rkey == 0) || (rkey_handle == NULL) || - (rkey_handle->mapped_addr == NULL)) { + if ((uct_rkey == 0) || (rkey_handle == NULL)) { + return; + } + + if ((rkey_handle->mapped_addr == NULL) && + !rkey_handle->free_unused_mapping) { return; } @@ -644,7 +654,8 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, rkey->d_bptr, rkey_handle->mapped_addr, rkey_handle->cu_dev, - uct_cuda_ipc_component.enable_remote_cache); + uct_cuda_ipc_component.enable_remote_cache && + !rkey_handle->free_unused_mapping); } static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component, diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 15ecc612383..b124bdc43ba 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -245,7 +245,7 @@ ucp_rkey_h test_ucp_mmap::mem_chunk::unpack(ucp_ep_h ep, ucp_md_map_t md_map) // Different MD map means different config index on proto v2 ASSERT_UCS_OK(ucp_ep_rkey_unpack_internal(ep, rkey_buffer, rkey_size, md_map, 0, - UCS_SYS_DEVICE_ID_UNKNOWN, + UCS_SYS_DEVICE_ID_UNKNOWN, 0, &rkey)); } @@ -497,7 +497,7 @@ void test_ucp_mmap::test_rkey_proto(ucp_mem_h memh) /* Unpack remote key buffer */ ucp_rkey_h rkey; status = ucp_ep_rkey_unpack_reachable(receiver().ep(), &rkey_buffer[0], - rkey_size, &rkey); + rkey_size, 0, &rkey); ASSERT_UCS_OK(status); /* Check rkey configuration and lanes distances */ diff --git a/test/mpi/test_rma_cuda.c b/test/mpi/test_rma_cuda.c index 36609a19821..3e9319f48a7 100644 --- a/test/mpi/test_rma_cuda.c +++ b/test/mpi/test_rma_cuda.c @@ -167,7 +167,7 @@ int main(int argc, char **argv) return -1; } - xfer_args.ucp = create_ucp(); + xfer_args.ucp = create_ucp(UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP); for (operation_idx = 0; operation_idx < ucs_static_array_size(operations); ++operation_idx) { diff --git a/test/mpi/test_ucp.c b/test/mpi/test_ucp.c index 22657b3ad5e..207c87de625 100644 --- a/test/mpi/test_ucp.c +++ b/test/mpi/test_ucp.c @@ -11,14 +11,14 @@ #include -static ucp_context_h create_ucp_context() +static ucp_context_h create_ucp_context(uint64_t features) { ucp_params_t params; ucp_config_t *config; ucp_context_h context; params.field_mask = UCP_PARAM_FIELD_FEATURES | UCP_PARAM_FIELD_NAME; - params.features = UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP; + params.features = features; params.name = "test_context"; UCX_CHECK(ucp_config_read(NULL, NULL, &config)); UCX_CHECK(ucp_init(¶ms, config, &context)); @@ -124,10 +124,10 @@ create_ucp_mem(void *buffer, size_t size, ucp_context_h context) return mem; } -ucp_t create_ucp() +ucp_t create_ucp(uint64_t features) { ucp_t ucp; - ucp.context = create_ucp_context(); + ucp.context = create_ucp_context(features); ucp.worker = create_ucp_worker(ucp.context); ucp.ep = create_ucp_ep(ucp.worker); return ucp; diff --git a/test/mpi/test_ucp.h b/test/mpi/test_ucp.h index bc679ab43b2..3f6a615163e 100644 --- a/test/mpi/test_ucp.h +++ b/test/mpi/test_ucp.h @@ -22,7 +22,7 @@ typedef struct { } rkey_t; -ucp_t create_ucp(); +ucp_t create_ucp(uint64_t features); void destroy_ucp(ucp_t); diff --git a/test/mpi/test_ucp_rkey_destroy.c b/test/mpi/test_ucp_rkey_destroy.c index 822f1a18f9a..8ce5bef1a89 100644 --- a/test/mpi/test_ucp_rkey_destroy.c +++ b/test/mpi/test_ucp_rkey_destroy.c @@ -13,6 +13,8 @@ #include "test_ucp.h" #include "test_ucx_check_def.h" +#include "ucp/api/device/ucp_host.h" + #include #include @@ -23,7 +25,28 @@ #define MPI_COMM_SIZE 2 #define SIZE (1024 * 1024 * 1024) -static int rank0(ucp_context_h ucp_context) +static void mpi_barrier(void) +{ + MPI_Request request; + + MPI_Ibarrier(MPI_COMM_WORLD, &request); + MPI_Wait(&request, MPI_STATUS_IGNORE); +} + +/* Progress the worker while waiting, to serve the peer's wireup messages */ +static void mpi_barrier_progress(ucp_worker_h worker) +{ + int done = 0; + MPI_Request request; + + MPI_Ibarrier(MPI_COMM_WORLD, &request); + while (!done) { + ucp_worker_progress(worker); + MPI_Test(&request, &done, MPI_STATUS_IGNORE); + } +} + +static int rank0(ucp_t ucp) { CUdeviceptr ptr; ucp_mem_h ucp_mem; @@ -32,11 +55,11 @@ static int rank0(ucp_context_h ucp_context) CUDA_CHECK(cuMemGetInfo(&free_bytes_before, &total_bytes)); CUDA_CHECK(cuMemAlloc(&ptr, SIZE)); - ucp_mem = send_rkey(1, (void*)ptr, SIZE, ucp_context); + ucp_mem = send_rkey(1, (void*)ptr, SIZE, ucp.context); - MPI_Barrier(MPI_COMM_WORLD); + mpi_barrier_progress(ucp.worker); - UCX_CHECK(ucp_mem_unmap(ucp_context, ucp_mem)); + UCX_CHECK(ucp_mem_unmap(ucp.context, ucp_mem)); CUDA_CHECK(cuMemFree(ptr)); CUDA_CHECK(cuCtxSynchronize()); @@ -47,17 +70,51 @@ static int rank0(ucp_context_h ucp_context) return (unreleased_memory == 0) ? EXIT_SUCCESS : EXIT_FAILURE; } -static void rank1(ucp_ep_h ucp_ep) +static void create_remote_mem_list(ucp_t ucp, const rkey_t *rkey, + ucp_device_remote_mem_list_h *mem_list_h) { + ucp_device_mem_list_elem_t elem; + ucp_device_mem_list_params_t params; + ucs_status_t status; + + elem.field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP | + UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | + UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY; + elem.ep = ucp.ep; + elem.remote_addr = rkey->remote_address; + elem.rkey = rkey->rkey; + + params.field_mask = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS | + UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE | + UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS; + params.elements = &elem; + params.element_size = sizeof(elem); + params.num_elements = 1; + + do { + ucp_worker_progress(ucp.worker); + status = ucp_device_remote_mem_list_create(¶ms, mem_list_h); + } while (status == UCS_ERR_NOT_CONNECTED); + + UCX_CHECK(status); +} + +static void rank1(ucp_t ucp) +{ + ucp_device_remote_mem_list_h mem_list_h; rkey_t rkey; void *ptr; - rkey = recv_rkey(0, ucp_ep); + rkey = recv_rkey(0, ucp.ep); UCX_CHECK(ucp_rkey_ptr(rkey.rkey, rkey.remote_address, &ptr)); + + create_remote_mem_list(ucp, &rkey, &mem_list_h); + ucp_device_mem_list_release(mem_list_h); + ucp_rkey_destroy(rkey.rkey); - MPI_Barrier(MPI_COMM_WORLD); + mpi_barrier(); } int main(int argc, char **argv) @@ -88,12 +145,12 @@ int main(int argc, char **argv) CUDA_CHECK(cuDevicePrimaryCtxRetain(&cu_ctx, cu_dev)); CUDA_CHECK(cuCtxSetCurrent(cu_ctx)); - ucp = create_ucp(); + ucp = create_ucp(UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP | UCP_FEATURE_DEVICE); if (rank == 0) { - exit_status = rank0(ucp.context); + exit_status = rank0(ucp); } else { - rank1(ucp.ep); + rank1(ucp); } MPI_Bcast(&exit_status, 1, MPI_INT, 0, MPI_COMM_WORLD);