prov/efa: dmabuf try / fallback logic

feat: Adding default dmabuf attempt and fallback logic for all efa_hmem_ifaces

Problem:
  - How make dmabuf usage default going forward and have fallback mechanism if dmabuf not supported

Solution:
  - Modified initial PR from @jiaxiyan at 6aa6708#diff-9b57a9410ed94ed1f1aea837412e68bbe9b49582edce813ba352fffb37dcc007
  - Added dmabuf_supported_by_device_b flag in efa_hmem_info structure in prov/efa/efa_hmem.h
  - Updated dmabuf_supported_by_device_b in each fi_hmem_iface type p2p_support fcn in prov/efa/efa_hmem.c
  - Removed per fi_hmem_iface type checks in prov/efa/efa_mr_reg_ibv_mr.c
  -

Testing:
  - Ran mpi perf tests on 2 nodes on p5en with dmabuf and fallback option hard set
  - Ran mpi perf tests on 16 nodes on p5en with dmabuf and fallback option hard set

Sim Issue:
  - N/A

Signed-off-by: Nick Mazzilli <[email protected]>

Author: nmazzilli3

prov/efa/src/efa_env.c

@@ -41,6 +41,13 @@ struct efa_env efa_env = {
 	.internal_rx_refill_threshold = 8,
 	.use_data_path_direct = true,
 	.implicit_av_size = 0,
+
+	/* HMEM DMABUF defaults */
+	.hmem_system_use_dmabuf = 0,
+	.hmem_cuda_use_dmabuf = 1,
+	.hmem_neuron_use_dmabuf = 1,
+	.hmem_synapseai_use_dmabuf = 1,
+	.hmem_rocr_use_dmabuf = 1,
 };
 
 /* @brief Read and store the FI_EFA_* environment variables.
@@ -141,6 +148,14 @@ void efa_env_param_get(void)
 	}
 	fi_param_get_bool(&efa_prov, "use_data_path_direct", &efa_env.use_data_path_direct);
 
+	/* HMEM DMABUF parameters */
+	fi_param_get_bool(&efa_prov, "hmem_system_use_dmabuf", &efa_env.hmem_system_use_dmabuf);
+	fi_param_get_bool(&efa_prov, "hmem_neuron_use_dmabuf", &efa_env.hmem_neuron_use_dmabuf);
+	fi_param_get_bool(&efa_prov, "hmem_synapseai_use_dmabuf", &efa_env.hmem_synapseai_use_dmabuf);
+	/* Use core libfabric parameters for CUDA and ROCR */
+	fi_param_get_bool(NULL, "hmem_cuda_use_dmabuf", &efa_env.hmem_cuda_use_dmabuf);
+	fi_param_get_bool(NULL, "hmem_rocr_use_dmabuf", &efa_env.hmem_rocr_use_dmabuf);
+
 	efa_fork_support_request_initialize();
 }
 
@@ -184,6 +199,15 @@ void efa_env_define()
 			"Set the maximum number of receive operations before the provider returns -FI_EAGAIN.");
 	fi_param_define(&efa_prov, "rx_copy_unexp", FI_PARAM_BOOL,
 			"Enables the use of a separate pool of bounce-buffers to copy unexpected messages out of the pre-posted receive buffers. (Default: 1)");
+
+	/* HMEM DMABUF environment variables */
+	fi_param_define(&efa_prov, "hmem_system_use_dmabuf", FI_PARAM_BOOL,
+			"Use DMABUF for system memory registration. (Default: 0)");
+	fi_param_define(&efa_prov, "hmem_neuron_use_dmabuf", FI_PARAM_BOOL,
+			"Use DMABUF for Neuron device memory registration. (Default: 1)");
+	fi_param_define(&efa_prov, "hmem_synapseai_use_dmabuf", FI_PARAM_BOOL,
+			"Use DMABUF for SynapseAI device memory registration. (Default: 1)");
+
 	fi_param_define(&efa_prov, "rx_copy_ooo", FI_PARAM_BOOL,
 			"Enables the use of a separate pool of bounce-buffers to copy out-of-order RTM packets out of the pre-posted receive buffers. (Default: 1)");
 	fi_param_define(&efa_prov, "max_timeout", FI_PARAM_INT,

prov/efa/src/efa_env.h

@@ -75,6 +75,13 @@ struct efa_env {
 	 * Value of 0 means there is no limit on the size.
 	 */
 	size_t implicit_av_size;
+
+	/* HMEM DMABUF settings */
+	int hmem_system_use_dmabuf;
+	int hmem_cuda_use_dmabuf;
+	int hmem_neuron_use_dmabuf;
+	int hmem_synapseai_use_dmabuf;
+	int hmem_rocr_use_dmabuf;
 };
 
 extern struct efa_env efa_env;

prov/efa/src/efa_hmem.c

@@ -7,6 +7,32 @@
 
 struct efa_hmem_info g_efa_hmem_info[OFI_HMEM_MAX];
 
+/**
+ * @brief Check if DMABUF is enabled for a specific HMEM interface
+ *
+ * This function checks the environment variables to determine if DMABUF
+ * should be used for the specified HMEM interface. It respects both
+ * EFA-specific and core libfabric environment variables.
+ *
+ * @param[in] iface The HMEM interface to check
+ * @return true if DMABUF is enabled for the interface, false otherwise
+ */
+bool efa_hmem_is_dmabuf_env_var_enabled(enum fi_hmem_iface iface)
+{
+	switch (iface) {
+	case FI_HMEM_SYSTEM:
+		return efa_env.hmem_system_use_dmabuf;
+	case FI_HMEM_CUDA:
+		return efa_env.hmem_cuda_use_dmabuf;
+	case FI_HMEM_NEURON:
+		return efa_env.hmem_neuron_use_dmabuf;
+	case FI_HMEM_SYNAPSEAI:
+		return efa_env.hmem_synapseai_use_dmabuf;
+	default:
+		return false;
+	}
+}
+
 #if HAVE_CUDA || HAVE_NEURON
 static size_t efa_max_eager_msg_size_with_largest_header() {
 	int mtu_size;
@@ -136,22 +162,28 @@ static inline void efa_hmem_info_check_p2p_support_cuda(struct efa_hmem_info *in
 	}
 
 #if HAVE_EFA_DMABUF_MR
-	ret = cuda_get_dmabuf_fd(ptr, len, &dmabuf_fd, &dmabuf_offset);
-	if (ret == FI_SUCCESS) {
-		ibv_mr = ibv_reg_dmabuf_mr(ibv_pd, dmabuf_offset,
-					   len, (uint64_t)ptr, dmabuf_fd, ibv_access);
-		(void)cuda_put_dmabuf_fd(dmabuf_fd);
-		if (!ibv_mr) {
+	if (efa_hmem_is_dmabuf_env_var_enabled(FI_HMEM_CUDA)) {
+		ret = cuda_get_dmabuf_fd(ptr, len, &dmabuf_fd, &dmabuf_offset);
+		if (ret == FI_SUCCESS) {
+			ibv_mr = ibv_reg_dmabuf_mr(ibv_pd, dmabuf_offset,
+						   len, (uint64_t)ptr, dmabuf_fd, ibv_access);
+			(void)cuda_put_dmabuf_fd(dmabuf_fd);
+			if (!ibv_mr) {
+				EFA_INFO(FI_LOG_CORE,
+					"Unable to register CUDA device buffer via dmabuf: %s. "
+					"Fall back to ibv_reg_mr\n", fi_strerror(-errno));
+				ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
+			} else {
+				info->dmabuf_supported_by_device_b = true;
+			}
+		} else {
 			EFA_INFO(FI_LOG_CORE,
-				"Unable to register CUDA device buffer via dmabuf: %s. "
-				"Fall back to ibv_reg_mr\n", fi_strerror(-errno));
+				"Unable to retrieve dmabuf fd of CUDA device buffer: %d. "
+				"Fall back to ibv_reg_mr\n", ret);
 			ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
 		}
-		info->dmabuf_supported_by_device_b = true;
 	} else {
-		EFA_INFO(FI_LOG_CORE,
-			"Unable to retrieve dmabuf fd of CUDA device buffer: %d. "
-			"Fall back to ibv_reg_mr\n", ret);
+		EFA_INFO(FI_LOG_CORE, "CUDA DMABUF disabled by environment variable\n");
 		ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
 	}
 #else
@@ -222,16 +254,28 @@ static inline void efa_hmem_info_check_p2p_support_neuron(struct efa_hmem_info *
 	}
 
 #if HAVE_EFA_DMABUF_MR
-	ret = neuron_get_dmabuf_fd(ptr, (uint64_t)len, &dmabuf_fd, &offset);
-	if (ret == FI_SUCCESS) {
-		ibv_mr = ibv_reg_dmabuf_mr(
-					ibv_pd, offset,
-					len, (uint64_t)ptr, dmabuf_fd, ibv_access);
-		info->dmabuf_supported_by_device_b = true;
-	} else if (ret == -FI_EOPNOTSUPP) {
-		EFA_INFO(FI_LOG_MR,
-			"Unable to retrieve dmabuf fd of Neuron device buffer, "
-			"Fall back to ibv_reg_mr\n");
+	if (efa_hmem_is_dmabuf_env_var_enabled(FI_HMEM_NEURON)) {
+		ret = neuron_get_dmabuf_fd(ptr, (uint64_t)len, &dmabuf_fd, &offset);
+		if (ret == FI_SUCCESS) {
+			ibv_mr = ibv_reg_dmabuf_mr(
+						ibv_pd, offset,
+						len, (uint64_t)ptr, dmabuf_fd, ibv_access);
+			if (ibv_mr) {
+				info->dmabuf_supported_by_device_b = true;
+			} else {
+				EFA_INFO(FI_LOG_CORE,
+					"Unable to register Neuron device buffer via dmabuf: %s. "
+					"Fall back to ibv_reg_mr\n", fi_strerror(-errno));
+				ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
+			}
+		} else {
+			EFA_INFO(FI_LOG_MR,
+				"Unable to retrieve dmabuf fd of Neuron device buffer: %d. "
+				"Fall back to ibv_reg_mr\n", ret);
+			ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
+		}
+	} else {
+		EFA_INFO(FI_LOG_CORE, "Neuron DMABUF disabled by environment variable\n");
 		ibv_mr = ibv_reg_mr(ibv_pd, ptr, len, ibv_access);
 	}
 #else
@@ -295,9 +339,12 @@ efa_hmem_info_init_iface(enum fi_hmem_iface iface)
 	info->min_read_msg_size = 0;
 	info->min_read_write_size = 0;
 
-	if (iface == FI_HMEM_SYNAPSEAI || iface == FI_HMEM_SYSTEM) {
+	if (iface == FI_HMEM_SYNAPSEAI) {
 		info->p2p_supported_by_device = true;
 		info->dmabuf_supported_by_device_b = true;
+	} else if(iface == FI_HMEM_SYSTEM) {
+		info->p2p_supported_by_device = true;
+		info->dmabuf_supported_by_device_b = false;
 	} else if (ofi_hmem_p2p_disabled()) {
 		info->p2p_supported_by_device = false;
 	} else {

prov/efa/src/efa_hmem.h

@@ -105,4 +105,5 @@ static inline int efa_copy_to_hmem(void *desc, void *dest, const void *buff, siz
 
 ssize_t efa_copy_from_hmem_iov(void **desc, char *buff, size_t buff_size, const struct iovec *hmem_iov, size_t iov_count);
 ssize_t efa_copy_to_hmem_iov(void **desc, struct iovec *hmem_iov, size_t iov_count, char *buff, size_t buff_size);
+bool efa_hmem_is_dmabuf_env_var_enabled(enum fi_hmem_iface iface);
 #endif

prov/efa/src/efa_mr.c

@@ -550,38 +550,44 @@ static struct ibv_mr *efa_mr_reg_ibv_mr(struct efa_mr *efa_mr,
                 &dmabuf_fd, &offset);
 
         if (ret == FI_SUCCESS) {
+            struct ibv_mr *dmabuf_mr;
             EFA_INFO(FI_LOG_MR,
                      "Registering dmabuf MR: fd=%d offset=%lu len=%zu\n",
                      dmabuf_fd, offset, mr_attr->mr_iov->iov_len);
 
-            return efa_mr_reg_ibv_dmabuf_mr(
+            dmabuf_mr = efa_mr_reg_ibv_dmabuf_mr(
                 efa_mr->domain->ibv_pd, offset,
                 mr_attr->mr_iov->iov_len,
                 (uint64_t)mr_attr->mr_iov->iov_base,
                 dmabuf_fd, access);
+            
+            /* Close the dmabuf file descriptor - it's no longer needed after registration */
+            (void) ofi_hmem_put_dmabuf_fd(mr_attr->iface, dmabuf_fd);
+            
+            if (dmabuf_mr) {
+                return dmabuf_mr;
+            }
+            
+            EFA_INFO(FI_LOG_MR, "DMABUF registration failed, returning NULL\n");
         }
 
         if (ret == -FI_EOPNOTSUPP || ret == -FI_ENOSYS) {
             EFA_WARN(FI_LOG_MR,
-                     "dmabuf not supported at runtime for iface=%d, disabling\n",
+                     "dmabuf not supported at runtime for iface=%d \n",
                      mr_attr->iface);
-			g_efa_hmem_info[mr_attr->iface].dmabuf_supported_by_device_b = false;
         } else {
             EFA_WARN(FI_LOG_MR,
                      "ofi_hmem_get_dmabuf_fd failed: ret=%d (%s)\n",
                      ret, fi_strerror(-ret));
         }
         /* fall through to ibv_reg_mr */
-    }
-
-    /* Fallback: plain ibv_reg_mr */
-    EFA_WARN(FI_LOG_MR,
-             "Fallback ibv_reg_mr: addr=%p len=%zu\n",
-             mr_attr->mr_iov->iov_base, mr_attr->mr_iov->iov_len);
-
-    return ibv_reg_mr(efa_mr->domain->ibv_pd,
+    } else {
+		return ibv_reg_mr(efa_mr->domain->ibv_pd,
                       (void *)mr_attr->mr_iov->iov_base,
                       mr_attr->mr_iov->iov_len, access);
+	}
+
+	return NULL;
 }
 
 #if HAVE_CUDA