[RFC PATCH v2 0/5] vfio/pci: Support ZONE_DEVICE-backed DMABUF Exports

[RFC PATCH v2 0/5] vfio/pci: Support ZONE_DEVICE-backed DMABUF Exports

[Pranjal Shrivastava]

Introduce ZONE_DEVICE backing for VFIO-exposed PCIe BARs via the DMABUF
subsystem. This series is based on Matt's on-going work on VFIO DMABUF
mmap [1].

Currently, kernel drivers can register their BARs with the P2PDMA
subsystem to enable high-performance, page-backed P2P DMA. However, when
a device is bound to vfio-pci, this capability is missing. This prevents
userspace drivers from performing zero-copy P2P DMA via standard POSIX APIs
(e.g., O_DIRECT) which require struct page metadata.

Based on feedback from v1, this series pivots away from system-wide P2P
registration. Instead, it integrates with the ongoing VFIO DMABUF-mmap
work [1], introducing on-demand ZONE_DEVICE allocation exclusively for
DMABUF exports. This allows DMABUFs to optionally register with ZONE_DEVICE
while laying the groundwork for userspace NFS clients and other storage
targets to perform zero-copy P2PDMA against VFIO-managed memory.

(Note: There's on-going work to support P2PDMA on NFS [2])

Design
======
The proposed design involves the following:

a) On-Demand ZONE_DEVICE Registration
A new UAPI flag (VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED) is introduced to
the VFIO_DEVICE_FEATURE_DMA_BUF ioctl to allow users to explicitly "opt-in"
to struct page backing on a per-dmabuf export basis. This ensures we only
allocate vmemmap memory when requested by the user (e.g., NFS + O_DIRECT).

b) Sticky Registration
In order to prevent vmemmap memory fragmentation, ZONE_DEVICE allocations
are "sticky" at the BAR level. Since the ZONE_DEVICE registration relies
on the pci_p2pdma_add_resource() API internally, it is guaranteed to
have the allocations stay until the vfio-pci driver is unbound.

Thus, once a DMABUF export requests page backing, struct pages are allocated
for the entire BAR (even if the DMABUF spans a smaller region within the BAR),
and this allocation survives DMABUF closures and device resets. 

If VM1 opts-in to ZONE_DEVICE, undergoes a reset, and the device is subsequently
assigned to VM2 without the opt-in flag, VFIO simply exports the standard DMABUF
and ignores the underlying struct pages.

c) Revocation Strategy
During a device reset or teardown, VFIO revokes the DMABUF. the design implements
a synchronous revocation fence that blocks indefinitely until all struct page 
refcounts drop to 1 (meaning the importer has fully released them) to prevent
DMA-after-free corruption.

d) DMABUF Interoperability
A custom .map_dma_buf handler is implemented to ensure a ZONE_DEVICE-backed DMABUF
can be used as a regular VFIO-exported DMABUF ensuring importers can still use
standard SG-table-based APIs seamlessly alongside page-based ones.

e) Concurrency
The page allocation and refcount initialization are serialized via the memory_lock
write semaphore to prevent concurrent fault races. Additionally, we call a
unmap_mapping_range() during revocation while holding memory_lock without triggering
a circular rmap deadlock (mmap_lock -> memory_lock -> i_mmap_rwsem). This is safely
avoided because pages allocated via devm_memremap_pages() do not have page->mapping
set, rendering them invisible to the rmap.

Call for Review & Design Trade-Offs
====================================
Please provide feedback on the sticky registration and revocation strategy.

I've evaluated the following and would appreciate guidance on these tradeoffs:

a) Sticky Registration vs. Ephemeral Teardown
Instead of making vmemmap allocations sticky across resets, we could free and 
re-allocate them per-session. While this would eliminate the need for our 
indefinite polling loop (as the devres teardown handles it natively) and
clean up struct pages that are not needed anymore after the fd closure.

This design opts for the sticky approach to avoid vmemmap fragmentation over the 
host's uptime. 

b) pci_p2pdma_add_resource vs. Open-Coded devm_memremap

This implementation relies on pci_p2pdma_add_resource(), which couples the
vmemmap lifecycle to the vfio-pci driver unbind event. Alternatively, we 
could open-code a devm_memremap_pages() implementation (similar to P2PDMA API)
directly within VFIO to gain finer-grained control over the teardown
(maybe something like vfio_p2pdma_add_resource() or something).

I've avoided that in this version to first gain consensus on the fragmentation
and revocation fence.

c) Interoperability
The propsed design assumes that an importer of a ZONE_DEVICE DMABUF may still
want to utilize standard DMABUF operations, and thus implemented the .map_dma_buf
op. An alternative would be to strictly enforce that ZONE_DEVICE DMABUFs only 
support page-backed usage, explicitly rejecting standard DMABUF operations.

[1] https://lore.kernel.org/all/20260715174737.15287-1-matt@ozlabs.org/
[2] https://lore.kernel.org/all/20260720150601.2702700-1-praan@google.com/

Thanks,
Praan

Pranjal Shrivastava (5):
  vfio: Add UAPI flag for ZONE_DEVICE-backed DMABUF exports
  vfio/pci: Implement ZONE_DEVICE registration for DMABUFs
  vfio/pci: Implement page-backed .map_dma_buf handler
  vfio/pci: Add .mmap handler for page-backed DMABUFs
  vfio/pci: Add revocation fence for ZONE_DEVICE DMABUFs

 drivers/vfio/pci/Kconfig           |  11 ++
 drivers/vfio/pci/vfio_pci_core.c   |  39 +++++-
 drivers/vfio/pci/vfio_pci_dmabuf.c | 213 ++++++++++++++++++++++++++++-
 drivers/vfio/pci/vfio_pci_priv.h   |   1 +
 include/linux/vfio_pci_core.h      |   4 +
 include/uapi/linux/vfio.h          |  10 +-
 6 files changed, 266 insertions(+), 12 deletions(-)

-- 
2.55.0.571.g244d577d93-goog

[RFC PATCH v2 1/5] vfio: Add UAPI flag for ZONE_DEVICE-backed DMABUF exports

[Pranjal Shrivastava]

Add an opt-in UAPI flag for exporting VFIO MMIO regions via dmabuf
backed by struct page metadata (via ZONE_DEVICE registration).

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 include/uapi/linux/vfio.h | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/include/uapi/linux/vfio.h b/include/uapi/linux/vfio.h
index cbe03e5d1009..3ae025b33e7f 100644
--- a/include/uapi/linux/vfio.h
+++ b/include/uapi/linux/vfio.h
@@ -1505,7 +1505,13 @@ struct vfio_device_feature_bus_master {
  * etc. offset/length specify a slice of the region to create the dmabuf from.
  * nr_ranges is the total number of (P2P DMA) ranges that comprise the dmabuf.
  *
- * flags should be 0.
+ * flags supports:
+ *   - VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED: Request the kernel to allocate
+ *     valid struct pages for the exported dmabuf range (via ZONE_DEVICE
+ *     registration). This allows the exported dmabuf fd to be used by
+ *     subsystems that rely on page metadata.
+ *     (NOTE: This consumes extra system RAM to store BAR_SIZE / PAGE_SIZE
+ *     struct pages).
  *
  * Return: The fd number on success, -1 and errno is set on failure.
  */
@@ -1516,6 +1522,8 @@ struct vfio_region_dma_range {
 	__u64 length;
 };
 
+#define VFIO_DMA_BUF_FLAG_ZONE_DEVICE_BACKED	(1 << 0)
+
 struct vfio_device_feature_dma_buf {
 	__u32	region_index;
 	__u32	open_flags;
-- 
2.55.0.571.g244d577d93-goog

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

[Pranjal Shrivastava]

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

[RFC PATCH v2 3/5] vfio/pci: Implement page-backed .map_dma_buf handler

[Pranjal Shrivastava]

When a DMABUF is backed by ZONE_DEVICE pages, the standard raw PFN
scatterlist builder (dma_buf_phys_vec_to_sgt) cannot be used because
it explicitly drops page metadata pointers. Implement a map_dma_buf
helper that loops through existing contiguous physical ranges and
generates SGL entries directly via sg_set_page while maps using the
standard dma_map_sgtable. Implement a corresponding .unmap_dma_buf as
well.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/vfio_pci_dmabuf.c | 81 ++++++++++++++++++++++++++++--
 1 file changed, 77 insertions(+), 4 deletions(-)

diff --git a/drivers/vfio/pci/vfio_pci_dmabuf.c b/drivers/vfio/pci/vfio_pci_dmabuf.c
index 900ac1851c9f..b936da3bcada 100644
--- a/drivers/vfio/pci/vfio_pci_dmabuf.c
+++ b/drivers/vfio/pci/vfio_pci_dmabuf.c
@@ -55,6 +55,7 @@ static int vfio_pci_dma_buf_mmap(struct dma_buf *dmabuf, struct vm_area_struct *
 	/* See comments in vfio_pci_core_mmap() re VM_ALLOW_ANY_UNCACHED. */
 	vm_flags_set(vma, VM_ALLOW_ANY_UNCACHED | VM_IO | VM_PFNMAP |
 		     VM_DONTEXPAND | VM_DONTDUMP);
+
 	vma->vm_private_data = priv;
 	vfio_pci_set_vma_ops(vma);
 
@@ -70,6 +71,66 @@ static void vfio_pci_dma_buf_done(struct kref *kref)
 	complete(&priv->comp);
 }
 
+/*
+ * For ZONE_DEVICE-backed DMABUFs, populate the scatterlist with struct page
+ * pointers so that dma_map_sgtable() can detect MEMORY_DEVICE_PCI_P2PDMA and
+ * perform peer-to-peer DMA mappings for importing devices.
+ */
+static struct sg_table *
+vfio_pci_dma_buf_map_page_backed(struct dma_buf_attachment *attachment,
+				 enum dma_data_direction dir)
+{
+	struct vfio_pci_dma_buf *priv = attachment->dmabuf->priv;
+	struct sg_table *sgt;
+	struct scatterlist *sgl;
+	unsigned int nents = 0;
+	int i, ret;
+
+	for (i = 0; i < priv->nr_ranges; i++) {
+		unsigned int added = DIV_ROUND_UP(priv->phys_vec[i].len,
+						 (UINT_MAX & PAGE_MASK));
+
+		if (check_add_overflow(nents, added, &nents))
+			return ERR_PTR(-EOVERFLOW);
+	}
+
+	sgt = kzalloc(sizeof(*sgt), GFP_KERNEL);
+	if (!sgt)
+		return ERR_PTR(-ENOMEM);
+
+	ret = sg_alloc_table(sgt, nents, GFP_KERNEL);
+	if (ret) {
+		kfree(sgt);
+		return ERR_PTR(ret);
+	}
+
+	sgl = sgt->sgl;
+	for (i = 0; i < priv->nr_ranges; i++) {
+		size_t range_len = priv->phys_vec[i].len;
+		unsigned long pfn = priv->phys_vec[i].paddr >> PAGE_SHIFT;
+
+		while (range_len > 0) {
+			unsigned int chunk_len = min_t(size_t, range_len, (UINT_MAX & PAGE_MASK));
+			struct page *page = pfn_to_page(pfn);
+
+			sg_set_page(sgl, page, chunk_len, 0);
+			sgl = sg_next(sgl);
+
+			range_len -= chunk_len;
+			pfn += chunk_len >> PAGE_SHIFT;
+		}
+	}
+
+	ret = dma_map_sgtable(attachment->dev, sgt, dir, 0);
+	if (ret) {
+		sg_free_table(sgt);
+		kfree(sgt);
+		return ERR_PTR(ret);
+	}
+
+	return sgt;
+}
+
 static struct sg_table *
 vfio_pci_dma_buf_map(struct dma_buf_attachment *attachment,
 		     enum dma_data_direction dir)
@@ -82,9 +143,14 @@ vfio_pci_dma_buf_map(struct dma_buf_attachment *attachment,
 	if (priv->status != VFIO_PCI_DMABUF_OK)
 		return ERR_PTR(-ENODEV);
 
-	ret = dma_buf_phys_vec_to_sgt(attachment, priv->provider,
-				      priv->phys_vec, priv->nr_ranges,
-				      priv->size, dir);
+	if (priv->zone_device_backed) {
+		ret = vfio_pci_dma_buf_map_page_backed(attachment, dir);
+	} else {
+		ret = dma_buf_phys_vec_to_sgt(attachment, priv->provider,
+					      priv->phys_vec, priv->nr_ranges,
+					      priv->size, dir);
+	}
+
 	if (IS_ERR(ret))
 		return ret;
 
@@ -100,7 +166,14 @@ static void vfio_pci_dma_buf_unmap(struct dma_buf_attachment *attachment,
 
 	dma_resv_assert_held(priv->dmabuf->resv);
 
-	dma_buf_free_sgt(attachment, sgt, dir);
+	if (priv->zone_device_backed) {
+		dma_unmap_sgtable(attachment->dev, sgt, dir, 0);
+		sg_free_table(sgt);
+		kfree(sgt);
+	} else {
+		dma_buf_free_sgt(attachment, sgt, dir);
+	}
+
 	kref_put(&priv->kref, vfio_pci_dma_buf_done);
 }
 
-- 
2.55.0.571.g244d577d93-goog

[RFC PATCH v2 4/5] vfio/pci: Add .mmap handler for page-backed DMABUFs

[Pranjal Shrivastava]

Implement mmap and a corresponding page fault handler for ZONE_DEVICE
backed DMABUFs. Set VM_MIXEDMAP for page-backed VMAs and introduce the
vfio_pci_vmf_insert_page() helper to insert struct page pointers into
the user PTEs.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/vfio_pci_core.c   | 39 +++++++++++++++++++++++++++---
 drivers/vfio/pci/vfio_pci_dmabuf.c | 22 +++++++++++++++--
 include/linux/vfio_pci_core.h      |  3 +++
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/drivers/vfio/pci/vfio_pci_core.c b/drivers/vfio/pci/vfio_pci_core.c
index 9ca99a5a61c3..f5c2912c3ffc 100644
--- a/drivers/vfio/pci/vfio_pci_core.c
+++ b/drivers/vfio/pci/vfio_pci_core.c
@@ -1790,6 +1790,34 @@ vm_fault_t vfio_pci_vmf_insert_pfn(struct vfio_pci_core_device *vdev,
 }
 EXPORT_SYMBOL_GPL(vfio_pci_vmf_insert_pfn);
 
+vm_fault_t vfio_pci_vmf_insert_page(struct vfio_pci_core_device *vdev,
+				    struct vm_fault *vmf,
+				    unsigned long pfn,
+				    unsigned int order)
+{
+	struct vm_area_struct *vma = vmf->vma;
+	struct page *page;
+
+	lockdep_assert_held_read(&vdev->memory_lock);
+
+	if (vdev->pm_runtime_engaged || !__vfio_pci_memory_enabled(vdev))
+		return VM_FAULT_SIGBUS;
+
+	/* vmf_insert_page only supports 0-order pages */
+	if (order > 0)
+		return VM_FAULT_FALLBACK;
+
+	if (WARN_ON_ONCE(!pfn_valid(pfn)))
+		return VM_FAULT_SIGBUS;
+
+	page = pfn_to_page(pfn);
+
+	if (vma->vm_flags & VM_WRITE)
+		return vmf_insert_page_mkwrite(vmf, page, false);
+
+	return vmf_insert_page(vma, vmf->address, page);
+}
+
 static vm_fault_t vfio_pci_mmap_huge_fault(struct vm_fault *vmf,
 					   unsigned int order)
 {
@@ -1877,11 +1905,14 @@ static vm_fault_t vfio_pci_mmap_huge_fault(struct vm_fault *vmf,
 							     vmf->address,
 							     order, &pfn);
 
-			if (pres == 0)
-				ret = vfio_pci_vmf_insert_pfn(vdev, vmf,
-							      pfn, order);
-			else if (pres == -EAGAIN)
+			if (pres == 0) {
+				if (priv->zone_device_backed)
+					ret = vfio_pci_vmf_insert_page(vdev, vmf, pfn, order);
+				else
+					ret = vfio_pci_vmf_insert_pfn(vdev, vmf, pfn, order);
+			} else if (pres == -EAGAIN) {
 				ret = VM_FAULT_FALLBACK;
+			}
 		}
 	}
 
diff --git a/drivers/vfio/pci/vfio_pci_dmabuf.c b/drivers/vfio/pci/vfio_pci_dmabuf.c
index b936da3bcada..b582e856ba7a 100644
--- a/drivers/vfio/pci/vfio_pci_dmabuf.c
+++ b/drivers/vfio/pci/vfio_pci_dmabuf.c
@@ -4,6 +4,7 @@
 #include <linux/dma-buf-mapping.h>
 #include <linux/pci-p2pdma.h>
 #include <linux/dma-resv.h>
+#include <linux/sched.h>
 #include <uapi/linux/dma-buf.h>
 
 #include "vfio_pci_priv.h"
@@ -53,8 +54,14 @@ static int vfio_pci_dma_buf_mmap(struct dma_buf *dmabuf, struct vm_area_struct *
 	vma->vm_page_prot = pgprot_decrypted(vma->vm_page_prot);
 
 	/* See comments in vfio_pci_core_mmap() re VM_ALLOW_ANY_UNCACHED. */
-	vm_flags_set(vma, VM_ALLOW_ANY_UNCACHED | VM_IO | VM_PFNMAP |
-		     VM_DONTEXPAND | VM_DONTDUMP);
+	if (priv->zone_device_backed) {
+		/* VM_MIXEDMAP is required for ZONE_DEVICE pages */
+		vm_flags_set(vma, VM_ALLOW_ANY_UNCACHED | VM_MIXEDMAP |
+			     VM_DONTEXPAND | VM_DONTDUMP);
+	} else {
+		vm_flags_set(vma, VM_ALLOW_ANY_UNCACHED | VM_IO | VM_PFNMAP |
+			     VM_DONTEXPAND | VM_DONTDUMP);
+	}
 
 	vma->vm_private_data = priv;
 	vfio_pci_set_vma_ops(vma);
@@ -521,6 +528,17 @@ static int vfio_pci_dma_buf_alloc_struct_pages(struct vfio_pci_core_device *vdev
 	if (ret) {
 		if (ret != -EEXIST)
 			return ret;
+	} else {
+		/* Initialise the refcount for the freshly allocated page. */
+		unsigned long pfn = pci_resource_start(pdev, bar_index) >> PAGE_SHIFT;
+		unsigned long npgs = pci_resource_len(pdev, bar_index) >> PAGE_SHIFT;
+
+		while (npgs--) {
+			set_page_count(pfn_to_page(pfn++), 1);
+			/* Yield the CPU periodically on large BARs to prevent soft lockups */
+			if (unlikely((npgs & 4095) == 0))
+				cond_resched();
+		}
 	}
 
 	vdev->p2p_struct_page_bars |= (1 << bar_index);
diff --git a/include/linux/vfio_pci_core.h b/include/linux/vfio_pci_core.h
index c28f06bae302..778671194b24 100644
--- a/include/linux/vfio_pci_core.h
+++ b/include/linux/vfio_pci_core.h
@@ -187,6 +187,9 @@ ssize_t vfio_pci_core_write(struct vfio_device *core_vdev, const char __user *bu
 vm_fault_t vfio_pci_vmf_insert_pfn(struct vfio_pci_core_device *vdev,
 				   struct vm_fault *vmf, unsigned long pfn,
 				   unsigned int order);
+vm_fault_t vfio_pci_vmf_insert_page(struct vfio_pci_core_device *vdev,
+				    struct vm_fault *vmf, unsigned long pfn,
+				    unsigned int order);
 int vfio_pci_core_mmap(struct vfio_device *core_vdev, struct vm_area_struct *vma);
 void vfio_pci_core_request(struct vfio_device *core_vdev, unsigned int count);
 int vfio_pci_core_match(struct vfio_device *core_vdev, char *buf);
-- 
2.55.0.571.g244d577d93-goog

[RFC PATCH v2 5/5] vfio/pci: Add revocation fence for ZONE_DEVICE DMABUFs

[Pranjal Shrivastava]

Implement a synchronization fence to safely revoke ZONE_DEVICE-backed
DMABUFs. Introduce a fence in vfio_pci_dma_buf_set_status(). The fence
waits for all struct page refcounts to drop to 1.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/vfio_pci_dmabuf.c | 55 ++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/drivers/vfio/pci/vfio_pci_dmabuf.c b/drivers/vfio/pci/vfio_pci_dmabuf.c
index b582e856ba7a..b4284b5cad03 100644
--- a/drivers/vfio/pci/vfio_pci_dmabuf.c
+++ b/drivers/vfio/pci/vfio_pci_dmabuf.c
@@ -4,6 +4,7 @@
 #include <linux/dma-buf-mapping.h>
 #include <linux/pci-p2pdma.h>
 #include <linux/dma-resv.h>
+#include <linux/iopoll.h>
 #include <linux/sched.h>
 #include <uapi/linux/dma-buf.h>
 
@@ -745,6 +746,44 @@ int vfio_pci_core_mmap_prep_dmabuf(struct vfio_pci_core_device *vdev,
 	return ret;
 }
 
+static void vfio_pci_zone_device_wait_fence(struct vfio_pci_dma_buf *priv)
+{
+	unsigned int i;
+
+	if (!priv->zone_device_backed)
+		return;
+
+	/*
+	 * Fence: Wait for any active references to the ZONE_DEVICE
+	 * pages to be dropped. A refcount of 1 represents the base
+	 * ownership.
+	 */
+	for (i = 0; i < priv->nr_ranges; i++) {
+		unsigned long pfn = priv->phys_vec[i].paddr >> PAGE_SHIFT;
+		unsigned long npgs = PAGE_ALIGN(priv->phys_vec[i].len) >> PAGE_SHIFT;
+
+		while (npgs--) {
+			struct page *page = pfn_to_page(pfn++);
+			int count, ret;
+
+			/*
+			 * Poll page_count() and block indefinitely until all
+			 * refs drop to avoid DMA-after-free.
+			 */
+			do {
+				ret = read_poll_timeout(page_count, count,
+							(count == 1),
+							1000, 10000000,
+							false, page);
+				if (ret)
+					dev_warn(&priv->vdev->pdev->dev,
+						 "Waiting for GUP pins to drop on PFN 0x%lx... (importer hung?)\n",
+						 pfn - 1);
+			} while (ret);
+		}
+	}
+}
+
 /* Set the DMABUF's revocation status (OK or temporarily/permanently revoked) */
 static void vfio_pci_dma_buf_set_status(struct vfio_pci_dma_buf *priv,
 					enum vfio_pci_dma_buf_status new_status)
@@ -779,8 +818,24 @@ static void vfio_pci_dma_buf_set_status(struct vfio_pci_dma_buf *priv,
 		dma_resv_unlock(priv->dmabuf->resv);
 		kref_put(&priv->kref, vfio_pci_dma_buf_done);
 		wait_for_completion(&priv->comp);
+
+		/*
+		 * Note: Rmap Deadlocks
+		 * unmap_mapping_range() is safe to call here within memory_lock
+		 * despite the VMA being VM_MIXEDMAP. Because our ZONE_DEVICE pages
+		 * are allocated via devm_memremap_pages(), page->mapping is never
+		 * set which makes them invisible to rmap.
+		 *
+		 * If this changes in the future, this call must be factored outside
+		 * the memory_lock to prevent a 3-way circular deadlock:
+		 * (mmap_lock -> memory_lock -> i_mmap_rwsem).
+		 */
 		unmap_mapping_range(priv->dmabuf->file->f_mapping,
 				    0, 0, true);
+
+		/* Wait for all page refs to drop if ZONE_DEVICE registered */
+		vfio_pci_zone_device_wait_fence(priv);
+
 		/*
 		 * Re-arm the registered kref reference and the
 		 * completion so the post-revoke state matches the
-- 
2.55.0.571.g244d577d93-goog