UCP/PROTO: Disable protocols that need memory type copy

src/ucp/am/eager_multi.c

@@ -48,8 +48,9 @@ ucp_am_eager_multi_bcopy_proto_probe(const ucp_proto_init_params_t *init_params)
         .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID
     };
 
-    if (!ucp_am_check_init_params(init_params, UCP_PROTO_AM_OP_ID_MASK,
-                                  UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
+    if (!ucp_am_check_init_params_avoid_copy(init_params,
+                                             UCP_PROTO_AM_OP_ID_MASK,
+                                             UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
         return;
     }
 
@@ -220,8 +221,9 @@ ucp_am_eager_multi_zcopy_proto_probe(const ucp_proto_init_params_t *init_params)
         .middle.tl_cap_flags = UCT_IFACE_FLAG_AM_ZCOPY
     };
 
-    if (!ucp_am_check_init_params(init_params, UCP_PROTO_AM_OP_ID_MASK,
-                                  UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
+    if (!ucp_am_check_init_params_avoid_copy(init_params,
+                                             UCP_PROTO_AM_OP_ID_MASK,
+                                             UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
         return;
     }
 

src/ucp/am/eager_single.c

@@ -116,8 +116,8 @@ ucp_am_eager_short_probe_common(const ucp_proto_init_params_t *init_params,
         .tl_cap_flags        = UCT_IFACE_FLAG_AM_SHORT
     };
 
-    if (!ucp_am_check_init_params(init_params, UCS_BIT(op_id),
-                                  UCP_PROTO_SELECT_OP_FLAG_AM_RNDV) ||
+    if (!ucp_am_check_init_params_avoid_copy(init_params, UCS_BIT(op_id),
+                                             UCP_PROTO_SELECT_OP_FLAG_AM_RNDV) ||
         !ucp_proto_is_short_supported(select_param)) {
         return;
     }
@@ -246,8 +246,8 @@ static void ucp_am_eager_single_bcopy_probe_common(
         .tl_cap_flags        = UCT_IFACE_FLAG_AM_BCOPY
     };
 
-    if (!ucp_am_check_init_params(init_params, UCS_BIT(op_id),
-                                  UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
+    if (!ucp_am_check_init_params_avoid_copy(init_params, UCS_BIT(op_id),
+                                             UCP_PROTO_SELECT_OP_FLAG_AM_RNDV)) {
         return;
     }
 
@@ -338,8 +338,8 @@ static void ucp_am_eager_single_zcopy_probe_common(
         .tl_cap_flags        = UCT_IFACE_FLAG_AM_ZCOPY
     };
 
-    if (!ucp_am_check_init_params(init_params, UCS_BIT(op_id),
-                                  UCP_PROTO_SELECT_OP_FLAG_AM_RNDV) ||
+    if (!ucp_am_check_init_params_avoid_copy(init_params, UCS_BIT(op_id),
+                                             UCP_PROTO_SELECT_OP_FLAG_AM_RNDV) ||
         (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG)) {
         return;
     }

src/ucp/am/ucp_am.inl

@@ -33,6 +33,14 @@ ucp_am_check_init_params(const ucp_proto_init_params_t *init_params,
              exclude_flags);
 }
 
+static UCS_F_ALWAYS_INLINE int
+ucp_am_check_init_params_avoid_copy(const ucp_proto_init_params_t *init_params,
+                                    uint64_t op_id_mask, uint16_t exclude_flags)
+{
+    return ucp_am_check_init_params(init_params, op_id_mask, exclude_flags) &&
+           !init_params->worker->context->config.ext.avoid_copy_mem_types;
+}
+
 static UCS_F_ALWAYS_INLINE int
 ucp_proto_config_is_am(const ucp_proto_config_t *proto_config)
 {

src/ucp/core/ucp_context.c

@@ -359,6 +359,10 @@ static ucs_config_field_t ucp_context_config_table[] = {
    ucs_offsetof(ucp_context_config_t, rndv_frag_mem_types),
    UCS_CONFIG_TYPE_BITMAP(ucs_memory_type_names)},
 
+  {"MEMTYPE_AVOID_COPY", "n",
+   "Avoid memory type copies when activated.\n",
+   ucs_offsetof(ucp_context_config_t, avoid_copy_mem_types), UCS_CONFIG_TYPE_BOOL},
+
   {"RNDV_PIPELINE_SEND_THRESH", "inf",
    "RNDV size threshold to enable sender side pipeline for mem type",
    ucs_offsetof(ucp_context_config_t, rndv_pipeline_send_thresh), UCS_CONFIG_TYPE_MEMUNITS},
@@ -1158,6 +1162,13 @@ static void ucp_add_tl_resource_if_enabled(
 
     if (ucp_is_resource_enabled(resource, config, aux_tls, &rsc_flags,
                                 dev_cfg_masks, tl_cfg_mask)) {
+        if (context->config.ext.avoid_copy_mem_types) {
+            if (!strcmp(resource->tl_name, "cuda_ipc")) {
+                ucs_debug("disabled cuda_ipc for memtype copy avoidance");
+                return;
+            }
+        }
+
         if ((resource->sys_device != UCS_SYS_DEVICE_ID_UNKNOWN) &&
             (resource->sys_device >= UCP_MAX_SYS_DEVICES)) {
             ucs_diag(UCT_TL_RESOURCE_DESC_FMT

src/ucp/core/ucp_context.h

@@ -90,6 +90,8 @@ typedef struct ucp_context_config {
     size_t                                 rndv_num_frags[UCS_MEMORY_TYPE_LAST];
     /** Memory types of fragments used for RNDV pipeline protocol */
     uint64_t                               rndv_frag_mem_types;
+    /** Avoid memtype copies when set to true */
+    int                                    avoid_copy_mem_types;
     /** RNDV pipeline send threshold */
     size_t                                 rndv_pipeline_send_thresh;
     /** Enabling 2-stage pipeline rndv protocol */

src/ucp/rma/amo_offload.c

@@ -149,6 +149,7 @@ static void ucp_proto_amo_probe(const ucp_proto_init_params_t *init_params,
                                 int is_memtype)
 {
     ucp_worker_h worker              = init_params->worker;
+    ucs_memory_type_t send_mem_type  = init_params->select_param->mem_type;
     ucs_memory_type_t reply_mem_type =
             init_params->select_param->op.reply.mem_type;
     ucp_proto_single_init_params_t params = {
@@ -181,6 +182,12 @@ static void ucp_proto_amo_probe(const ucp_proto_init_params_t *init_params,
         return;
     }
 
+    if ((!UCP_MEM_IS_ACCESSIBLE_FROM_CPU(send_mem_type) ||
+         !UCP_MEM_IS_ACCESSIBLE_FROM_CPU(reply_mem_type)) &&
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
+        return;
+    }
+
     if (op_id != UCP_OP_ID_AMO_POST) {
         params.super.flags |= UCP_PROTO_COMMON_INIT_FLAG_RESPONSE;
         if (!UCP_MEM_IS_ACCESSIBLE_FROM_CPU(reply_mem_type) &&

src/ucp/rma/amo_sw.c

@@ -431,6 +431,10 @@ static void ucp_proto_amo_sw_probe(const ucp_proto_init_params_t *init_params,
     const ucp_ep_config_key_lane_t *lane_config;
     const uct_iface_attr_t *iface_attr;
 
+    if (init_params->worker->context->config.ext.avoid_copy_mem_types) {
+        return;
+    }
+
     /* If the endpoint has device atomic lanes, it means the target worker
        expects only device atomics, so we cannot use SW atomics. */
     ucs_carray_for_each(lane_config, ep_config_key->lanes,
@@ -485,7 +489,8 @@ ucp_proto_amo_sw_fetch_probe(const ucp_proto_init_params_t *init_params)
     if (!ucp_proto_init_check_op(init_params,
                                  UCS_BIT(UCP_OP_ID_AMO_FETCH) |
                                  UCS_BIT(UCP_OP_ID_AMO_CSWAP)) ||
-        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG)) {
+        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rma/get_am.c

@@ -100,7 +100,8 @@ ucp_proto_get_am_bcopy_probe(const ucp_proto_init_params_t *init_params)
         .tl_cap_flags        = UCT_IFACE_FLAG_AM_BCOPY
     };
 
-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_GET))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_GET)) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rma/get_offload.c

@@ -108,7 +108,8 @@ ucp_proto_get_offload_bcopy_probe(const ucp_proto_init_params_t *init_params)
     };
 
     if ((init_params->select_param->dt_class != UCP_DATATYPE_CONTIG) ||
-        !ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_GET))) {
+        !ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_GET)) ||
+        context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rma/put_am.c

@@ -108,7 +108,8 @@ ucp_proto_put_am_bcopy_probe(const ucp_proto_init_params_t *init_params)
         .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID
     };
 
-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT)) ||
+        context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rma/put_offload.c

@@ -76,7 +76,8 @@ ucp_proto_put_offload_short_probe(const ucp_proto_init_params_t *init_params)
     };
 
     if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT)) ||
-        !ucp_proto_is_short_supported(init_params->select_param)) {
+        !ucp_proto_is_short_supported(init_params->select_param) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -178,7 +179,8 @@ ucp_proto_put_offload_bcopy_probe(const ucp_proto_init_params_t *init_params)
         .opt_align_offs      = UCP_PROTO_COMMON_OFFSET_INVALID
     };
 
-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_PUT)) ||
+        context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rndv/rndv_am.c

@@ -16,7 +16,8 @@ static void ucp_rndv_am_probe_common(ucp_proto_multi_init_params_t *params)
     ucp_context_h context = params->super.super.worker->context;
 
     if (!ucp_proto_rndv_op_check(&params->super.super, UCP_OP_ID_RNDV_SEND,
-                                 0)) {
+                                 0) ||
+        context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/rndv/rndv_mtype.inl

@@ -34,7 +34,8 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params,
 
     if ((init_params->select_param->dt_class != UCP_DATATYPE_CONTIG) ||
         (ucp_proto_rndv_mtype_ep(worker, frag_mem_type, mem_type) == NULL) ||
-        !ucp_proto_init_check_op(init_params, UCP_PROTO_RNDV_OP_ID_MASK)) {
+        !ucp_proto_init_check_op(init_params, UCP_PROTO_RNDV_OP_ID_MASK) ||
+        context->config.ext.avoid_copy_mem_types) {
         return UCS_ERR_UNSUPPORTED;
     }
 

src/ucp/rndv/rndv_ppln.c

@@ -321,7 +321,8 @@ static size_t ucp_proto_rndv_ppln_pack_ack(void *dest, void *arg)
 static void
 ucp_proto_rndv_send_ppln_probe(const ucp_proto_init_params_t *init_params)
 {
-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_RNDV_SEND))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_RNDV_SEND)) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -356,7 +357,8 @@ ucp_proto_t ucp_rndv_send_ppln_proto = {
 static void
 ucp_proto_rndv_recv_ppln_probe(const ucp_proto_init_params_t *init_params)
 {
-    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_RNDV_RECV))) {
+    if (!ucp_proto_init_check_op(init_params, UCS_BIT(UCP_OP_ID_RNDV_RECV)) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/stream/stream_multi.c

@@ -26,7 +26,8 @@ static void
 ucp_stream_multi_common_probe(const ucp_proto_multi_init_params_t *params)
 {
     if (!ucp_proto_init_check_op(&params->super.super,
-                                 UCS_BIT(UCP_OP_ID_STREAM_SEND))) {
+                                 UCS_BIT(UCP_OP_ID_STREAM_SEND)) ||
+        params->super.super.worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/tag/eager_multi.c

@@ -35,7 +35,8 @@ ucp_proto_eager_multi_probe_common(ucp_proto_multi_init_params_t *params,
 {
     ucp_context_config_t *context_config;
 
-    if (!ucp_tag_eager_check_op_id(&params->super.super, op_id, 0)) {
+    if (!ucp_tag_eager_check_op_id(&params->super.super, op_id, 0) ||
+        params->super.super.worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/tag/eager_single.c

@@ -74,7 +74,8 @@ ucp_proto_eager_short_probe(const ucp_proto_init_params_t *init_params)
 
     /* AM based proto can not be used if tag offload lane configured */
     if (!ucp_tag_eager_check_op_id(init_params, UCP_OP_ID_TAG_SEND, 0) ||
-        !ucp_proto_is_short_supported(select_param)) {
+        !ucp_proto_is_short_supported(select_param) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -147,7 +148,8 @@ ucp_proto_eager_bcopy_single_probe(const ucp_proto_init_params_t *init_params)
     };
 
     /* AM based proto can not be used if tag offload lane configured */
-    if (!ucp_tag_eager_check_op_id(init_params, UCP_OP_ID_TAG_SEND, 0)) {
+    if (!ucp_tag_eager_check_op_id(init_params, UCP_OP_ID_TAG_SEND, 0) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -197,7 +199,8 @@ ucp_proto_eager_zcopy_single_probe(const ucp_proto_init_params_t *init_params)
 
     /* AM based proto can not be used if tag offload lane configured */
     if (!ucp_tag_eager_check_op_id(init_params, UCP_OP_ID_TAG_SEND, 0) ||
-        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG)) {
+        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/tag/offload/eager.c

@@ -71,7 +71,8 @@ static void ucp_proto_eager_tag_offload_short_probe(
     };
 
     if (!ucp_tag_eager_check_op_id(init_params, UCP_OP_ID_TAG_SEND, 1) ||
-        !ucp_proto_is_short_supported(select_param)) {
+        !ucp_proto_is_short_supported(select_param) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -146,7 +147,8 @@ static void ucp_proto_eager_tag_offload_bcopy_probe_common(
     };
 
     /* offload proto can not be used if no tag offload lane configured */
-    if (!ucp_tag_eager_check_op_id(init_params, op_id, 1)) {
+    if (!ucp_tag_eager_check_op_id(init_params, op_id, 1) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 
@@ -259,7 +261,8 @@ static void ucp_proto_eager_tag_offload_zcopy_probe_common(
 
     /* offload proto can not be used if no tag offload lane configured */
     if (!ucp_tag_eager_check_op_id(init_params, op_id, 1) ||
-        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG)) {
+        (init_params->select_param->dt_class != UCP_DATATYPE_CONTIG) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }
 

src/ucp/tag/offload/rndv.c

@@ -192,7 +192,8 @@ ucp_tag_rndv_offload_sw_proto_probe(const ucp_proto_init_params_t *init_params)
     ucp_proto_rndv_ctrl_priv_t rpriv;
 
     if (!ucp_tag_rndv_check_op_id(init_params) ||
-        !ucp_ep_config_key_has_tag_lane(init_params->ep_config_key)) {
+        !ucp_ep_config_key_has_tag_lane(init_params->ep_config_key) ||
+        init_params->worker->context->config.ext.avoid_copy_mem_types) {
         return;
     }