[PATCH 2/3] iommufd: Use iommu_unmap_sparse() to unpin domain pages

[Fred Griffoul <griffoul@gmail.com>]

From: Fred Griffoul <fgriffo@amazon.co.uk>

iommufd keeps no copy of the PFNs it pinned for an area, so
iopt_area_unpin_domain() reads each one back from the domain with
iommu_iova_to_phys(), one CPU page at a time, and then unmaps the range
in a second walk. With 4 KiB pages, a 1 GiB IOPTE costs 262144 lookups
before the unmap starts.

Call iommu_unmap_sparse() instead, and add the PFNs of each IOPTE it
removes to the pfn_batch, so one walk both unmaps the range and
collects the PFNs. The last IOPTE removed can hold pages beyond the
ones this pass unpins. The loop already handles that case: it keeps
those pages in the batch and carries them over to the next pass.

Domains that do not use iommupt keep the per-page read-back, and so do
areas whose IOVA or memory starts or ends inside a CPU page. Such areas
exist only when the smallest page size of every domain in the IOAS is
smaller than PAGE_SIZE.

Signed-off-by: Fred Griffoul <fgriffo@amazon.co.uk>
---
 drivers/iommu/iommufd/pages.c | 98 +++++++++++++++++++++++++++++++++--
 1 file changed, 93 insertions(+), 5 deletions(-)

diff --git a/drivers/iommu/iommufd/pages.c b/drivers/iommu/iommufd/pages.c
index 2f6153108add7..c2a04a8845cc8 100644
--- a/drivers/iommu/iommufd/pages.c
+++ b/drivers/iommu/iommufd/pages.c
@@ -468,6 +468,79 @@ static void batch_from_domain_continue(struct pfn_batch *batch,
 	batch->array_size = array_size;
 }
 
+struct batch_unmap {
+	struct pfn_batch *batch;
+	unsigned long first_iova;
+	unsigned long last_iova;
+};
+
+static bool batch_record_unmapped(void *priv, unsigned long iova,
+				  phys_addr_t paddr, size_t size)
+{
+	struct batch_unmap *unmap = priv;
+	unsigned long first, end;
+
+	/* Stop after the entry that holds the last page to unpin */
+	if (iova > unmap->last_iova)
+		return false;
+
+	/*
+	 * Count each CPU page once, from the entry that holds its first byte.
+	 * Pages before first_iova belong to an access and stay pinned.
+	 */
+	first = max(ALIGN(iova, PAGE_SIZE), unmap->first_iova);
+	end = ALIGN(iova + size, PAGE_SIZE);
+	if (first >= end)
+		return true;
+	return batch_add_pfn_num(unmap->batch, PHYS_PFN(paddr + (first - iova)),
+				 (end - first) >> PAGE_SHIFT, BATCH_CPU_MEMORY);
+}
+
+/* True if the area starts and ends on CPU pages, in both IOVA and memory */
+static bool iopt_area_pages_aligned(struct iopt_area *area)
+{
+	return !area->page_offset &&
+	       PAGE_ALIGNED(iopt_area_iova(area) |
+			    (iopt_area_last_iova(area) + 1));
+}
+
+/*
+ * Unmap from unmapped_index and collect the PFNs from start_index on into the
+ * batch in one walk. The walk ends after the entry that holds last_index, which
+ * can reach past last_index, or before an entry the batch has no room for.
+ * Returns the number of indexes unmapped, 0 if the caller has to read the PFNs
+ * back with batch_from_domain() instead.
+ */
+static unsigned long batch_unmap_from_domain(struct pfn_batch *batch,
+					     struct iommu_domain *domain,
+					     struct iopt_area *area,
+					     unsigned long unmapped_index,
+					     unsigned long start_index,
+					     unsigned long last_index,
+					     unsigned long real_last_index)
+{
+	unsigned long start = max(start_index, unmapped_index);
+	struct batch_unmap unmap = {
+		.batch = batch,
+		.first_iova = iopt_area_index_to_iova(area, start),
+		.last_iova = iopt_area_index_to_iova_last(area, last_index),
+	};
+	unsigned long iova = iopt_area_index_to_iova(area, unmapped_index);
+	unsigned long last =
+		iopt_area_index_to_iova_last(area, real_last_index);
+	size_t unmapped;
+
+	if (!iopt_area_pages_aligned(area))
+		return 0;
+
+	/* The range runs to the area end, batch_record_unmapped() stops it */
+	unmapped = iommu_unmap_sparse(domain, iova, last - iova + 1,
+				      batch_record_unmapped, &unmap);
+	if (IS_ENABLED(CONFIG_IOMMUFD_TEST))
+		WARN_ON(unmapped % PAGE_SIZE);
+	return unmapped >> PAGE_SHIFT;
+}
+
 /*
  * This is part of the VFIO compatibility support for VFIO_TYPE1_IOMMU. That
  * mode permits splitting a mapped area up, and then one of the splits is
@@ -1689,19 +1762,33 @@ iopt_area_unpin_domain(struct pfn_batch *batch, struct iopt_area *area,
 {
 	while (start_index <= last_index) {
 		unsigned long batch_last_index;
+		bool read_ahead = true;
 
 		if (*unmapped_end_index <= last_index) {
 			unsigned long start =
 				max(start_index, *unmapped_end_index);
+			unsigned long unmapped, batch_end;
 
 			if (IS_ENABLED(CONFIG_IOMMUFD_TEST) &&
 			    batch->total_pfns)
 				WARN_ON(*unmapped_end_index -
 						batch->total_pfns !=
 					start_index);
-			batch_from_domain(batch, domain, area, start,
-					  last_index);
-			batch_last_index = start_index + batch->total_pfns - 1;
+			unmapped = batch_unmap_from_domain(batch, domain, area,
+							   *unmapped_end_index,
+							   start_index,
+							   last_index,
+							   real_last_index);
+			if (unmapped) {
+				/* The walk ended on an entry boundary */
+				*unmapped_end_index += unmapped;
+				read_ahead = false;
+			} else {
+				batch_from_domain(batch, domain, area, start,
+						  last_index);
+			}
+			batch_end = start_index + batch->total_pfns - 1;
+			batch_last_index = min(batch_end, last_index);
 		} else {
 			batch_last_index = last_index;
 		}
@@ -1716,9 +1803,10 @@ iopt_area_unpin_domain(struct pfn_batch *batch, struct iopt_area *area,
 		 * middle of the domains we need to keep reading pfns until we
 		 * find a cut point to do the unmap. The pfns we read are
 		 * carried over and either skipped or integrated into the next
-		 * batch.
+		 * batch. batch_unmap_from_domain() already ends on an entry
+		 * boundary and needs no read-ahead.
 		 */
-		if (batch_last_index == last_index &&
+		if (read_ahead && batch_last_index == last_index &&
 		    last_index != real_last_index)
 			batch_from_domain_continue(batch, domain, area,
 						   last_index + 1,
-- 
2.47.3