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

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

Hide file tree

Showing 6 changed files with 60 additions and 2 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
@@ -10,6 +10,7 @@ void test_efa_mr_reg_counters(struct efa_resource **state)
     size_t mr_size = 64;
     char *buf;
     struct fid_mr *mr;
+    struct efa_mr *efa_mr;
 
     efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);
 
@@ -27,6 +28,48 @@ void test_efa_mr_reg_counters(struct efa_resource **state)
     assert_true(efa_domain->ibv_mr_reg_ct == 1);
     assert_true(efa_domain->ibv_mr_reg_sz == mr_size);
 
+    efa_mr = container_of(mr, struct efa_mr, mr_fid);
+    /* FI_MR_DMABUF flag was not set, so GDRCopy should be registered if available */
+    if (cuda_is_gdrcopy_enabled())
+        assert_true(efa_mr->peer.flags & OFI_HMEM_DATA_DEV_REG_HANDLE);
+
+    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);
+
+    free(buf);
+}
+
+
+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;
+    char *buf;
+    struct fid_mr *mr;
+    struct efa_mr *efa_mr;
+
+    efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_DIRECT_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);
+
+    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);
+
+    efa_mr = container_of(mr, struct efa_mr, mr_fid);
+    /* Buffer should not be registered with GDRCopy in the efa-direct path */
+    assert_false(efa_mr->peer.flags & OFI_HMEM_DATA_DEV_REG_HANDLE);
+
     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);

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_mr_reg_counters, 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
@@ -288,6 +288,11 @@ 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_mr_reg_counters();
+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)
 {