TL/UCP: force enablement of dynamic segments

openucx · Feb 21, 2024 · 845d045 · 845d045
1 parent cc71ef8
commit 845d045
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c
@@ -189,6 +189,11 @@ ucc_config_field_t ucc_tl_ucp_lib_config_table[] = {
      ucc_offsetof(ucc_tl_ucp_lib_config_t, use_reordering),
      UCC_CONFIG_TYPE_BOOL},
 
+    {"USE_DYNAMIC_SEGMENTS", "n",
+     "Use dynamic segments in TL UCP for onesided collectives",
+     ucc_offsetof(ucc_tl_ucp_lib_config_t, use_dynamic_segments),
+     UCC_CONFIG_TYPE_BOOL},
+
     {NULL}};
 
 static ucs_config_field_t ucc_tl_ucp_context_config_table[] = {

diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h
@@ -74,6 +74,7 @@ typedef struct ucc_tl_ucp_lib_config {
     uint32_t                 alltoallv_hybrid_pairwise_num_posts;
     ucc_ternary_auto_value_t use_topo;
     int                      use_reordering;
+    int                      use_dynamic_segments;
 } ucc_tl_ucp_lib_config_t;
 
 typedef struct ucc_tl_ucp_context_config {

diff --git a/src/components/tl/ucp/tl_ucp_coll.c b/src/components/tl/ucp/tl_ucp_coll.c
@@ -257,11 +257,12 @@ ucc_status_t ucc_tl_ucp_coll_dynamic_segments(ucc_coll_args_t   *coll_args,
                                               ucc_tl_ucp_task_t *task)
 {
     ucc_tl_ucp_team_t    *tl_team        = UCC_TL_UCP_TASK_TEAM(task);
+    ucc_tl_ucp_lib_t     *tl_lib         = UCC_TL_UCP_TEAM_LIB(tl_team);
     ucc_tl_ucp_context_t *ctx            = UCC_TL_UCP_TEAM_CTX(tl_team);
     int                   i              = 0;
     ucc_status_t          status;
 
-    if (coll_args->mem_map.n_segments > 0) {
+    if (tl_lib->cfg.use_dynamic_segments && coll_args->mem_map.n_segments > 0) {
         int                   starting_index = ctx->n_dynrinfo_segs;
         size_t                seg_pack_size  = 0;
         size_t               *global_size    = NULL;

diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c
@@ -163,8 +163,12 @@ 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 | UCP_FEATURE_RMA |
-                          UCP_FEATURE_AMO64 | UCP_FEATURE_EXPORTED_MEMH;
+    ucp_params.features = UCP_FEATURE_TAG | UCP_FEATURE_AM;
+    if ((params->params.mask & UCC_CONTEXT_PARAM_FIELD_MEM_PARAMS) ||
+         lib->cfg.use_dynamic_segments) {
+        ucp_params.features |= UCP_FEATURE_RMA | UCP_FEATURE_AMO64 |
+                               UCP_FEATURE_EXPORTED_MEMH;
+    }
     ucp_params.tag_sender_mask = UCC_TL_UCP_TAG_SENDER_MASK;
     ucp_params.name = "UCC_UCP_CONTEXT";