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

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

[Rik van Riel]

follow_page_mask() reports the size of the page it found with a page_mask:
a bitmask of the enclosing naturally-aligned huge page. That form can only
describe a power-of-two block aligned to its own size, so a PTE-mapped
large folio (mTHP) is walked one page at a time even though it maps a
contiguous run, while the PMD/PUD-mapped cases already return their whole
mapping in one step.

Patch 1 replaces the page_mask output with nr_pages (a plain count) and
adds an @end argument bounding the walk, with no functional change: the
huge PMD/PUD paths report the same stride as before, just as a count.

Patch 2 uses the new interface to batch the PTE case: follow_pte_batch()
counts 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, then follow_page_pte() hands the whole run back in
one call.

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.

Both patches apply standalone against mm-unstable; no other series is
required.

Rik van Riel (2):
  mm/gup: pass an end address to follow_page_mask() and return a page
    count
  mm/gup: batch contiguous PTE-mapped large folios in follow_page_mask()

 mm/gup.c | 139 ++++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 97 insertions(+), 42 deletions(-)

base-commit: 7368ccdeff50c27809d3a8276c85bae7e402c96c
---
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: linux-mm@kvack.org
To: linux-kernel@vger.kernel.org
Cc: Dave Hansen <dave.hansen@linux.intel.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Suren Baghdasaryan <surenb@google.com>
Cc: Lorenzo Stoakes <ljs@kernel.org>
Cc: Vlastimil Babka <vbabka@kernel.org>
Cc: David Hildenbrand <david@kernel.org>
Cc: "Liam R. Howlett" <liam@infradead.org>
Cc: Mike Rapoport <rppt@kernel.org>
Cc: Michal Hocko <mhocko@suse.com>
Cc: Jason Gunthorpe <jgg@ziepe.ca>
Cc: John Hubbard <jhubbard@nvidia.com>
Cc: Peter Xu <peterx@redhat.com>
Cc: Matthew Wilcox <willy@infradead.org>
Cc: Usama Arif <usamaarif642@gmail.com>
Cc: Rik van Riel <riel@surriel.com>
Cc: kernel-team@meta.com
-- 
2.53.0-Meta

[PATCH 1/2] mm/gup: pass an end address to follow_page_mask() and return a page count

[Rik van Riel]

follow_page_mask() reports how many pages the returned page covers with a
page_mask: a bitmask of the enclosing naturally-aligned huge page. Callers
turn that into a stride, which assumes the run is a power-of-two block
aligned to its own size.

That form cannot describe an arbitrary contiguous run, which a later patch
needs to batch PTE-mapped large folios.

Replace the page_mask output with nr_pages, the number of contiguous pages
the returned page covers starting at @address, and add an @end argument
bounding how far the walk looks, so a small request does not scan an entire
large folio.

The huge PMD and PUD paths report the same stride as before, now as a count
clamped to @end. __get_user_pages() uses the count directly.

No functional change intended.

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

diff --git a/mm/gup.c b/mm/gup.c
index 99902c15703b..3437cd3407d5 100644
--- a/mm/gup.c
+++ b/mm/gup.c
@@ -647,8 +647,9 @@ static inline bool can_follow_write_pud(pud_t pud, struct page *page,
 }
 
 static struct page *follow_huge_pud(struct vm_area_struct *vma,
-				    unsigned long addr, pud_t *pudp,
-				    int flags, unsigned long *page_mask)
+				    unsigned long addr, unsigned long end,
+				    pud_t *pudp, int flags,
+				    unsigned long *nr_pages)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	struct page *page;
@@ -672,10 +673,13 @@ static struct page *follow_huge_pud(struct vm_area_struct *vma,
 		return ERR_PTR(-EMLINK);
 
 	ret = try_grab_folio(page_folio(page), 1, flags);
-	if (ret)
+	if (ret) {
 		page = ERR_PTR(ret);
-	else
-		*page_mask = HPAGE_PUD_NR - 1;
+	} else {
+		unsigned long off = (addr & ~PUD_MASK) >> PAGE_SHIFT;
+
+		*nr_pages = min(HPAGE_PUD_NR - off, (end - addr) >> PAGE_SHIFT);
+	}
 
 	return page;
 }
@@ -699,9 +703,9 @@ static inline bool can_follow_write_pmd(pmd_t pmd, struct page *page,
 }
 
 static struct page *follow_huge_pmd(struct vm_area_struct *vma,
-				    unsigned long addr, pmd_t *pmd,
-				    unsigned int flags,
-				    unsigned long *page_mask)
+				    unsigned long addr, unsigned long end,
+				    pmd_t *pmd, unsigned int flags,
+				    unsigned long *nr_pages)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	pmd_t pmdval = *pmd;
@@ -738,23 +742,25 @@ static struct page *follow_huge_pmd(struct vm_area_struct *vma,
 #endif	/* CONFIG_TRANSPARENT_HUGEPAGE */
 
 	page += (addr & ~HPAGE_PMD_MASK) >> PAGE_SHIFT;
-	*page_mask = HPAGE_PMD_NR - 1;
+	*nr_pages = min(HPAGE_PMD_NR - ((addr & ~HPAGE_PMD_MASK) >> PAGE_SHIFT),
+			(end - addr) >> PAGE_SHIFT);
 
 	return page;
 }
 
 #else  /* CONFIG_PGTABLE_HAS_HUGE_LEAVES */
 static struct page *follow_huge_pud(struct vm_area_struct *vma,
-				    unsigned long addr, pud_t *pudp,
-				    int flags, unsigned long *page_mask)
+				    unsigned long addr, unsigned long end,
+				    pud_t *pudp, int flags,
+				    unsigned long *nr_pages)
 {
 	return NULL;
 }
 
 static struct page *follow_huge_pmd(struct vm_area_struct *vma,
-				    unsigned long addr, pmd_t *pmd,
-				    unsigned int flags,
-				    unsigned long *page_mask)
+				    unsigned long addr, unsigned long end,
+				    pmd_t *pmd, unsigned int flags,
+				    unsigned long *nr_pages)
 {
 	return NULL;
 }
@@ -800,7 +806,8 @@ static inline bool can_follow_write_pte(pte_t pte, struct page *page,
 }
 
 static struct page *follow_page_pte(struct vm_area_struct *vma,
-		unsigned long address, pmd_t *pmd, unsigned int flags)
+		unsigned long address, unsigned long end, pmd_t *pmd,
+		unsigned int flags, unsigned long *nr_pages)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	struct folio *folio;
@@ -885,6 +892,7 @@ static struct page *follow_page_pte(struct vm_area_struct *vma,
 		 */
 		folio_mark_accessed(folio);
 	}
+
 out:
 	pte_unmap_unlock(ptep, ptl);
 	return page;
@@ -896,9 +904,9 @@ static struct page *follow_page_pte(struct vm_area_struct *vma,
 }
 
 static struct page *follow_pmd_mask(struct vm_area_struct *vma,
-				    unsigned long address, pud_t *pudp,
-				    unsigned int flags,
-				    unsigned long *page_mask)
+				    unsigned long address, unsigned long end,
+				    pud_t *pudp, unsigned int flags,
+				    unsigned long *nr_pages)
 {
 	pmd_t *pmd, pmdval;
 	spinlock_t *ptl;
@@ -912,7 +920,7 @@ static struct page *follow_pmd_mask(struct vm_area_struct *vma,
 	if (!pmd_present(pmdval))
 		return no_page_table(vma, flags, address);
 	if (likely(!pmd_leaf(pmdval)))
-		return follow_page_pte(vma, address, pmd, flags);
+		return follow_page_pte(vma, address, end, pmd, flags, nr_pages);
 
 	if (pmd_protnone(pmdval) && !gup_can_follow_protnone(vma, flags))
 		return no_page_table(vma, flags, address);
@@ -925,24 +933,24 @@ static struct page *follow_pmd_mask(struct vm_area_struct *vma,
 	}
 	if (unlikely(!pmd_leaf(pmdval))) {
 		spin_unlock(ptl);
-		return follow_page_pte(vma, address, pmd, flags);
+		return follow_page_pte(vma, address, end, pmd, flags, nr_pages);
 	}
 	if (pmd_trans_huge(pmdval) && (flags & FOLL_SPLIT_PMD)) {
 		spin_unlock(ptl);
 		split_huge_pmd(vma, pmd, address);
 		/* If pmd was left empty, stuff a page table in there quickly */
 		return pte_alloc(mm, pmd) ? ERR_PTR(-ENOMEM) :
-			follow_page_pte(vma, address, pmd, flags);
+			follow_page_pte(vma, address, end, pmd, flags, nr_pages);
 	}
-	page = follow_huge_pmd(vma, address, pmd, flags, page_mask);
+	page = follow_huge_pmd(vma, address, end, pmd, flags, nr_pages);
 	spin_unlock(ptl);
 	return page;
 }
 
 static struct page *follow_pud_mask(struct vm_area_struct *vma,
-				    unsigned long address, p4d_t *p4dp,
-				    unsigned int flags,
-				    unsigned long *page_mask)
+				    unsigned long address, unsigned long end,
+				    p4d_t *p4dp, unsigned int flags,
+				    unsigned long *nr_pages)
 {
 	pud_t *pudp, pud;
 	spinlock_t *ptl;
@@ -955,7 +963,7 @@ static struct page *follow_pud_mask(struct vm_area_struct *vma,
 		return no_page_table(vma, flags, address);
 	if (pud_leaf(pud)) {
 		ptl = pud_lock(mm, pudp);
-		page = follow_huge_pud(vma, address, pudp, flags, page_mask);
+		page = follow_huge_pud(vma, address, end, pudp, flags, nr_pages);
 		spin_unlock(ptl);
 		if (page)
 			return page;
@@ -964,13 +972,13 @@ static struct page *follow_pud_mask(struct vm_area_struct *vma,
 	if (unlikely(pud_bad(pud)))
 		return no_page_table(vma, flags, address);
 
-	return follow_pmd_mask(vma, address, pudp, flags, page_mask);
+	return follow_pmd_mask(vma, address, end, pudp, flags, nr_pages);
 }
 
 static struct page *follow_p4d_mask(struct vm_area_struct *vma,
-				    unsigned long address, pgd_t *pgdp,
-				    unsigned int flags,
-				    unsigned long *page_mask)
+				    unsigned long address, unsigned long end,
+				    pgd_t *pgdp, unsigned int flags,
+				    unsigned long *nr_pages)
 {
 	p4d_t *p4dp, p4d;
 
@@ -981,15 +989,16 @@ static struct page *follow_p4d_mask(struct vm_area_struct *vma,
 	if (!p4d_present(p4d) || p4d_bad(p4d))
 		return no_page_table(vma, flags, address);
 
-	return follow_pud_mask(vma, address, p4dp, flags, page_mask);
+	return follow_pud_mask(vma, address, end, p4dp, flags, nr_pages);
 }
 
 /**
  * follow_page_mask - look up a page descriptor from a user-virtual address
  * @vma: vm_area_struct mapping @address
  * @address: virtual address to look up
+ * @end: virtual address at which to stop batching contiguous pages
  * @flags: flags modifying lookup behaviour
- * @page_mask: a pointer to output page_mask
+ * @nr_pages: output; number of contiguous pages the caller can read
  *
  * @flags can have FOLL_ flags set, defined in <linux/mm.h>
  *
@@ -998,15 +1007,17 @@ static struct page *follow_p4d_mask(struct vm_area_struct *vma,
  * trigger a fault with FAULT_FLAG_UNSHARE set. Note that unsharing is only
  * relevant with FOLL_PIN and !FOLL_WRITE.
  *
- * On output, @page_mask is set according to the size of the page.
+ * On output, @nr_pages holds how many contiguous pages the folio that includes
+ * the returned page has mapped into this process, so the caller can advance
+ * over a large folio in one step.
  *
  * Return: the mapped (struct page *), %NULL if no mapping exists, or
  * an error pointer if there is a mapping to something not represented
  * by a page descriptor (see also vm_normal_page()).
  */
 static struct page *follow_page_mask(struct vm_area_struct *vma,
-			      unsigned long address, unsigned int flags,
-			      unsigned long *page_mask)
+			      unsigned long address, unsigned long end,
+			      unsigned int flags, unsigned long *nr_pages)
 {
 	pgd_t *pgd;
 	struct mm_struct *mm = vma->vm_mm;
@@ -1014,13 +1025,13 @@ static struct page *follow_page_mask(struct vm_area_struct *vma,
 
 	vma_pgtable_walk_begin(vma);
 
-	*page_mask = 0;
+	*nr_pages = 1;
 	pgd = pgd_offset(mm, address);
 
 	if (pgd_none(*pgd) || unlikely(pgd_bad(*pgd)))
 		page = no_page_table(vma, flags, address);
 	else
-		page = follow_p4d_mask(vma, address, pgd, flags, page_mask);
+		page = follow_p4d_mask(vma, address, end, pgd, flags, nr_pages);
 
 	vma_pgtable_walk_end(vma);
 
@@ -1358,7 +1369,6 @@ static long __get_user_pages(struct mm_struct *mm,
 {
 	long ret = 0, i = 0;
 	struct vm_area_struct *vma = NULL;
-	unsigned long page_mask = 0;
 
 	if (!nr_pages)
 		return 0;
@@ -1373,7 +1383,7 @@ static long __get_user_pages(struct mm_struct *mm,
 
 	do {
 		struct page *page;
-		unsigned int page_increm;
+		unsigned long page_increm;
 
 		/* first iteration or cross vma bound */
 		if (!vma || start >= vma->vm_end) {
@@ -1400,7 +1410,7 @@ static long __get_user_pages(struct mm_struct *mm,
 						pages ? &page : NULL);
 				if (ret)
 					goto out;
-				page_mask = 0;
+				page_increm = 1;
 				goto next_page;
 			}
 
@@ -1423,7 +1433,8 @@ static long __get_user_pages(struct mm_struct *mm,
 		}
 		cond_resched();
 
-		page = follow_page_mask(vma, start, gup_flags, &page_mask);
+		page = follow_page_mask(vma, start, start + nr_pages * PAGE_SIZE,
+					gup_flags, &page_increm);
 		if (!page || PTR_ERR(page) == -EMLINK) {
 			ret = faultin_page(vma, start, gup_flags,
 					   PTR_ERR(page) == -EMLINK, locked);
@@ -1456,7 +1467,6 @@ static long __get_user_pages(struct mm_struct *mm,
 			goto out;
 		}
 next_page:
-		page_increm = 1 + (~(start >> PAGE_SHIFT) & page_mask);
 		if (page_increm > nr_pages)
 			page_increm = nr_pages;
 
-- 
2.53.0-Meta

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

[Rik van Riel]

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