[PATCH V1] accel/amdxdna: Remap BO pages instead of using vm_insert_pages

[Lizhi Hou <lizhi.hou@amd.com>]

VM_PFNMAP is set for shmem BO to prevent pin_user_pages() from pining BO
memory. Replacing it with VM_MIXEDMAP breaks this behavior.

Keep the VM_PFNMAP mapping established by drm_gem_shmem_mmap() and
populate its page tables with remap_pfn_range() immediately after
creating the mapping.

Fixes: df4d5352063a ("accel/amdxdna: Fix unsafe use of handle_mm_fault()")
Signed-off-by: Lizhi Hou <lizhi.hou@amd.com>
---
 drivers/accel/amdxdna/amdxdna_gem.c | 157 +++++++++++-----------------
 1 file changed, 62 insertions(+), 95 deletions(-)

diff --git a/drivers/accel/amdxdna/amdxdna_gem.c b/drivers/accel/amdxdna/amdxdna_gem.c
index 46b9f6c960d7..09e48c794e1b 100644
--- a/drivers/accel/amdxdna/amdxdna_gem.c
+++ b/drivers/accel/amdxdna/amdxdna_gem.c
@@ -564,11 +564,51 @@ static void amdxdna_mark_mapp_invalid(struct amdxdna_gem_obj *abo,
 	up_write(&xdna->notifier_lock);
 }
 
+static int amdxdna_remap_pages(struct amdxdna_gem_obj *abo,
+			       struct vm_area_struct *vma)
+{
+	unsigned long num_pages = vma_pages(vma);
+	struct page **pages = abo->base.pages;
+	unsigned long i, nr;
+	int ret;
+
+	for (i = 0; i < num_pages; i += nr) {
+		unsigned long pfn = page_to_pfn(pages[i]);
+
+		for (nr = 1; i + nr < num_pages; nr++) {
+			if (page_to_pfn(pages[i + nr]) != pfn + nr)
+				break;
+		}
+
+		ret = remap_pfn_range(vma, vma->vm_start + (i << PAGE_SHIFT), pfn,
+				      nr << PAGE_SHIFT, vma->vm_page_prot);
+		if (ret) {
+			zap_special_vma_range(vma, vma->vm_start, i << PAGE_SHIFT);
+			return ret;
+		}
+	}
+
+	/*
+	 * The page tables are populated upfront, thus the CPU never takes a
+	 * write fault and drm_gem_shmem_vm_ops.pfn_mkwrite() never marks the
+	 * folios dirty. The device writes to the pages through the IOMMU
+	 * without any notification either. And the PTEs are special in a
+	 * VM_PFNMAP VMA, so unmapping does not transfer their dirty bits
+	 * back to the folios. Mark the pages dirty when they are released,
+	 * the same way drm_gem_shmem_vmap() does.
+	 */
+	dma_resv_lock(to_gobj(abo)->resv, NULL);
+	abo->base.pages_mark_dirty_on_put = true;
+	abo->base.pages_mark_accessed_on_put = true;
+	dma_resv_unlock(to_gobj(abo)->resv);
+
+	return 0;
+}
+
 static int amdxdna_insert_pages(struct amdxdna_gem_obj *abo,
 				struct vm_area_struct *vma)
 {
 	struct amdxdna_dev *xdna = to_xdna_dev(to_gobj(abo)->dev);
-	unsigned long num_pages = vma_pages(vma);
 	int ret;
 
 	/*
@@ -587,7 +627,8 @@ static int amdxdna_insert_pages(struct amdxdna_gem_obj *abo,
 			return ret;
 		}
 
-		amdxdna_mark_mapp_invalid(abo, vma);
+		if (!(vma->vm_flags & VM_PFNMAP))
+			amdxdna_mark_mapp_invalid(abo, vma);
 
 		/* Drop the reference drm_gem_mmap_obj() acquired.*/
 		drm_gem_object_put(to_gobj(abo));
@@ -600,10 +641,9 @@ static int amdxdna_insert_pages(struct amdxdna_gem_obj *abo,
 		return ret;
 	}
 
-	vm_flags_mod(vma, VM_MIXEDMAP, VM_PFNMAP);
-	ret = vm_insert_pages(vma, vma->vm_start, abo->base.pages, &num_pages);
+	ret = amdxdna_remap_pages(abo, vma);
 	if (ret) {
-		XDNA_ERR(xdna, "Failed to insert pages %d", ret);
+		XDNA_ERR(xdna, "Failed to remap pages %d", ret);
 		dma_resv_lock(to_gobj(abo)->resv, NULL);
 		drm_gem_shmem_put_pages_locked(&abo->base);
 		dma_resv_unlock(to_gobj(abo)->resv);
@@ -641,99 +681,21 @@ static int amdxdna_gem_obj_mmap(struct drm_gem_object *gobj,
 	return ret;
 }
 
-/*
- * VM operations for amdxdna shmem VMAs that use VM_MIXEDMAP.
- *
- * drm_gem_shmem_vm_ops cannot be used on VM_MIXEDMAP VMAs because its fault
- * handler calls vmf_insert_pfn() → vmf_insert_pfn_prot() which contains:
- *   BUG_ON((vma->vm_flags & VM_MIXEDMAP) && pfn_valid(pfn))
- * All amdxdna shmem pages are ordinary struct pages so pfn_valid() is always
- * true, making the combination fatal.
- *
- * These ops use vmf_insert_page() (struct-page based) instead, which is the
- * correct API for VM_MIXEDMAP VMAs backed by real struct pages.  The open and
- * close handlers replicate drm_gem_shmem_vm_open/close using only exported
- * symbols.
- */
-static vm_fault_t amdxdna_gem_mixedmap_fault(struct vm_fault *vmf)
-{
-	struct vm_area_struct *vma = vmf->vma;
-	struct drm_gem_object *gobj = vma->vm_private_data;
-	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
-	loff_t num_pages = gobj->size >> PAGE_SHIFT;
-	vm_fault_t ret = VM_FAULT_SIGBUS;
-	pgoff_t page_offset;
-	struct page *page;
-
-	/*
-	 * Partial free of vma is unexpected. Otherwise, the wrong page
-	 * will be faulted in and the user application may crash itself.
-	 */
-	page_offset = linear_page_delta(vma, vmf->address);
-
-	dma_resv_lock(gobj->resv, NULL);
-
-	if (!shmem->pages || shmem->madv < 0 || page_offset >= num_pages)
-		goto out;
-
-	page = shmem->pages[page_offset];
-	if (WARN_ON_ONCE(!page))
-		goto out;
-
-	/*
-	 * Use vmf_insert_page() (struct-page path) not vmf_insert_pfn()
-	 * (PFN path) because this VMA carries VM_MIXEDMAP.
-	 */
-	ret = vmf_insert_page(vma, vmf->address, page);
-	if (ret == VM_FAULT_NOPAGE)
-		folio_mark_accessed(page_folio(page));
-
-out:
-	dma_resv_unlock(gobj->resv);
-	return ret;
-}
-
-static void amdxdna_gem_mixedmap_vm_open(struct vm_area_struct *vma)
-{
-	struct drm_gem_object *gobj = vma->vm_private_data;
-	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
-
-	/*
-	 * Bump pages_use_count so the page array stays alive for the new
-	 * mapping copy created by fork().  Mirrors drm_gem_shmem_vm_open().
-	 */
-	dma_resv_lock(gobj->resv, NULL);
-	drm_WARN_ON_ONCE(gobj->dev, !refcount_inc_not_zero(&shmem->pages_use_count));
-	dma_resv_unlock(gobj->resv);
-
-	drm_gem_vm_open(vma);
-}
-
-static void amdxdna_gem_mixedmap_vm_close(struct vm_area_struct *vma)
-{
-	struct drm_gem_object *gobj = vma->vm_private_data;
-	struct drm_gem_shmem_object *shmem = to_drm_gem_shmem_obj(gobj);
-
-	dma_resv_lock(gobj->resv, NULL);
-	drm_gem_shmem_put_pages_locked(shmem);
-	dma_resv_unlock(gobj->resv);
-
-	drm_gem_vm_close(vma);
-}
-
-static const struct vm_operations_struct amdxdna_gem_mixedmap_vm_ops = {
-	.fault  = amdxdna_gem_mixedmap_fault,
-	.open   = amdxdna_gem_mixedmap_vm_open,
-	.close  = amdxdna_gem_mixedmap_vm_close,
-};
-
 static int amdxdna_gem_dmabuf_mmap(struct dma_buf *dma_buf, struct vm_area_struct *vma)
 {
 	struct drm_gem_object *gobj = dma_buf->priv;
 	struct amdxdna_gem_obj *abo = to_xdna_obj(gobj);
 	int ret;
 
-	vma->vm_ops = &amdxdna_gem_mixedmap_vm_ops;
+	/*
+	 * vm_pgoff is the page offset within the dma-buf, see dma_buf_mmap().
+	 * drm_gem_shmem_vm_ops indexes the BO pages by the page offset within
+	 * the VMA. Thus the mapping has to start at the beginning of the BO.
+	 */
+	if (vma->vm_pgoff)
+		return -EINVAL;
+
+	vma->vm_ops = &drm_gem_shmem_vm_ops;
 	vma->vm_private_data = gobj;
 
 	drm_gem_object_get(gobj);
@@ -741,11 +703,16 @@ static int amdxdna_gem_dmabuf_mmap(struct dma_buf *dma_buf, struct vm_area_struc
 	if (ret)
 		goto put_obj;
 
-	/* The buffer is based on memory pages. Fix the flag. */
-	vm_flags_mod(vma, VM_MIXEDMAP, VM_PFNMAP);
+	ret = amdxdna_remap_pages(abo, vma);
+	if (ret)
+		goto put_pages;
 
 	return 0;
 
+put_pages:
+	dma_resv_lock(gobj->resv, NULL);
+	drm_gem_shmem_put_pages_locked(&abo->base);
+	dma_resv_unlock(gobj->resv);
 put_obj:
 	drm_gem_object_put(gobj);
 	return ret;
@@ -1043,7 +1010,7 @@ static const struct drm_gem_object_funcs amdxdna_gem_shmem_funcs = {
 	.vmap = amdxdna_gem_obj_vmap,
 	.vunmap = amdxdna_gem_obj_vunmap,
 	.mmap = amdxdna_gem_obj_mmap,
-	.vm_ops = &amdxdna_gem_mixedmap_vm_ops,
+	.vm_ops = &drm_gem_shmem_vm_ops,
 	.export = amdxdna_gem_prime_export,
 };
 
-- 
2.34.1