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

[Pranjal Shrivastava <praan@google.com>]

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