Batch unmap of uffd-wp file folios

pte_install_uffd_wp_if_needed() has grown too large for mm_inline.h. Move it to memory.c.

While at it, convert the comment to kerneldoc and rename the local arguments from pte/pteval to ptep/pte so the pointer and pte value are easier to distinguish.

Signed-off-by Dev Jain <[email protected]>

---
 include/linux/mm.h        |  2 ++
 include/linux/mm_inline.h | 53 ----------------------------------
 mm/memory.c               | 60 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 53 deletions(-)

include/linux/mm.h +0 -0

index 550fb92957d18..a71341c44655e 100644

Quoted 22 lines · Dev Jain

pte_install_uffd_wp_if_needed() has grown too large for mm_inline.h. Move it to memory.c.

While at it, convert the comment to kerneldoc and rename the local arguments from pte/pteval to ptep/pte so the pointer and pte value are easier to distinguish.

Signed-off-by Dev Jain <[email protected]>

---
 include/linux/mm.h        |  2 ++
 include/linux/mm_inline.h | 53 ----------------------------------
 mm/memory.c               | 60 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 53 deletions(-)

include/linux/mm.h +0 -0

index 550fb92957d18..a71341c44655e 100644
@@ -5406,4 +5406,6 @@ void map_anon_folio_pte_nopf(struct folio *folio, pte_t *pte,
 		struct vm_area_struct *vma, unsigned long addr,
 		bool uffd_wp);
 

include/linux/mm.h void map_anon_folio_pte_nopf(struct folio *folio, pte_t *pte, · line 5410

Quoted 2 lines

include/linux/mm.h +2 -0

+bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
+		unsigned long addr, pte_t *ptep, pte_t pte);

This could probably go to mm/internal.h instead?

Apart from that LGTM.

Enable batch setting of uffd-wp ptes.

The code paths passing nr > 1 to zap_install_uffd_wp_if_needed() produce that nr through either folio_pte_batch or swap_pte_batch, therefore batching is correct:

1) all ptes belong to the same type of VMA (anonymous or non-anonymous, wp-armed or non-wp-armed)

2) all ptes being marked with uffd-wp or all being not marked (same is the case with the pte_swp_uffd_wp_any check)

3) uffd_supports_wp_marker() is independent of the function parameters

Note that we will have to use set_pte_at() in a loop instead of set_ptes() since the latter cannot handle present->non-present conversion for nr_pages > 1.

Rename the function to cond_install_uffd_wp_ptes.

Signed-off-by Dev Jain <[email protected]>

---
To handle nonpresent->nonpresent transition in the ptes, we can have a set_nonpresent_ptes() (in my unmap series) : if !softleaf_has_pfn(), use set the same pte value to all ptep's in the patch. if softleaf_has_pfn(), then add a softleaf_next_pfn() to construct the next softleaf, and pte_next_softleaf() to call softleaf_next_pfn() and preserve the wp bit, s-d bit, etc from the previous pte.

 include/linux/mm.h |  6 +++--
 mm/memory.c        | 64 +++++++++++++++++-----------------------------
 mm/rmap.c          |  2 +-
 3 files changed, 29 insertions(+), 43 deletions(-)

include/linux/mm.h +0 -0

index a71341c44655e..94e0a92bc70b1 100644

include/linux/mm.h +1 -0

+
 #endif /* _LINUX_MM_H */

mm/memory.c +18 -16

index 98b3ace15cef2..5d2b567b383d4 100644
@@ -1676,27 +1676,29 @@ static inline bool zap_drop_markers(struct zap_details *details)
 }
 
 /**
- * pte_install_uffd_wp_if_needed - install uffd-wp marker after clearing a PTE
- * @vma: The VMA the page is mapped into.
- * @addr: Address the page is mapped at.
- * @ptep: Page table pointer for this entry.
+ * cond_install_uffd_wp_ptes - install uffd-wp markers after clearing PTEs
+ * @vma: The VMA the pages are mapped into.
+ * @addr: Address the first page of this batch is mapped at.
+ * @ptep: Page table pointer for the first entry of this batch.
  * @pte: Old value of the entry pointed to by @ptep.
+ * @nr_ptes: Number of entries to install.
  *
- * If the PTE was write-protected by uffd-wp in any form, arm a special PTE
- * to replace a none PTE. NOTE! This should only be called when the PTE is
- * already cleared so we will never accidentally replace something valuable.
- * Meanwhile none PTEs also mean we are not demoting the PTE so a TLB flush is
- * not needed. E.g., when the PTE was cleared, the caller should have taken care
- * of the TLB flush.
+ * If the PTEs were write-protected by uffd-wp in any form, arm special
+ * PTEs to replace none PTEs. NOTE! This should only be called when the PTEs
+ * are already cleared so we will never accidentally replace something
+ * valuable. Meanwhile none PTEs also mean we are not demoting the PTEs so a
+ * TLB flush is not needed. E.g., when PTEs were cleared, the caller should
+ * have taken care of the TLB flush.
  *
- * Must be called with the page table lock held so that no thread will see the
- * none PTE, and if they see it, they'll fault and serialize at the page table
- * lock.
+ * Must be called with the page table lock held so that no thread will see
+ * the none PTEs, and if they see them, they'll fault and serialize at the
+ * page table lock.
  *
- * Returns true if an uffd-wp PTE was installed, false otherwise.
+ * Returns true if uffd-wp PTEs were installed, false otherwise.
  */
-bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
-				   unsigned long addr, pte_t *ptep, pte_t pte)
+bool cond_install_uffd_wp_ptes(struct vm_area_struct *vma,
+		unsigned long addr, pte_t *ptep, pte_t pte,
+		unsigned long nr_ptes)
 {
 	bool arm_uffd_pte = false;
 

@@ -1726,13 +1728,14 @@ bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
 	if (unlikely(pte_swp_uffd_any(pte)))
 		arm_uffd_pte = true;

mm/memory.c +2 -21

-	return false;
+	return true;
 }
 
 /*

@@ -1746,29 +1749,10 @@ zap_install_uffd_wp_if_needed(struct vm_area_struct *vma,
 			      unsigned long addr, pte_t *pte, int nr,
 			      struct zap_details *details, pte_t pteval)
 {
-	bool was_installed = false;
-
-	if (!uffd_supports_wp_marker())
-		return false;
-
-	/* Zap on anonymous always means dropping everything */
-	if (vma_is_anonymous(vma))
-		return false;
-
 	if (zap_drop_markers(details))
 		return false;
 
-	for (;;) {
-		/* the PFN in the PTE is irrelevant. */
-		if (pte_install_uffd_wp_if_needed(vma, addr, pte, pteval))
-			was_installed = true;
-		if (--nr == 0)
-			break;
-		pte++;
-		addr += PAGE_SIZE;
-	}
-
-	return was_installed;
+	return cond_install_uffd_wp_ptes(vma, addr, pte, pteval, nr);
 }
 
 static __always_inline void zap_present_folio_ptes(struct mmu_gather *tlb,

mm/rmap.c +0 -0

index ad820fe86f7d8..2f938d0ac6953 100644

Quoted 42 lines · Dev Jain

Enable batch setting of uffd-wp ptes.

The code paths passing nr > 1 to zap_install_uffd_wp_if_needed() produce that nr through either folio_pte_batch or swap_pte_batch, therefore batching is correct:

1) all ptes belong to the same type of VMA (anonymous or non-anonymous, wp-armed or non-wp-armed)

2) all ptes being marked with uffd-wp or all being not marked (same is the case with the pte_swp_uffd_wp_any check)

3) uffd_supports_wp_marker() is independent of the function parameters

Note that we will have to use set_pte_at() in a loop instead of set_ptes() since the latter cannot handle present->non-present conversion for nr_pages > 1.

Rename the function to cond_install_uffd_wp_ptes.

Signed-off-by Dev Jain <[email protected]>

---
To handle nonpresent->nonpresent transition in the ptes, we can have a set_nonpresent_ptes() (in my unmap series) : if !softleaf_has_pfn(), use set the same pte value to all ptep's in the patch. if softleaf_has_pfn(), then add a softleaf_next_pfn() to construct the next softleaf, and pte_next_softleaf() to call softleaf_next_pfn() and preserve the wp bit, s-d bit, etc from the previous pte.

 include/linux/mm.h |  6 +++--
 mm/memory.c        | 64 +++++++++++++++++-----------------------------
 mm/rmap.c          |  2 +-
 3 files changed, 29 insertions(+), 43 deletions(-)

include/linux/mm.h +0 -0

index a71341c44655e..94e0a92bc70b1 100644
@@ -5406,6 +5406,8 @@ void map_anon_folio_pte_nopf(struct folio *folio, pte_t *pte,
 		struct vm_area_struct *vma, unsigned long addr,
 		bool uffd_wp);
 

include/linux/mm.h void map_anon_folio_pte_nopf(struct folio *folio, pte_t *pte, · line 5411

Quoted 5 lines

include/linux/mm.h +3 -2

-bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
-		unsigned long addr, pte_t *ptep, pte_t pte);
+bool cond_install_uffd_wp_ptes(struct vm_area_struct *vma, unsigned long addr,
+			       pte_t *ptep, pte_t pte,
+			       unsigned long nr_ptes);

Two tab ...

Quoted 57 lines

include/linux/mm.h +1 -0

+
 #endif /* _LINUX_MM_H */

mm/memory.c +18 -16

index 98b3ace15cef2..5d2b567b383d4 100644
@@ -1676,27 +1676,29 @@ static inline bool zap_drop_markers(struct zap_details *details)
 }
 
 /**
- * pte_install_uffd_wp_if_needed - install uffd-wp marker after clearing a PTE
- * @vma: The VMA the page is mapped into.
- * @addr: Address the page is mapped at.
- * @ptep: Page table pointer for this entry.
+ * cond_install_uffd_wp_ptes - install uffd-wp markers after clearing PTEs
+ * @vma: The VMA the pages are mapped into.
+ * @addr: Address the first page of this batch is mapped at.
+ * @ptep: Page table pointer for the first entry of this batch.
  * @pte: Old value of the entry pointed to by @ptep.
+ * @nr_ptes: Number of entries to install.
  *
- * If the PTE was write-protected by uffd-wp in any form, arm a special PTE
- * to replace a none PTE. NOTE! This should only be called when the PTE is
- * already cleared so we will never accidentally replace something valuable.
- * Meanwhile none PTEs also mean we are not demoting the PTE so a TLB flush is
- * not needed. E.g., when the PTE was cleared, the caller should have taken care
- * of the TLB flush.
+ * If the PTEs were write-protected by uffd-wp in any form, arm special
+ * PTEs to replace none PTEs. NOTE! This should only be called when the PTEs
+ * are already cleared so we will never accidentally replace something
+ * valuable. Meanwhile none PTEs also mean we are not demoting the PTEs so a
+ * TLB flush is not needed. E.g., when PTEs were cleared, the caller should
+ * have taken care of the TLB flush.
  *
- * Must be called with the page table lock held so that no thread will see the
- * none PTE, and if they see it, they'll fault and serialize at the page table
- * lock.
+ * Must be called with the page table lock held so that no thread will see
+ * the none PTEs, and if they see them, they'll fault and serialize at the
+ * page table lock.
  *
- * Returns true if an uffd-wp PTE was installed, false otherwise.
+ * Returns true if uffd-wp PTEs were installed, false otherwise.
  */
-bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
-				   unsigned long addr, pte_t *ptep, pte_t pte)
+bool cond_install_uffd_wp_ptes(struct vm_area_struct *vma,
+		unsigned long addr, pte_t *ptep, pte_t pte,
+		unsigned long nr_ptes)
 {
 	bool arm_uffd_pte = false;
 

@@ -1726,13 +1728,14 @@ bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma,
 	if (unlikely(pte_swp_uffd_any(pte)))
 		arm_uffd_pte = true;
 

mm/memory.c bool pte_install_uffd_wp_if_needed(struct vm_area_struct *vma, · line 1733

Quoted 9 lines

mm/memory.c +4 -3

-	if (unlikely(arm_uffd_pte)) {
+	if (likely(!arm_uffd_pte))
+		return false;
+
+	for (unsigned long i = 0; i < nr_ptes; ++i, ++ptep, addr += PAGE_SIZE)
 		set_pte_at(vma->vm_mm, addr, ptep,
 			   make_pte_marker(PTE_MARKER_UFFD_WP));
-		return true;
-	}

Can we keep the "for (;;)" style of iterating PTEs that we use elsewhere?

Quoted 43 lines

mm/memory.c +2 -21

-	return false;
+	return true;
 }
 
 /*

@@ -1746,29 +1749,10 @@ zap_install_uffd_wp_if_needed(struct vm_area_struct *vma,
 			      unsigned long addr, pte_t *pte, int nr,
 			      struct zap_details *details, pte_t pteval)
 {
-	bool was_installed = false;
-
-	if (!uffd_supports_wp_marker())
-		return false;
-
-	/* Zap on anonymous always means dropping everything */
-	if (vma_is_anonymous(vma))
-		return false;
-
 	if (zap_drop_markers(details))
 		return false;
 
-	for (;;) {
-		/* the PFN in the PTE is irrelevant. */
-		if (pte_install_uffd_wp_if_needed(vma, addr, pte, pteval))
-			was_installed = true;
-		if (--nr == 0)
-			break;
-		pte++;
-		addr += PAGE_SIZE;
-	}
-
-	return was_installed;
+	return cond_install_uffd_wp_ptes(vma, addr, pte, pteval, nr);
 }
 
 static __always_inline void zap_present_folio_ptes(struct mmu_gather *tlb,

mm/rmap.c +0 -0

index ad820fe86f7d8..2f938d0ac6953 100644

mm/rmap.c static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma, · line 2348

Quoted 5 lines

mm/rmap.c +1 -1

@@ -2345,7 +2345,7 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		 * we may want to replace a none pte with a marker pte if
 		 * it's file-backed, so we don't lose the tracking info.
 		 */
-		pte_install_uffd_wp_if_needed(vma, address, pvmw.pte, pteval);
+		cond_install_uffd_wp_ptes(vma, address, pvmw.pte, pteval, 1);

Was about to ask whether we should provide a wrapper, but the next patch converts this case.

Commit a67fe41e214f ("mm: rmap: support batched unmapping for file large folios") extended batched unmapping for file folios. That also required making install_uffd_wp_pte_if_needed() support batching, but that was left out for the time being, and correctness was maintained by stopping batching in case the VMA the folio belongs to is marked uffd-wp.

Now that we have a batched version called cond_install_uffd_wp_ptes, simply call that. folio_unmap_pte_batch() ensures that the original state of the ptes is either all uffd or all non-uffd, so we maintain correctness.

If uffd-wp bit is there, we have the following transitions of ptes after unmapping, for a file folio:

present -> uffd-wp marker

We must ensure that these ptes are not reprocessed by the while loop - if the batch length is less than the number of pages in the folio, then we must skip over this batch.

The page_vma_mapped_walk API ensures this - check_pte() will return true only if any of [pvmw->pfn, pvmw->pfn + nr_pages) is mapped by the pte. There is no pfn underlying a uffd-wp marker pte, so check_pte returns false and we keep skipping until we hit a present entry, which is where we want to batch from next.

Note that for the uffd-rwp case, we already do not make a uffd marker in cond_install_uffd_wp_ptes.

Signed-off-by Dev Jain <[email protected]>

---
Dropped David's ACK due to the uffd-rwp changes.

 mm/rmap.c | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

mm/rmap.c +0 -3

index 2f938d0ac6953..ac88673600fa0 100644
@@ -1965,9 +1965,6 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 	if (pte_unused(pte))
 		return 1;
 
-	if (userfaultfd_protected(vma))
-		return 1;
-
 	/*
 	 * If unmap fails, we need to restore the ptes. To avoid accidentally
 	 * upgrading write permissions for ptes that were not originally

Currently, batched unmapping is supported if:

1) folio is a file folio, not belonging to uffd-wp VMA
2) folio is anonymous and not swapbacked (lazyfree), not belonging to uffd-wp VMA

So the cases which are not supported are

1) folio belonging to uffd-wp VMA
2) folio is anonymous and swapbacked

It is easy to see that this adds a lot of cognitive load while reading try_to_unmap_one - we need to remember throughout whether nr_pages == 1 or > 1.

The uffd-wp handling in try_to_unmap_one is regarding preserving the uffd-wp state for file folios via pte_install_uffd_wp_if_needed (for anon folio, we handle that while constructing the swap pte).

Stop special casing on uffd-wp VMAs by simply adding batching support to pte_install_uffd_wp_if_needed.

Quoted 19 lines · On Wed, 15 Jul 2026 11:18:33 +0000 Dev Jain <[email protected]> wrote:

Currently, batched unmapping is supported if:

1) folio is a file folio, not belonging to uffd-wp VMA
2) folio is anonymous and not swapbacked (lazyfree), not belonging to uffd-wp VMA

So the cases which are not supported are

1) folio belonging to uffd-wp VMA
2) folio is anonymous and swapbacked

It is easy to see that this adds a lot of cognitive load while reading try_to_unmap_one - we need to remember throughout whether nr_pages == 1 or > 1.

The uffd-wp handling in try_to_unmap_one is regarding preserving the uffd-wp state for file folios via pte_install_uffd_wp_if_needed (for anon folio, we handle that while constructing the swap pte).

Quoted 2 lines

Stop special casing on uffd-wp VMAs by simply adding batching support to pte_install_uffd_wp_if_needed.

Thanks, I'll await reviewer input on this.

Sashiko might have found a pre-existing issue:

	https://sashiko.dev/#/patchset/[email protected]