From 845d045c3d711239dc977e2ee92cdbedad9cdacd Mon Sep 17 00:00:00 2001
From: ferrol aderholdt <faderholdt@nvidia.com>
Date: Wed, 21 Feb 2024 13:01:36 -0800
Subject: [PATCH] TL/UCP: force enablement of dynamic segments

---
 src/components/tl/ucp/tl_ucp.c         | 5 +++++
 src/components/tl/ucp/tl_ucp.h         | 1 +
 src/components/tl/ucp/tl_ucp_coll.c    | 3 ++-
 src/components/tl/ucp/tl_ucp_context.c | 8 ++++++--
 4 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c
index 72586dc5de..038c3ba4ca 100644
--- 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
index fc2d4acb07..67b3d73d4d 100644
--- 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
index 3a318d20cf..c25f1319a7 100644
--- 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
index 77efbba97f..c6ec78cf4e 100644
--- 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";