[PATCH 2/2] mm/gup: batch contiguous PTE-mapped large folios in follow_page_mask()

[Rik van Riel <riel@surriel.com>]

follow_page_mask() returns one page per call for a PTE-mapped large folio,
so __get_user_pages() re-walks the page tables for every page of an mTHP
even though the folio maps a contiguous run. The huge PMD and PUD paths
already return the whole mapping in one step.

Report the contiguous run for the PTE case too. follow_pte_batch() uses
folio_pte_batch_flags() to count consecutive present PTEs that map
consecutive pages of the same folio with a uniform write bit, bounded by
the page table, @end, the VMA, and the folio itself.

Keep the per-PTE guarantees that follow_page_pte() makes for the head page.
folio_pte_batch_flags() with FPB_RESPECT_WRITE stops the run at a change in
the write bit, so the whole run matches the head.

A writable run is safe for any access: a writable anon page is exclusive,
so gup_must_unshare() cannot fire, and FOLL_WRITE is satisfied.

A read-only run is batched only for a plain read, since FOLL_WRITE would
need a COW fault per page and FOLL_PIN would need a per-page
gup_must_unshare() check.

Measured with mm/gup_test.c (PIN_LONGTERM_BENCHMARK, the slow
pin_user_pages() path) on a 256 MB MADV_HUGEPAGE anonymous region in a
4 CPU VM, median get time over 16 iterations. Each folio size was confirmed
through the per-size anon_fault_alloc counters (4096 folios for 64 kB, 128
for 2 MB):

  gup_test -L -m 256 -n 65536 -r 16 -t
                        before      after
  64 kB mTHP            3140 us      412 us   (7.6x)
  2 MB THP (control)      78 us       76 us
  4 kB base (control)   3010 us     3042 us

The PMD-mapped 2 MB THP already returns the whole mapping in one step, so
it stays fast and unchanged. The 4 kB baseline shows the per-page walk cost
that the 64 kB case paid before this change; only the PTE-mapped large
folio case improves.

Assisted-by: Claude:claude-opus-4.8
Signed-off-by: Rik van Riel <riel@surriel.com>
---
 mm/gup.c | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/mm/gup.c b/mm/gup.c
index 3437cd3407d5..8c6ad1ee7ccc 100644
--- a/mm/gup.c
+++ b/mm/gup.c
@@ -805,6 +805,43 @@ static inline bool can_follow_write_pte(pte_t pte, struct page *page,
 	return !userfaultfd_pte_wp(vma, pte);
 }
 
+/*
+ * Count the pages, starting at @address and bounded by @end, that a PTE-mapped
+ * large @folio maps contiguously and that can be returned together with the
+ * page at @address: consecutive present PTEs mapping consecutive pages of
+ * @folio with a uniform write bit, within this VMA and a single page table.
+ * Returns at least 1.
+ *
+ * gup_must_unshare() and the write-fault check are per PTE. A writable run is
+ * always safe: a writable anon page is exclusive, and FOLL_WRITE is satisfied.
+ * A read-only run is only safe for a plain read; FOLL_WRITE would need a COW
+ * fault per page and FOLL_PIN would need a per-page gup_must_unshare() check,
+ * so those fall back to a single page.
+ */
+static unsigned long follow_pte_batch(struct vm_area_struct *vma,
+		unsigned long address, unsigned long end, struct folio *folio,
+		struct page *page, pte_t *ptep, pte_t pte, unsigned int flags)
+{
+	pte_t batch_pte = pte;
+	unsigned long max;
+
+	if (!pte_write(pte) && (flags & (FOLL_WRITE | FOLL_PIN)))
+		return 1;
+
+	/*
+	 * folio_pte_batch_flags() scans forward from @ptep, so the run must
+	 * stay within this page table: bound it by the PMD as well as @end and
+	 * the VMA, since a large folio can be PTE-mapped across a PMD boundary.
+	 */
+	max = min((pmd_addr_end(address, end) - address) >> PAGE_SHIFT,
+		  (vma->vm_end - address) >> PAGE_SHIFT);
+	if (max <= 1)
+		return 1;
+
+	return folio_pte_batch_flags(folio, vma, ptep, &batch_pte, max,
+				     FPB_RESPECT_WRITE);
+}
+
 static struct page *follow_page_pte(struct vm_area_struct *vma,
 		unsigned long address, unsigned long end, pmd_t *pmd,
 		unsigned int flags, unsigned long *nr_pages)
@@ -893,6 +930,14 @@ static struct page *follow_page_pte(struct vm_area_struct *vma,
 		folio_mark_accessed(folio);
 	}
 
+	/*
+	 * A PTE-mapped large folio can be handed back as a contiguous batch,
+	 * so the caller advances over the whole run in one step instead of
+	 * walking the page tables for every page.
+	 */
+	if (folio_test_large(folio))
+		*nr_pages = follow_pte_batch(vma, address, end, folio, page,
+					     ptep, pte, flags);
 out:
 	pte_unmap_unlock(ptep, ptl);
 	return page;
-- 
2.53.0-Meta