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;
@@ -118,6 +144,8 @@ static inline void efa_hmem_info_check_p2p_support_cuda(struct efa_hmem_info *in
 	int dmabuf_fd;
 	uint64_t dmabuf_offset;
 
+	info->dmabuf_supported_by_device_b = false;
+
 	cuda_ret = ofi_cudaMalloc(&ptr, len);
 	if (cuda_ret != cudaSuccess) {
 		info->initialized = false;
@@ -134,21 +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);
 		}
 	} 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
@@ -193,6 +228,8 @@ static inline void efa_hmem_info_check_p2p_support_neuron(struct efa_hmem_info *
 	uint64_t offset;
 	int ret;
 
+	info->dmabuf_supported_by_device_b = false;
+
 	if (g_efa_selected_device_list[0].device_caps & EFADV_DEVICE_ATTR_CAPS_RDMA_READ) {
 		ibv_access |= IBV_ACCESS_REMOTE_READ;
 	}
@@ -217,15 +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);
-	} 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
@@ -284,9 +334,17 @@ efa_hmem_info_init_iface(enum fi_hmem_iface iface)
 	}
 
 	info->initialized = true;
+	info->max_medium_msg_size = 0;
+	info->runt_size = 0;
+	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

@@ -24,6 +24,7 @@ static const enum fi_hmem_iface efa_hmem_ifaces[] = {
 struct efa_hmem_info {
 	bool initialized; 	/* do we support it at all */
 	bool p2p_supported_by_device;	/* do we support p2p with this device */
+	bool dmabuf_supported_by_device_b;	/* do we support dmabuf with this device */
 
 	size_t max_medium_msg_size;
 	size_t runt_size;
@@ -104,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

@@ -506,77 +506,88 @@ struct ibv_mr *efa_mr_reg_ibv_dmabuf_mr(struct ibv_pd *pd, uint64_t offset,
  * @param flags flags in fi_mr_reg/fi_mr_regattr
  * @return struct ibv_mr* the ptr to the registered MR
  */
-static struct ibv_mr *efa_mr_reg_ibv_mr(struct efa_mr *efa_mr, struct fi_mr_attr *mr_attr,
-					int access, const uint64_t flags)
+static struct ibv_mr *efa_mr_reg_ibv_mr(struct efa_mr *efa_mr,
+                                        struct fi_mr_attr *mr_attr,
+                                        int access, const uint64_t flags)
 {
-	int dmabuf_fd;
-	uint64_t offset;
-	int ret;
-	struct ibv_mr *ibv_mr;
-
-	if (flags & FI_MR_DMABUF)
-		return efa_mr_reg_ibv_dmabuf_mr(
-			efa_mr->domain->ibv_pd,
-			mr_attr->dmabuf->offset,
-			mr_attr->dmabuf->len,
-			(uintptr_t) mr_attr->dmabuf->base_addr + mr_attr->dmabuf->offset,
-			mr_attr->dmabuf->fd,
-			access
-		);
-
-	if (efa_mr_is_synapseai(efa_mr)) {
-		ret = ofi_hmem_get_dmabuf_fd(efa_mr->peer.iface,
-					     mr_attr->mr_iov->iov_base,
-					     (uint64_t) mr_attr->mr_iov->iov_len,
-					     &dmabuf_fd, &offset);
-		if (ret != FI_SUCCESS) {
-			EFA_WARN(FI_LOG_MR, "Unable to get dmabuf fd for Gaudi device buffer \n");
-			return NULL;
-		}
-		return 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);
-	}
-
-	/*
-	 * TODO: need such fallback for cuda as well when
-	 * FI_CUDA_API_PERMITTED is true
-	 */
-	if (efa_mr_is_neuron(efa_mr)) {
-		ret = ofi_hmem_get_dmabuf_fd(
-				efa_mr->peer.iface,
-				mr_attr->mr_iov->iov_base,
-				mr_attr->mr_iov->iov_len,
-				&dmabuf_fd,
-				&offset);
-
-		if (ret == FI_SUCCESS) {
-			/* Success => invoke ibv_reg_dmabuf_mr */
-			ibv_mr = efa_mr_reg_ibv_dmabuf_mr(
-					efa_mr->domain->ibv_pd, 0,
-					mr_attr->mr_iov->iov_len,
-					(uint64_t)mr_attr->mr_iov->iov_base,
-					dmabuf_fd, access);
-			(void) ofi_hmem_put_dmabuf_fd(efa_mr->peer.iface, dmabuf_fd);
-			return ibv_mr;
-		} else if (ret == -FI_EOPNOTSUPP) {
-			/* Protocol not available => fallback */
-			EFA_INFO(FI_LOG_MR,
-				"Unable to get dmabuf fd for Neuron device buffer, "
-				"Fall back to ibv_reg_mr\n");
-			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;
+    int dmabuf_fd;
+    uint64_t offset;
+    int ret;
+
+    /* Explicit dmabuf registration */
+    if (flags & FI_MR_DMABUF) {
+        if (!mr_attr->dmabuf) {
+            EFA_WARN(FI_LOG_MR, "FI_MR_DMABUF set but mr_attr->dmabuf == NULL\n");
+            return NULL;
+        }
+        if (!g_efa_hmem_info[mr_attr->iface].dmabuf_supported_by_device_b) {
+            EFA_WARN(FI_LOG_MR,
+                     "Requested FI_MR_DMABUF, but dmabuf not supported for iface=%d\n",
+                     mr_attr->iface);
+            return NULL;
+        }
+
+        EFA_INFO(FI_LOG_MR,
+                 "FI_MR_DMABUF: fd=%d offset=%lu len=%zu\n",
+                 mr_attr->dmabuf->fd, mr_attr->dmabuf->offset,
+                 mr_attr->dmabuf->len);
+
+        return efa_mr_reg_ibv_dmabuf_mr(
+            efa_mr->domain->ibv_pd,
+            mr_attr->dmabuf->offset,
+            mr_attr->dmabuf->len,
+            (uintptr_t) mr_attr->dmabuf->base_addr + mr_attr->dmabuf->offset,
+            mr_attr->dmabuf->fd,
+            access);
+    }
+
+    /* Implicit VA path with dmabuf-first */
+    if (g_efa_hmem_info[mr_attr->iface].dmabuf_supported_by_device_b) {
+        ret = ofi_hmem_get_dmabuf_fd(
+                mr_attr->iface,
+                mr_attr->mr_iov->iov_base,
+                (uint64_t) mr_attr->mr_iov->iov_len,
+                &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);
+
+            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 \n",
+                     mr_attr->iface);
+        } 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 */
+    } 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 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