prov/efa: Do GDRCopy registrations only in the EFA RDM path

Signed-off-by: Sai Sunku <sunkusa@amazon.com>
sunkuamzn · Feb 12, 2025 · 645b4ad · 645b4ad
1 parent 45c3522
commit 645b4ad
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 13 deletions.
diff --git a/prov/efa/src/efa_domain.c b/prov/efa/src/efa_domain.c
@@ -114,6 +114,9 @@ static int efa_domain_init_rdm(struct efa_domain *efa_domain, struct fi_info *in
 	int err;
 	bool enable_shm = efa_env.enable_shm_transfer;
 
+	assert(EFA_INFO_TYPE_IS_RDM(info));
+	efa_domain->rdm_ep = true;
+
 	/* App provided hints supercede environmental variables.
 	 *
 	 * Using the shm provider comes with some overheads, so avoid

diff --git a/prov/efa/src/efa_domain.h b/prov/efa/src/efa_domain.h
@@ -33,6 +33,7 @@ struct efa_domain {
 	size_t ibv_mr_reg_sz;
 
 	/* Only valid for RDM EP type */
+	bool			rdm_ep;		/* Set to true for RDM domain. False otherwise. */
 	uint64_t		rdm_mode;
 	size_t			rdm_cq_size;
 	/* number of rdma-read messages in flight */
@@ -46,7 +47,6 @@ struct efa_domain {
 	struct dlist_entry peer_backoff_list;
 	/* list of #efa_rdm_peer that will retry posting handshake pkt */
 	struct dlist_entry handshake_queued_peer_list;
-
 };
 
 extern struct dlist_entry g_efa_domain_list;

diff --git a/prov/efa/src/efa_mr.c b/prov/efa/src/efa_mr.c
@@ -221,7 +221,8 @@ static int efa_mr_hmem_setup(struct efa_mr *efa_mr,
 		efa_mr->needs_sync = true;
 		efa_mr->peer.device.cuda = attr->device.cuda;
 
-		if (!(flags & FI_MR_DMABUF) && cuda_is_gdrcopy_enabled()) {
+		/* Only attempt GDRCopy registrations for efa rdm path */
+		if (efa_mr->domain->rdm_ep && !(flags & FI_MR_DMABUF) && cuda_is_gdrcopy_enabled()) {
 			mr_iov = *attr->mr_iov;
 			err = ofi_hmem_dev_register(FI_HMEM_CUDA, mr_iov.iov_base, mr_iov.iov_len,
 						    (uint64_t *)&efa_mr->peer.hmem_data);

diff --git a/prov/efa/test/efa_unit_test_mr.c b/prov/efa/test/efa_unit_test_mr.c
@@ -3,33 +3,158 @@
 
 #include "efa_unit_tests.h"
 
-void test_efa_mr_reg_counters(struct efa_resource **state)
+static void test_efa_mr_counters(struct efa_domain *efa_domain, struct fid_mr *mr, int mr_reg_count, int mr_reg_size, bool gdrcopy_flag)
+{
+    struct efa_mr *efa_mr;
+
+    assert_int_equal(efa_domain->ibv_mr_reg_ct, mr_reg_count);
+    assert_int_equal(efa_domain->ibv_mr_reg_sz, mr_reg_size);
+
+    if (mr) {
+        efa_mr = container_of(mr, struct efa_mr, mr_fid);
+        if (cuda_is_gdrcopy_enabled()) {
+            if (gdrcopy_flag)
+                assert_true(efa_mr->peer.flags & OFI_HMEM_DATA_DEV_REG_HANDLE);
+            else
+                assert_false(efa_mr->peer.flags & OFI_HMEM_DATA_DEV_REG_HANDLE);
+        }
+    }
+}
+
+void test_efa_rdm_mr_reg_host_memory(struct efa_resource **state)
 {
     struct efa_resource *resource = *state;
     struct efa_domain *efa_domain;
     size_t mr_size = 64;
-    char *buf;
-    struct fid_mr *mr;
+    void *buf;
+    struct fid_mr *mr = NULL;
 
     efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);
 
     efa_domain = container_of(resource->domain, struct efa_domain, util_domain.domain_fid);
-    assert_true(efa_domain->ibv_mr_reg_ct == 0);
-    assert_true(efa_domain->ibv_mr_reg_sz == 0);
-
 
     buf = malloc(mr_size);
     assert_non_null(buf);
 
+    test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
     assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
             FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL), 0);
 
-    assert_true(efa_domain->ibv_mr_reg_ct == 1);
-    assert_true(efa_domain->ibv_mr_reg_sz == mr_size);
+    /* No GDRCopy registration for host memory */
+    test_efa_mr_counters(efa_domain, mr, 1, mr_size, false);
 
     assert_int_equal(fi_close(&mr->fid), 0);
-    assert_true(efa_domain->ibv_mr_reg_ct == 0);
-    assert_true(efa_domain->ibv_mr_reg_sz == 0);
-
+    test_efa_mr_counters(efa_domain, mr, 0, 0, false);
     free(buf);
 }
+
+#if HAVE_CUDA
+void test_efa_rdm_mr_reg_cuda_memory(struct efa_resource **state)
+{
+    struct efa_resource *resource = *state;
+    struct efa_domain *efa_domain;
+    size_t mr_size = 64;
+    void *buf;
+    struct fid_mr *mr = NULL;
+    struct fi_mr_attr mr_reg_attr = {0};
+    struct iovec iovec;
+    int err;
+
+    if (hmem_ops[FI_HMEM_CUDA].initialized) {
+	    resource->hints = efa_unit_test_alloc_hints(ep_type, EFA_FABRIC_NAME);
+        resource->hints->caps |= FI_HMEM;
+		efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM, FI_VERSION(2, 0),
+							    resource->hints, true, true);
+
+        efa_domain = container_of(resource->domain, struct efa_domain, util_domain.domain_fid);
+        test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
+        err = ofi_cudaMalloc(&buf, mr_size);
+        assert_int_equal(err, 0);
+        assert_non_null(buf);
+
+        mr_reg_attr.access = FI_SEND | FI_RECV;
+        mr_reg_attr.iface = FI_HMEM_CUDA;
+        iovec.iov_base = buf;
+        iovec.iov_len = mr_size;
+        mr_reg_attr.mr_iov = &iovec;
+        mr_reg_attr.iov_count = 1;
+
+        err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
+        assert_int_equal(err, 0);
+
+        /* FI_MR_DMABUF flag was not set, so GDRCopy should be registered if available */
+        test_efa_mr_counters(efa_domain, mr, 1, mr_size, true);
+
+        assert_int_equal(fi_close(&mr->fid), 0);
+        test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
+        err = ofi_cudaFree(buf);
+        assert_int_equal(err, 0);
+    }
+}
+#else
+void test_efa_rdm_mr_reg_cuda_memory(struct efa_resource **state)
+{
+    skip();
+}
+#endif
+
+#if HAVE_CUDA
+void test_efa_direct_mr_reg_no_gdrcopy(struct efa_resource **state)
+{
+    struct efa_resource *resource = *state;
+    struct efa_domain *efa_domain;
+    size_t mr_size = 64;
+    void *buf;
+    struct fid_mr *mr = NULL;
+    struct fi_mr_attr mr_reg_attr = {0};
+    struct iovec iovec;
+    int err;
+
+    if (hmem_ops[FI_HMEM_CUDA].initialized) {
+	    resource->hints = efa_unit_test_alloc_hints(ep_type, EFA_DIRECT_FABRIC_NAME);
+        resource->hints->caps |= FI_HMEM;
+		efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM, FI_VERSION(2, 0),
+							    resource->hints, true, true);
+
+        efa_domain = container_of(resource->domain, struct efa_domain, util_domain.domain_fid);
+        test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
+        err = ofi_cudaMalloc(&buf, mr_size);
+        assert_int_equal(err, 0);
+        assert_non_null(buf);
+
+        mr_reg_attr.access = FI_SEND | FI_RECV;
+        mr_reg_attr.iface = FI_HMEM_CUDA;
+        iovec.iov_base = buf;
+        iovec.iov_len = mr_size;
+        mr_reg_attr.mr_iov = &iovec;
+        mr_reg_attr.iov_count = 1;
+
+        err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
+        assert_int_equal(err, 0);
+
+        /* FI_MR_DMABUF flag was not set, so GDRCopy should be registered if available */
+        test_efa_mr_counters(efa_domain, mr, 1, mr_size, true);
+
+        assert_int_equal(fi_close(&mr->fid), 0);
+        test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
+        err = ofi_cudaFree(buf);
+        assert_int_equal(err, 0);
+        efa_domain = container_of(resource->domain, struct efa_domain, util_domain.domain_fid);
+        test_efa_mr_counters(efa_domain, mr, 0, 0, false);
+
+        err = ofi_cudaMalloc(buf, mr_size);
+        assert_non_null(buf);
+        assert_int_equal(err, 0);
+    }
+}
+#else
+void test_efa_direct_mr_reg_no_gdrcopy(struct efa_resource **state)
+{
+    skip();
+}
+#endif
diff --git a/prov/efa/test/efa_unit_tests.c b/prov/efa/test/efa_unit_tests.c
@@ -266,6 +266,12 @@ int main(void)
 		cmocka_unit_test_setup_teardown(test_efa_ep_bind_and_enable, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_cntr_ibv_cq_poll_list_same_tx_rx_cq_single_ep, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_cntr_ibv_cq_poll_list_separate_tx_rx_cq_single_ep, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+
+		/* begin efa_unit_test_mr.c */
+		cmocka_unit_test_setup_teardown(test_efa_rdm_mr_reg_host_memory, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		cmocka_unit_test_setup_teardown(test_efa_rdm_mr_reg_cuda_memory, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		cmocka_unit_test_setup_teardown(test_efa_direct_mr_reg_no_gdrcopy, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		/* end efa_unit_test_mr.c */
 	};
 
 	cmocka_set_message_output(CM_OUTPUT_XML);

diff --git a/prov/efa/test/efa_unit_tests.h b/prov/efa/test/efa_unit_tests.h
@@ -164,6 +164,7 @@ void test_ibv_cq_ex_read_ignore_removed_peer();
 /* begin efa_unit_test_info.c */
 void test_info_open_ep_with_wrong_info();
 void test_info_rdm_attributes();
+void test_info_rdm_attributes_fork_support();
 void test_info_dgram_attributes();
 void test_info_direct_attributes();
 void test_info_direct_hmem_support_p2p();
@@ -288,6 +289,12 @@ void test_efa_ep_bind_and_enable();
 void test_efa_cntr_ibv_cq_poll_list_same_tx_rx_cq_single_ep();
 void test_efa_cntr_ibv_cq_poll_list_separate_tx_rx_cq_single_ep();
 
+/* begin efa_unit_test_mr.c */
+void test_efa_rdm_mr_reg_host_memory();
+void test_efa_rdm_mr_reg_cuda_memory();
+void test_efa_direct_mr_reg_no_gdrcopy();
+/* end efa_unit_test_mr.c */
+
 static inline
 int efa_unit_test_get_dlist_length(struct dlist_entry *head)
 {