prov/efa: fix missing mem_desc and iface initialization in non-p2p path

When p2p is not supported for CUDA/ROCR memory, the non-p2p path in
efa_rdm_mr_reg_impl() was missing two critical initializations:
efa_mr->mr_fid.mem_desc was not set (causing fi_mr_desc() to return
NULL/invalid pointers) and efa_mr_hmem_setup() was not called
(leaving efa_mr->iface uninitialized). Fix by adding the missing
mem_desc assignment, calling efa_mr_hmem_setup() in the non-p2p path,
making efa_mr_hmem_setup() non-static, and adding a unit test to
verify both fields are properly initialized when p2p is disabled.

Signed-off-by: Shi Jin <sjina@amazon.com>

Author: shijin-aws

prov/efa/src/efa_mr.c

@@ -89,9 +89,9 @@ int efa_mr_regattr_validate(struct fid *fid, const struct fi_mr_attr *attr,
  *
  * @return FI_SUCCESS or negative FI error code
  */
-static int efa_mr_hmem_setup(struct efa_mr *efa_mr,
-                             const struct fi_mr_attr *attr,
-							 uint64_t flags)
+int efa_mr_hmem_setup(struct efa_mr *efa_mr,
+		       const struct fi_mr_attr *attr,
+		       uint64_t flags)
 {
 	if (attr->iface == FI_HMEM_SYSTEM) {
 		efa_mr->iface = FI_HMEM_SYSTEM;

prov/efa/src/efa_mr.h

@@ -18,6 +18,7 @@ struct efa_mr {
 
 int efa_mr_reg_impl(struct efa_mr *efa_mr, uint64_t flags, const struct fi_mr_attr *mr_attr);
 int efa_mr_dereg_impl(struct efa_mr *efa_mr);
+int efa_mr_hmem_setup(struct efa_mr *efa_mr, const struct fi_mr_attr *attr, uint64_t flags);
 int efa_mr_validate_regattr(struct fid *fid, const struct fi_mr_attr *attr, uint64_t flags);
 
 #define EFA_MR_ATTR_INIT_SYSTEM(iov, count, access, offset, requested_key, context) \

prov/efa/src/rdm/efa_rdm_mr.c

@@ -497,7 +497,11 @@ static int efa_rdm_mr_reg_impl(struct efa_rdm_mr *efa_rdm_mr, uint64_t flags,
 	 */
 	if ((mr_attr->iface == FI_HMEM_CUDA || mr_attr->iface == FI_HMEM_ROCR)
 		&& !g_efa_hmem_info[mr_attr->iface].p2p_supported_by_device) {
+		ret = efa_mr_hmem_setup(&efa_rdm_mr->efa_mr, mr_attr, flags);
+		if (ret)
+			return ret;
 		efa_rdm_mr->efa_mr.mr_fid.key = efa_rdm_mr_non_p2p_keygen();
+		efa_rdm_mr->efa_mr.mr_fid.mem_desc = &efa_rdm_mr->efa_mr;
 	} else {
 		/* base mr registration (ibv mr), must be called the first before RDM specific fields are setup */
 		ret = efa_mr_reg_impl(&efa_rdm_mr->efa_mr, flags, mr_attr);

prov/efa/test/efa_unit_test_mr.c

@@ -1356,3 +1356,97 @@ void test_efa_rdm_mr_cache_reference_counting(struct efa_resource **state)
 	free(buf);
 #endif
 }
+
+#if HAVE_CUDA
+/**
+ * @brief Test RDM MR registration for CUDA memory in non-p2p path
+ *
+ * This test verifies that when p2p is not supported, the non-p2p path
+ * properly initializes the MR structure, including the mem_desc field
+ * and iface field. It uses fi_mr_desc to verify the descriptor is correct.
+ */
+void test_efa_rdm_mr_reg_cuda_memory_non_p2p(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;
+	struct efa_rdm_mr *efa_rdm_mr;
+	void *desc;
+	int err, baseline_ct, baseline_sz;
+
+	if (!g_efa_hmem_info[FI_HMEM_CUDA].initialized) {
+		skip();
+		return;
+	}
+
+	resource->hints = efa_unit_test_alloc_hints_hmem(
+		FI_EP_RDM, EFA_FABRIC_NAME);
+	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);
+
+	/* Mock p2p as not supported to force non-p2p path */
+	g_efa_hmem_info[FI_HMEM_CUDA].p2p_supported_by_device = false;
+
+	/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
+	baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
+	baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);
+
+	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);
+	assert_non_null(mr);
+
+	/* Verify this is an RDM MR */
+	efa_rdm_mr = container_of(mr, struct efa_rdm_mr, efa_mr.mr_fid);
+	assert_non_null(efa_rdm_mr);
+
+	/* Test fi_mr_desc returns the correct descriptor */
+	desc = fi_mr_desc(mr);
+	assert_non_null(desc);
+	assert_ptr_equal(desc, &efa_rdm_mr->efa_mr);
+
+	/* Verify iface is properly set (second bug fix) */
+	assert_int_equal(efa_rdm_mr->efa_mr.iface, FI_HMEM_CUDA);
+
+	/* Verify key is set (non-p2p path should generate proprietary key) */
+	assert_true(fi_mr_key(mr) != FI_KEY_NOTAVAIL);
+	assert_true(fi_mr_key(mr) > UINT32_MAX); /* Non-p2p keys should be > UINT32_MAX */
+
+	/* Since we're in non-p2p path, no ibv_mr should be registered */
+	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct), baseline_ct);
+	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz), baseline_sz);
+
+	/* Verify RDM-specific fields are properly initialized */
+	assert_true(efa_rdm_mr->inserted_to_mr_map);
+	assert_true(efa_rdm_mr->needs_sync); /* CUDA memory should need sync */
+
+	/* Cleanup */
+	assert_int_equal(fi_close(&mr->fid), 0);
+	err = ofi_cudaFree(buf);
+	assert_int_equal(err, 0);
+}
+#else
+void test_efa_rdm_mr_reg_cuda_memory_non_p2p(struct efa_resource **state)
+{
+	skip();
+}
+#endif

prov/efa/test/efa_unit_tests.h

@@ -505,6 +505,7 @@ void test_efa_rdm_mr_reg_host_memory();
 void test_efa_rdm_mr_reg_host_memory_no_mr_local();
 void test_efa_rdm_mr_reg_host_memory_overlapping_buffers();
 void test_efa_rdm_mr_reg_cuda_memory();
+void test_efa_rdm_mr_reg_cuda_memory_non_p2p();
 void test_efa_direct_mr_reg_cuda_memory();
 void test_efa_direct_mr_reg_rdma_read_not_supported();
 void test_efa_direct_mr_reg_rdma_write_not_supported();