UCP/PROTO: Add option to force ZCOPY (#11289)

tvegas1 · web-flow · commit fd3bc3a9be79 · 2026-05-04T17:47:37.000+02:00
diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c
@@ -471,6 +471,12 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "lane without waiting for remote completion.",
    ucs_offsetof(ucp_context_config_t, rndv_put_force_flush), UCS_CONFIG_TYPE_BOOL},
 
+  {"PROTO_EMULATION_ENABLE", "y",
+   "When set to 'no', emulation protocols for put and get are disabled. If no native\n"
+   "zero-copy RMA protocol exist for the memory type pair, RMA requests will be\n"
+   "cancelled.",
+   ucs_offsetof(ucp_context_config_t, proto_emulation_enable), UCS_CONFIG_TYPE_BOOL},
+
   {"SA_DATA_VERSION", "v2",
    "Defines the minimal header version the client will use for establishing\n"
    "client/server connection",
diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h
@@ -167,6 +167,9 @@ typedef struct ucp_context_config {
     uint64_t                               reg_whole_alloc_bitmap;
     /** Always use flush operation in rendezvous put */
     int                                    rndv_put_force_flush;
+    /** Allow RMA emulation protocols. When disabled, provide an explicit error
+      * if no suitable proto is found */
+    int                                    proto_emulation_enable;
     /** Maximum size of mem type direct rndv*/
     size_t                                 rndv_memtype_direct_size;
     /** UCP sockaddr private data format version */
diff --git a/src/ucp/proto/proto_reconfig.c b/src/ucp/proto/proto_reconfig.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2021. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2021-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -15,6 +15,7 @@
 
 #include <ucp/am/ucp_am.inl>
 #include <ucp/core/ucp_worker.inl>
+#include <ucs/memory/memory_type.h>
 #include <ucs/sys/math.h>
 
 
@@ -43,6 +44,37 @@ static void ucp_proto_reconfig_abort(ucp_request_t *req, ucs_status_t status)
     ucp_request_complete_send(req, status);
 }
 
+static int
+ucp_proto_reconfig_report_no_rma_emulation_no_proto(ucp_request_t *req,
+                                                    ucp_ep_h ep)
+{
+    ucp_operation_id_t op_id;
+    ucs_memory_type_t local_mem_type, remote_mem_type;
+
+    if (ep->worker->context->config.ext.proto_emulation_enable) {
+        return 0;
+    }
+
+    op_id = ucp_proto_select_op_id(&req->send.proto_config->select_param);
+    if (((op_id != UCP_OP_ID_PUT) && (op_id != UCP_OP_ID_GET))) {
+        return 0;
+    }
+
+    local_mem_type  = req->send.proto_config->select_param.mem_type;
+    remote_mem_type = req->send.rma.rkey->mem_type;
+
+    ucs_error("No zero-copy protocol found for %s %s %s %s, %zu bytes. "
+              "Please check for proper GPU and/or HCA support, or set "
+              "UCX_PROTO_EMULATION_ENABLE=y to proceed by allowing slower "
+              "software emulation.",
+              (op_id == UCP_OP_ID_PUT) ? "put from" : "get into",
+              ucs_memory_type_names[local_mem_type],
+              (op_id == UCP_OP_ID_PUT) ? "to" : "from",
+              ucs_memory_type_names[remote_mem_type],
+              req->send.state.dt_iter.length);
+    return 1;
+}
+
 static ucs_status_t ucp_proto_reconfig_progress(uct_pending_req_t *self)
 {
     ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct);
@@ -52,6 +84,11 @@ static ucs_status_t ucp_proto_reconfig_progress(uct_pending_req_t *self)
 
     /* This protocol should not be selected for valid and connected endpoint */
     if (ep->flags & UCP_EP_FLAG_REMOTE_CONNECTED) {
+        if (ucp_proto_reconfig_report_no_rma_emulation_no_proto(req, ep)) {
+            ucp_proto_request_abort(req, UCS_ERR_CANCELED);
+            return UCS_OK;
+        }
+
         ucp_ep_config_name(ep->worker, req->send.proto_config->ep_cfg_index,
                            &strb);
         ucs_string_buffer_appendf(&strb, " | ");
diff --git a/src/ucp/rma/get_am.c b/src/ucp/rma/get_am.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -111,6 +111,10 @@ ucp_proto_get_am_bcopy_probe(const ucp_proto_init_params_t *init_params)
         return;
     }
 
+    if (!init_params->worker->context->config.ext.proto_emulation_enable) {
+        return;
+    }
+
     ucp_proto_single_probe(&params);
 }
 
diff --git a/src/ucp/rma/put_am.c b/src/ucp/rma/put_am.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -119,6 +119,10 @@ ucp_proto_put_am_bcopy_probe(const ucp_proto_init_params_t *init_params)
         return;
     }
 
+    if (!init_params->worker->context->config.ext.proto_emulation_enable) {
+        return;
+    }
+
     ucp_proto_multi_probe(&params);
 }
 
diff --git a/test/gtest/ucp/test_ucp_rma.cc b/test/gtest/ucp/test_ucp_rma.cc
@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2015. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 * Copyright (c) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
@@ -53,6 +53,21 @@ class test_ucp_rma : public test_ucp_memheap {
         }
     }
 
+    ucs_status_ptr_t do_put(size_t size, void *expected_data, ucp_mem_h memh,
+                            void *target_ptr, ucp_rkey_h rkey, void *arg)
+    {
+        ucs_memory_type_t *mem_types = reinterpret_cast<ucs_memory_type_t*>(
+                arg);
+        mem_buffer::pattern_fill(expected_data, size, ucs::rand(),
+                                 mem_types[0]);
+
+        ucp_request_param_t param;
+        request_param_init(&param, memh);
+
+        return ucp_put_nbx(sender().ep(), expected_data, size,
+                           (uintptr_t)target_ptr, rkey, &param);
+    }
+
     void put_b(size_t size, void *expected_data, ucp_mem_h memh,
                void *target_ptr, ucp_rkey_h rkey, void *arg)
     {
@@ -76,6 +91,16 @@ class test_ucp_rma : public test_ucp_memheap {
                    rkey, arg);
     }
 
+    ucs_status_ptr_t do_get(size_t size, void *expected_data, ucp_mem_h memh,
+                            void *target_ptr, ucp_rkey_h rkey)
+    {
+        ucp_request_param_t param;
+        request_param_init(&param, memh);
+
+        return ucp_get_nbx(sender().ep(), expected_data, size,
+                           (uintptr_t)target_ptr, rkey, &param);
+    }
+
     void get_b(size_t size, void *expected_data, ucp_mem_h memh,
                void *target_ptr, ucp_rkey_h rkey, void *arg)
     {
@@ -118,7 +143,6 @@ class test_ucp_rma : public test_ucp_memheap {
                 ucs::supported_mem_type_pairs();
 
         for (size_t i = 0; i < pairs.size(); ++i) {
-
             /* Memory type put/get is fully supported only with new protocols */
             if (!is_proto_enabled() && (!UCP_MEM_IS_HOST(pairs[i][0]) ||
                                         !UCP_MEM_IS_HOST(pairs[i][1]))) {
@@ -209,19 +233,6 @@ class test_ucp_rma : public test_ucp_memheap {
         param->memh          = memh;
     }
 
-    ucs_status_ptr_t do_put(size_t size, void *expected_data, ucp_mem_h memh,
-                            void *target_ptr, ucp_rkey_h rkey, void *arg)
-    {
-        ucs_memory_type_t *mem_types = reinterpret_cast<ucs_memory_type_t*>(arg);
-        mem_buffer::pattern_fill(expected_data, size, ucs::rand(), mem_types[0]);
-
-        ucp_request_param_t param;
-        request_param_init(&param, memh);
-
-        return ucp_put_nbx(sender().ep(), expected_data, size,
-                           (uintptr_t)target_ptr, rkey, &param);
-    }
-
     ucs_status_ptr_t do_put_iov(size_t size, void *expected_data,
                                 ucp_request_param_t *param, void *target_ptr,
                                 ucp_rkey_h rkey, ucp_dt_iov_t *iov,
@@ -241,16 +252,6 @@ class test_ucp_rma : public test_ucp_memheap {
                            rkey, param);
     }
 
-    ucs_status_ptr_t do_get(size_t size, void *expected_data, ucp_mem_h memh,
-                            void *target_ptr, ucp_rkey_h rkey)
-    {
-        ucp_request_param_t param;
-        request_param_init(&param, memh);
-
-        return ucp_get_nbx(sender().ep(), expected_data, size,
-                           (uintptr_t)target_ptr, rkey, &param);
-    }
-
     ucs_status_ptr_t do_get_iov(size_t size, void *expected_data,
                                 ucp_request_param_t *param, void *target_ptr,
                                 ucp_rkey_h rkey, ucp_dt_iov_t *iov,
@@ -345,6 +346,114 @@ UCS_TEST_P(test_ucp_rma, proto_disabled_unsupported, "PROTO_ENABLE=n")
 UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_rma)
 
 
+class test_ucp_proto_emulation_enable : public test_ucp_rma {
+public:
+    static constexpr size_t SMALL_SIZE = 8;
+    static constexpr size_t BIG_SIZE   = 512 * UCS_KBYTE;
+
+    static void get_test_variants(std::vector<ucp_test_variant> &variants)
+    {
+        add_variant_with_value(variants, UCP_FEATURE_RMA, 0, "");
+    }
+
+    test_ucp_proto_emulation_enable()
+    {
+        modify_config("PROTO_EMULATION_ENABLE", "n");
+        modify_config("IB_TX_INLINE_RESP", "0", SETENV_IF_NOT_EXIST);
+    }
+
+protected:
+    enum rma_op_t {
+        RMA_OP_PUT,
+        RMA_OP_GET
+    };
+
+    void run_expect_canceled(rma_op_t op, size_t size)
+    {
+        mapped_buffer rbuf(size * 2, receiver());
+        ucs::handle<ucp_rkey_h> rkey = rbuf.rkey(sender());
+        mem_buffer lbuf(size, UCS_MEMORY_TYPE_HOST);
+        ucs_memory_type_t mem_types[] = {UCS_MEMORY_TYPE_HOST,
+                                         UCS_MEMORY_TYPE_HOST};
+
+        scoped_log_handler slh(wrap_errors_logger);
+
+        ucs_status_ptr_t req;
+        if (op == RMA_OP_PUT) {
+            req = do_put(size, lbuf.ptr(), NULL, rbuf.ptr(), rkey.get(),
+                         mem_types);
+        } else {
+            req = do_get(size, lbuf.ptr(), NULL, rbuf.ptr(), rkey.get());
+        }
+        ucs_status_t status = request_wait(req);
+        EXPECT_EQ(UCS_ERR_CANCELED, status)
+                << (op == RMA_OP_PUT ? "put" : "get") << " should be canceled";
+
+        /* Verify PROTO_EMULATION_ENABLE error message was logged */
+        bool found_rma_msg = false;
+        for (const auto &err : m_errors) {
+            if (err.find("set UCX_PROTO_EMULATION_ENABLE=y to proceed") !=
+                std::string::npos) {
+                found_rma_msg = true;
+                break;
+            }
+        }
+        EXPECT_TRUE(found_rma_msg) << "Expected error message with "
+                                      "UCX_PROTO_EMULATION_ENABLE=y advice";
+    }
+
+    void test_forced_message_sizes(send_func_t send_func)
+    {
+        for (const auto &pair : ucs::supported_mem_type_pairs()) {
+            if (check_reg_mem_types(sender(), pair[0]) &&
+                check_reg_mem_types(sender(), pair[1])) {
+                test_message_sizes(send_func, SMALL_SIZE, BIG_SIZE, pair[0],
+                                   pair[1], 0);
+            }
+        }
+    }
+};
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, no_zcopy_proto_fails_put_small,
+           "PROTOS=put/am/*,get/am/*,reconfig")
+{
+    run_expect_canceled(RMA_OP_PUT, SMALL_SIZE);
+}
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, no_zcopy_proto_fails_put_big,
+           "PROTOS=put/am/*,get/am/*,reconfig")
+{
+    run_expect_canceled(RMA_OP_PUT, BIG_SIZE);
+}
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, no_zcopy_proto_fails_get_small,
+           "PROTOS=put/am/*,get/am/*,reconfig")
+{
+    run_expect_canceled(RMA_OP_GET, SMALL_SIZE);
+}
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, no_zcopy_proto_fails_get_big,
+           "PROTOS=put/am/*,get/am/*,reconfig")
+{
+    run_expect_canceled(RMA_OP_GET, BIG_SIZE);
+}
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, get_zcopy_forced_success,
+           "PROTOS=get/bcopy,get/zcopy,reconfig")
+{
+    test_forced_message_sizes(static_cast<send_func_t>(&test_ucp_rma::get_b));
+}
+
+UCS_TEST_P(test_ucp_proto_emulation_enable, put_zcopy_forced_success,
+           "PROTOS=put/offload/*,reconfig")
+{
+    test_forced_message_sizes(static_cast<send_func_t>(&test_ucp_rma::put_b));
+}
+
+UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_proto_emulation_enable, ib,
+                                        "ib")
+
+
 class test_ucp_rma_reg : public test_ucp_rma {
 public:
     static void get_test_variants(std::vector<ucp_test_variant>& variants) {

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`/**`
`2`		`- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.`
	`2`	`+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.`
`3`	`3`	`*`
`4`	`4`	`* See file LICENSE for terms.`
`5`	`5`	`*/`
`@@ -111,6 +111,10 @@ ucp_proto_get_am_bcopy_probe(const ucp_proto_init_params_t *init_params)`
`111`	`111`	`return;`
`112`	`112`	`}`
`113`	`113`
	`114`	`+ if (!init_params->worker->context->config.ext.proto_emulation_enable) {`
	`115`	`+ return;`
	`116`	`+ }`
	`117`	`+`
`114`	`118`	`ucp_proto_single_probe(&params);`
`115`	`119`	`}`
`116`	`120`