Refactor dmabuf reg

Signed-off-by: Jessie Yang <[email protected]>
jiaxiyan · May 29, 2024 · ceda8cb · ceda8cb
1 parent f58d5f9
commit ceda8cb
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 49 deletions.
diff --git a/prov/efa/src/efa_hmem.c b/prov/efa/src/efa_hmem.c
@@ -97,6 +97,8 @@ static int efa_domain_hmem_info_init_system(struct efa_domain *efa_domain)
 	info->p2p_disabled_by_user = false;
 	info->p2p_required_by_impl = false;
 	info->p2p_supported_by_device = true;
+	info->is_dmabuf_supported = false;
+
 	efa_domain_hmem_info_init_protocol_thresholds(efa_domain, FI_HMEM_SYSTEM);
 	return 0;
 }
@@ -137,6 +139,7 @@ static int efa_domain_hmem_info_init_cuda(struct efa_domain *efa_domain)
 
 	info->initialized = true;
 	info->p2p_disabled_by_user = false;
+	info->is_dmabuf_supported = false;
 
 	/* If user is using libfabric API 1.18 or later, by default EFA provider is permitted to
 	 * use CUDA library to support CUDA memory, therefore p2p is not required.
@@ -156,6 +159,8 @@ static int efa_domain_hmem_info_init_cuda(struct efa_domain *efa_domain)
 				"Unable to register CUDA device buffer via dmabuf: %s. "
 				"Fall back to ibv_reg_mr\n", fi_strerror(-errno));
 			ibv_mr = ibv_reg_mr(g_device_list[0].ibv_pd, ptr, len, ibv_access);
+		} else {
+			info->is_dmabuf_supported = true;
 		}
 	} else {
 		EFA_INFO(FI_LOG_DOMAIN,
@@ -247,13 +252,15 @@ static int efa_domain_hmem_info_init_neuron(struct efa_domain *efa_domain)
 	info->p2p_disabled_by_user = false;
 	/* Neuron currently requires P2P */
 	info->p2p_required_by_impl = true;
+	info->is_dmabuf_supported = false;
 
 #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(
 					g_device_list[0].ibv_pd, offset,
 					len, (uint64_t)ptr, dmabuf_fd, ibv_access);
+		info->is_dmabuf_supported = true;
 	} else if (ret == -FI_ENOPROTOOPT) {
 		EFA_INFO(FI_LOG_MR,
 			"Unable to retrieve dmabuf fd of Neuron device buffer, "
@@ -325,6 +332,7 @@ static int efa_domain_hmem_info_init_synapseai(struct efa_domain *efa_domain)
 	/* SynapseAI currently requires P2P */
 	info->p2p_required_by_impl = true;
 	info->p2p_supported_by_device = true;
+	info->is_dmabuf_supported = true;
 	efa_domain_hmem_info_init_protocol_thresholds(efa_domain, FI_HMEM_SYNAPSEAI);
 
 	/*  Only the long read protocol is supported */

diff --git a/prov/efa/src/efa_hmem.h b/prov/efa/src/efa_hmem.h
@@ -26,6 +26,7 @@ struct efa_hmem_info {
 	bool p2p_disabled_by_user;	/* Did the user disable p2p via FI_OPT_FI_HMEM_P2P? */
 	bool p2p_required_by_impl;	/* Is p2p required for this interface? */
 	bool p2p_supported_by_device;	/* do we support p2p with this device */
+	bool is_dmabuf_supported;
 
 	size_t max_intra_eager_size; /* Maximum message size to use eager protocol for intra-node */
 	size_t max_medium_msg_size;

diff --git a/prov/efa/src/efa_mr.c b/prov/efa/src/efa_mr.c
@@ -511,7 +511,10 @@ struct ibv_mr *efa_mr_reg_ibv_dmabuf_mr(struct ibv_pd *pd, uint64_t offset,
 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)
 {
-	if (flags & FI_MR_DMABUF)
+	if (flags & FI_MR_DMABUF) {
+		assert(efa_mr->domain->hmem_info[mr_attr->iface].is_dmabuf_supported);
+		EFA_INFO(FI_LOG_MR, "FI_MR_DMABUF is set. Registering dmabuf mr with 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,
@@ -520,64 +523,32 @@ static struct ibv_mr *efa_mr_reg_ibv_mr(struct efa_mr *efa_mr, struct fi_mr_attr
 			mr_attr->dmabuf->fd,
 			access
 		);
+	}
 
-	/*
-	 * TODO: remove the synapseai and neuron blocks by onboarding the
-	 * ofi_hmem_get_dmabuf_fd API.
-	 */
-#if HAVE_SYNAPSEAI
-	if (efa_mr_is_synapseai(efa_mr)) {
-		int dmabuf_fd;
-		uint64_t offset;
-		int ret;
-
-		ret = synapseai_get_dmabuf_fd(mr_attr->mr_iov->iov_base,
-						(uint64_t) mr_attr->mr_iov->iov_len,
-						&dmabuf_fd, &offset);
+	int dmabuf_fd;
+	uint64_t offset;
+	int ret;
+
+	if (efa_mr->domain->hmem_info[mr_attr->iface].is_dmabuf_supported) {
+		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) {
-			EFA_WARN(FI_LOG_MR, "Unable to get dmabuf fd for Gaudi device buffer \n");
+			EFA_WARN(FI_LOG_MR, "Unable to get dmabuf fd for device buffer \n");
 			return NULL;
 		}
+		EFA_INFO(FI_LOG_MR, "Registering dmabuf mr with fd %d, offset: %lu, len: %zu, \n", 
+				 dmabuf_fd, offset, mr_attr->mr_iov->iov_len);
 		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);
 	}
-#endif
-
-#if HAVE_NEURON
-	if (efa_mr_is_neuron(efa_mr)) {
-		int dmabuf_fd;
-		uint64_t offset;
-		int ret;
-
-		ret = neuron_get_dmabuf_fd(
-				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 */
-			return 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);
-		} else if (ret == -FI_ENOPROTOOPT) {
-			/* Protocol not availabe => 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;
-	}
-#endif
 
+	assert(efa_mr->domain->hmem_info[mr_attr->iface].p2p_supported_by_device);
+	EFA_INFO(FI_LOG_MR, "Registering memory with 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);