UCP/RNDV: Use local memory detection for RTR_REQ responder - v1.23.x

src/ucp/rma/rma_rndv.c

@@ -599,11 +599,8 @@ ucs_status_t ucp_rma_rndv_process_rts(ucp_worker_h worker,
     recv_req->recv.worker        = worker;
     recv_req->recv.op_attr       = 0;
     recv_req->recv.remote_req_id = rts->super.sreq.req_id;
-    ucp_memory_detect(worker->context, (void*)(uintptr_t)rts->address,
-                      rts->super.size, &mem_info);
-    ucs_assertv(mem_info.type == rts->mem_type, "detected mem_type=%s rts=%s",
-                ucs_memory_type_names[mem_info.type],
-                ucs_memory_type_names[rts->mem_type]);
+    ucp_proto_rndv_memory_detect(worker->context, rts->address, rts->super.size,
+                                 rts->mem_type, rts->sys_dev, &mem_info);
     ucp_datatype_iter_init_contig(&recv_req->recv.dt_iter,
                                   (void*)(uintptr_t)rts->address,
                                   rts->super.size, &mem_info);

src/ucp/rndv/proto_rndv.c

@@ -1059,23 +1059,19 @@ ucp_proto_rndv_rtr_req_sreq_init(ucp_ep_h ep, ucp_request_t *req,
                                  const ucp_rndv_rtr_req_hdr_t *rtr_req)
 {
     const ucp_rndv_rtr_hdr_t *rtr = &rtr_req->super;
-    ucp_memory_info_t mem_info    = {
-        .type    = rtr_req->mem_type,
-        .sys_dev = rtr_req->sys_dev
-    };
-    ucp_memory_info_t local_mem_info;
+    ucp_memory_info_t mem_info;
 
-    ucp_memory_detect(ep->worker->context, (void*)(uintptr_t)rtr_req->address,
-                      rtr->size, &local_mem_info);
-    mem_info.flags = local_mem_info.flags;
+    ucp_proto_rndv_memory_detect(ep->worker->context, rtr_req->address,
+                                 rtr->size, rtr_req->mem_type, rtr_req->sys_dev,
+                                 &mem_info);
 
     ucp_proto_request_send_init(req, ep,
                                 UCP_REQUEST_FLAG_RNDV_SEND_INTERNAL);
     ucp_request_set_callback(req, send.cb,
                              ucp_proto_rndv_rtr_req_send_complete);
     req->send.buffer              = (void*)(uintptr_t)rtr_req->address;
     req->send.length              = rtr->size;
-    req->send.mem_type            = rtr_req->mem_type;
+    req->send.mem_type            = mem_info.type;
     req->send.rndv.remote_req_id  = rtr->rreq_id;
     req->send.rndv.rkey           = NULL;
     req->send.rndv.remote_address = rtr_req->address;

src/ucp/rndv/proto_rndv.inl

@@ -583,4 +583,27 @@ ucp_proto_rndv_recv_complete(ucp_request_t *req)
     return ucp_proto_rndv_recv_complete_status(req, rreq->status);
 }
 
+/* Detect the local buffer at a peer-provided address, and check it against the
+ * peer-provided mem_type and sys_dev when known */
+static UCS_F_ALWAYS_INLINE void
+ucp_proto_rndv_memory_detect(ucp_context_h context, uint64_t address,
+                             size_t length,
+                             ucs_memory_type_t UCS_V_UNUSED mem_type,
+                             ucs_sys_device_t UCS_V_UNUSED sys_dev,
+                             ucp_memory_info_t *mem_info)
+{
+    ucp_memory_detect(context, (void*)(uintptr_t)address, length, mem_info);
+
+    ucs_assertv(mem_info->type == mem_type,
+                "detected mem_type=%s doesn't match peer provided mem_type=%s",
+                ucs_memory_type_names[mem_info->type],
+                ucs_memory_type_names[mem_type]);
+
+    ucs_assertv((sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) ||
+                        (sys_dev == mem_info->sys_dev),
+                "detected sys_dev=%s doesn't match peer provided sys_dev=%s",
+                ucs_topo_sys_device_get_name(mem_info->sys_dev),
+                ucs_topo_sys_device_get_name(sys_dev));
+}
+
 #endif

test/gtest/ucp/test_ucp_rma.cc

@@ -13,6 +13,7 @@ extern "C" {
 #include <ucp/core/ucp_mm.h> /* for UCP_MEM_IS_ACCESSIBLE_FROM_CPU */
 #include <ucp/core/ucp_ep.inl>
 #include <ucp/core/ucp_rkey.h>
+#include <ucp/core/ucp_rkey.inl>
 #include <ucp/proto/proto_multi.h>
 #include <ucs/sys/sys.h>
 #include <uct/api/v2/uct_v2.h>
@@ -588,6 +589,59 @@ UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_rma_rndv_gpu_nic, rcx_cuda,
                               "rc_x,cuda_copy")
 
 
+/* Without GPUDirect RDMA no NIC registers the responder's CUDA buffer, so the
+ * peer's rkey has no sys_dev, and the responder must detect the buffer's GPU
+ * locally to keep its host-staged rndv/put/mtype on the NICs assigned to it */
+class test_ucp_rma_rndv_gpu_nic_no_gdr : public test_ucp_rma_rndv_gpu_nic {
+public:
+    static void get_test_variants(std::vector<ucp_test_variant> &variants)
+    {
+        add_variant_values(variants, test_ucp_rma_gpu_nic::get_test_variants,
+                           UCS_BIT(RNDV_SCHEME_PUT_PPLN), gpu_nic_rndv_schemes);
+    }
+
+    void init() override
+    {
+        m_env.push_back(new ucs::scoped_setenv("UCX_IB_GPU_DIRECT_RDMA", "n"));
+        test_ucp_rma_rndv_gpu_nic::init();
+    }
+
+    /* RTR_REQ carries this rkey's sys_dev, which must stay unknown so the
+     * responder has to detect its buffer's GPU locally */
+    void get_b_expect_empty_rkey(size_t size, void *expected_data,
+                                 ucp_mem_h memh, void *target_ptr,
+                                 ucp_rkey_h rkey, void *arg)
+    {
+        EXPECT_EQ(0, rkey->md_map);
+        EXPECT_EQ(UCS_SYS_DEVICE_ID_UNKNOWN,
+                  ucp_rkey_config(sender().worker(), rkey)->key.sys_dev);
+        get_b(size, expected_data, memh, target_ptr, rkey, arg);
+    }
+};
+
+UCS_TEST_P(test_ucp_rma_rndv_gpu_nic_no_gdr, get_blocking)
+{
+    if (has_cuda_net_md()) {
+        UCS_TEST_SKIP_R("a network memory domain can register CUDA memory");
+    }
+
+    if (!is_buffer_gpu_assigned(receiver().ucph())) {
+        UCS_TEST_SKIP_R("no nic is assigned to the test buffers' gpu");
+    }
+
+    /* The initiator's buffer is host memory, so the put can land */
+    test_message_sizes(
+            static_cast<send_func_t>(
+                    &test_ucp_rma_rndv_gpu_nic_no_gdr::get_b_expect_empty_rkey),
+            128, PPLN_FRAG_SIZE, UCS_MEMORY_TYPE_HOST, UCS_MEMORY_TYPE_CUDA, 0);
+    expect_assigned_lanes(entities(), "rndv/put/mtype");
+}
+
+/* cuda_ipc is reachable intra-node and would resolve the sys_dev */
+UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_rma_rndv_gpu_nic_no_gdr, rcx_cuda,
+                              "rc_x,cuda_copy")
+
+
 class test_ucp_rma_rndv : public test_ucp_rma {
 public:
     static constexpr size_t SIZE = 512 * UCS_KBYTE;