TL/UCP: disable proto that need memory type copy

openucx · Feb 24, 2025 · 0fa8262 · 0fa8262
1 parent bc996dd
commit 0fa8262
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -246,6 +246,12 @@ static ucs_config_field_t ucc_tl_ucp_context_config_table[] = {
      ucc_offsetof(ucc_tl_ucp_context_config_t, service_throttling_thresh),
      UCC_CONFIG_TYPE_UINT},
 
+    {"MEMTYPE_AVOID_COPY", "n",
+     "Avoid memory type copies when activated.",
+     ucc_offsetof(ucc_tl_ucp_context_config_t, avoid_copy_mem_types),
+     UCC_CONFIG_TYPE_BOOL,
+    },
+
     {NULL}};
 
 UCC_CLASS_DEFINE_NEW_FUNC(ucc_tl_ucp_lib_t, ucc_base_lib_t,

diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -87,6 +87,7 @@ typedef struct ucc_tl_ucp_context_config {
     uint32_t                pre_reg_mem;
     uint32_t                service_worker;
     uint32_t                service_throttling_thresh;
+    int                     avoid_copy_mem_types;
 } ucc_tl_ucp_context_config_t;
 
 typedef ucc_tl_ucp_lib_config_t ucc_tl_ucp_team_config_t;

diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -163,6 +163,11 @@ UCC_CLASS_INIT_FUNC(ucc_tl_ucp_context_t,
     prefix[strlen(prefix) - 1] = '\0';
     UCP_CHECK(ucp_config_read(prefix, NULL, &ucp_config),
               "failed to read ucp configuration", err_cfg_read, self);
+    if (tl_ucp_config->avoid_copy_mem_types) {
+        UCP_CHECK(ucp_config_modify(ucp_config, "MEMTYPE_AVOID_COPY", "yes"),
+                  "failed to set memtype avoid copy option for UCX",
+                  err_cfg, self);
+    }
 
     ucp_params.field_mask =
         UCP_PARAM_FIELD_FEATURES | UCP_PARAM_FIELD_TAG_SENDER_MASK | UCP_PARAM_FIELD_NAME;