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

[Pranjal Shrivastava <praan@google.com>]

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