[RFC PATCH v2 2/5] vfio/pci: Implement ZONE_DEVICE registration for DMABUFs

[Pranjal Shrivastava <praan@google.com>]

When a DMABUF is requested with the VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED
flag is requested during DMABUF export, invoke pci_p2pdma_add_resource()
to register the target BAR with the ZONE_DEVICE framework. This allocates
the underlying struct page metadata. Add internal state tracking for
ZONE_DEVICE-backed DMABUF exports. Add p2p_struct_page_bars, a bitmask to
track PCI BARs that have been registered. Add a zone_device_backed flag
to track whether a specific DMABUF contains page-backed memory.

Introduce a new Kconfig option, VFIO_PCI_DMABUF_ZONE_DEVICE to handle
the dependency on PCI_P2PDMA without clobbering existing dependencies.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/Kconfig           | 11 ++++++
 drivers/vfio/pci/vfio_pci_dmabuf.c | 55 +++++++++++++++++++++++++++++-
 drivers/vfio/pci/vfio_pci_priv.h   |  1 +
 include/linux/vfio_pci_core.h      |  1 +
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/drivers/vfio/pci/Kconfig b/drivers/vfio/pci/Kconfig
index 67a2ae1fbc04..c9cd04a16297 100644
--- a/drivers/vfio/pci/Kconfig
+++ b/drivers/vfio/pci/Kconfig
@@ -61,6 +61,17 @@ config VFIO_PCI_DMABUF
 	def_bool y if PCI_P2PDMA
 	depends on VFIO_PCI_CORE
 
+config VFIO_PCI_DMABUF_ZONE_DEVICE
+	bool "VFIO PCI DMABUF ZONE_DEVICE page-backed export support"
+	depends on VFIO_PCI_DMABUF
+	depends on PCI_P2PDMA
+	help
+	  Say Y here to enable optional struct page backing (ZONE_DEVICE)
+	  for VFIO exported DMABUFs. This is required to support peer-to-peer
+	  (P2P) DMA transactions with subsystems that rely on page metadata.
+
+	  If unsure, say N.
+
 source "drivers/vfio/pci/mlx5/Kconfig"
 
 source "drivers/vfio/pci/ism/Kconfig"
diff --git a/drivers/vfio/pci/vfio_pci_dmabuf.c b/drivers/vfio/pci/vfio_pci_dmabuf.c
index 36bf07530840..900ac1851c9f 100644
--- a/drivers/vfio/pci/vfio_pci_dmabuf.c
+++ b/drivers/vfio/pci/vfio_pci_dmabuf.c
@@ -427,6 +427,41 @@ static int validate_dmabuf_input(struct vfio_device_feature_dma_buf *dma_buf,
 	return 0;
 }
 
+#ifdef CONFIG_VFIO_PCI_DMABUF_ZONE_DEVICE
+static int vfio_pci_dma_buf_alloc_struct_pages(struct vfio_pci_core_device *vdev,
+					       struct vfio_device_feature_dma_buf *dma_buf)
+{
+	struct pci_dev *pdev = vdev->pdev;
+	u32 bar_index = dma_buf->region_index;
+	int ret;
+
+	if (vdev->p2p_struct_page_bars & (1 << bar_index))
+		return 0;
+
+	/*
+	 * Allocate vmemmap (struct pages) for the ENTIRE BAR, even if the
+	 * specific DMABUF only exports a partial slice of it. This prevents
+	 * vmemmap fragmentation and ensures that any subsequent slice exports
+	 * from the same BAR get struct page backing instantly.
+	 */
+	ret = pci_p2pdma_add_resource(pdev, bar_index, 0, 0);
+	if (ret) {
+		if (ret != -EEXIST)
+			return ret;
+	}
+
+	vdev->p2p_struct_page_bars |= (1 << bar_index);
+
+	return 0;
+}
+#else
+static inline int vfio_pci_dma_buf_alloc_struct_pages(struct vfio_pci_core_device *vdev,
+						      struct vfio_device_feature_dma_buf *dma_buf)
+{
+	return -EOPNOTSUPP;
+}
+#endif
+
 int vfio_pci_core_feature_dma_buf(struct vfio_pci_core_device *vdev, u32 flags,
 				  struct vfio_device_feature_dma_buf __user *arg,
 				  size_t argsz)
@@ -448,7 +483,8 @@ int vfio_pci_core_feature_dma_buf(struct vfio_pci_core_device *vdev, u32 flags,
 	if (copy_from_user(&get_dma_buf, arg, sizeof(get_dma_buf)))
 		return -EFAULT;
 
-	if (!get_dma_buf.nr_ranges || get_dma_buf.flags)
+	if (!get_dma_buf.nr_ranges ||
+	    (get_dma_buf.flags & ~VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED))
 		return -EINVAL;
 
 	/*
@@ -468,6 +504,21 @@ int vfio_pci_core_feature_dma_buf(struct vfio_pci_core_device *vdev, u32 flags,
 	if (ret)
 		goto err_free_ranges;
 
+	if (get_dma_buf.flags & VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED) {
+		/*
+		 * Serialize the allocation and the page initialization.
+		 * Holding the memory_lock here prevents a race where a
+		 * concurrent ioctl bypasses the allocation + init, handing
+		 * a DMABUF to the user before the ZONE_DEVICE registration
+		 * is completed successfully.
+		 */
+		down_write(&vdev->memory_lock);
+		ret = vfio_pci_dma_buf_alloc_struct_pages(vdev, &get_dma_buf);
+		up_write(&vdev->memory_lock);
+		if (ret)
+			goto err_free_ranges;
+	}
+
 	priv = kzalloc_obj(*priv);
 	if (!priv) {
 		ret = -ENOMEM;
@@ -480,6 +531,8 @@ int vfio_pci_core_feature_dma_buf(struct vfio_pci_core_device *vdev, u32 flags,
 	}
 
 	priv->vdev = vdev;
+	if (get_dma_buf.flags & VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED)
+		priv->zone_device_backed = 1;
 	priv->nr_ranges = get_dma_buf.nr_ranges;
 	priv->size = length;
 	ret = vdev->pci_ops->get_dmabuf_phys(vdev, &priv->provider,
diff --git a/drivers/vfio/pci/vfio_pci_priv.h b/drivers/vfio/pci/vfio_pci_priv.h
index 8741abd04461..d0bf5c118793 100644
--- a/drivers/vfio/pci/vfio_pci_priv.h
+++ b/drivers/vfio/pci/vfio_pci_priv.h
@@ -42,6 +42,7 @@ struct vfio_pci_dma_buf {
 	struct completion comp;
 	unsigned long vma_pgoff_adjust;
 	enum vfio_pci_dma_buf_status status;
+	u8 zone_device_backed : 1;
 };
 
 bool vfio_pci_intx_mask(struct vfio_pci_core_device *vdev);
diff --git a/include/linux/vfio_pci_core.h b/include/linux/vfio_pci_core.h
index e2b4252e7c3f..c28f06bae302 100644
--- a/include/linux/vfio_pci_core.h
+++ b/include/linux/vfio_pci_core.h
@@ -135,6 +135,7 @@ struct vfio_pci_core_device {
 	bool			pm_runtime_engaged;
 	bool			sriov_active;
 	bool			zap_bars_on_revoke;
+	u8			p2p_struct_page_bars;
 	struct pci_saved_state	*pci_saved_state;
 	struct pci_saved_state	*pm_save;
 	int			ioeventfds_nr;
-- 
2.55.0.571.g244d577d93-goog