[RFC PATCH v3 8/8] mm/gup: batch contiguous same-folio PTEs into one refcount grab

[Rik van Riel]

follow_page_pte() now walks a whole page table in one call, but still
resolves and commits each PTE on its own, so a PTE-mapped large folio
(mTHP) pays one try_grab_folio() per subpage.

Add follow_pte_batch() and give follow_page_pte_commit() a run length.
Each run starts with a full per-PTE resolve, then follow_pte_batch()
finds how many more PTEs extend it with a pte_same() scan rather than a
re-derivation of each page, and the run is committed with one refcount
grab. Runs stop at folio boundaries, so plain base pages are unaffected.

Gather the dirty bits from all the PTEs in a batch, in order to mark
the folio dirty for FOLL_WRITE.

Same benchmark as the previous change, with before again taken at the
base of the series:

  gup_test -L -m 256 -n 65536 -r 16 -t
                        before      after
  64 kB mTHP           2946 us      231 us  (12.8x)
  4 kB base (control)  2753 us     1198 us   (2.3x)
  2 MB THP (control)     72 us       71 us

The 4 kB column is the previous change's 2.3x, unchanged: base pages
share no folio, so there is nothing to batch. 64 kB mTHP goes from that
change's 2.2x to 12.8x, a further 5.9x from committing a run with one
refcount grab instead of one per subpage.

Suggested-by: David Hildenbrand <david@xxxxxxxxxx>
Assisted-by: Claude:claude-opus-4-8
Signed-off-by: Rik van Riel <riel@xxxxxxxxxxx>
---
 mm/gup.c | 58 ++++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 50 insertions(+), 8 deletions(-)

diff --git a/mm/gup.c b/mm/gup.c
index c4233a7b8a48..106806634b3c 100644
--- a/mm/gup.c
+++ b/mm/gup.c
@@ -833,12 +833,13 @@ static inline bool can_follow_write_pte(pte_t pte, struct page *page,
  */
 static long follow_page_pte_commit(struct vm_area_struct *vma,
 		unsigned long address, struct folio *folio, struct page *page,
-		pte_t pte, unsigned int flags, struct page **pages)
+		pte_t pte, unsigned long nr, unsigned int flags,
+		struct page **pages)
 {
 	long ret;
 
 	/* try_grab_folio() does nothing unless FOLL_GET or FOLL_PIN is set. */
-	ret = try_grab_folio(folio, 1, flags);
+	ret = try_grab_folio(folio, nr, flags);
 	if (unlikely(ret))
 		return ret;
 
@@ -850,7 +851,7 @@ static long follow_page_pte_commit(struct vm_area_struct *vma,
 	if (flags & FOLL_PIN) {
 		ret = arch_make_folio_accessible(folio);
 		if (ret) {
-			gup_put_folio(folio, 1, flags);
+			gup_put_folio(folio, nr, flags);
 			return ret;
 		}
 	}
@@ -866,7 +867,7 @@ static long follow_page_pte_commit(struct vm_area_struct *vma,
 		folio_mark_accessed(folio);
 	}
 
-	gup_fill_pages(vma, address, page, 1, pages);
+	gup_fill_pages(vma, address, page, nr, pages);
 
 	return 0;
 }
@@ -916,6 +917,35 @@ static long follow_one_pte(struct vm_area_struct *vma, unsigned long address,
 	return 0;
 }
 
+/*
+ * Return how many PTEs map consecutive pages of the same folio and can be
+ * committed as one run. Always at least 1.
+ *
+ * The write-fault and unshare checks in follow_one_pte() are per PTE, but a
+ * writable run needs no repeat: a writable anon page is exclusive. A read-only
+ * run under FOLL_WRITE or FOLL_PIN does need the per-page check, so it stays
+ * one page at a time.
+ */
+static unsigned long follow_pte_batch(struct vm_area_struct *vma,
+		unsigned long address, unsigned long walk_end,
+		struct folio *folio, pte_t *ptep, pte_t *batch_pte, unsigned int flags)
+{
+	unsigned long max;
+
+	if (!folio_test_large(folio))
+		return 1;
+	if (!pte_write(*batch_pte) && (flags & (FOLL_WRITE | FOLL_PIN)))
+		return 1;
+
+	max = (walk_end - address) >> PAGE_SHIFT;
+	if (max <= 1)
+		return 1;
+
+	/* Merge young/dirty across batch so folio_mark_dirty sees any dirty. */
+	return folio_pte_batch_flags(folio, vma, ptep, batch_pte, max,
+				     FPB_RESPECT_WRITE | FPB_MERGE_YOUNG_DIRTY);
+}
+
 /*
  * Walk the PTEs from the start address to the end of this page table or VMA,
  * whichever comes first, and commit every page found.
@@ -947,12 +977,24 @@ static long follow_page_pte(struct vm_area_struct *vma,
 
 		ret = follow_one_pte(vma, address, ptep, pte, flags, &page);
 		if (!ret && page) {
-			ret = follow_page_pte_commit(vma, address,
-						     page_folio(page), page,
-						     pte, flags,
+			struct folio *folio = page_folio(page);
+			unsigned long batch;
+
+			pte_t batch_pte = pte;
+
+			batch = follow_pte_batch(vma, address, walk_end, folio,
+						 ptep, &batch_pte, flags);
+			ret = follow_page_pte_commit(vma, address, folio, page,
+						     batch_pte, batch, flags,
 						     pages ? pages + nr : NULL);
 			if (!ret) {
-				nr++;
+				nr += batch;
+				/*
+				 * The loop's own increment covers one PTE; skip
+				 * the rest of the batch.
+				 */
+				ptep += batch - 1;
+				address += (batch - 1) * PAGE_SIZE;
 				continue;
 			}
 		}
-- 
2.55.0