From 89bf554d003dc93c1a654ba1b9443454104bed4b Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Wed, 16 Oct 2024 10:29:48 -0700 Subject: [PATCH 01/14] API: add explicit memory mapping --- src/ucc/api/ucc.h | 80 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 3257c79896..2243d622c5 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2227,5 +2227,85 @@ ucc_status_t ucc_ee_wait(ucc_ee_h ee, ucc_ev_t *ev); */ ucc_status_t ucc_collective_triggered_post(ucc_ee_h ee, ucc_ev_t *ee_event); +typedef void * ucc_mem_map_mem_h; + +/** + * @ingroup UCC_DATATYPE + */ +enum ucc_mem_map_flags { + UCC_MEM_MAP_EXPORT = 0, /*!< Indicate @ref ucc_mem_map() should export + memory handles from TLs used by context */ + UCC_MEM_MAP_IMPORT = 1 /*!< Indicate @ref ucc_mem_map() should import + memory handles from user memory handle */ +}; + +/** + * @ingroup UCC_CONTEXT + * @brief Routine registers or maps memory for use in future collectives. + * + * This routine maps memory a user-specified memory segment with a @ref + * ucc_context_t. The segment is considered "mapped" with the context until + * either the user calls @ref ucc_mem_unmap or @ref ucc_context_destroy(). A + * handle to the mapped memory is provided in memh. If the flag + * UCC_MEM_MAP_EXPORT is used, the memory will be mapped and memory + * handles from TLs will be generated and stored in the memh. If the flag + * UCC_MEM_MAP_IMPORT is used, the user must provide a valid memh, otherwise + * behavior is undefined. + * + * @params [in] context Context mapped memory is associated with + * @params [in] flags flags dictating the behavior of the routine + * @params [in] params parameters indicating the address and length of + * memory to map + * @params [inout] *memh Handle for the registered memory + * + * @return Error code as defined by @ref ucc_status_t. + */ + +ucc_status_t ucc_mem_map(ucc_context_t context, + ucc_mem_map_flags flags, + ucc_mem_map_params params, + ucc_mem_map_mem_h *memh); + +/** + * @ingroup UCC_CONTEXT + * @brief Routine exchanges mapped memory with peers to enable future + * collectives. + * + * This routine performs an nonblocking exchange of mapped memory with peers + * within the same context on which it is mapped. This is a collective routine + * and must be called by all members of the ucc_context on which the memory + * is mapped. + * + * @params [in] *memh Handle of the registered memory + * + * @return Error code as defined by @ref ucc_status_t. + */ +ucc_status_t ucc_mem_exchange_post(ucc_mem_map_mem_h *memh); + +/** + * @ingroup UCC_CONTEXT + * @brief Routine tests for the completion of a memory exchange on a context. + * + * This routine tests for the completion of a memory exchange on a context. + * + * @params [in] *memh Handle of the registered memory + * + * @return Error code as defined by @ref ucc_status_t. + */ +ucc_status_t ucc_mem_exchange_test(ucc_mem_map_mem_h *memh); + +/** + * @ingroup UCC_CONTEXT + * @brief Routine unmaps memory from a context + * + * This routine unmaps memory and all resources associated with the memory + * from a context. The memh object is freed and cannot be reused. + * + * @params [in] *memh Handle of the registered memory + * + * @return Error code as defined by @ref ucc_status_t. + */ +ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh); + END_C_DECLS #endif From e94a0eb525d9cb25a7737513462ecc64af5834e6 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 1 Nov 2024 10:12:43 -0700 Subject: [PATCH 02/14] REVIEW: address feedback --- src/ucc/api/ucc.h | 29 ----------------------------- 1 file changed, 29 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 2243d622c5..d678dedec6 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2265,35 +2265,6 @@ ucc_status_t ucc_mem_map(ucc_context_t context, ucc_mem_map_flags flags, ucc_mem_map_params params, ucc_mem_map_mem_h *memh); - -/** - * @ingroup UCC_CONTEXT - * @brief Routine exchanges mapped memory with peers to enable future - * collectives. - * - * This routine performs an nonblocking exchange of mapped memory with peers - * within the same context on which it is mapped. This is a collective routine - * and must be called by all members of the ucc_context on which the memory - * is mapped. - * - * @params [in] *memh Handle of the registered memory - * - * @return Error code as defined by @ref ucc_status_t. - */ -ucc_status_t ucc_mem_exchange_post(ucc_mem_map_mem_h *memh); - -/** - * @ingroup UCC_CONTEXT - * @brief Routine tests for the completion of a memory exchange on a context. - * - * This routine tests for the completion of a memory exchange on a context. - * - * @params [in] *memh Handle of the registered memory - * - * @return Error code as defined by @ref ucc_status_t. - */ -ucc_status_t ucc_mem_exchange_test(ucc_mem_map_mem_h *memh); - /** * @ingroup UCC_CONTEXT * @brief Routine unmaps memory from a context From 1aa5d5c92eb6cfa9ff38ddd89d16b7d31ec9139e Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 1 Nov 2024 10:30:17 -0700 Subject: [PATCH 03/14] CODESTYLE: fix code style --- src/ucc/api/ucc.h | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index d678dedec6..50ae12417e 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2227,30 +2227,29 @@ ucc_status_t ucc_ee_wait(ucc_ee_h ee, ucc_ev_t *ev); */ ucc_status_t ucc_collective_triggered_post(ucc_ee_h ee, ucc_ev_t *ee_event); -typedef void * ucc_mem_map_mem_h; +typedef void *ucc_mem_map_mem_h; /** * @ingroup UCC_DATATYPE */ -enum ucc_mem_map_flags { +typedef enum { UCC_MEM_MAP_EXPORT = 0, /*!< Indicate @ref ucc_mem_map() should export memory handles from TLs used by context */ - UCC_MEM_MAP_IMPORT = 1 /*!< Indicate @ref ucc_mem_map() should import + UCC_MEM_MAP_IMPORT = 1 /*!< Indicate @ref ucc_mem_map() should import memory handles from user memory handle */ -}; +} ucc_mem_map_flags_t; /** * @ingroup UCC_CONTEXT * @brief Routine registers or maps memory for use in future collectives. * - * This routine maps memory a user-specified memory segment with a @ref - * ucc_context_t. The segment is considered "mapped" with the context until - * either the user calls @ref ucc_mem_unmap or @ref ucc_context_destroy(). A - * handle to the mapped memory is provided in memh. If the flag - * UCC_MEM_MAP_EXPORT is used, the memory will be mapped and memory - * handles from TLs will be generated and stored in the memh. If the flag - * UCC_MEM_MAP_IMPORT is used, the user must provide a valid memh, otherwise - * behavior is undefined. + * This routine maps a user-specified memory segment with a ucc_context_h. The + * segment is considered "mapped" with the context until either the user calls + * @ref ucc_mem_unmap or @ref ucc_context_destroy(). A handle to the mapped + * memory is provided in memh. If the flag UCC_MEM_MAP_EXPORT is used, the + * memory will be mapped and memory handles from TLs will be generated and + * stored in the memh. If the flag UCC_MEM_MAP_IMPORT is used, the user must + * provide a valid memh, otherwise behavior is undefined. * * @params [in] context Context mapped memory is associated with * @params [in] flags flags dictating the behavior of the routine @@ -2261,10 +2260,8 @@ enum ucc_mem_map_flags { * @return Error code as defined by @ref ucc_status_t. */ -ucc_status_t ucc_mem_map(ucc_context_t context, - ucc_mem_map_flags flags, - ucc_mem_map_params params, - ucc_mem_map_mem_h *memh); +ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, + ucc_mem_map_params_t params, ucc_mem_map_mem_h *memh); /** * @ingroup UCC_CONTEXT * @brief Routine unmaps memory from a context From fe8f872001391fe4c231d1861c526b2689d915bd Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 1 Nov 2024 11:01:54 -0700 Subject: [PATCH 04/14] REVIEW: remove ref --- src/ucc/api/ucc.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 50ae12417e..0b9e034075 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2233,9 +2233,9 @@ typedef void *ucc_mem_map_mem_h; * @ingroup UCC_DATATYPE */ typedef enum { - UCC_MEM_MAP_EXPORT = 0, /*!< Indicate @ref ucc_mem_map() should export + UCC_MEM_MAP_EXPORT = 0, /*!< Indicate ucc_mem_map() should export memory handles from TLs used by context */ - UCC_MEM_MAP_IMPORT = 1 /*!< Indicate @ref ucc_mem_map() should import + UCC_MEM_MAP_IMPORT = 1 /*!< Indicate ucc_mem_map() should import memory handles from user memory handle */ } ucc_mem_map_flags_t; @@ -2245,7 +2245,7 @@ typedef enum { * * This routine maps a user-specified memory segment with a ucc_context_h. The * segment is considered "mapped" with the context until either the user calls - * @ref ucc_mem_unmap or @ref ucc_context_destroy(). A handle to the mapped + * ucc_mem_unmap or ucc_context_destroy(). A handle to the mapped * memory is provided in memh. If the flag UCC_MEM_MAP_EXPORT is used, the * memory will be mapped and memory handles from TLs will be generated and * stored in the memh. If the flag UCC_MEM_MAP_IMPORT is used, the user must @@ -2257,7 +2257,7 @@ typedef enum { * memory to map * @params [inout] *memh Handle for the registered memory * - * @return Error code as defined by @ref ucc_status_t. + * @return Error code as defined by ucc_status_t. */ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, @@ -2271,7 +2271,7 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, * * @params [in] *memh Handle of the registered memory * - * @return Error code as defined by @ref ucc_status_t. + * @return Error code as defined by ucc_status_t. */ ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh); From 3765beca0963d79973140182ce9acebc4d32b7df Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 1 Nov 2024 11:20:03 -0700 Subject: [PATCH 05/14] REVIEW: address feedback --- src/ucc/api/ucc.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 0b9e034075..28ce28aa8f 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2251,11 +2251,11 @@ typedef enum { * stored in the memh. If the flag UCC_MEM_MAP_IMPORT is used, the user must * provide a valid memh, otherwise behavior is undefined. * - * @params [in] context Context mapped memory is associated with - * @params [in] flags flags dictating the behavior of the routine - * @params [in] params parameters indicating the address and length of + * @param [in] context Context mapped memory is associated with + * @param [in] flags flags dictating the behavior of the routine + * @param [in] params parameters indicating the address and length of * memory to map - * @params [inout] *memh Handle for the registered memory + * @param [inout] *memh Handle for the registered memory * * @return Error code as defined by ucc_status_t. */ @@ -2269,7 +2269,7 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, * This routine unmaps memory and all resources associated with the memory * from a context. The memh object is freed and cannot be reused. * - * @params [in] *memh Handle of the registered memory + * @param [in] *memh Handle of the registered memory * * @return Error code as defined by ucc_status_t. */ From 887d0a24483e22f1db8a2d7f47fbf7712e652af8 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Wed, 11 Dec 2024 09:49:06 -0800 Subject: [PATCH 06/14] REVIEW: address feedback --- src/ucc/api/ucc.h | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 28ce28aa8f..711c66e658 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2244,8 +2244,9 @@ typedef enum { * @brief Routine registers or maps memory for use in future collectives. * * This routine maps a user-specified memory segment with a ucc_context_h. The - * segment is considered "mapped" with the context until either the user calls - * ucc_mem_unmap or ucc_context_destroy(). A handle to the mapped + * segment is considered "mapped" with the context until the user calls + * ucc_mem_unmap. It is the user's responsibility to unmap all mapped segments + * prior to calling ucc_context_destroy(). A handle to the mapped * memory is provided in memh. If the flag UCC_MEM_MAP_EXPORT is used, the * memory will be mapped and memory handles from TLs will be generated and * stored in the memh. If the flag UCC_MEM_MAP_IMPORT is used, the user must @@ -2262,12 +2263,14 @@ typedef enum { ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, ucc_mem_map_params_t params, ucc_mem_map_mem_h *memh); + /** * @ingroup UCC_CONTEXT * @brief Routine unmaps memory from a context * - * This routine unmaps memory and all resources associated with the memory - * from a context. The memh object is freed and cannot be reused. + * This is a collective routine that unmaps memory and all resources + * associated with the memory from a context. The memh object is freed and + * cannot be reused. * * @param [in] *memh Handle of the registered memory * From 51a8118bddac3b8158ad2d61d50bd3b82140f849 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 13 Dec 2024 11:07:55 -0800 Subject: [PATCH 07/14] REVIEW: address feedback --- src/ucc/api/ucc.h | 41 +++++++++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 711c66e658..a73c62661c 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -1805,9 +1805,12 @@ enum ucc_coll_args_field { UCC_COLL_ARGS_FIELD_TAG = UCC_BIT(1), UCC_COLL_ARGS_FIELD_CB = UCC_BIT(2), UCC_COLL_ARGS_FIELD_GLOBAL_WORK_BUFFER = UCC_BIT(3), - UCC_COLL_ARGS_FIELD_ACTIVE_SET = UCC_BIT(4) + UCC_COLL_ARGS_FIELD_ACTIVE_SET = UCC_BIT(4), + UCC_COLL_ARGS_FIELD_MEM_MAP_MEMH = UCC_BIT(5) }; +typedef void *ucc_mem_map_mem_h; + /** * @ingroup UCC_COLLECTIVES * @@ -1880,6 +1883,26 @@ typedef struct ucc_coll_args { int64_t stride; uint64_t size; } active_set; + union { + ucc_mem_map_mem_h local_memh; /*!< Memory handle of locally + mapped memory for the src + buffer */ + ucc_mem_map_mem_h *global_memh; /*!< Array of memory handles + for the src buffer + corresponding to all + participating processes + in the collective */ + } src_memh; + union { + ucc_mem_map_mem_h local_memh; /*!< Memory handle of locally + mapped memory for the dst + buffer */ + ucc_mem_map_mem_h *global_memh; /*!< Array of memory handles + for the dst buffer + corresponding to all + participating processes + in the collective */ + } dst_memh; } ucc_coll_args_t; /** @@ -2227,8 +2250,6 @@ ucc_status_t ucc_ee_wait(ucc_ee_h ee, ucc_ev_t *ev); */ ucc_status_t ucc_collective_triggered_post(ucc_ee_h ee, ucc_ev_t *ee_event); -typedef void *ucc_mem_map_mem_h; - /** * @ingroup UCC_DATATYPE */ @@ -2243,12 +2264,12 @@ typedef enum { * @ingroup UCC_CONTEXT * @brief Routine registers or maps memory for use in future collectives. * - * This routine maps a user-specified memory segment with a ucc_context_h. The - * segment is considered "mapped" with the context until the user calls - * ucc_mem_unmap. It is the user's responsibility to unmap all mapped segments - * prior to calling ucc_context_destroy(). A handle to the mapped - * memory is provided in memh. If the flag UCC_MEM_MAP_EXPORT is used, the - * memory will be mapped and memory handles from TLs will be generated and + * This local routine maps a user-specified memory segment with a + * ucc_context_h. The segment is considered "mapped" with the context until + * the user calls ucc_mem_unmap. It is the user's responsibility to unmap all + * mapped segments prior to calling ucc_context_destroy(). A handle to the + * mapped memory is provided in memh. If the flag UCC_MEM_MAP_EXPORT is used, + * the memory will be mapped and memory handles from TLs will be generated and * stored in the memh. If the flag UCC_MEM_MAP_IMPORT is used, the user must * provide a valid memh, otherwise behavior is undefined. * @@ -2268,7 +2289,7 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, * @ingroup UCC_CONTEXT * @brief Routine unmaps memory from a context * - * This is a collective routine that unmaps memory and all resources + * This is a local routine that unmaps memory and all resources * associated with the memory from a context. The memh object is freed and * cannot be reused. * From 9d6f2a53d27cd3b49632d8e89a000ba9118c1b19 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Thu, 9 Jan 2025 14:23:27 -0800 Subject: [PATCH 08/14] API: update mem_map with exchange size --- src/ucc/api/ucc.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index a73c62661c..854e320676 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2277,13 +2277,15 @@ typedef enum { * @param [in] flags flags dictating the behavior of the routine * @param [in] params parameters indicating the address and length of * memory to map + * @param [out] memh_size Size of memory handle when exported * @param [inout] *memh Handle for the registered memory * * @return Error code as defined by ucc_status_t. */ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, - ucc_mem_map_params_t params, ucc_mem_map_mem_h *memh); + ucc_mem_map_params_t *params, size_t *memh_size, + ucc_mem_map_mem_h *memh); /** * @ingroup UCC_CONTEXT From c2ebac7ec626ce45ae61cba39499850955ec9e32 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Thu, 9 Jan 2025 14:26:08 -0800 Subject: [PATCH 09/14] CORE: Add support for ucc_mem_map/unmap CL/BASIC: add support for ucc_mem_map/unmap CL/HIER: add support for ucc_mem_map/unmap CL/DOCA_UROM: add support for ucc_mem_map/unmap TL/CUDA: add support for ucc_mem_map/unmap TL/MLX5: add support for ucc_mem_map/unmap TL/NCCL: add support for ucc_mem_map/unmap TL/RCCL: add support for ucc_mem_map/unmap TL/SELF: add support for ucc_mem_map/unmap TL/SHARP: add support for ucc_mem_map/unmap --- src/components/base/ucc_base_iface.h | 8 + src/components/cl/basic/cl_basic.c | 7 + src/components/cl/basic/cl_basic_context.c | 19 ++ src/components/cl/doca_urom/cl_doca_urom.c | 10 + .../cl/doca_urom/cl_doca_urom_context.c | 19 ++ src/components/cl/hier/cl_hier.c | 10 + src/components/cl/hier/cl_hier_context.c | 19 ++ src/components/tl/cuda/tl_cuda.c | 9 + src/components/tl/cuda/tl_cuda_context.c | 18 ++ src/components/tl/mlx5/tl_mlx5.c | 9 + src/components/tl/mlx5/tl_mlx5_context.c | 18 ++ src/components/tl/nccl/tl_nccl.c | 9 + src/components/tl/nccl/tl_nccl_context.c | 18 ++ src/components/tl/rccl/tl_rccl.c | 9 + src/components/tl/rccl/tl_rccl_context.c | 18 ++ src/components/tl/self/tl_self.c | 9 + src/components/tl/self/tl_self_context.c | 18 ++ src/components/tl/sharp/tl_sharp.c | 9 + src/components/tl/sharp/tl_sharp_context.c | 18 ++ src/core/ucc_context.c | 180 ++++++++++++++++++ src/core/ucc_context.h | 22 +++ 21 files changed, 456 insertions(+) diff --git a/src/components/base/ucc_base_iface.h b/src/components/base/ucc_base_iface.h index 1a637274a2..7c3e074e7a 100644 --- a/src/components/base/ucc_base_iface.h +++ b/src/components/base/ucc_base_iface.h @@ -123,6 +123,11 @@ typedef struct ucc_base_context_iface { void (*destroy)(ucc_base_context_t *ctx); ucc_status_t (*get_attr)(const ucc_base_context_t *context, ucc_base_ctx_attr_t *attr); + ucc_status_t (*mem_map)(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, void *tl_h); + ucc_status_t (*mem_unmap)(const ucc_base_context_t *context, int type, void *tl_h); + ucc_status_t (*memh_pack)(const ucc_base_context_t *context, void *memh, + void **pack_buffer); } ucc_base_context_iface_t; @@ -233,6 +238,9 @@ typedef struct ucc_base_coll_alg_info { .super.context.destroy = \ UCC_CLASS_DELETE_FUNC_NAME(ucc_##_f##_name##_context_t), \ .super.context.get_attr = ucc_##_f##_name##_get_context_attr, \ + .super.context.mem_map = ucc_##_f##_name##_mem_map, \ + .super.context.mem_unmap = ucc_##_f##_name##_mem_unmap, \ + .super.context.memh_pack = ucc_##_f##_name##_memh_pack, \ .super.team.create_post = \ UCC_CLASS_NEW_FUNC_NAME(ucc_##_f##_name##_team_t), \ .super.team.create_test = ucc_##_f##_name##_team_create_test, \ diff --git a/src/components/cl/basic/cl_basic.c b/src/components/cl/basic/cl_basic.c index ac3e344250..e31808bb3f 100644 --- a/src/components/cl/basic/cl_basic.c +++ b/src/components/cl/basic/cl_basic.c @@ -10,6 +10,13 @@ ucc_status_t ucc_cl_basic_get_lib_attr(const ucc_base_lib_t *lib, ucc_base_lib_attr_t *base_attr); ucc_status_t ucc_cl_basic_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_cl_basic_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h); +ucc_status_t ucc_cl_basic_mem_unmap(const ucc_base_context_t *context, int type, + void *tl_h); +ucc_status_t ucc_cl_basic_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer); ucc_status_t ucc_cl_basic_get_lib_properties(ucc_base_lib_properties_t *prop); diff --git a/src/components/cl/basic/cl_basic_context.c b/src/components/cl/basic/cl_basic_context.c index e4546a5ba5..3dc3a608ff 100644 --- a/src/components/cl/basic/cl_basic_context.c +++ b/src/components/cl/basic/cl_basic_context.c @@ -61,6 +61,25 @@ UCC_CLASS_CLEANUP_FUNC(ucc_cl_basic_context_t) ucc_free(self->super.tl_ctxs); } +ucc_status_t ucc_cl_basic_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_basic_mem_unmap(const ucc_base_context_t *context, int type, + void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_basic_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} + UCC_CLASS_DEFINE(ucc_cl_basic_context_t, ucc_cl_context_t); ucc_status_t diff --git a/src/components/cl/doca_urom/cl_doca_urom.c b/src/components/cl/doca_urom/cl_doca_urom.c index ae09833074..c9f4df875e 100644 --- a/src/components/cl/doca_urom/cl_doca_urom.c +++ b/src/components/cl/doca_urom/cl_doca_urom.c @@ -13,6 +13,16 @@ ucc_status_t ucc_cl_doca_urom_get_lib_attr(const ucc_base_lib_t *lib, ucc_status_t ucc_cl_doca_urom_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_cl_doca_urom_mem_map(const ucc_base_context_t *context, + int type, void *address, size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_cl_doca_urom_mem_unmap(const ucc_base_context_t *context, + int type, void *tl_h); + +ucc_status_t ucc_cl_doca_urom_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer); + ucc_status_t ucc_cl_doca_urom_get_lib_properties(ucc_base_lib_properties_t *prop); static ucc_config_field_t ucc_cl_doca_urom_lib_config_table[] = { diff --git a/src/components/cl/doca_urom/cl_doca_urom_context.c b/src/components/cl/doca_urom/cl_doca_urom_context.c index 58b3cc9ed8..eafa8b0aec 100644 --- a/src/components/cl/doca_urom/cl_doca_urom_context.c +++ b/src/components/cl/doca_urom/cl_doca_urom_context.c @@ -556,3 +556,22 @@ ucc_cl_doca_urom_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_clear(attr); return UCC_OK; } + +ucc_status_t ucc_cl_doca_urom_mem_map(const ucc_base_context_t *context, + int type, void *address, size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_doca_urom_mem_unmap(const ucc_base_context_t *context, + int type, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_doca_urom_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/components/cl/hier/cl_hier.c b/src/components/cl/hier/cl_hier.c index edbb469d78..071ce325a4 100644 --- a/src/components/cl/hier/cl_hier.c +++ b/src/components/cl/hier/cl_hier.c @@ -18,6 +18,16 @@ ucc_status_t ucc_cl_hier_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_cl_hier_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_cl_hier_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h); + +ucc_status_t ucc_cl_hier_mem_unmap(const ucc_base_context_t *context, int type, + void *tl_h); + +ucc_status_t ucc_cl_hier_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer); + static ucc_config_field_t ucc_cl_hier_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_cl_hier_lib_config_t, super), UCC_CONFIG_TYPE_TABLE(ucc_cl_lib_config_table)}, diff --git a/src/components/cl/hier/cl_hier_context.c b/src/components/cl/hier/cl_hier_context.c index 4f7f219426..44bb0636f5 100644 --- a/src/components/cl/hier/cl_hier_context.c +++ b/src/components/cl/hier/cl_hier_context.c @@ -91,3 +91,22 @@ ucc_cl_hier_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ attr->topo_required = 1; return UCC_OK; } + +ucc_status_t ucc_cl_hier_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_hier_mem_unmap(const ucc_base_context_t *context, int type, + void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_cl_hier_memh_pack(const ucc_base_context_t *context, + void *memh, void **packed_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/components/tl/cuda/tl_cuda.c b/src/components/tl/cuda/tl_cuda.c index 98dccf26bf..5bfb8d0da5 100644 --- a/src/components/tl/cuda/tl_cuda.c +++ b/src/components/tl/cuda/tl_cuda.c @@ -64,6 +64,15 @@ static ucs_config_field_t ucc_tl_cuda_context_config_table[] = { ucc_status_t ucc_tl_cuda_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_cuda_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + UCC_CLASS_DEFINE_NEW_FUNC(ucc_tl_cuda_context_t, ucc_base_context_t, const ucc_base_context_params_t *, const ucc_base_config_t *); diff --git a/src/components/tl/cuda/tl_cuda_context.c b/src/components/tl/cuda/tl_cuda_context.c index 4d89029680..b3c231c548 100644 --- a/src/components/tl/cuda/tl_cuda_context.c +++ b/src/components/tl/cuda/tl_cuda_context.c @@ -76,6 +76,24 @@ UCC_CLASS_INIT_FUNC(ucc_tl_cuda_context_t, return status; } +ucc_status_t ucc_tl_cuda_mem_map(const ucc_base_context_t *context, + void *address, + size_t len, + void *memh) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} + +ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, void *memh) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} + +ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} + UCC_CLASS_CLEANUP_FUNC(ucc_tl_cuda_context_t) { tl_debug(self->super.super.lib, "finalizing tl context: %p", self); diff --git a/src/components/tl/mlx5/tl_mlx5.c b/src/components/tl/mlx5/tl_mlx5.c index 5cdd6c51a1..50f961b468 100644 --- a/src/components/tl/mlx5/tl_mlx5.c +++ b/src/components/tl/mlx5/tl_mlx5.c @@ -11,6 +11,15 @@ ucc_status_t ucc_tl_mlx5_get_lib_attr(const ucc_base_lib_t *lib, ucc_status_t ucc_tl_mlx5_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t * base_attr); +ucc_status_t ucc_tl_mlx5_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + ucc_status_t ucc_tl_mlx5_get_lib_properties(ucc_base_lib_properties_t *prop); static ucc_config_field_t ucc_tl_mlx5_lib_config_table[] = { diff --git a/src/components/tl/mlx5/tl_mlx5_context.c b/src/components/tl/mlx5/tl_mlx5_context.c index 1011fe72c0..c1958c3cf3 100644 --- a/src/components/tl/mlx5/tl_mlx5_context.c +++ b/src/components/tl/mlx5/tl_mlx5_context.c @@ -302,3 +302,21 @@ ucc_status_t ucc_tl_mlx5_context_create_epilog(ucc_base_context_t *context) { return ucc_tl_mlx5_context_ib_ctx_pd_setup(context); } + +ucc_status_t ucc_tl_mlx5_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} + +ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} + +ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_IMPLEMENTED; +} diff --git a/src/components/tl/nccl/tl_nccl.c b/src/components/tl/nccl/tl_nccl.c index 46fdcff8e3..4c82db7574 100644 --- a/src/components/tl/nccl/tl_nccl.c +++ b/src/components/tl/nccl/tl_nccl.c @@ -16,6 +16,15 @@ ucc_status_t ucc_tl_nccl_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_tl_nccl_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_nccl_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + static ucc_config_field_t ucc_tl_nccl_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_nccl_lib_config_t, super), UCC_CONFIG_TYPE_TABLE(ucc_tl_lib_config_table)}, diff --git a/src/components/tl/nccl/tl_nccl_context.c b/src/components/tl/nccl/tl_nccl_context.c index 4b5254f32b..909bc16e0f 100644 --- a/src/components/tl/nccl/tl_nccl_context.c +++ b/src/components/tl/nccl/tl_nccl_context.c @@ -223,3 +223,21 @@ ucc_tl_nccl_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ ucc_base_ctx_attr_clear(attr); return UCC_OK; } + +ucc_status_t ucc_tl_nccl_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/components/tl/rccl/tl_rccl.c b/src/components/tl/rccl/tl_rccl.c index a6b44fc062..e12ab28251 100644 --- a/src/components/tl/rccl/tl_rccl.c +++ b/src/components/tl/rccl/tl_rccl.c @@ -17,6 +17,15 @@ ucc_status_t ucc_tl_rccl_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_tl_rccl_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_rccl_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + static ucc_config_field_t ucc_tl_rccl_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_rccl_lib_config_t, super), UCC_CONFIG_TYPE_TABLE(ucc_tl_lib_config_table)}, diff --git a/src/components/tl/rccl/tl_rccl_context.c b/src/components/tl/rccl/tl_rccl_context.c index d4e94e73b0..9dc247e56c 100644 --- a/src/components/tl/rccl/tl_rccl_context.c +++ b/src/components/tl/rccl/tl_rccl_context.c @@ -120,3 +120,21 @@ ucc_tl_rccl_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ ucc_base_ctx_attr_clear(attr); return UCC_OK; } + +ucc_status_t ucc_tl_rccl_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/components/tl/self/tl_self.c b/src/components/tl/self/tl_self.c index 80ba2c0af8..334c0409ea 100644 --- a/src/components/tl/self/tl_self.c +++ b/src/components/tl/self/tl_self.c @@ -16,6 +16,15 @@ ucc_status_t ucc_tl_self_get_lib_attr(const ucc_base_lib_t *lib, ucc_status_t ucc_tl_self_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_self_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + ucc_status_t ucc_tl_self_get_lib_properties(ucc_base_lib_properties_t *prop); static ucc_config_field_t ucc_tl_self_lib_config_table[] = { diff --git a/src/components/tl/self/tl_self_context.c b/src/components/tl/self/tl_self_context.c index d8dfb08fe6..4b21c5601b 100644 --- a/src/components/tl/self/tl_self_context.c +++ b/src/components/tl/self/tl_self_context.c @@ -49,3 +49,21 @@ ucc_tl_self_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ ucc_base_ctx_attr_clear(attr); return UCC_OK; } + +ucc_status_t ucc_tl_self_mem_map(const ucc_base_context_t *context, + int type, void *address, + size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/components/tl/sharp/tl_sharp.c b/src/components/tl/sharp/tl_sharp.c index 464ef50478..705f1aa3dd 100644 --- a/src/components/tl/sharp/tl_sharp.c +++ b/src/components/tl/sharp/tl_sharp.c @@ -15,6 +15,15 @@ ucc_status_t ucc_tl_sharp_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_tl_sharp_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_sharp_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h); + +ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + +ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + static ucc_config_field_t ucc_tl_sharp_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_sharp_lib_config_t, super), UCC_CONFIG_TYPE_TABLE(ucc_tl_lib_config_table)}, diff --git a/src/components/tl/sharp/tl_sharp_context.c b/src/components/tl/sharp/tl_sharp_context.c index 42d10f8d87..6a0b58c53b 100644 --- a/src/components/tl/sharp/tl_sharp_context.c +++ b/src/components/tl/sharp/tl_sharp_context.c @@ -511,3 +511,21 @@ ucc_status_t ucc_tl_sharp_get_context_attr(const ucc_base_context_t *context, /* attr->topo_required = 1; return UCC_OK; } + +ucc_status_t ucc_tl_sharp_mem_map(const ucc_base_context_t *context, int type, + void *address, + size_t len, + void *memh, void *tl_h) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + return UCC_ERR_NOT_SUPPORTED; +} + +ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + return UCC_ERR_NOT_SUPPORTED; +} diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index 13e51246b8..1a2bf9bcca 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1131,3 +1131,183 @@ ucc_status_t ucc_context_get_attr(ucc_context_t *context, return status; } + +ucc_status_t ucc_mem_map_import(ucc_context_h context, + ucc_mem_map_params_t *params, size_t *memh_size, + ucc_mem_map_mem_h *memh) +{ + ucc_context_t *ctx = (ucc_context_t *)context; + ucc_status_t status = UCC_OK; + int i; + ucc_mem_map_memh_t *local_memh; + ucc_tl_lib_t *tl_lib; + + if (!memh) { + ucc_error("cannot import NULL memory handle"); + return UCC_ERR_INVALID_PARAM; + } + + /* FIXME: it's legal for the user to export and immediately import a handle, + * is this an issue? */ + local_memh = (ucc_mem_map_memh_t *)memh; + + /* memh should have been used in exchanges or from a remote process, addresses, etc likely garbage. fix it */ + local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( + ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh"); + for (i = 0; i < ctx->n_tl_ctx; i++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + /* FIXME: i don't think this will work properly for more than 1 TL */ + status = tl_lib->iface->context.mem_map( + (const ucc_base_context_t *)ctx, UCC_MEM_MAP_TYPE_GLOBAL, params->segments[0].address, params->segments[0].len, + local_memh, &local_memh->tl_h[i]); + if (status != UCC_OK) { + ucc_error("failed to import mem map memh"); + return status; + } + } + + local_memh->type = UCC_MEM_MAP_TYPE_GLOBAL; + return status; +} + +ucc_status_t ucc_mem_map_export(ucc_context_h context, + ucc_mem_map_params_t *params, size_t *memh_size, + ucc_mem_map_mem_h *memh) +{ + ucc_context_t *ctx = (ucc_context_t *)context; + size_t total_pack_size = 0; + ucc_mem_map_memh_t *local_memh; + ucc_mem_map_memh_t *exported_memh; + void **packed_buffers; + ucc_status_t status; + ucc_tl_lib_t *tl_lib; + size_t offset; + int i; + + local_memh = (ucc_mem_map_memh_t *)ucc_calloc( + 1, sizeof(ucc_mem_map_memh_t), "local memh"); + if (!local_memh) { + ucc_error("failed to allocate a local memory handle"); + return UCC_ERR_NO_MEMORY; + } + + local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( + ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh"); + packed_buffers = + (void **)ucc_malloc(sizeof(void *) * ctx->n_tl_ctx, "packed buffers"); + + /* map all the memory */ + for (i = 0; i < ctx->n_tl_ctx; i++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + /* always treat as a local mem handle */ + status = tl_lib->iface->context.mem_map( + (const ucc_base_context_t *)context, UCC_MEM_MAP_TYPE_LOCAL, params->segments[0].address, params->segments[0].len, + local_memh, &local_memh->tl_h[i]); + if (status != UCC_OK) { + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("failed to map memory"); + goto failed_mem_map; + } + if (status == UCC_ERR_NOT_IMPLEMENTED || + status == UCC_ERR_NOT_SUPPORTED) { + /* either not implemented or not supported, set memh to null */ + local_memh->tl_h[i].packed_size = 0; + } + } + } + + /* now pack all the memories */ + for (i = 0; i < ctx->n_tl_ctx; i++) { + if (local_memh->tl_h[i].packed_size > 0) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + /* tl should set packed_size, allocate buffer, pack memh */ + status = tl_lib->iface->context.memh_pack( + (const ucc_base_context_t *)context, &local_memh->tl_h[i], &packed_buffers[i]); + if (status != UCC_OK) { + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("failed to map memory"); + goto failed_pack; + } + } + total_pack_size += local_memh->tl_h[i].packed_size; + } + } + + // allocate exported memh, copy items over + exported_memh = (ucc_mem_map_memh_t *)ucc_calloc( + 1, sizeof(ucc_mem_map_memh_t) + total_pack_size, "exported memh"); + if (!exported_memh) { + ucc_error("failed to allocate handle for exported buffers"); + return UCC_ERR_NO_MEMORY; + } + + // copying + exported_memh->tl_h = local_memh->tl_h; + offset = 0; + for (i = 0; i < ctx->n_tl_ctx; i++) { + if (local_memh->tl_h[i].packed_size == 0) { + continue; + } + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + &local_memh->tl_h[i].packed_size, sizeof(size_t)); + offset += sizeof(size_t); + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + packed_buffers[i], local_memh->tl_h[i].packed_size); + ucc_free(packed_buffers[i]); + offset += local_memh->tl_h[i].packed_size; + } + ucc_free(local_memh); + exported_memh->context = context; + *memh = (ucc_mem_map_memh_t *)exported_memh; + return UCC_OK; +failed_pack: +failed_mem_map: + for (int j = 0; j < i; j++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, UCC_MEM_MAP_TYPE_LOCAL, &local_memh->tl_h[j]); + } + return status; +} + +ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, + ucc_mem_map_params_t *params, size_t *memh_size, + ucc_mem_map_mem_h *memh) +{ + // check if flags is import / export + if (flags == UCC_MEM_MAP_IMPORT) { + // set map type to global + return ucc_mem_map_import(context, params, memh_size, memh); + } else { + // set map type to local + return ucc_mem_map_export(context, params, memh_size, memh); + } + return UCC_OK; +} + +ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh) +{ + ucc_context_t *ctx; + ucc_status_t status; + ucc_mem_map_memh_t *lmemh; + ucc_tl_lib_t *tl_lib; + int i; + + if (!memh) { + ucc_warn("unable to free NULL memh"); + return UCC_ERR_INVALID_PARAM; + } + + /* it could be global or local type */ + lmemh = (ucc_mem_map_memh_t *)memh; + ctx = lmemh->context; + + for (i = 0; i < ctx->n_tl_ctx; i++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + status = tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, lmemh->type, &lmemh->tl_h[i]); + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("we had an error"); + return status; + } + } + return UCC_OK; +} diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index 3944d5675b..909fed5446 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -94,6 +94,28 @@ typedef struct ucc_context_config { uint32_t throttle_progress; } ucc_context_config_t; +typedef enum { + UCC_MEM_MAP_TYPE_LOCAL, + UCC_MEM_MAP_TYPE_GLOBAL +} ucc_mem_map_type_t; + +typedef struct ucc_mem_map_tl_t { + size_t packed_size; + void *tl_data; /* tl specific data */ +} ucc_mem_map_tl_t; + +typedef struct ucc_mem_map_memh_t { + ucc_mem_map_type_t type; + ucc_context_h context; + void *address; + size_t len; + /* rank of exporting process */ + ucc_rank_t my_ctx_rank; + /* handles for each tl */ + ucc_mem_map_tl_t *tl_h; + char pack_buffer[0]; +} ucc_mem_map_memh_t; + /* Internal function for context creation that takes explicit pointer for proc_info */ ucc_status_t ucc_context_create_proc_info(ucc_lib_h lib, From 93547eab6a938ca602f9d5ab5f6b7bcf44361826 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Thu, 9 Jan 2025 14:45:51 -0800 Subject: [PATCH 10/14] TL/UCP: add support for ucc_mem_map/unmap TL/UCP: add support for communication with memmap REVIEW: various fixes TL/UCP: Enable multiple packed buffers --- src/components/tl/ucp/alltoall/alltoall.c | 7 + .../tl/ucp/alltoall/alltoall_onesided.c | 11 +- .../tl/ucp/alltoallv/alltoallv_onesided.c | 12 +- src/components/tl/ucp/tl_ucp.c | 11 ++ src/components/tl/ucp/tl_ucp.h | 7 + src/components/tl/ucp/tl_ucp_context.c | 172 ++++++++++++++++++ src/components/tl/ucp/tl_ucp_sendrecv.h | 103 ++++++++++- src/core/ucc_context.c | 63 +++++-- src/core/ucc_context.h | 2 + src/ucc/api/ucc.h | 3 +- 10 files changed, 361 insertions(+), 30 deletions(-) diff --git a/src/components/tl/ucp/alltoall/alltoall.c b/src/components/tl/ucp/alltoall/alltoall.c index 3803d96426..8ad4fa8486 100644 --- a/src/components/tl/ucp/alltoall/alltoall.c +++ b/src/components/tl/ucp/alltoall/alltoall.c @@ -95,6 +95,13 @@ ucc_status_t ucc_tl_ucp_alltoall_onesided_init(ucc_base_coll_args_t *coll_args, goto out; } } + if (!(coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_SRC_MEMH)) { + coll_args->args.src_memh.global_memh = NULL; + } + if (!(coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH)) { + coll_args->args.dst_memh.global_memh = NULL; + } + task = ucc_tl_ucp_init_task(coll_args, team); *task_h = &task->super; task->super.post = ucc_tl_ucp_alltoall_onesided_start; diff --git a/src/components/tl/ucp/alltoall/alltoall_onesided.c b/src/components/tl/ucp/alltoall/alltoall_onesided.c index 856b392534..649f7c010e 100644 --- a/src/components/tl/ucp/alltoall/alltoall_onesided.c +++ b/src/components/tl/ucp/alltoall/alltoall_onesided.c @@ -24,22 +24,25 @@ ucc_status_t ucc_tl_ucp_alltoall_onesided_start(ucc_coll_task_t *ctask) ucc_rank_t gsize = UCC_TL_TEAM_SIZE(team); ucc_rank_t start = (grank + 1) % gsize; long * pSync = TASK_ARGS(task).global_work_buffer; + ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; + ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; ucc_rank_t peer; ucc_tl_ucp_task_reset(task, UCC_INPROGRESS); /* TODO: change when support for library-based work buffers is complete */ nelems = (nelems / gsize) * ucc_dt_size(TASK_ARGS(task).src.info.datatype); dest = dest + grank * nelems; + UCPCHECK_GOTO(ucc_tl_ucp_put_nb((void *)(src + start * nelems), - (void *)dest, nelems, start, team, task), + (void *)dest, nelems, start, src_memh[start], dst_memh[start], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, team), task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, src_memh[start], dst_memh[start], team), task, out); for (peer = (start + 1) % gsize; peer != start; peer = (peer + 1) % gsize) { UCPCHECK_GOTO(ucc_tl_ucp_put_nb((void *)(src + peer * nelems), - (void *)dest, nelems, peer, team, task), + (void *)dest, nelems, peer, src_memh[peer], dst_memh[peer], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, team), task, + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], dst_memh[peer], team), task, out); } diff --git a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c index bb6fa14b3e..decfd470a5 100644 --- a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c +++ b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c @@ -24,6 +24,8 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_start(ucc_coll_task_t *ctask) ucc_aint_t *d_disp = TASK_ARGS(task).dst.info_v.displacements; size_t sdt_size = ucc_dt_size(TASK_ARGS(task).src.info_v.datatype); size_t rdt_size = ucc_dt_size(TASK_ARGS(task).dst.info_v.datatype); + ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; + ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; ucc_rank_t peer; size_t sd_disp, dd_disp, data_size; @@ -46,9 +48,9 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_start(ucc_coll_task_t *ctask) UCPCHECK_GOTO(ucc_tl_ucp_put_nb(PTR_OFFSET(src, sd_disp), PTR_OFFSET(dest, dd_disp), - data_size, peer, team, task), + data_size, peer, src_memh[peer], dst_memh[peer], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, team), task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], dst_memh[peer], team), task, out); } return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super); out: @@ -93,6 +95,12 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_init(ucc_base_coll_args_t *coll_args, goto out; } } + if (!(coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_SRC_MEMH)) { + coll_args->args.src_memh.global_memh = NULL; + } + if (!(coll_args->args.mask & UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH)) { + coll_args->args.dst_memh.global_memh = NULL; + } task = ucc_tl_ucp_init_task(coll_args, team); *task_h = &task->super; diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c index 7db99bdaf2..2d508da50b 100644 --- a/src/components/tl/ucp/tl_ucp.c +++ b/src/components/tl/ucp/tl_ucp.c @@ -32,6 +32,17 @@ ucc_status_t ucc_tl_ucp_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_tl_ucp_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); +ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, + int type, + void *address, + size_t len, + void *memh, + void *tl_h); + +ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); + +ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, void *memh); + ucc_config_field_t ucc_tl_ucp_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_ucp_lib_config_t, super), UCC_CONFIG_TYPE_TABLE(ucc_tl_lib_config_table)}, diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h index 3c439f4ae5..59a063040f 100644 --- a/src/components/tl/ucp/tl_ucp.h +++ b/src/components/tl/ucp/tl_ucp.h @@ -107,6 +107,13 @@ typedef struct ucc_tl_ucp_remote_info { size_t packed_key_len; } ucc_tl_ucp_remote_info_t; +typedef struct ucc_tl_ucp_memh_data { + ucc_tl_ucp_remote_info_t rinfo; + void *packed_memh; + size_t packed_memh_len; + ucp_rkey_h rkey; +} ucc_tl_ucp_memh_data_t; + typedef struct ucc_tl_ucp_worker { ucp_context_h ucp_context; ucp_worker_h ucp_worker; diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c index 2a945bd4f0..7440d263bf 100644 --- a/src/components/tl/ucp/tl_ucp_context.c +++ b/src/components/tl/ucp/tl_ucp_context.c @@ -555,6 +555,178 @@ static void ucc_tl_ucp_ctx_remote_pack_data(ucc_tl_ucp_context_t *ctx, } } +ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, + int type, + void *address, + size_t len, + void *memh, + void *tl_h) +{ + ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); + ucc_status_t ucc_status = UCC_OK; + ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *)tl_h; + ucc_tl_ucp_memh_data_t *m_data = (ucc_tl_ucp_memh_data_t *)p_memh->tl_data; + ucp_mem_map_params_t mmap_params; + ucp_mem_h mh = NULL; + ucs_status_t status; + ucp_memh_pack_params_t pack_params; + + // basically an import here + if (type == UCC_MEM_MAP_TYPE_GLOBAL) { + // m_data is lost in the exchange, make a new one + m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); + if (!m_data) { + tl_error(ctx->super.super.lib, "failed to allocate tl data"); + return UCC_ERR_NO_MEMORY; + } + p_memh->tl_data = m_data; +#if defined(__aarch64__) + ucc_mem_map_memh_t * l_memh = (ucc_mem_map_memh_t *)memh; + /* unpack here */ + size_t *key_size = (size_t *)l_memh->pack_buffer; + void *packed_memh = PTR_OFFSET(l_memh->pack_buffer, sizeof(size_t) * 2 + *key_size); + mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_EXPORTED_MEMH_BUFFER; + mmap_params.exported_memh_buffer = packed_memh; + + status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, &mh); + if (UCS_OK != status) { + if (status != UCS_ERR_UNREACHABLE) { + tl_error(ctx->super.super.lib, + "ucp_mem_map failed with error code: %s", ucs_status_string(status)); + ucc_status = ucs_status_to_ucc_status(status); + return ucc_status; + } else { + tl_debug(ctx->super.super.lib, + "ucp_mem_map could not map exported memory handle"); + ucc_status = UCC_OK; // this is still OK + } + } else { + m_data->rinfo.mem_h = mh; + // the rest of the data is garbage. fix it + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; + } +#endif + } else { + if (!m_data) { + m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); + if (!m_data) { + tl_error(ctx->super.super.lib, "failed to allocate TL/UCP specific data"); + return UCC_ERR_NO_MEMORY; + } + p_memh->tl_data = m_data; + } + // local, we need to map it here + if (m_data->rinfo.mem_h == NULL) { + mmap_params.field_mask = + UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; + mmap_params.address = address; + mmap_params.length = len; + + status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, &mh); + if (UCS_OK != status) { + tl_error(ctx->super.super.lib, + "ucp_mem_map failed with error code: %d", status); + ucc_status = ucs_status_to_ucc_status(status); + } + } + m_data->rinfo.mem_h = mh; + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; + + // export memh + pack_params.field_mask = UCP_MEMH_PACK_PARAM_FIELD_FLAGS; + pack_params.flags = UCP_MEMH_PACK_FLAG_EXPORT; + + status = ucp_memh_pack(mh, &pack_params, &m_data->packed_memh, &m_data->packed_memh_len); + if (status != UCS_OK) { + // we don't support memory pack, or it failed + tl_debug("ucp_memh_pack() returned error %s", ucs_status_string(status)); + m_data->packed_memh = 0; + m_data->packed_memh_len = 0; + } + // pack rkey + status = ucp_rkey_pack(ctx->worker.ucp_context, mh, &m_data->rinfo.packed_key, &m_data->rinfo.packed_key_len); + if (status != UCS_OK) { + tl_error("unable to pack rkey with error %s", ucs_status_string(status)); + } + p_memh->packed_size = m_data->packed_memh_len + m_data->rinfo.packed_key_len; + } + + return ucc_status; +} + +ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +{ + ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); + ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *) memh; + ucc_tl_ucp_memh_data_t *data; + ucs_status_t status; + + if (!p_memh) { + return UCC_OK; + } + + data = (ucc_tl_ucp_memh_data_t *)p_memh->tl_data; + + if (type == UCC_MEM_MAP_TYPE_LOCAL) { + status = ucp_mem_unmap(ctx->worker.ucp_context, data->rinfo.mem_h); + if (status != UCS_OK) { + tl_error(ctx->super.super.lib, "ucp_mem_unmap failed with error code %d", status); + return ucs_status_to_ucc_status(status); + } + } else if (type == UCC_MEM_MAP_TYPE_GLOBAL) { + // need to free rkeys (data->rkey) , packed memh (data->packed_memh) + if (data->packed_memh) { + ucc_free(data->packed_memh); + } + if (data->rinfo.packed_key) { + ucp_rkey_buffer_release(data->rinfo.packed_key); + } + if (data->rkey) { + ucp_rkey_destroy(data->rkey); + } + } else { + ucc_error("Unknown type entered"); + return UCC_ERR_INVALID_PARAM; + } + return UCC_OK; +} + +ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +{ + ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *) memh; + ucc_tl_ucp_memh_data_t *data = p_memh->tl_data; + void *packed_buffer; + size_t *key_size; + size_t *memh_size; + + if (!data) { + return UCC_OK; + } + /* + * data order + * + * packed_key_size | packed_memh_size | packed_key | packed_memh + */ + packed_buffer = ucc_malloc(sizeof(size_t) * 2 + data->packed_memh_len + data->rinfo.packed_key_len, + "packed buffer"); + if (!packed_buffer) { + ucc_error("failed to allocate a packed buffer of size %lu", data->packed_memh_len + data->rinfo.packed_key_len); + return UCC_ERR_NO_MEMORY; + } + key_size = packed_buffer; + *key_size = data->rinfo.packed_key_len; + memh_size = PTR_OFFSET(packed_buffer, sizeof(size_t)); + *memh_size = data->packed_memh_len; + memcpy(PTR_OFFSET(packed_buffer, sizeof(size_t) * 2), data->rinfo.packed_key, *key_size); + memcpy(PTR_OFFSET(packed_buffer, sizeof(size_t) * 2 + data->rinfo.packed_key_len), data->packed_memh, *memh_size); + + p_memh->packed_size = sizeof(size_t) * 2 + data->packed_memh_len + data->rinfo.packed_key_len; + *pack_buffer = packed_buffer; + return UCC_OK; +} + ucc_status_t ucc_tl_ucp_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *attr) { diff --git a/src/components/tl/ucp/tl_ucp_sendrecv.h b/src/components/tl/ucp/tl_ucp_sendrecv.h index 4d1bc84440..fa02790abe 100644 --- a/src/components/tl/ucp/tl_ucp_sendrecv.h +++ b/src/components/tl/ucp/tl_ucp_sendrecv.h @@ -225,10 +225,64 @@ static inline ucc_status_t ucc_tl_ucp_send_nz(void *buffer, size_t msglen, dest_group_rank, team, task); } +static inline ucc_status_t find_tl_index(ucc_mem_map_mem_h map_memh, int *tl_index) +{ + ucc_mem_map_memh_t *memh = (ucc_mem_map_memh_t *)map_memh; + int i = 0; + + for (; i < memh->num_tls; i++) { + if (strncmp(memh->tl_h[i].tl_name, "ucp", 3) == 0) { + *tl_index = i; + return UCC_OK; + } + } + return UCC_ERR_NOT_FOUND; +} + +static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_t *rva, + ucp_rkey_h *rkey, int tl_index, ucc_mem_map_mem_h map_memh) +{ + // check if src_memh or dest_memh have segment + ucc_mem_map_memh_t *memh = map_memh; + ucc_tl_ucp_memh_data_t *tl_data = (ucc_tl_ucp_memh_data_t *)memh->tl_h[tl_index].tl_data; + uint64_t base, end; + ucs_status_t ucs_status; + int i; + size_t offset; + + base = (uint64_t)memh->address; + end = base + memh->len; + + if ((uint64_t)va >= base && (uint64_t)va < end) { + *rva = (uint64_t)PTR_OFFSET(memh->address, ((uint64_t)va - (uint64_t)memh->address)); + if (NULL == tl_data->rkey) { + offset = 0; + /* find pack location for tl */ + for (i = 0; i < tl_index; i++) { + size_t *p = PTR_OFFSET(memh->pack_buffer, offset); + if (p[0] == tl_index) { + break; + } + offset += p[1]; + } + ucs_status = + ucp_ep_rkey_unpack(*ep, PTR_OFFSET(memh->pack_buffer, offset + sizeof(size_t) * 4), + &tl_data->rkey); + if (UCS_OK != ucs_status) { + return ucs_status_to_ucc_status(ucs_status); + } + } + *rkey = tl_data->rkey; + /* FIXME: packed memh? */ + return UCC_OK; + } + return UCC_ERR_NOT_FOUND; +} + static inline ucc_status_t ucc_tl_ucp_resolve_p2p_by_va(ucc_tl_ucp_team_t *team, void *va, ucp_ep_h *ep, ucc_rank_t peer, uint64_t *rva, ucp_rkey_h *rkey, - int *segment) + int *segment, ucc_mem_map_mem_h src_memh, ucc_mem_map_mem_h dest_memh) { ucc_tl_ucp_context_t *ctx = UCC_TL_UCP_TEAM_CTX(team); ptrdiff_t key_offset = 0; @@ -264,6 +318,41 @@ ucc_tl_ucp_resolve_p2p_by_va(ucc_tl_ucp_team_t *team, void *va, ucp_ep_h *ep, key_offset += key_sizes[i]; } if (ucc_unlikely(0 > *segment)) { + // check if src_memh or dest_memh have segment + int tl_index = 0; + ucc_status_t status; + + if (src_memh) { + //status = UCC_OK; + status = find_tl_index(src_memh, &tl_index); + if (status == UCC_ERR_NOT_FOUND) { + tl_error(UCC_TL_TEAM_LIB(team), + "attempt to perform one-sided operation with malformed mem map handle"); + return status; + } + + status = ucc_tl_ucp_check_memh(ep, va, rva, rkey, tl_index, src_memh); + if (status == UCC_OK) { + return UCC_OK; + } + } + + if (dest_memh) { + if (tl_index == -1) { + status = find_tl_index(dest_memh, &tl_index); + if (status == UCC_ERR_NOT_FOUND) { + tl_error(UCC_TL_TEAM_LIB(team), + "attempt to perform one-sided operation with malformed mem map handle"); + return status; + } + } + + status = ucc_tl_ucp_check_memh(ep, va, rva, rkey, tl_index, dest_memh); + if (status == UCC_OK) { + return UCC_OK; + } + } + tl_error(UCC_TL_TEAM_LIB(team), "attempt to perform one-sided operation on non-registered memory %p", va); return UCC_ERR_NOT_FOUND; @@ -322,6 +411,8 @@ static inline ucc_status_t ucc_tl_ucp_ep_flush(ucc_rank_t dest_group_rank, static inline ucc_status_t ucc_tl_ucp_put_nb(void *buffer, void *target, size_t msglen, ucc_rank_t dest_group_rank, + ucc_mem_map_mem_h src_memh, + ucc_mem_map_mem_h dest_memh, ucc_tl_ucp_team_t *team, ucc_tl_ucp_task_t *task) { @@ -339,7 +430,7 @@ static inline ucc_status_t ucc_tl_ucp_put_nb(void *buffer, void *target, } status = ucc_tl_ucp_resolve_p2p_by_va(team, target, &ep, dest_group_rank, - &rva, &rkey, &segment); + &rva, &rkey, &segment, src_memh, dest_memh); if (ucc_unlikely(UCC_OK != status)) { return status; } @@ -365,6 +456,8 @@ static inline ucc_status_t ucc_tl_ucp_put_nb(void *buffer, void *target, static inline ucc_status_t ucc_tl_ucp_get_nb(void *buffer, void *target, size_t msglen, ucc_rank_t dest_group_rank, + ucc_mem_map_mem_h src_memh, + ucc_mem_map_mem_h dest_memh, ucc_tl_ucp_team_t *team, ucc_tl_ucp_task_t *task) { @@ -382,7 +475,7 @@ static inline ucc_status_t ucc_tl_ucp_get_nb(void *buffer, void *target, } status = ucc_tl_ucp_resolve_p2p_by_va(team, target, &ep, dest_group_rank, - &rva, &rkey, &segment); + &rva, &rkey, &segment, src_memh, dest_memh); if (ucc_unlikely(UCC_OK != status)) { return status; } @@ -408,6 +501,8 @@ static inline ucc_status_t ucc_tl_ucp_get_nb(void *buffer, void *target, static inline ucc_status_t ucc_tl_ucp_atomic_inc(void * target, ucc_rank_t dest_group_rank, + ucc_mem_map_mem_h src_memh, + ucc_mem_map_mem_h dest_memh, ucc_tl_ucp_team_t *team) { ucp_request_param_t req_param = {0}; @@ -425,7 +520,7 @@ static inline ucc_status_t ucc_tl_ucp_atomic_inc(void * target, } status = ucc_tl_ucp_resolve_p2p_by_va(team, target, &ep, dest_group_rank, - &rva, &rkey, &segment); + &rva, &rkey, &segment, src_memh, dest_memh); if (ucc_unlikely(UCC_OK != status)) { return status; } diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index 1a2bf9bcca..5586a85fbd 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1138,6 +1138,7 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, { ucc_context_t *ctx = (ucc_context_t *)context; ucc_status_t status = UCC_OK; + ucc_config_names_array_t *tls = &ctx->all_tls; int i; ucc_mem_map_memh_t *local_memh; ucc_tl_lib_t *tl_lib; @@ -1146,10 +1147,14 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, ucc_error("cannot import NULL memory handle"); return UCC_ERR_INVALID_PARAM; } + if (!params) { + ucc_error("params cannot be NULL"); + return UCC_ERR_INVALID_PARAM; + } /* FIXME: it's legal for the user to export and immediately import a handle, * is this an issue? */ - local_memh = (ucc_mem_map_memh_t *)memh; + local_memh = *memh; /* memh should have been used in exchanges or from a remote process, addresses, etc likely garbage. fix it */ local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( @@ -1158,15 +1163,18 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* FIXME: i don't think this will work properly for more than 1 TL */ status = tl_lib->iface->context.mem_map( - (const ucc_base_context_t *)ctx, UCC_MEM_MAP_TYPE_GLOBAL, params->segments[0].address, params->segments[0].len, + (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_GLOBAL, params->segments[0].address, params->segments[0].len, local_memh, &local_memh->tl_h[i]); - if (status != UCC_OK) { - ucc_error("failed to import mem map memh"); + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("failed to import mem map memh %d", status); return status; } + strncpy(local_memh->tl_h[i].tl_name, tls->names[i], 8); } - local_memh->type = UCC_MEM_MAP_TYPE_GLOBAL; + /* fix context as it will be incorrect on a different system */ + local_memh->context = ctx; + return status; } @@ -1183,6 +1191,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, ucc_tl_lib_t *tl_lib; size_t offset; int i; + ucc_config_names_array_t *tls = &ctx->all_tls; local_memh = (ucc_mem_map_memh_t *)ucc_calloc( 1, sizeof(ucc_mem_map_memh_t), "local memh"); @@ -1194,14 +1203,14 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh"); packed_buffers = - (void **)ucc_malloc(sizeof(void *) * ctx->n_tl_ctx, "packed buffers"); + (void **)ucc_calloc(ctx->n_tl_ctx, sizeof(void *), "packed buffers"); /* map all the memory */ for (i = 0; i < ctx->n_tl_ctx; i++) { tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* always treat as a local mem handle */ status = tl_lib->iface->context.mem_map( - (const ucc_base_context_t *)context, UCC_MEM_MAP_TYPE_LOCAL, params->segments[0].address, params->segments[0].len, + (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_LOCAL, params->segments[0].address, params->segments[0].len, local_memh, &local_memh->tl_h[i]); if (status != UCC_OK) { if (status < UCC_ERR_NOT_IMPLEMENTED) { @@ -1222,7 +1231,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* tl should set packed_size, allocate buffer, pack memh */ status = tl_lib->iface->context.memh_pack( - (const ucc_base_context_t *)context, &local_memh->tl_h[i], &packed_buffers[i]); + (const ucc_base_context_t *)ctx->tl_ctx[i], &local_memh->tl_h[i], &packed_buffers[i]); if (status != UCC_OK) { if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("failed to map memory"); @@ -1233,21 +1242,24 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, } } - // allocate exported memh, copy items over + /* allocate exported memh, copy items over */ exported_memh = (ucc_mem_map_memh_t *)ucc_calloc( - 1, sizeof(ucc_mem_map_memh_t) + total_pack_size, "exported memh"); + 1, sizeof(ucc_mem_map_memh_t) + total_pack_size + 2 * sizeof(size_t) * ctx->n_tl_ctx, "exported memh"); if (!exported_memh) { ucc_error("failed to allocate handle for exported buffers"); return UCC_ERR_NO_MEMORY; } - // copying + /* copying */ exported_memh->tl_h = local_memh->tl_h; - offset = 0; - for (i = 0; i < ctx->n_tl_ctx; i++) { + for (i = 0, offset = 0; i < ctx->n_tl_ctx; i++) { + uint64_t tl_index = i; if (local_memh->tl_h[i].packed_size == 0) { continue; } + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + &tl_index, sizeof(size_t)); + offset += sizeof(size_t); memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), &local_memh->tl_h[i].packed_size, sizeof(size_t)); offset += sizeof(size_t); @@ -1255,10 +1267,19 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, packed_buffers[i], local_memh->tl_h[i].packed_size); ucc_free(packed_buffers[i]); offset += local_memh->tl_h[i].packed_size; + + // copy name information for look ups later + strncpy(exported_memh->tl_h[i].tl_name, tls->names[i], 8); } ucc_free(local_memh); - exported_memh->context = context; - *memh = (ucc_mem_map_memh_t *)exported_memh; + exported_memh->type = UCC_MEM_MAP_TYPE_LOCAL; + exported_memh->context = ctx; + exported_memh->address = params->segments[0].address; + exported_memh->len = params->segments[0].len; + exported_memh->my_ctx_rank = ctx->rank; + exported_memh->num_tls = ctx->n_tl_ctx; + *memh = exported_memh; + *memh_size = total_pack_size; return UCC_OK; failed_pack: failed_mem_map: @@ -1273,6 +1294,11 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { + if (params->n_segments > 1) { + ucc_error("UCC only supports one mapping per call"); + return UCC_ERR_INVALID_PARAM; + } + // check if flags is import / export if (flags == UCC_MEM_MAP_IMPORT) { // set map type to global @@ -1296,14 +1322,13 @@ ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh) ucc_warn("unable to free NULL memh"); return UCC_ERR_INVALID_PARAM; } - /* it could be global or local type */ - lmemh = (ucc_mem_map_memh_t *)memh; - ctx = lmemh->context; + lmemh = (ucc_mem_map_memh_t *)*memh; + ctx = (ucc_context_t *)lmemh->context; for (i = 0; i < ctx->n_tl_ctx; i++) { tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); - status = tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, lmemh->type, &lmemh->tl_h[i]); + status = tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx->tl_ctx[i], lmemh->type, &lmemh->tl_h[i]); if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("we had an error"); return status; diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index 909fed5446..f8e255257d 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -101,6 +101,7 @@ typedef enum { typedef struct ucc_mem_map_tl_t { size_t packed_size; + char tl_name[8]; // typically less than 4 letters void *tl_data; /* tl specific data */ } ucc_mem_map_tl_t; @@ -113,6 +114,7 @@ typedef struct ucc_mem_map_memh_t { ucc_rank_t my_ctx_rank; /* handles for each tl */ ucc_mem_map_tl_t *tl_h; + int num_tls; char pack_buffer[0]; } ucc_mem_map_memh_t; diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 854e320676..8e194c7dcf 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -1806,7 +1806,8 @@ enum ucc_coll_args_field { UCC_COLL_ARGS_FIELD_CB = UCC_BIT(2), UCC_COLL_ARGS_FIELD_GLOBAL_WORK_BUFFER = UCC_BIT(3), UCC_COLL_ARGS_FIELD_ACTIVE_SET = UCC_BIT(4), - UCC_COLL_ARGS_FIELD_MEM_MAP_MEMH = UCC_BIT(5) + UCC_COLL_ARGS_FIELD_MEM_MAP_SRC_MEMH = UCC_BIT(5), + UCC_COLL_ARGS_FIELD_MEM_MAP_DST_MEMH = UCC_BIT(5), }; typedef void *ucc_mem_map_mem_h; From 4a2f9c6c0336d241e6b622d557f4b68dba691926 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Wed, 15 Jan 2025 16:26:46 -0800 Subject: [PATCH 11/14] CODESTYLE: fix code style --- src/components/tl/cuda/tl_cuda.c | 11 +- src/components/tl/cuda/tl_cuda_context.c | 10 +- src/components/tl/mlx5/tl_mlx5.c | 11 +- src/components/tl/mlx5/tl_mlx5_context.c | 11 +- src/components/tl/nccl/tl_nccl.c | 11 +- src/components/tl/nccl/tl_nccl_context.c | 11 +- src/components/tl/rccl/tl_rccl.c | 11 +- src/components/tl/rccl/tl_rccl_context.c | 11 +- src/components/tl/self/tl_self.c | 11 +- src/components/tl/self/tl_self_context.c | 13 +- src/components/tl/sharp/tl_sharp.c | 11 +- src/components/tl/sharp/tl_sharp_context.c | 11 +- .../tl/ucp/alltoall/alltoall_onesided.c | 45 ++-- .../tl/ucp/alltoallv/alltoallv_onesided.c | 10 +- src/components/tl/ucp/tl_ucp.c | 13 +- src/components/tl/ucp/tl_ucp.h | 6 +- src/components/tl/ucp/tl_ucp_context.c | 204 ++++++++++-------- src/components/tl/ucp/tl_ucp_sendrecv.h | 1 - src/core/ucc_context.c | 86 ++++---- src/core/ucc_context.h | 4 +- 20 files changed, 275 insertions(+), 227 deletions(-) diff --git a/src/components/tl/cuda/tl_cuda.c b/src/components/tl/cuda/tl_cuda.c index 5bfb8d0da5..2e5ade2077 100644 --- a/src/components/tl/cuda/tl_cuda.c +++ b/src/components/tl/cuda/tl_cuda.c @@ -65,13 +65,14 @@ ucc_status_t ucc_tl_cuda_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); ucc_status_t ucc_tl_cuda_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); UCC_CLASS_DEFINE_NEW_FUNC(ucc_tl_cuda_context_t, ucc_base_context_t, const ucc_base_context_params_t *, diff --git a/src/components/tl/cuda/tl_cuda_context.c b/src/components/tl/cuda/tl_cuda_context.c index b3c231c548..e6859a1471 100644 --- a/src/components/tl/cuda/tl_cuda_context.c +++ b/src/components/tl/cuda/tl_cuda_context.c @@ -77,19 +77,19 @@ UCC_CLASS_INIT_FUNC(ucc_tl_cuda_context_t, } ucc_status_t ucc_tl_cuda_mem_map(const ucc_base_context_t *context, - void *address, - size_t len, - void *memh) + void *address, size_t len, void *memh) { return UCC_ERR_NOT_IMPLEMENTED; } -ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, void *memh) +ucc_status_t ucc_tl_cuda_mem_unmap(const ucc_base_context_t *context, + void *memh) { return UCC_ERR_NOT_IMPLEMENTED; } -ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_cuda_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_IMPLEMENTED; } diff --git a/src/components/tl/mlx5/tl_mlx5.c b/src/components/tl/mlx5/tl_mlx5.c index 50f961b468..190aef3833 100644 --- a/src/components/tl/mlx5/tl_mlx5.c +++ b/src/components/tl/mlx5/tl_mlx5.c @@ -12,13 +12,14 @@ ucc_status_t ucc_tl_mlx5_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t * base_attr); ucc_status_t ucc_tl_mlx5_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); ucc_status_t ucc_tl_mlx5_get_lib_properties(ucc_base_lib_properties_t *prop); diff --git a/src/components/tl/mlx5/tl_mlx5_context.c b/src/components/tl/mlx5/tl_mlx5_context.c index c1958c3cf3..1fdf3f1093 100644 --- a/src/components/tl/mlx5/tl_mlx5_context.c +++ b/src/components/tl/mlx5/tl_mlx5_context.c @@ -304,19 +304,20 @@ ucc_status_t ucc_tl_mlx5_context_create_epilog(ucc_base_context_t *context) } ucc_status_t ucc_tl_mlx5_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h) + void *address, size_t len, void *memh, + void *tl_h) { return UCC_ERR_NOT_IMPLEMENTED; } -ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_mlx5_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { return UCC_ERR_NOT_IMPLEMENTED; } -ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_mlx5_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_IMPLEMENTED; } diff --git a/src/components/tl/nccl/tl_nccl.c b/src/components/tl/nccl/tl_nccl.c index 4c82db7574..2016a7a70c 100644 --- a/src/components/tl/nccl/tl_nccl.c +++ b/src/components/tl/nccl/tl_nccl.c @@ -17,13 +17,14 @@ ucc_status_t ucc_tl_nccl_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); ucc_status_t ucc_tl_nccl_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); static ucc_config_field_t ucc_tl_nccl_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_nccl_lib_config_t, super), diff --git a/src/components/tl/nccl/tl_nccl_context.c b/src/components/tl/nccl/tl_nccl_context.c index 909bc16e0f..14396d7e1c 100644 --- a/src/components/tl/nccl/tl_nccl_context.c +++ b/src/components/tl/nccl/tl_nccl_context.c @@ -225,19 +225,20 @@ ucc_tl_nccl_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ } ucc_status_t ucc_tl_nccl_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h) + void *address, size_t len, void *memh, + void *tl_h) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_nccl_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_nccl_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_SUPPORTED; } diff --git a/src/components/tl/rccl/tl_rccl.c b/src/components/tl/rccl/tl_rccl.c index e12ab28251..857732148c 100644 --- a/src/components/tl/rccl/tl_rccl.c +++ b/src/components/tl/rccl/tl_rccl.c @@ -18,13 +18,14 @@ ucc_status_t ucc_tl_rccl_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); ucc_status_t ucc_tl_rccl_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); static ucc_config_field_t ucc_tl_rccl_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_rccl_lib_config_t, super), diff --git a/src/components/tl/rccl/tl_rccl_context.c b/src/components/tl/rccl/tl_rccl_context.c index 9dc247e56c..36c2cf6184 100644 --- a/src/components/tl/rccl/tl_rccl_context.c +++ b/src/components/tl/rccl/tl_rccl_context.c @@ -122,19 +122,20 @@ ucc_tl_rccl_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ } ucc_status_t ucc_tl_rccl_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h) + void *address, size_t len, void *memh, + void *tl_h) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_rccl_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_rccl_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_SUPPORTED; } diff --git a/src/components/tl/self/tl_self.c b/src/components/tl/self/tl_self.c index 334c0409ea..ca4e6be6f3 100644 --- a/src/components/tl/self/tl_self.c +++ b/src/components/tl/self/tl_self.c @@ -17,13 +17,14 @@ ucc_status_t ucc_tl_self_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); ucc_status_t ucc_tl_self_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); ucc_status_t ucc_tl_self_get_lib_properties(ucc_base_lib_properties_t *prop); diff --git a/src/components/tl/self/tl_self_context.c b/src/components/tl/self/tl_self_context.c index 4b21c5601b..6c11f3a428 100644 --- a/src/components/tl/self/tl_self_context.c +++ b/src/components/tl/self/tl_self_context.c @@ -50,20 +50,21 @@ ucc_tl_self_get_context_attr(const ucc_base_context_t *context, /* NOLINT */ return UCC_OK; } -ucc_status_t ucc_tl_self_mem_map(const ucc_base_context_t *context, - int type, void *address, - size_t len, - void *memh, void *tl_h) +ucc_status_t ucc_tl_self_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_self_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_self_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_SUPPORTED; } diff --git a/src/components/tl/sharp/tl_sharp.c b/src/components/tl/sharp/tl_sharp.c index 705f1aa3dd..cea76856f2 100644 --- a/src/components/tl/sharp/tl_sharp.c +++ b/src/components/tl/sharp/tl_sharp.c @@ -16,13 +16,14 @@ ucc_status_t ucc_tl_sharp_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); ucc_status_t ucc_tl_sharp_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h); + void *address, size_t len, void *memh, + void *tl_h); -ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); -ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer); static ucc_config_field_t ucc_tl_sharp_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_sharp_lib_config_t, super), diff --git a/src/components/tl/sharp/tl_sharp_context.c b/src/components/tl/sharp/tl_sharp_context.c index 6a0b58c53b..8b23cbbe15 100644 --- a/src/components/tl/sharp/tl_sharp_context.c +++ b/src/components/tl/sharp/tl_sharp_context.c @@ -513,19 +513,20 @@ ucc_status_t ucc_tl_sharp_get_context_attr(const ucc_base_context_t *context, /* } ucc_status_t ucc_tl_sharp_mem_map(const ucc_base_context_t *context, int type, - void *address, - size_t len, - void *memh, void *tl_h) + void *address, size_t len, void *memh, + void *tl_h) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_sharp_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { return UCC_ERR_NOT_SUPPORTED; } -ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_sharp_memh_pack(const ucc_base_context_t *context, + void *memh, void **pack_buffer) { return UCC_ERR_NOT_SUPPORTED; } diff --git a/src/components/tl/ucp/alltoall/alltoall_onesided.c b/src/components/tl/ucp/alltoall/alltoall_onesided.c index 649f7c010e..3a6666908b 100644 --- a/src/components/tl/ucp/alltoall/alltoall_onesided.c +++ b/src/components/tl/ucp/alltoall/alltoall_onesided.c @@ -15,17 +15,17 @@ void ucc_tl_ucp_alltoall_onesided_progress(ucc_coll_task_t *ctask); ucc_status_t ucc_tl_ucp_alltoall_onesided_start(ucc_coll_task_t *ctask) { - ucc_tl_ucp_task_t *task = ucc_derived_of(ctask, ucc_tl_ucp_task_t); - ucc_tl_ucp_team_t *team = TASK_TEAM(task); - ptrdiff_t src = (ptrdiff_t)TASK_ARGS(task).src.info.buffer; - ptrdiff_t dest = (ptrdiff_t)TASK_ARGS(task).dst.info.buffer; - size_t nelems = TASK_ARGS(task).src.info.count; - ucc_rank_t grank = UCC_TL_TEAM_RANK(team); - ucc_rank_t gsize = UCC_TL_TEAM_SIZE(team); - ucc_rank_t start = (grank + 1) % gsize; - long * pSync = TASK_ARGS(task).global_work_buffer; - ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; - ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; + ucc_tl_ucp_task_t *task = ucc_derived_of(ctask, ucc_tl_ucp_task_t); + ucc_tl_ucp_team_t *team = TASK_TEAM(task); + ptrdiff_t src = (ptrdiff_t)TASK_ARGS(task).src.info.buffer; + ptrdiff_t dest = (ptrdiff_t)TASK_ARGS(task).dst.info.buffer; + size_t nelems = TASK_ARGS(task).src.info.count; + ucc_rank_t grank = UCC_TL_TEAM_RANK(team); + ucc_rank_t gsize = UCC_TL_TEAM_SIZE(team); + ucc_rank_t start = (grank + 1) % gsize; + long *pSync = TASK_ARGS(task).global_work_buffer; + ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; + ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; ucc_rank_t peer; ucc_tl_ucp_task_reset(task, UCC_INPROGRESS); @@ -33,17 +33,22 @@ ucc_status_t ucc_tl_ucp_alltoall_onesided_start(ucc_coll_task_t *ctask) nelems = (nelems / gsize) * ucc_dt_size(TASK_ARGS(task).src.info.datatype); dest = dest + grank * nelems; - UCPCHECK_GOTO(ucc_tl_ucp_put_nb((void *)(src + start * nelems), - (void *)dest, nelems, start, src_memh[start], dst_memh[start], team, task), - task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, src_memh[start], dst_memh[start], team), task, out); + UCPCHECK_GOTO( + ucc_tl_ucp_put_nb((void *)(src + start * nelems), (void *)dest, nelems, + start, src_memh[start], dst_memh[start], team, task), + task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, src_memh[start], + dst_memh[start], team), + task, out); for (peer = (start + 1) % gsize; peer != start; peer = (peer + 1) % gsize) { - UCPCHECK_GOTO(ucc_tl_ucp_put_nb((void *)(src + peer * nelems), - (void *)dest, nelems, peer, src_memh[peer], dst_memh[peer], team, task), - task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], dst_memh[peer], team), task, - out); + UCPCHECK_GOTO(ucc_tl_ucp_put_nb( + (void *)(src + peer * nelems), (void *)dest, nelems, + peer, src_memh[peer], dst_memh[peer], team, task), + task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], + dst_memh[peer], team), + task, out); } return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super); diff --git a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c index decfd470a5..ffbec5ee9a 100644 --- a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c +++ b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c @@ -24,8 +24,8 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_start(ucc_coll_task_t *ctask) ucc_aint_t *d_disp = TASK_ARGS(task).dst.info_v.displacements; size_t sdt_size = ucc_dt_size(TASK_ARGS(task).src.info_v.datatype); size_t rdt_size = ucc_dt_size(TASK_ARGS(task).dst.info_v.datatype); - ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; - ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; + ucc_mem_map_mem_h *src_memh = TASK_ARGS(task).src_memh.global_memh; + ucc_mem_map_mem_h *dst_memh = TASK_ARGS(task).dst_memh.global_memh; ucc_rank_t peer; size_t sd_disp, dd_disp, data_size; @@ -48,9 +48,11 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_start(ucc_coll_task_t *ctask) UCPCHECK_GOTO(ucc_tl_ucp_put_nb(PTR_OFFSET(src, sd_disp), PTR_OFFSET(dest, dd_disp), - data_size, peer, src_memh[peer], dst_memh[peer], team, task), + data_size, peer, src_memh[peer], + dst_memh[peer], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], dst_memh[peer], team), task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], + dst_memh[peer], team), task, out); } return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super); out: diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c index 2d508da50b..b10d7e56cf 100644 --- a/src/components/tl/ucp/tl_ucp.c +++ b/src/components/tl/ucp/tl_ucp.c @@ -32,16 +32,15 @@ ucc_status_t ucc_tl_ucp_get_lib_properties(ucc_base_lib_properties_t *prop); ucc_status_t ucc_tl_ucp_get_context_attr(const ucc_base_context_t *context, ucc_base_ctx_attr_t *base_attr); -ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, - int type, - void *address, - size_t len, - void *memh, +ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, void *tl_h); -ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer); +ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, + void **pack_buffer); -ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, void *memh); +ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, + void *memh); ucc_config_field_t ucc_tl_ucp_lib_config_table[] = { {"", "", NULL, ucc_offsetof(ucc_tl_ucp_lib_config_t, super), diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h index 59a063040f..83658a62d5 100644 --- a/src/components/tl/ucp/tl_ucp.h +++ b/src/components/tl/ucp/tl_ucp.h @@ -109,9 +109,9 @@ typedef struct ucc_tl_ucp_remote_info { typedef struct ucc_tl_ucp_memh_data { ucc_tl_ucp_remote_info_t rinfo; - void *packed_memh; - size_t packed_memh_len; - ucp_rkey_h rkey; + void *packed_memh; + size_t packed_memh_len; + ucp_rkey_h rkey; } ucc_tl_ucp_memh_data_t; typedef struct ucc_tl_ucp_worker { diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c index 7440d263bf..74dc39b74c 100644 --- a/src/components/tl/ucp/tl_ucp_context.c +++ b/src/components/tl/ucp/tl_ucp_context.c @@ -555,113 +555,127 @@ static void ucc_tl_ucp_ctx_remote_pack_data(ucc_tl_ucp_context_t *ctx, } } -ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, - int type, - void *address, - size_t len, - void *memh, - void *tl_h) +ucc_status_t ucc_tl_ucp_mem_map_memhbuf(ucc_tl_ucp_context_t *ctx, + void *pack_buffer, ucp_mem_h *mh) { - ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); - ucc_status_t ucc_status = UCC_OK; - ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *)tl_h; - ucc_tl_ucp_memh_data_t *m_data = (ucc_tl_ucp_memh_data_t *)p_memh->tl_data; ucp_mem_map_params_t mmap_params; - ucp_mem_h mh = NULL; ucs_status_t status; - ucp_memh_pack_params_t pack_params; - // basically an import here - if (type == UCC_MEM_MAP_TYPE_GLOBAL) { - // m_data is lost in the exchange, make a new one + *mh = NULL; + /* unpack here */ + size_t *key_size = (size_t *)pack_buffer; + void *packed_memh = PTR_OFFSET(pack_buffer, sizeof(size_t) * 2 + *key_size); + + mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_EXPORTED_MEMH_BUFFER; + mmap_params.exported_memh_buffer = packed_memh; + + status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, mh); + if (UCS_OK != status) { + if (status != UCS_ERR_UNREACHABLE) { + tl_error(ctx->super.super.lib, + "ucp_mem_map failed with error code: %s", + ucs_status_string(status)); + return ucs_status_to_ucc_status(status); + } else { + tl_debug(ctx->super.super.lib, + "ucp_mem_map could not map exported memory handle"); + } + } + return UCC_OK; +} + +ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, + void *address, size_t len, void *memh, + void *tl_h) +{ + ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); + ucc_status_t ucc_status = UCC_OK; + ucc_mem_map_tl_t *p_memh = (ucc_mem_map_tl_t *)tl_h; + ucc_tl_ucp_memh_data_t *m_data = (ucc_tl_ucp_memh_data_t *)p_memh->tl_data; + ucp_mem_h mh = NULL; + ucc_mem_map_memh_t *l_memh = (ucc_mem_map_memh_t *)memh; + size_t offset = 0; + ucp_mem_map_params_t mmap_params; + ucs_status_t status; + ucp_memh_pack_params_t pack_params; + + if (type == UCC_MEM_MAP_TYPE_GLOBAL || !m_data) { + /* either we are importing or m_data is null */ m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); if (!m_data) { tl_error(ctx->super.super.lib, "failed to allocate tl data"); return UCC_ERR_NO_MEMORY; } p_memh->tl_data = m_data; -#if defined(__aarch64__) - ucc_mem_map_memh_t * l_memh = (ucc_mem_map_memh_t *)memh; - /* unpack here */ - size_t *key_size = (size_t *)l_memh->pack_buffer; - void *packed_memh = PTR_OFFSET(l_memh->pack_buffer, sizeof(size_t) * 2 + *key_size); - mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_EXPORTED_MEMH_BUFFER; - mmap_params.exported_memh_buffer = packed_memh; - + } + if (type == UCC_MEM_MAP_TYPE_LOCAL) { + mmap_params.field_mask = + UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; + mmap_params.address = address; + mmap_params.length = len; + status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, &mh); if (UCS_OK != status) { - if (status != UCS_ERR_UNREACHABLE) { - tl_error(ctx->super.super.lib, - "ucp_mem_map failed with error code: %s", ucs_status_string(status)); - ucc_status = ucs_status_to_ucc_status(status); - return ucc_status; - } else { - tl_debug(ctx->super.super.lib, - "ucp_mem_map could not map exported memory handle"); - ucc_status = UCC_OK; // this is still OK - } - } else { - m_data->rinfo.mem_h = mh; - // the rest of the data is garbage. fix it - m_data->rinfo.va_base = address; - m_data->rinfo.len = len; + tl_error(ctx->super.super.lib, + "ucp_mem_map failed with error code: %d", status); + ucc_status = ucs_status_to_ucc_status(status); } -#endif } else { - if (!m_data) { - m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); - if (!m_data) { - tl_error(ctx->super.super.lib, "failed to allocate TL/UCP specific data"); - return UCC_ERR_NO_MEMORY; + for (int i = 0; i < l_memh->num_tls; i++) { + size_t *p = (size_t *)PTR_OFFSET(l_memh->pack_buffer, offset); + + if (tl_h == (void *)&l_memh->tl_h[i]) { + break; } - p_memh->tl_data = m_data; - } - // local, we need to map it here - if (m_data->rinfo.mem_h == NULL) { - mmap_params.field_mask = - UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; - mmap_params.address = address; - mmap_params.length = len; - - status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, &mh); - if (UCS_OK != status) { - tl_error(ctx->super.super.lib, - "ucp_mem_map failed with error code: %d", status); - ucc_status = ucs_status_to_ucc_status(status); + /* this is not the index, skip this section of buffer if exists */ + if (p[0] == i) { + offset += p[1]; } } - m_data->rinfo.mem_h = mh; - m_data->rinfo.va_base = address; - m_data->rinfo.len = len; - // export memh - pack_params.field_mask = UCP_MEMH_PACK_PARAM_FIELD_FLAGS; - pack_params.flags = UCP_MEMH_PACK_FLAG_EXPORT; + ucc_status = ucc_tl_ucp_mem_map_memhbuf( + ctx, PTR_OFFSET(l_memh->pack_buffer, offset), &mh); + if (ucc_status != UCC_OK) { + tl_error(ctx->super.super.lib, "ucp_mem_map failed to map memh"); + return ucc_status; + } + } + m_data->rinfo.mem_h = mh; + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; - status = ucp_memh_pack(mh, &pack_params, &m_data->packed_memh, &m_data->packed_memh_len); + if (type == UCC_MEM_MAP_TYPE_LOCAL) { + status = ucp_memh_pack(mh, &pack_params, &m_data->packed_memh, + &m_data->packed_memh_len); if (status != UCS_OK) { // we don't support memory pack, or it failed - tl_debug("ucp_memh_pack() returned error %s", ucs_status_string(status)); - m_data->packed_memh = 0; + tl_debug(ctx->super.super.lib, "ucp_memh_pack() returned error %s", + ucs_status_string(status)); + m_data->packed_memh = 0; m_data->packed_memh_len = 0; } // pack rkey - status = ucp_rkey_pack(ctx->worker.ucp_context, mh, &m_data->rinfo.packed_key, &m_data->rinfo.packed_key_len); + status = ucp_rkey_pack(ctx->worker.ucp_context, mh, + &m_data->rinfo.packed_key, + &m_data->rinfo.packed_key_len); if (status != UCS_OK) { - tl_error("unable to pack rkey with error %s", ucs_status_string(status)); + tl_error(ctx->super.super.lib, "unable to pack rkey with error %s", + ucs_status_string(status)); } - p_memh->packed_size = m_data->packed_memh_len + m_data->rinfo.packed_key_len; + p_memh->packed_size = + m_data->packed_memh_len + m_data->rinfo.packed_key_len; } return ucc_status; } -ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, void *memh) +ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, + void *memh) { - ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); - ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *) memh; + ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); + ucc_mem_map_tl_t *p_memh = (ucc_mem_map_tl_t *)memh; ucc_tl_ucp_memh_data_t *data; - ucs_status_t status; + ucs_status_t status; if (!p_memh) { return UCC_OK; @@ -672,7 +686,8 @@ ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, v if (type == UCC_MEM_MAP_TYPE_LOCAL) { status = ucp_mem_unmap(ctx->worker.ucp_context, data->rinfo.mem_h); if (status != UCS_OK) { - tl_error(ctx->super.super.lib, "ucp_mem_unmap failed with error code %d", status); + tl_error(ctx->super.super.lib, + "ucp_mem_unmap failed with error code %d", status); return ucs_status_to_ucc_status(status); } } else if (type == UCC_MEM_MAP_TYPE_GLOBAL) { @@ -693,13 +708,15 @@ ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, v return UCC_OK; } -ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, void **pack_buffer) +ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, + void **pack_buffer) { - ucc_mem_map_tl_t * p_memh = (ucc_mem_map_tl_t *) memh; - ucc_tl_ucp_memh_data_t *data = p_memh->tl_data; - void *packed_buffer; - size_t *key_size; - size_t *memh_size; + ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); + ucc_mem_map_tl_t *p_memh = (ucc_mem_map_tl_t *)memh; + ucc_tl_ucp_memh_data_t *data = p_memh->tl_data; + void *packed_buffer; + size_t *key_size; + size_t *memh_size; if (!data) { return UCC_OK; @@ -709,20 +726,27 @@ ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context, void *memh, * * packed_key_size | packed_memh_size | packed_key | packed_memh */ - packed_buffer = ucc_malloc(sizeof(size_t) * 2 + data->packed_memh_len + data->rinfo.packed_key_len, + packed_buffer = ucc_malloc(sizeof(size_t) * 2 + data->packed_memh_len + + data->rinfo.packed_key_len, "packed buffer"); if (!packed_buffer) { - ucc_error("failed to allocate a packed buffer of size %lu", data->packed_memh_len + data->rinfo.packed_key_len); + tl_error(ctx->super.super.lib, + "failed to allocate a packed buffer of size %lu", + data->packed_memh_len + data->rinfo.packed_key_len); return UCC_ERR_NO_MEMORY; } - key_size = packed_buffer; - *key_size = data->rinfo.packed_key_len; - memh_size = PTR_OFFSET(packed_buffer, sizeof(size_t)); + key_size = packed_buffer; + *key_size = data->rinfo.packed_key_len; + memh_size = PTR_OFFSET(packed_buffer, sizeof(size_t)); *memh_size = data->packed_memh_len; - memcpy(PTR_OFFSET(packed_buffer, sizeof(size_t) * 2), data->rinfo.packed_key, *key_size); - memcpy(PTR_OFFSET(packed_buffer, sizeof(size_t) * 2 + data->rinfo.packed_key_len), data->packed_memh, *memh_size); - - p_memh->packed_size = sizeof(size_t) * 2 + data->packed_memh_len + data->rinfo.packed_key_len; + memcpy(PTR_OFFSET(packed_buffer, sizeof(size_t) * 2), + data->rinfo.packed_key, *key_size); + memcpy(PTR_OFFSET(packed_buffer, + sizeof(size_t) * 2 + data->rinfo.packed_key_len), + data->packed_memh, *memh_size); + + p_memh->packed_size = + sizeof(size_t) * 2 + data->packed_memh_len + data->rinfo.packed_key_len; *pack_buffer = packed_buffer; return UCC_OK; } diff --git a/src/components/tl/ucp/tl_ucp_sendrecv.h b/src/components/tl/ucp/tl_ucp_sendrecv.h index fa02790abe..bcdb43cb3b 100644 --- a/src/components/tl/ucp/tl_ucp_sendrecv.h +++ b/src/components/tl/ucp/tl_ucp_sendrecv.h @@ -273,7 +273,6 @@ static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_ } } *rkey = tl_data->rkey; - /* FIXME: packed memh? */ return UCC_OK; } return UCC_ERR_NOT_FOUND; diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index 5586a85fbd..d63c144d60 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1132,16 +1132,16 @@ ucc_status_t ucc_context_get_attr(ucc_context_t *context, return status; } -ucc_status_t ucc_mem_map_import(ucc_context_h context, +ucc_status_t ucc_mem_map_import(ucc_context_h context, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { - ucc_context_t *ctx = (ucc_context_t *)context; - ucc_status_t status = UCC_OK; - ucc_config_names_array_t *tls = &ctx->all_tls; - int i; - ucc_mem_map_memh_t *local_memh; - ucc_tl_lib_t *tl_lib; + ucc_context_t *ctx = (ucc_context_t *)context; + ucc_status_t status = UCC_OK; + ucc_config_names_array_t *tls = &ctx->all_tls; + int i; + ucc_mem_map_memh_t *local_memh; + ucc_tl_lib_t *tl_lib; if (!memh) { ucc_error("cannot import NULL memory handle"); @@ -1163,8 +1163,9 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* FIXME: i don't think this will work properly for more than 1 TL */ status = tl_lib->iface->context.mem_map( - (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_GLOBAL, params->segments[0].address, params->segments[0].len, - local_memh, &local_memh->tl_h[i]); + (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_GLOBAL, + params->segments[0].address, params->segments[0].len, local_memh, + &local_memh->tl_h[i]); if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("failed to import mem map memh %d", status); return status; @@ -1174,27 +1175,28 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, local_memh->type = UCC_MEM_MAP_TYPE_GLOBAL; /* fix context as it will be incorrect on a different system */ local_memh->context = ctx; + *memh_size = 0; return status; } -ucc_status_t ucc_mem_map_export(ucc_context_h context, +ucc_status_t ucc_mem_map_export(ucc_context_h context, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { - ucc_context_t *ctx = (ucc_context_t *)context; - size_t total_pack_size = 0; - ucc_mem_map_memh_t *local_memh; - ucc_mem_map_memh_t *exported_memh; - void **packed_buffers; - ucc_status_t status; - ucc_tl_lib_t *tl_lib; - size_t offset; - int i; + ucc_context_t *ctx = (ucc_context_t *)context; + size_t total_pack_size = 0; + ucc_mem_map_memh_t *local_memh; + ucc_mem_map_memh_t *exported_memh; + void **packed_buffers; + ucc_status_t status; + ucc_tl_lib_t *tl_lib; + size_t offset; + int i; ucc_config_names_array_t *tls = &ctx->all_tls; - local_memh = (ucc_mem_map_memh_t *)ucc_calloc( - 1, sizeof(ucc_mem_map_memh_t), "local memh"); + local_memh = (ucc_mem_map_memh_t *)ucc_calloc(1, sizeof(ucc_mem_map_memh_t), + "local memh"); if (!local_memh) { ucc_error("failed to allocate a local memory handle"); return UCC_ERR_NO_MEMORY; @@ -1210,8 +1212,9 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* always treat as a local mem handle */ status = tl_lib->iface->context.mem_map( - (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_LOCAL, params->segments[0].address, params->segments[0].len, - local_memh, &local_memh->tl_h[i]); + (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_LOCAL, + params->segments[0].address, params->segments[0].len, local_memh, + &local_memh->tl_h[i]); if (status != UCC_OK) { if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("failed to map memory"); @@ -1231,7 +1234,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* tl should set packed_size, allocate buffer, pack memh */ status = tl_lib->iface->context.memh_pack( - (const ucc_base_context_t *)ctx->tl_ctx[i], &local_memh->tl_h[i], &packed_buffers[i]); + (const ucc_base_context_t *)ctx->tl_ctx[i], + &local_memh->tl_h[i], &packed_buffers[i]); if (status != UCC_OK) { if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("failed to map memory"); @@ -1244,7 +1248,10 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, /* allocate exported memh, copy items over */ exported_memh = (ucc_mem_map_memh_t *)ucc_calloc( - 1, sizeof(ucc_mem_map_memh_t) + total_pack_size + 2 * sizeof(size_t) * ctx->n_tl_ctx, "exported memh"); + 1, + sizeof(ucc_mem_map_memh_t) + total_pack_size + + 2 * sizeof(size_t) * ctx->n_tl_ctx, + "exported memh"); if (!exported_memh) { ucc_error("failed to allocate handle for exported buffers"); return UCC_ERR_NO_MEMORY; @@ -1257,8 +1264,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, if (local_memh->tl_h[i].packed_size == 0) { continue; } - memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), - &tl_index, sizeof(size_t)); + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), &tl_index, + sizeof(size_t)); offset += sizeof(size_t); memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), &local_memh->tl_h[i].packed_size, sizeof(size_t)); @@ -1272,20 +1279,22 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, strncpy(exported_memh->tl_h[i].tl_name, tls->names[i], 8); } ucc_free(local_memh); - exported_memh->type = UCC_MEM_MAP_TYPE_LOCAL; - exported_memh->context = ctx; - exported_memh->address = params->segments[0].address; - exported_memh->len = params->segments[0].len; + exported_memh->type = UCC_MEM_MAP_TYPE_LOCAL; + exported_memh->context = ctx; + exported_memh->address = params->segments[0].address; + exported_memh->len = params->segments[0].len; exported_memh->my_ctx_rank = ctx->rank; - exported_memh->num_tls = ctx->n_tl_ctx; - *memh = exported_memh; - *memh_size = total_pack_size; + exported_memh->num_tls = ctx->n_tl_ctx; + *memh = exported_memh; + *memh_size = total_pack_size; return UCC_OK; failed_pack: failed_mem_map: for (int j = 0; j < i; j++) { tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); - tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, UCC_MEM_MAP_TYPE_LOCAL, &local_memh->tl_h[j]); + tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, + UCC_MEM_MAP_TYPE_LOCAL, + &local_memh->tl_h[j]); } return status; } @@ -1307,7 +1316,6 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, // set map type to local return ucc_mem_map_export(context, params, memh_size, memh); } - return UCC_OK; } ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh) @@ -1324,11 +1332,13 @@ ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh) } /* it could be global or local type */ lmemh = (ucc_mem_map_memh_t *)*memh; - ctx = (ucc_context_t *)lmemh->context; + ctx = (ucc_context_t *)lmemh->context; for (i = 0; i < ctx->n_tl_ctx; i++) { tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); - status = tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx->tl_ctx[i], lmemh->type, &lmemh->tl_h[i]); + status = tl_lib->iface->context.mem_unmap( + (const ucc_base_context_t *)ctx->tl_ctx[i], lmemh->type, + &lmemh->tl_h[i]); if (status < UCC_ERR_NOT_IMPLEMENTED) { ucc_error("we had an error"); return status; diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index f8e255257d..916c87478a 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -101,7 +101,7 @@ typedef enum { typedef struct ucc_mem_map_tl_t { size_t packed_size; - char tl_name[8]; // typically less than 4 letters + char tl_name[8]; void *tl_data; /* tl specific data */ } ucc_mem_map_tl_t; @@ -110,9 +110,7 @@ typedef struct ucc_mem_map_memh_t { ucc_context_h context; void *address; size_t len; - /* rank of exporting process */ ucc_rank_t my_ctx_rank; - /* handles for each tl */ ucc_mem_map_tl_t *tl_h; int num_tls; char pack_buffer[0]; From 94523157cbf1ae2fb410ae5ac52fcbd61d0cd8cc Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Wed, 15 Jan 2025 16:44:53 -0800 Subject: [PATCH 12/14] REVIEW: address failures --- src/core/ucc_context.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index 916c87478a..23a583c5bd 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -13,6 +13,8 @@ #include "utils/ucc_proc_info.h" #include "components/topo/ucc_topo.h" +#define UCC_MEM_MAP_TL_NAME_LEN 8 + typedef struct ucc_lib_info ucc_lib_info_t; typedef struct ucc_cl_context ucc_cl_context_t; typedef struct ucc_tl_context ucc_tl_context_t; @@ -101,7 +103,7 @@ typedef enum { typedef struct ucc_mem_map_tl_t { size_t packed_size; - char tl_name[8]; + char tl_name[UCC_MEM_MAP_TL_NAME_LEN]; void *tl_data; /* tl specific data */ } ucc_mem_map_tl_t; From 033942ffd964434b286690f839887cc47dc84ffe Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 14 Feb 2025 10:15:25 -0800 Subject: [PATCH 13/14] REVIEW: various fixes for DPU --- .../tl/ucp/alltoall/alltoall_onesided.c | 8 +- src/components/tl/ucp/tl_ucp_context.c | 47 ++++-- src/components/tl/ucp/tl_ucp_sendrecv.h | 37 ++++- src/core/ucc_context.c | 140 ++++++++++++++++-- src/core/ucc_context.h | 4 +- src/ucc/api/ucc.h | 4 +- 6 files changed, 207 insertions(+), 33 deletions(-) diff --git a/src/components/tl/ucp/alltoall/alltoall_onesided.c b/src/components/tl/ucp/alltoall/alltoall_onesided.c index 3a6666908b..8d5b57b2e5 100644 --- a/src/components/tl/ucp/alltoall/alltoall_onesided.c +++ b/src/components/tl/ucp/alltoall/alltoall_onesided.c @@ -35,18 +35,18 @@ ucc_status_t ucc_tl_ucp_alltoall_onesided_start(ucc_coll_task_t *ctask) UCPCHECK_GOTO( ucc_tl_ucp_put_nb((void *)(src + start * nelems), (void *)dest, nelems, - start, src_memh[start], dst_memh[start], team, task), + start, *src_memh, dst_memh[start], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, src_memh[start], + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, start, *src_memh, dst_memh[start], team), task, out); for (peer = (start + 1) % gsize; peer != start; peer = (peer + 1) % gsize) { UCPCHECK_GOTO(ucc_tl_ucp_put_nb( (void *)(src + peer * nelems), (void *)dest, nelems, - peer, src_memh[peer], dst_memh[peer], team, task), + peer, *src_memh, dst_memh[peer], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, *src_memh, dst_memh[peer], team), task, out); } diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c index 74dc39b74c..e48ab11f26 100644 --- a/src/components/tl/ucp/tl_ucp_context.c +++ b/src/components/tl/ucp/tl_ucp_context.c @@ -166,10 +166,8 @@ UCC_CLASS_INIT_FUNC(ucc_tl_ucp_context_t, ucp_params.field_mask = UCP_PARAM_FIELD_FEATURES | UCP_PARAM_FIELD_TAG_SENDER_MASK | UCP_PARAM_FIELD_NAME; - ucp_params.features = UCP_FEATURE_TAG | UCP_FEATURE_AM; - if (params->params.mask & UCC_CONTEXT_PARAM_FIELD_MEM_PARAMS) { - ucp_params.features |= UCP_FEATURE_RMA | UCP_FEATURE_AMO64; - } + ucp_params.features = UCP_FEATURE_TAG | UCP_FEATURE_AM | UCP_FEATURE_EXPORTED_MEMH | + UCP_FEATURE_RMA | UCP_FEATURE_AMO64; ucp_params.tag_sender_mask = UCC_TL_UCP_TAG_SENDER_MASK; ucp_params.name = "UCC_UCP_CONTEXT"; @@ -563,14 +561,17 @@ ucc_status_t ucc_tl_ucp_mem_map_memhbuf(ucc_tl_ucp_context_t *ctx, *mh = NULL; /* unpack here */ - size_t *key_size = (size_t *)pack_buffer; - void *packed_memh = PTR_OFFSET(pack_buffer, sizeof(size_t) * 2 + *key_size); + size_t *key_size = (size_t *)PTR_OFFSET(pack_buffer, sizeof(size_t) * 2); + void *packed_memh = PTR_OFFSET(pack_buffer, sizeof(size_t) * 4 + *key_size); mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_EXPORTED_MEMH_BUFFER; mmap_params.exported_memh_buffer = packed_memh; status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, mh); if (UCS_OK != status) { + if (status == UCS_ERR_INVALID_PARAM) { + return UCC_OK; //it's ok. + } if (status != UCS_ERR_UNREACHABLE) { tl_error(ctx->super.super.lib, "ucp_mem_map failed with error code: %s", @@ -599,7 +600,8 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, ucs_status_t status; ucp_memh_pack_params_t pack_params; - if (type == UCC_MEM_MAP_TYPE_GLOBAL || !m_data) { + /* technically, only need to do this on import */ + if (type == UCC_MEM_MAP_TYPE_GLOBAL || type == UCC_MEM_MAP_TYPE_DPU_IMPORT || !m_data) { /* either we are importing or m_data is null */ m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); if (!m_data) { @@ -608,7 +610,7 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, } p_memh->tl_data = m_data; } - if (type == UCC_MEM_MAP_TYPE_LOCAL) { + if (type == UCC_MEM_MAP_TYPE_LOCAL) { /* this is an export */ mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; mmap_params.address = address; @@ -620,7 +622,10 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, "ucp_mem_map failed with error code: %d", status); ucc_status = ucs_status_to_ucc_status(status); } - } else { + m_data->rinfo.mem_h = mh; + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; + } else if (type == UCC_MEM_MAP_TYPE_DPU_IMPORT) { for (int i = 0; i < l_memh->num_tls; i++) { size_t *p = (size_t *)PTR_OFFSET(l_memh->pack_buffer, offset); @@ -639,10 +644,10 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, tl_error(ctx->super.super.lib, "ucp_mem_map failed to map memh"); return ucc_status; } + m_data->rinfo.mem_h = mh; /* used for xgvmi */ + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; } - m_data->rinfo.mem_h = mh; - m_data->rinfo.va_base = address; - m_data->rinfo.len = len; if (type == UCC_MEM_MAP_TYPE_LOCAL) { status = ucp_memh_pack(mh, &pack_params, &m_data->packed_memh, @@ -654,6 +659,7 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, m_data->packed_memh = 0; m_data->packed_memh_len = 0; } + // pack rkey status = ucp_rkey_pack(ctx->worker.ucp_context, mh, &m_data->rinfo.packed_key, @@ -662,9 +668,22 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, tl_error(ctx->super.super.lib, "unable to pack rkey with error %s", ucs_status_string(status)); } - p_memh->packed_size = - m_data->packed_memh_len + m_data->rinfo.packed_key_len; + } + if (type == UCC_MEM_MAP_TYPE_DPU_EXPORT) { + // pack rkey + status = ucp_rkey_pack(ctx->worker.ucp_context, m_data->rinfo.mem_h, + &m_data->rinfo.packed_key, + &m_data->rinfo.packed_key_len); + if (status != UCS_OK) { + tl_error(ctx->super.super.lib, "unable to pack rkey with error %s", + ucs_status_string(status)); + } + /* memh is already there.. reuse it */ + l_memh->type = UCC_MEM_MAP_TYPE_DPU_EXPORT; + } + p_memh->packed_size = + m_data->packed_memh_len + m_data->rinfo.packed_key_len; return ucc_status; } diff --git a/src/components/tl/ucp/tl_ucp_sendrecv.h b/src/components/tl/ucp/tl_ucp_sendrecv.h index bcdb43cb3b..d3195ffc7b 100644 --- a/src/components/tl/ucp/tl_ucp_sendrecv.h +++ b/src/components/tl/ucp/tl_ucp_sendrecv.h @@ -239,6 +239,24 @@ static inline ucc_status_t find_tl_index(ucc_mem_map_mem_h map_memh, int *tl_ind return UCC_ERR_NOT_FOUND; } +static inline ucc_status_t ucc_tl_ucp_get_memh(ucc_tl_ucp_team_t *team, ucc_mem_map_mem_h map_memh, void **ucp_memh) +{ + ucc_mem_map_memh_t *memh = map_memh; + ucc_tl_ucp_memh_data_t *tl_data;// = (ucc_tl_ucp_memh_data_t *)memh->tl_h[tl_index].tl_data; + int tl_index = 0; + ucc_status_t status; + + status = find_tl_index(memh, &tl_index); + if (status == UCC_ERR_NOT_FOUND) { + tl_error(UCC_TL_TEAM_LIB(team), + "attempt to perform one-sided operation with malformed mem map handle"); + return status; + } + tl_data = (ucc_tl_ucp_memh_data_t *)memh->tl_h[tl_index].tl_data; + *ucp_memh = tl_data->rinfo.mem_h; + return UCC_OK; +} + static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_t *rva, ucp_rkey_h *rkey, int tl_index, ucc_mem_map_mem_h map_memh) { @@ -249,6 +267,7 @@ static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_ ucs_status_t ucs_status; int i; size_t offset; + uint64_t *key_size; base = (uint64_t)memh->address; end = base + memh->len; @@ -271,6 +290,9 @@ static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_ if (UCS_OK != ucs_status) { return ucs_status_to_ucc_status(ucs_status); } + /* if they don't have a key, they don't have a memh */ + key_size = (uint64_t *)PTR_OFFSET(memh->pack_buffer, offset + sizeof(size_t) * 2); + tl_data->packed_memh = PTR_OFFSET(memh->pack_buffer, offset + sizeof(size_t) * 4 + *key_size); } *rkey = tl_data->rkey; return UCC_OK; @@ -322,7 +344,6 @@ ucc_tl_ucp_resolve_p2p_by_va(ucc_tl_ucp_team_t *team, void *va, ucp_ep_h *ep, ucc_status_t status; if (src_memh) { - //status = UCC_OK; status = find_tl_index(src_memh, &tl_index); if (status == UCC_ERR_NOT_FOUND) { tl_error(UCC_TL_TEAM_LIB(team), @@ -332,6 +353,7 @@ ucc_tl_ucp_resolve_p2p_by_va(ucc_tl_ucp_team_t *team, void *va, ucp_ep_h *ep, status = ucc_tl_ucp_check_memh(ep, va, rva, rkey, tl_index, src_memh); if (status == UCC_OK) { + printf("found %d va %p rva %lx\n", peer, va, *rva); return UCC_OK; } } @@ -422,22 +444,33 @@ static inline ucc_status_t ucc_tl_ucp_put_nb(void *buffer, void *target, ucs_status_ptr_t ucp_status; ucc_status_t status; ucp_ep_h ep; + void *ucp_memh = NULL; status = ucc_tl_ucp_get_ep(team, dest_group_rank, &ep); if (ucc_unlikely(UCC_OK != status)) { return status; } + status = ucc_tl_ucp_get_memh(team, src_memh, &ucp_memh); + if (ucc_unlikely(UCC_OK != status)) { + return status; + } + status = ucc_tl_ucp_resolve_p2p_by_va(team, target, &ep, dest_group_rank, &rva, &rkey, &segment, src_memh, dest_memh); if (ucc_unlikely(UCC_OK != status)) { return status; } + req_param.op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_USER_DATA; req_param.cb.send = ucc_tl_ucp_put_completion_cb; req_param.user_data = (void *)task; + if (ucp_memh) { + req_param.op_attr_mask |= UCP_OP_ATTR_FIELD_MEMH; + req_param.memh = ucp_memh; + } ucp_status = ucp_put_nbx(ep, buffer, msglen, rva, rkey, &req_param); @@ -467,6 +500,7 @@ static inline ucc_status_t ucc_tl_ucp_get_nb(void *buffer, void *target, ucs_status_ptr_t ucp_status; ucc_status_t status; ucp_ep_h ep; + //void *packed_memh; status = ucc_tl_ucp_get_ep(team, dest_group_rank, &ep); if (ucc_unlikely(UCC_OK != status)) { @@ -512,6 +546,7 @@ static inline ucc_status_t ucc_tl_ucp_atomic_inc(void * target, ucs_status_ptr_t ucp_status; ucc_status_t status; ucp_ep_h ep; +// void *packed_memh; status = ucc_tl_ucp_get_ep(team, dest_group_rank, &ep); if (ucc_unlikely(UCC_OK != status)) { diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index d63c144d60..9bdd895ddc 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1132,13 +1132,14 @@ ucc_status_t ucc_context_get_attr(ucc_context_t *context, return status; } -ucc_status_t ucc_mem_map_import(ucc_context_h context, +ucc_status_t ucc_mem_map_import(ucc_context_h context, ucc_mem_map_flags_t flags, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { ucc_context_t *ctx = (ucc_context_t *)context; ucc_status_t status = UCC_OK; ucc_config_names_array_t *tls = &ctx->all_tls; + ucc_mem_map_type_t type = (flags == UCC_MEM_MAP_IMPORT) ? UCC_MEM_MAP_TYPE_GLOBAL : UCC_MEM_MAP_TYPE_DPU_IMPORT; int i; ucc_mem_map_memh_t *local_memh; ucc_tl_lib_t *tl_lib; @@ -1155,7 +1156,6 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, /* FIXME: it's legal for the user to export and immediately import a handle, * is this an issue? */ local_memh = *memh; - /* memh should have been used in exchanges or from a remote process, addresses, etc likely garbage. fix it */ local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh"); @@ -1163,7 +1163,7 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); /* FIXME: i don't think this will work properly for more than 1 TL */ status = tl_lib->iface->context.mem_map( - (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_GLOBAL, + (const ucc_base_context_t *)ctx->tl_ctx[i], type, params->segments[0].address, params->segments[0].len, local_memh, &local_memh->tl_h[i]); if (status < UCC_ERR_NOT_IMPLEMENTED) { @@ -1172,7 +1172,7 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, } strncpy(local_memh->tl_h[i].tl_name, tls->names[i], 8); } - local_memh->type = UCC_MEM_MAP_TYPE_GLOBAL; + local_memh->type = type; /* fix context as it will be incorrect on a different system */ local_memh->context = ctx; *memh_size = 0; @@ -1286,7 +1286,117 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, exported_memh->my_ctx_rank = ctx->rank; exported_memh->num_tls = ctx->n_tl_ctx; *memh = exported_memh; - *memh_size = total_pack_size; + *memh_size = sizeof(ucc_mem_map_memh_t) + total_pack_size + + 2 * sizeof(size_t) * ctx->n_tl_ctx; + return UCC_OK; +failed_pack: +failed_mem_map: + for (int j = 0; j < i; j++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx, + UCC_MEM_MAP_TYPE_LOCAL, + &local_memh->tl_h[j]); + } + return status; +} + +ucc_status_t ucc_mem_map_dpu_export(ucc_context_h context, + ucc_mem_map_params_t *params, size_t *memh_size, + ucc_mem_map_mem_h *memh) +{ + ucc_context_t *ctx = (ucc_context_t *)context; + size_t total_pack_size = 0; + ucc_mem_map_memh_t *local_memh = *memh; + ucc_mem_map_memh_t *exported_memh; + void **packed_buffers; + ucc_status_t status; + ucc_tl_lib_t *tl_lib; + size_t offset; + int i; + ucc_config_names_array_t *tls = &ctx->all_tls; + + packed_buffers = + (void **)ucc_calloc(ctx->n_tl_ctx, sizeof(void *), "packed buffers"); + + /* map all the memory */ + for (i = 0; i < ctx->n_tl_ctx; i++) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + /* always treat as a local mem handle */ + status = tl_lib->iface->context.mem_map( + (const ucc_base_context_t *)ctx->tl_ctx[i], UCC_MEM_MAP_TYPE_DPU_EXPORT, + params->segments[0].address, params->segments[0].len, local_memh, + &local_memh->tl_h[i]); + if (status != UCC_OK) { + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("failed to map memory"); + goto failed_mem_map; + } + if (status == UCC_ERR_NOT_IMPLEMENTED || + status == UCC_ERR_NOT_SUPPORTED) { + /* either not implemented or not supported, set memh to null */ + local_memh->tl_h[i].packed_size = 0; + } + } + } + + /* now pack all the memories */ + for (i = 0; i < ctx->n_tl_ctx; i++) { + if (local_memh->tl_h[i].packed_size > 0) { + tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); + /* tl should set packed_size, allocate buffer, pack memh */ + status = tl_lib->iface->context.memh_pack( + (const ucc_base_context_t *)ctx->tl_ctx[i], + &local_memh->tl_h[i], &packed_buffers[i]); + if (status != UCC_OK) { + if (status < UCC_ERR_NOT_IMPLEMENTED) { + ucc_error("failed to map memory"); + goto failed_pack; + } + } + total_pack_size += local_memh->tl_h[i].packed_size; + } + } + /* allocate exported memh, copy items over */ + exported_memh = (ucc_mem_map_memh_t *)ucc_calloc( + 1, sizeof(ucc_mem_map_memh_t) + total_pack_size + + 2 * sizeof(size_t) * ctx->n_tl_ctx, + "exported memh"); + if (!exported_memh) { + ucc_error("failed to allocate handle for exported buffers"); + return UCC_ERR_NO_MEMORY; + } + exported_memh = local_memh; + + /* copying */ + exported_memh->tl_h = local_memh->tl_h; + for (i = 0, offset = 0; i < ctx->n_tl_ctx; i++) { + uint64_t tl_index = i; + if (local_memh->tl_h[i].packed_size == 0) { + continue; + } + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), &tl_index, + sizeof(size_t)); + offset += sizeof(size_t); + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + &exported_memh->tl_h[i].packed_size, sizeof(size_t)); + offset += sizeof(size_t); + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + packed_buffers[i], exported_memh->tl_h[i].packed_size); + ucc_free(packed_buffers[i]); + offset += exported_memh->tl_h[i].packed_size; + + // copy name information for look ups later + strncpy(exported_memh->tl_h[i].tl_name, tls->names[i], 8); + } + exported_memh->type = UCC_MEM_MAP_TYPE_DPU_EXPORT; + exported_memh->context = ctx; + exported_memh->address = local_memh->address; //params->segments[0].address; + exported_memh->len = local_memh->len; //params->segments[0].len; + exported_memh->my_ctx_rank = ctx->rank; + exported_memh->num_tls = ctx->n_tl_ctx; + *memh = exported_memh; + *memh_size = sizeof(ucc_mem_map_memh_t) + total_pack_size + + 2 * sizeof(size_t) * ctx->n_tl_ctx; return UCC_OK; failed_pack: failed_mem_map: @@ -1303,16 +1413,22 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_flags_t flags, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { - if (params->n_segments > 1) { - ucc_error("UCC only supports one mapping per call"); - return UCC_ERR_INVALID_PARAM; - } - // check if flags is import / export - if (flags == UCC_MEM_MAP_IMPORT) { + if (flags == UCC_MEM_MAP_IMPORT || flags == UCC_MEM_MAP_IMPORT_OFFLOAD) { // set map type to global - return ucc_mem_map_import(context, params, memh_size, memh); + return ucc_mem_map_import(context, flags, params, memh_size, memh); + } else if (flags == UCC_MEM_MAP_EXPORT_OFFLOAD) { + /* similar to an export */ + if (params->n_segments > 1) { + ucc_error("UCC only supports one mapping per call"); + return UCC_ERR_INVALID_PARAM; + } + return ucc_mem_map_dpu_export(context, params, memh_size, memh); } else { + if (params->n_segments > 1) { + ucc_error("UCC only supports one mapping per call"); + return UCC_ERR_INVALID_PARAM; + } // set map type to local return ucc_mem_map_export(context, params, memh_size, memh); } diff --git a/src/core/ucc_context.h b/src/core/ucc_context.h index 23a583c5bd..e70b46659b 100644 --- a/src/core/ucc_context.h +++ b/src/core/ucc_context.h @@ -98,7 +98,9 @@ typedef struct ucc_context_config { typedef enum { UCC_MEM_MAP_TYPE_LOCAL, - UCC_MEM_MAP_TYPE_GLOBAL + UCC_MEM_MAP_TYPE_GLOBAL, + UCC_MEM_MAP_TYPE_DPU_IMPORT, /* special case */ + UCC_MEM_MAP_TYPE_DPU_EXPORT } ucc_mem_map_type_t; typedef struct ucc_mem_map_tl_t { diff --git a/src/ucc/api/ucc.h b/src/ucc/api/ucc.h index 8e194c7dcf..e8342cbfe2 100644 --- a/src/ucc/api/ucc.h +++ b/src/ucc/api/ucc.h @@ -2257,8 +2257,10 @@ ucc_status_t ucc_collective_triggered_post(ucc_ee_h ee, ucc_ev_t *ee_event); typedef enum { UCC_MEM_MAP_EXPORT = 0, /*!< Indicate ucc_mem_map() should export memory handles from TLs used by context */ - UCC_MEM_MAP_IMPORT = 1 /*!< Indicate ucc_mem_map() should import + UCC_MEM_MAP_IMPORT = 1, /*!< Indicate ucc_mem_map() should import memory handles from user memory handle */ + UCC_MEM_MAP_EXPORT_OFFLOAD = 2, + UCC_MEM_MAP_IMPORT_OFFLOAD = 3 } ucc_mem_map_flags_t; /** From 4f039a22f2d3ae415a3b5693cd53b10eb7728c84 Mon Sep 17 00:00:00 2001 From: ferrol aderholdt Date: Fri, 14 Feb 2025 16:16:06 -0800 Subject: [PATCH 14/14] fixes --- .../tl/ucp/alltoallv/alltoallv_onesided.c | 5 +- src/components/tl/ucp/tl_ucp.h | 18 +++ src/components/tl/ucp/tl_ucp_context.c | 129 ++++++++++++++++-- src/components/tl/ucp/tl_ucp_sendrecv.h | 1 - src/core/ucc_context.c | 23 ++-- 5 files changed, 148 insertions(+), 28 deletions(-) diff --git a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c index ffbec5ee9a..fa92011c8b 100644 --- a/src/components/tl/ucp/alltoallv/alltoallv_onesided.c +++ b/src/components/tl/ucp/alltoallv/alltoallv_onesided.c @@ -48,11 +48,10 @@ ucc_status_t ucc_tl_ucp_alltoallv_onesided_start(ucc_coll_task_t *ctask) UCPCHECK_GOTO(ucc_tl_ucp_put_nb(PTR_OFFSET(src, sd_disp), PTR_OFFSET(dest, dd_disp), - data_size, peer, src_memh[peer], + data_size, peer, *src_memh, dst_memh[peer], team, task), task, out); - UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, src_memh[peer], - dst_memh[peer], team), task, out); + UCPCHECK_GOTO(ucc_tl_ucp_atomic_inc(pSync, peer, *src_memh, dst_memh[peer], team), task, out); } return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super); out: diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h index 83658a62d5..eccf9c5c0e 100644 --- a/src/components/tl/ucp/tl_ucp.h +++ b/src/components/tl/ucp/tl_ucp.h @@ -197,6 +197,24 @@ extern ucc_config_field_t ucc_tl_ucp_lib_config_table[]; #define UCC_TL_UCP_REMOTE_RKEY(_ctx, _rank, _seg) \ ((_ctx)->rkeys[_rank * _ctx->n_rinfo_segs + _seg]) +#define UCC_TL_UCP_MEMH_TL_HEADERS 2 + +#define UCC_TL_UCP_MEMH_TL_PACKED_HEADERS 2 + +#define UCC_TL_UCP_MEMH_TL_HEADER_SIZE \ + (sizeof(size_t) * UCC_TL_UCP_MEMH_TL_HEADERS) + +#define UCC_TL_UCP_MEMH_TL_PACKED_HEADER_SIZE \ + (sizeof(size_t) * \ + (UCC_TL_UCP_MEMH_TL_HEADERS + UCC_TL_UCP_MEMH_TL_PACKED_HEADERS)) + +#define UCC_TL_UCP_MEMH_TL_KEY_SIZE(buffer) \ + (*(size_t *)(PTR_OFFSET(buffer, UCC_TL_UCP_MEMH_TL_HEADER_SIZE))) + +#define UCC_TL_UCP_MEMH_TL_PACKED_MEMH(buffer) \ + (PTR_OFFSET(buffer, UCC_TL_UCP_MEMH_TL_PACKED_HEADER_SIZE + \ + UCC_TL_UCP_MEMH_TL_KEY_SIZE(buffer))) + extern ucs_memory_type_t ucc_memtype_to_ucs[UCC_MEMORY_TYPE_LAST+1]; void ucc_tl_ucp_pre_register_mem(ucc_tl_ucp_team_t *team, void *addr, diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c index e48ab11f26..12878b21a8 100644 --- a/src/components/tl/ucp/tl_ucp_context.c +++ b/src/components/tl/ucp/tl_ucp_context.c @@ -558,11 +558,13 @@ ucc_status_t ucc_tl_ucp_mem_map_memhbuf(ucc_tl_ucp_context_t *ctx, { ucp_mem_map_params_t mmap_params; ucs_status_t status; + void *packed_memh; *mh = NULL; /* unpack here */ - size_t *key_size = (size_t *)PTR_OFFSET(pack_buffer, sizeof(size_t) * 2); - void *packed_memh = PTR_OFFSET(pack_buffer, sizeof(size_t) * 4 + *key_size); + packed_memh = UCC_TL_UCP_MEMH_TL_PACKED_MEMH(pack_buffer); +// size_t *key_size = (size_t *)PTR_OFFSET(pack_buffer, sizeof(size_t) * 2); +// void *packed_memh = PTR_OFFSET(pack_buffer, sizeof(size_t) * 4 + *key_size); mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_EXPORTED_MEMH_BUFFER; mmap_params.exported_memh_buffer = packed_memh; @@ -585,23 +587,113 @@ ucc_status_t ucc_tl_ucp_mem_map_memhbuf(ucc_tl_ucp_context_t *ctx, return UCC_OK; } +ucc_status_t ucc_tl_ucp_mem_map_export(ucc_tl_ucp_context_t *ctx, void *address, + size_t len, + int type, + ucc_tl_ucp_memh_data_t *m_data) +{ + ucc_status_t ucc_status = UCC_OK; + ucp_mem_h mh; + ucp_mem_map_params_t mmap_params; + ucp_memh_pack_params_t pack_params; + ucs_status_t status; + + if (type == UCC_MEM_MAP_TYPE_LOCAL) { + mmap_params.field_mask = + UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; + mmap_params.address = address; + mmap_params.length = len; + + status = ucp_mem_map(ctx->worker.ucp_context, &mmap_params, &mh); + if (UCS_OK != status) { + tl_error(ctx->super.super.lib, "ucp_mem_map failed with error code: %d", + status); + ucc_status = ucs_status_to_ucc_status(status); + return ucc_status; + } + m_data->rinfo.mem_h = mh; + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; + + status = ucp_memh_pack(mh, &pack_params, &m_data->packed_memh, + &m_data->packed_memh_len); + if (status != UCS_OK) { + // we don't support memory pack, or it failed + tl_debug(ctx->super.super.lib, "ucp_memh_pack() returned error %s", + ucs_status_string(status)); + m_data->packed_memh = 0; + m_data->packed_memh_len = 0; + } + } + + // pack rkey + status = + ucp_rkey_pack(ctx->worker.ucp_context, mh, &m_data->rinfo.packed_key, + &m_data->rinfo.packed_key_len); + if (status != UCS_OK) { + tl_error(ctx->super.super.lib, "unable to pack rkey with error %s", + ucs_status_string(status)); + ucp_mem_unmap(ctx->worker.ucp_context, mh); + ucc_status = ucs_status_to_ucc_status(status); + return ucc_status; + } + + return ucc_status; +} + +ucc_status_t ucc_tl_ucp_mem_map_dpu_import(ucc_tl_ucp_context_t *ctx, + void *address, size_t len, + ucc_tl_ucp_memh_data_t *m_data, + ucc_mem_map_memh_t *l_memh, + void *tl_h) +{ + ucp_mem_h mh; + ucc_status_t ucc_status; + size_t offset = 0; + + for (int i = 0; i < l_memh->num_tls; i++) { + size_t *p = (size_t *)PTR_OFFSET(l_memh->pack_buffer, offset); + + if (tl_h == (void *)&l_memh->tl_h[i]) { + break; + } + /* this is not the index, skip this section of buffer if exists */ + if (p[0] == i) { + offset += p[1]; + } + } + + ucc_status = ucc_tl_ucp_mem_map_memhbuf( + ctx, PTR_OFFSET(l_memh->pack_buffer, offset), &mh); + if (ucc_status != UCC_OK) { + tl_error(ctx->super.super.lib, "ucp_mem_map failed to map memh"); + return ucc_status; + } + m_data->rinfo.mem_h = mh; /* used for xgvmi */ + m_data->rinfo.va_base = address; + m_data->rinfo.len = len; + + return UCC_OK; +} + ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, void *address, size_t len, void *memh, void *tl_h) { ucc_tl_ucp_context_t *ctx = ucc_derived_of(context, ucc_tl_ucp_context_t); - ucc_status_t ucc_status = UCC_OK; - ucc_mem_map_tl_t *p_memh = (ucc_mem_map_tl_t *)tl_h; + ucc_mem_map_tl_t *p_memh = (ucc_mem_map_tl_t *)tl_h; ucc_tl_ucp_memh_data_t *m_data = (ucc_tl_ucp_memh_data_t *)p_memh->tl_data; - ucp_mem_h mh = NULL; ucc_mem_map_memh_t *l_memh = (ucc_mem_map_memh_t *)memh; - size_t offset = 0; - ucp_mem_map_params_t mmap_params; - ucs_status_t status; - ucp_memh_pack_params_t pack_params; +// size_t offset = 0; +// ucp_mem_h mh = NULL; + ucc_status_t ucc_status = UCC_OK; +// ucp_mem_map_params_t mmap_params; +// ucs_status_t status; +// ucp_memh_pack_params_t pack_params; /* technically, only need to do this on import */ - if (type == UCC_MEM_MAP_TYPE_GLOBAL || type == UCC_MEM_MAP_TYPE_DPU_IMPORT || !m_data) { + if (type == UCC_MEM_MAP_TYPE_GLOBAL || + type == UCC_MEM_MAP_TYPE_DPU_IMPORT || !m_data) { /* either we are importing or m_data is null */ m_data = ucc_calloc(1, sizeof(ucc_tl_ucp_memh_data_t), "tl data"); if (!m_data) { @@ -610,6 +702,19 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, } p_memh->tl_data = m_data; } + + if (type == UCC_MEM_MAP_TYPE_LOCAL || type == UCC_MEM_MAP_TYPE_DPU_EXPORT) { + ucc_status = ucc_tl_ucp_mem_map_export(ctx, address, len, type, m_data); + } else if (type == UCC_MEM_MAP_TYPE_DPU_IMPORT) { + ucc_status = ucc_tl_ucp_mem_map_dpu_import(ctx, address, len, m_data, l_memh, tl_h); + } + + p_memh->packed_size = + m_data->packed_memh_len + m_data->rinfo.packed_key_len; + + return ucc_status; +} +#if 0 if (type == UCC_MEM_MAP_TYPE_LOCAL) { /* this is an export */ mmap_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH; @@ -668,7 +773,6 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, tl_error(ctx->super.super.lib, "unable to pack rkey with error %s", ucs_status_string(status)); } - } if (type == UCC_MEM_MAP_TYPE_DPU_EXPORT) { // pack rkey @@ -683,10 +787,11 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, int type, l_memh->type = UCC_MEM_MAP_TYPE_DPU_EXPORT; } p_memh->packed_size = - m_data->packed_memh_len + m_data->rinfo.packed_key_len; + m_data->packed_memh_len + m_data->rinfo.packed_key_len; return ucc_status; } +#endif ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, int type, void *memh) diff --git a/src/components/tl/ucp/tl_ucp_sendrecv.h b/src/components/tl/ucp/tl_ucp_sendrecv.h index d3195ffc7b..a3f30b05f4 100644 --- a/src/components/tl/ucp/tl_ucp_sendrecv.h +++ b/src/components/tl/ucp/tl_ucp_sendrecv.h @@ -546,7 +546,6 @@ static inline ucc_status_t ucc_tl_ucp_atomic_inc(void * target, ucs_status_ptr_t ucp_status; ucc_status_t status; ucp_ep_h ep; -// void *packed_memh; status = ucc_tl_ucp_get_ep(team, dest_group_rank, &ep); if (ucc_unlikely(UCC_OK != status)) { diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index 9bdd895ddc..cc04b65be2 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1132,14 +1132,17 @@ ucc_status_t ucc_context_get_attr(ucc_context_t *context, return status; } -ucc_status_t ucc_mem_map_import(ucc_context_h context, ucc_mem_map_flags_t flags, +ucc_status_t ucc_mem_map_import(ucc_context_h context, + ucc_mem_map_flags_t flags, ucc_mem_map_params_t *params, size_t *memh_size, ucc_mem_map_mem_h *memh) { ucc_context_t *ctx = (ucc_context_t *)context; ucc_status_t status = UCC_OK; ucc_config_names_array_t *tls = &ctx->all_tls; - ucc_mem_map_type_t type = (flags == UCC_MEM_MAP_IMPORT) ? UCC_MEM_MAP_TYPE_GLOBAL : UCC_MEM_MAP_TYPE_DPU_IMPORT; + ucc_mem_map_type_t type = (flags == UCC_MEM_MAP_IMPORT) + ? UCC_MEM_MAP_TYPE_GLOBAL + : UCC_MEM_MAP_TYPE_DPU_IMPORT; int i; ucc_mem_map_memh_t *local_memh; ucc_tl_lib_t *tl_lib; @@ -1153,15 +1156,12 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, ucc_mem_map_flags return UCC_ERR_INVALID_PARAM; } - /* FIXME: it's legal for the user to export and immediately import a handle, - * is this an issue? */ local_memh = *memh; /* memh should have been used in exchanges or from a remote process, addresses, etc likely garbage. fix it */ local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc( ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh"); for (i = 0; i < ctx->n_tl_ctx; i++) { tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t); - /* FIXME: i don't think this will work properly for more than 1 TL */ status = tl_lib->iface->context.mem_map( (const ucc_base_context_t *)ctx->tl_ctx[i], type, params->segments[0].address, params->segments[0].len, local_memh, @@ -1185,6 +1185,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, ucc_mem_map_mem_h *memh) { ucc_context_t *ctx = (ucc_context_t *)context; + ucc_config_names_array_t *tls = &ctx->all_tls; size_t total_pack_size = 0; ucc_mem_map_memh_t *local_memh; ucc_mem_map_memh_t *exported_memh; @@ -1193,7 +1194,6 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, ucc_tl_lib_t *tl_lib; size_t offset; int i; - ucc_config_names_array_t *tls = &ctx->all_tls; local_memh = (ucc_mem_map_memh_t *)ucc_calloc(1, sizeof(ucc_mem_map_memh_t), "local memh"); @@ -1286,8 +1286,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, exported_memh->my_ctx_rank = ctx->rank; exported_memh->num_tls = ctx->n_tl_ctx; *memh = exported_memh; - *memh_size = sizeof(ucc_mem_map_memh_t) + total_pack_size + - 2 * sizeof(size_t) * ctx->n_tl_ctx; + *memh_size = sizeof(ucc_mem_map_memh_t) + offset;// total_pack_size + 2 * sizeof(size_t) * ctx->n_tl_ctx; return UCC_OK; failed_pack: failed_mem_map: @@ -1390,13 +1389,13 @@ ucc_status_t ucc_mem_map_dpu_export(ucc_context_h context, } exported_memh->type = UCC_MEM_MAP_TYPE_DPU_EXPORT; exported_memh->context = ctx; - exported_memh->address = local_memh->address; //params->segments[0].address; - exported_memh->len = local_memh->len; //params->segments[0].len; + exported_memh->address = local_memh->address; + exported_memh->len = local_memh->len; exported_memh->my_ctx_rank = ctx->rank; exported_memh->num_tls = ctx->n_tl_ctx; *memh = exported_memh; - *memh_size = sizeof(ucc_mem_map_memh_t) + total_pack_size + - 2 * sizeof(size_t) * ctx->n_tl_ctx; + *memh_size = sizeof(ucc_mem_map_memh_t) + offset;//total_pack_size + 2 * sizeof(size_t) * ctx->n_tl_ctx; + ucc_free(local_memh); return UCC_OK; failed_pack: failed_mem_map: