Add support for ucc_mem_map and ucc_mem_unmap

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,          \

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);
 

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

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[] = {

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;
+}

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)},

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;
+}

src/components/tl/cuda/tl_cuda.c

@@ -64,6 +64,16 @@ 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 *);

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);

src/components/tl/mlx5/tl_mlx5.c

@@ -11,6 +11,16 @@ 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[] = {

src/components/tl/mlx5/tl_mlx5_context.c

@@ -302,3 +302,22 @@ 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;
+}

src/components/tl/nccl/tl_nccl.c

@@ -16,6 +16,16 @@ 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)},

src/components/tl/nccl/tl_nccl_context.c

@@ -223,3 +223,22 @@ 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;
+}

src/components/tl/rccl/tl_rccl.c

@@ -17,6 +17,16 @@ 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)},

src/components/tl/rccl/tl_rccl_context.c

@@ -120,3 +120,22 @@ 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;
+}

src/components/tl/self/tl_self.c

@@ -16,6 +16,16 @@ 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[] = {

src/components/tl/self/tl_self_context.c

@@ -49,3 +49,22 @@ 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;
+}

src/components/tl/sharp/tl_sharp.c

@@ -15,6 +15,16 @@ 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)},

src/components/tl/sharp/tl_sharp_context.c

@@ -511,3 +511,22 @@ 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;
+}

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;