mm/rmap: Refactor try_to_unmap_one

try_to_unmap() receives hugetlb folios only from the hwpoison path. hugetlb_update_hwpoison() sets the hugetlb folio's head-page hwpoison bit, and page_vma_mapped_walk() reports the hugetlb mapping at the head PFN, so the previous PageHWPoison(subpage) check happened to work for hugetlb.

For non-hugetlb folios, unmap_poisoned_folio() currently rejects large folios before calling try_to_unmap(). Hence it is always the case that if try_to_unmap_one() handles an hwpoisoned folio, then the head page is marked with the poison bit.

Therefore, convert the poisoned subpage checks to folio_test_hwpoison().

No functional change intended, except that, while at it, convert VM_BUG_* to VM_WARN_*.

Acked-by David Hildenbrand (Arm) <[email protected]>

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

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

mm/rmap.c +4 -3

index 9939400e77c79..dbb077f8443e3 100644
@@ -2122,10 +2122,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			bool anon = folio_test_anon(folio);
 
 			/*
-			 * The try_to_unmap() is only passed a hugetlb page
-			 * in the case where the hugetlb page is poisoned.
+			 * The try_to_unmap() is only passed a hugetlb folio
+			 * in the case where the hugetlb folio contains a
+			 * poisoned page.
 			 */
-			VM_BUG_ON_PAGE(!PageHWPoison(subpage), subpage);
+			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
 			/*
 			 * huge_pmd_unshare may unmap an entire PMD page.
 			 * There is no way of knowing exactly which PMDs may

Quoted 45 lines · Dev Jain

try_to_unmap() receives hugetlb folios only from the hwpoison path. hugetlb_update_hwpoison() sets the hugetlb folio's head-page hwpoison bit, and page_vma_mapped_walk() reports the hugetlb mapping at the head PFN, so the previous PageHWPoison(subpage) check happened to work for hugetlb.

For non-hugetlb folios, unmap_poisoned_folio() currently rejects large folios before calling try_to_unmap(). Hence it is always the case that if try_to_unmap_one() handles an hwpoisoned folio, then the head page is marked with the poison bit.

Therefore, convert the poisoned subpage checks to folio_test_hwpoison().

No functional change intended, except that, while at it, convert VM_BUG_* to VM_WARN_*.

Acked-by David Hildenbrand (Arm) <[email protected]>

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

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

mm/rmap.c +4 -3

index 9939400e77c79..dbb077f8443e3 100644
@@ -2122,10 +2122,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			bool anon = folio_test_anon(folio);
 
 			/*
-			 * The try_to_unmap() is only passed a hugetlb page
-			 * in the case where the hugetlb page is poisoned.
+			 * The try_to_unmap() is only passed a hugetlb folio
+			 * in the case where the hugetlb folio contains a
+			 * poisoned page.
 			 */
-			VM_BUG_ON_PAGE(!PageHWPoison(subpage), subpage);
+			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
 			/*
 			 * huge_pmd_unshare may unmap an entire PMD page.
 			 * There is no way of knowing exactly which PMDs may

@@ -2204,7 +2205,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		/* Update high watermark before we lower rss */
 		update_hiwater_rss(mm);
 

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

Quoted 5 lines

mm/rmap.c +4 -1

-		if (PageHWPoison(subpage) && (flags & TTU_HWPOISON)) {
+		/*
+		 * With TTU_HWPOISON, we only expect small folios or hugetlb
+		 * folios here for now.
+		 */

Is this comment helpful given that you change that in patch #2, where you forget to update the comment? :)

try_to_unmap() receives hugetlb folios only from the hwpoison path.

hugetlb_update_hwpoison() sets the hugetlb folio's head-page hwpoison bit, and page_vma_mapped_walk() reports the hugetlb mapping at the head PFN, so the previous PageHWPoison(subpage) check happened to work for hugetlb.

For non-hugetlb folios, unmap_poisoned_folio() currently rejects large folios before calling try_to_unmap(). Hence it is always the case that if try_to_unmap_one() handles an hwpoisoned folio, then the head page is marked with the poison bit.

Therefore, convert the poisoned subpage checks to folio_test_hwpoison().

No functional change intended, except that, while at it, convert VM_BUG_* to VM_WARN_*.

Acked-by David Hildenbrand (Arm) <[email protected]>

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

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

mm/rmap.c +4 -3

index 9939400e77c79..dbb077f8443e3 100644
@@ -2122,10 +2122,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			bool anon = folio_test_anon(folio);
 
 			/*
-			 * The try_to_unmap() is only passed a hugetlb page
-			 * in the case where the hugetlb page is poisoned.
+			 * The try_to_unmap() is only passed a hugetlb folio
+			 * in the case where the hugetlb folio contains a
+			 * poisoned page.
 			 */
-			VM_BUG_ON_PAGE(!PageHWPoison(subpage), subpage);
+			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
 			/*
 			 * huge_pmd_unshare may unmap an entire PMD page.
 			 * There is no way of knowing exactly which PMDs may

mm/rmap.c +1 -1

+		if (folio_test_hwpoison(folio) && (flags & TTU_HWPOISON)) {
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
 			if (folio_test_hugetlb(folio)) {
 				hugetlb_count_sub(folio_nr_pages(folio), mm);
--

2.43.0

Quoted 1 line · Dev Jain

try_to_unmap() receives hugetlb folios only from the hwpoison path.

Horrifying :)

Quoted 10 lines

hugetlb_update_hwpoison() sets the hugetlb folio's head-page hwpoison bit, and page_vma_mapped_walk() reports the hugetlb mapping at the head PFN, so the previous PageHWPoison(subpage) check happened to work for hugetlb.

For non-hugetlb folios, unmap_poisoned_folio() currently rejects large folios before calling try_to_unmap(). Hence it is always the case that if try_to_unmap_one() handles an hwpoisoned folio, then the head page is marked with the poison bit.

Quoted 7 lines

Therefore, convert the poisoned subpage checks to folio_test_hwpoison().

No functional change intended, except that, while at it, convert VM_BUG_* to VM_WARN_*.

Acked-by David Hildenbrand (Arm) <[email protected]>

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

Small nit below but generally LGTM so:

Reviewed-by Lorenzo Stoakes (ARM) <[email protected]>

Quoted 27 lines

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

mm/rmap.c +4 -3

index 9939400e77c79..dbb077f8443e3 100644
@@ -2122,10 +2122,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			bool anon = folio_test_anon(folio);
 
 			/*
-			 * The try_to_unmap() is only passed a hugetlb page
-			 * in the case where the hugetlb page is poisoned.
+			 * The try_to_unmap() is only passed a hugetlb folio
+			 * in the case where the hugetlb folio contains a
+			 * poisoned page.
 			 */
-			VM_BUG_ON_PAGE(!PageHWPoison(subpage), subpage);
+			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
 			/*
 			 * huge_pmd_unshare may unmap an entire PMD page.
 			 * There is no way of knowing exactly which PMDs may

@@ -2204,7 +2205,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		/* Update high watermark before we lower rss */
 		update_hiwater_rss(mm);
 

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

Quoted 5 lines

mm/rmap.c +4 -1

-		if (PageHWPoison(subpage) && (flags & TTU_HWPOISON)) {
+		/*
+		 * With TTU_HWPOISON, we only expect small folios or hugetlb
+		 * folios here for now.
+		 */

I think probably we can drop this? It seems implied by the code (David comment similarly I think? :)

Quoted 8 lines

mm/rmap.c +1 -1

+		if (folio_test_hwpoison(folio) && (flags & TTU_HWPOISON)) {
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
 			if (folio_test_hugetlb(folio)) {
 				hugetlb_count_sub(folio_nr_pages(folio), mm);
--

2.43.0

Cheers, Lorenzo

Quoted 54 lines · Lorenzo Stoakes (ARM)

Quoted 1 line · Dev Jain

try_to_unmap() receives hugetlb folios only from the hwpoison path.

Horrifying :)

Quoted 10 lines

hugetlb_update_hwpoison() sets the hugetlb folio's head-page hwpoison bit, and page_vma_mapped_walk() reports the hugetlb mapping at the head PFN, so the previous PageHWPoison(subpage) check happened to work for hugetlb.

For non-hugetlb folios, unmap_poisoned_folio() currently rejects large folios before calling try_to_unmap(). Hence it is always the case that if try_to_unmap_one() handles an hwpoisoned folio, then the head page is marked with the poison bit.

Quoted 7 lines

Therefore, convert the poisoned subpage checks to folio_test_hwpoison().

No functional change intended, except that, while at it, convert VM_BUG_* to VM_WARN_*.

Acked-by David Hildenbrand (Arm) <[email protected]>

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

Small nit below but generally LGTM so:

Reviewed-by Lorenzo Stoakes (ARM) <[email protected]>

Quoted 27 lines

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

mm/rmap.c +4 -3

index 9939400e77c79..dbb077f8443e3 100644
@@ -2122,10 +2122,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			bool anon = folio_test_anon(folio);
 
 			/*
-			 * The try_to_unmap() is only passed a hugetlb page
-			 * in the case where the hugetlb page is poisoned.
+			 * The try_to_unmap() is only passed a hugetlb folio
+			 * in the case where the hugetlb folio contains a
+			 * poisoned page.
 			 */
-			VM_BUG_ON_PAGE(!PageHWPoison(subpage), subpage);
+			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
 			/*
 			 * huge_pmd_unshare may unmap an entire PMD page.
 			 * There is no way of knowing exactly which PMDs may

@@ -2204,7 +2205,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		/* Update high watermark before we lower rss */
 		update_hiwater_rss(mm);
 

Quoted 8 lines

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

Quoted 5 lines

mm/rmap.c +4 -1

-		if (PageHWPoison(subpage) && (flags & TTU_HWPOISON)) {
+		/*
+		 * With TTU_HWPOISON, we only expect small folios or hugetlb
+		 * folios here for now.
+		 */

I think probably we can drop this? It seems implied by the code (David comment similarly I think? :)

I had posted a v3 here:
https://lore.kernel.org/all/[email protected]/

In that I had fixed this comment to drop "or hugetlb folios ..." . But I would still like to retain the "only expect small folios" comment, it adds cognitive load when I am reading try_to_unmap_one and then have to go read hwpoison code to figure that large folios would be split before coming here.

Quoted 12 lines

Quoted 10 lines

mm/rmap.c +1 -1

+		if (folio_test_hwpoison(folio) && (flags & TTU_HWPOISON)) {
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
 			if (folio_test_hugetlb(folio)) {
 				hugetlb_count_sub(folio_nr_pages(folio), mm);
--

2.43.0

Cheers, Lorenzo

Simplify try_to_unmap_one() by separating the hugetlb parts into try_to_unmap_hugetlb_one().

To understand the correctness of the refactoring, the following points are noted:

1. try_to_unmap() is called for hugetlb folios only when they are hwpoisoned.

2. A hugetlb VMA cannot be mlocked.

3. page_vma_mapped_walk() returns at most one hugetlb mapping in a VMA, and that mapping points at the head PFN.

4. We won't ever process a softleaf entry that encodes a hugetlb folio; hugetlb folios are never swapped out, migration entries will be skipped (PVMW_MIGRATION not passed), and device-exclusive does not work for hugetlb.

5. The hwpoison entry is constructed from the poisoned folio, just as in the pre-refactor code. Any previous uffd-wp state is deliberately not preserved for the hwpoison entry.

6. TTU_HWPOISON is always present; for it to not be present, either the folio has to be in swapcache, or mapping_can_writeback() is true (see unmap_poisoned_folio), none of which is true for hugetlb folios.

7. Hugetlb uses separate counters from normal rss counters, therefore update_highwater_rss() need not be called.

While at it:

- Change VM_BUG_* to VM_WARN_*.

- Do not declare variables which are only used once.

- Constify some variables.

- Add some more VM_WARN_* to assert some invariants.

Except the above 4 points, no functional change intended.

Suggested-by David Hildenbrand (Arm) <[email protected]>

Acked-by David Hildenbrand (Arm) <[email protected]>

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

---

 include/linux/hugetlb.h |   1 +
 mm/rmap.c               | 178 ++++++++++++++++++++++------------------
 2 files changed, 97 insertions(+), 82 deletions(-)

include/linux/hugetlb.h +1 -0

index 4115076e4922a..bf7e163e3779d 100644
@@ -1271,6 +1271,7 @@ static inline void hugetlb_count_sub(long l, struct mm_struct *mm)
 }
 
 pte_t huge_ptep_get(struct mm_struct *mm, unsigned long addr, pte_t *ptep);
+unsigned long huge_pte_dirty(pte_t pte);
 
 static inline pte_t huge_ptep_clear_flush(struct vm_area_struct *vma,
 					  unsigned long addr, pte_t *ptep)

mm/rmap.c +96 -82

index dbb077f8443e3..cf02b315a2ead 100644
@@ -1978,6 +1978,96 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 				     FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY);
 }
 
+static bool try_to_unmap_hugetlb_one(struct folio *folio,
+		struct vm_area_struct *vma, unsigned long address, void *arg)
+{
+	DEFINE_FOLIO_VMA_WALK(pvmw, folio, vma, address, 0);
+	const unsigned long hsz = huge_page_size(hstate_vma(vma));
+	const enum ttu_flags flags = (enum ttu_flags)(long)arg;
+	struct mm_struct *mm = vma->vm_mm;
+	struct mmu_notifier_range range;
+	bool ret = true;
+	pte_t pteval;
+
+	/*
+	 * The try_to_unmap() is only passed a hugetlb folio in the case
+	 * where the hugetlb folio is poisoned.
+	 */
+	VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
+	VM_WARN_ON_ONCE(!(flags & TTU_HWPOISON));
+
+	range.end = vma_address_end(&pvmw);
+	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, vma->vm_mm,
+				address, range.end);
+	adjust_range_if_pmd_sharing_possible(vma, &range.start, &range.end);
+	mmu_notifier_invalidate_range_start(&range);
+
+	/* There is only a single mapping in a VMA. */
+	if (!page_vma_mapped_walk(&pvmw))
+		goto range_end;
+
+	VM_WARN_ON_ONCE(address != pvmw.address);
+
+	pteval = huge_ptep_get(mm, address, pvmw.pte);
+	VM_WARN_ON_ONCE(!pte_present(pteval));
+	VM_WARN_ON_ONCE(pte_pfn(pteval) != folio_pfn(folio));
+
+	/*
+	 * huge_pmd_unshare may unmap an entire PMD page. There is no way of
+	 * knowing exactly which PMDs may be cached for this mm, so we must
+	 * flush them all. start/end were already adjusted above to cover this
+	 * range.
+	 */
+	flush_cache_range(vma, range.start, range.end);
+
+	/*
+	 * To call huge_pmd_unshare, i_mmap_rwsem must be held in write mode.
+	 * Caller needs to explicitly do this outside rmap routines.
+	 *
+	 * We also must hold hugetlb vma_lock in write mode. Lock order dictates
+	 * acquiring vma_lock BEFORE i_mmap_rwsem. We can only try lock here and
+	 * fail if unsuccessful.
+	 */
+	if (!folio_test_anon(folio)) {
+		struct mmu_gather tlb;
+
+		VM_WARN_ON(!(flags & TTU_RMAP_LOCKED));
+		if (!hugetlb_vma_trylock_write(vma)) {
+			ret = false;
+			goto walk_done;
+		}
+
+		tlb_gather_mmu_vma(&tlb, vma);
+		if (huge_pmd_unshare(&tlb, vma, address, pvmw.pte)) {
+			hugetlb_vma_unlock_write(vma);
+			huge_pmd_unshare_flush(&tlb, vma);
+			tlb_finish_mmu(&tlb);
+			/*
+			 * The PMD table was unmapped, consequently unmapping
+			 * the folio.
+			 */
+			goto walk_done;
+		}
+		hugetlb_vma_unlock_write(vma);
+		tlb_finish_mmu(&tlb);
+	}
+	pteval = huge_ptep_clear_flush(vma, address, pvmw.pte);
+	if (huge_pte_dirty(pteval))
+		folio_mark_dirty(folio);
+
+	pteval = swp_entry_to_pte(make_hwpoison_entry(folio_page(folio, 0)));
+	hugetlb_count_sub(folio_nr_pages(folio), mm);
+	set_huge_pte_at(mm, address, pvmw.pte, pteval, hsz);
+	hugetlb_remove_rmap(folio);
+	folio_put_refs(folio, 1);
+
+walk_done:
+	page_vma_mapped_walk_done(&pvmw);
+range_end:
+	mmu_notifier_invalidate_range_end(&range);
+	return ret;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument
  */

@@ -1993,7 +2083,6 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 	enum ttu_flags flags = (enum ttu_flags)(long)arg;
 	unsigned long nr_pages = 1, end_addr;
 	unsigned long pfn;
-	unsigned long hsz = 0;
 	int ptes = 0;
 
 	/*

@@ -2007,8 +2096,6 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 
 	/*
 	 * For THP, we have to assume the worse case ie pmd for invalidation.
-	 * For hugetlb, it could be much worse if we need to do pud
-	 * invalidation in the case of pmd sharing.
 	 *
 	 * Note that the folio can not be freed in this function as call of
 	 * try_to_unmap() must hold a reference on the folio.

@@ -2016,17 +2103,6 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 	range.end = vma_address_end(&pvmw);
 	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, vma->vm_mm,
 				address, range.end);
-	if (folio_test_hugetlb(folio)) {
-		/*
-		 * If sharing is possible, start and end will be adjusted
-		 * accordingly.
-		 */
-		adjust_range_if_pmd_sharing_possible(vma, &range.start,
-						     &range.end);
-
-		/* We need the huge page size for set_huge_pte_at() */
-		hsz = huge_page_size(hstate_vma(vma));
-	}
 	mmu_notifier_invalidate_range_start(&range);
 
 	while (page_vma_mapped_walk(&pvmw)) {

@@ -2111,66 +2187,13 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			const softleaf_t entry = softleaf_from_pte(pteval);
 
 			pfn = softleaf_to_pfn(entry);
-			VM_WARN_ON_FOLIO(folio_test_hugetlb(folio), folio);
 		}
 
 		subpage = folio_page(folio, pfn - folio_pfn(folio));
 		anon_exclusive = folio_test_anon(folio) &&
 				 PageAnonExclusive(subpage);
 
-		if (folio_test_hugetlb(folio)) {
-			bool anon = folio_test_anon(folio);
-
-			/*
-			 * The try_to_unmap() is only passed a hugetlb folio
-			 * in the case where the hugetlb folio contains a
-			 * poisoned page.
-			 */
-			VM_WARN_ON_FOLIO(!folio_test_hwpoison(folio), folio);
-			/*
-			 * huge_pmd_unshare may unmap an entire PMD page.
-			 * There is no way of knowing exactly which PMDs may
-			 * be cached for this mm, so we must flush them all.
-			 * start/end were already adjusted above to cover this
-			 * range.
-			 */
-			flush_cache_range(vma, range.start, range.end);
-
-			/*
-			 * To call huge_pmd_unshare, i_mmap_rwsem must be
-			 * held in write mode.  Caller needs to explicitly
-			 * do this outside rmap routines.
-			 *
-			 * We also must hold hugetlb vma_lock in write mode.
-			 * Lock order dictates acquiring vma_lock BEFORE
-			 * i_mmap_rwsem.  We can only try lock here and fail
-			 * if unsuccessful.
-			 */
-			if (!anon) {
-				struct mmu_gather tlb;
-
-				VM_BUG_ON(!(flags & TTU_RMAP_LOCKED));
-				if (!hugetlb_vma_trylock_write(vma))
-					goto walk_abort;
-
-				tlb_gather_mmu_vma(&tlb, vma);
-				if (huge_pmd_unshare(&tlb, vma, address, pvmw.pte)) {
-					hugetlb_vma_unlock_write(vma);
-					huge_pmd_unshare_flush(&tlb, vma);
-					tlb_finish_mmu(&tlb);
-					/*
-					 * The PMD table was unmapped,
-					 * consequently unmapping the folio.
-					 */
-					goto walk_done;
-				}
-				hugetlb_vma_unlock_write(vma);
-				tlb_finish_mmu(&tlb);
-			}
-			pteval = huge_ptep_clear_flush(vma, address, pvmw.pte);
-			if (pte_dirty(pteval))
-				folio_mark_dirty(folio);
-		} else if (likely(pte_present(pteval))) {
+		if (likely(pte_present(pteval))) {
 			nr_pages = folio_unmap_pte_batch(folio, &pvmw, flags, pteval);
 			end_addr = address + nr_pages * PAGE_SIZE;
 			flush_cache_range(vma, address, end_addr);

@@ -2211,14 +2234,8 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 		 */
 		if (folio_test_hwpoison(folio) && (flags & TTU_HWPOISON)) {
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
-			if (folio_test_hugetlb(folio)) {
-				hugetlb_count_sub(folio_nr_pages(folio), mm);
-				set_huge_pte_at(mm, address, pvmw.pte, pteval,
-						hsz);
-			} else {
-				dec_mm_counter(mm, mm_counter(folio));
-				set_pte_at(mm, address, pvmw.pte, pteval);
-			}
+			dec_mm_counter(mm, mm_counter(folio));
+			set_pte_at(mm, address, pvmw.pte, pteval);
 		} else if (likely(pte_present(pteval)) && pte_unused(pteval) &&
 			   !userfaultfd_armed(vma)) {
 			/*

@@ -2346,11 +2363,7 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			add_mm_counter(mm, mm_counter_file(folio), -nr_pages);
 		}
 discard:
-		if (unlikely(folio_test_hugetlb(folio))) {
-			hugetlb_remove_rmap(folio);
-		} else {
-			folio_remove_rmap_ptes(folio, subpage, nr_pages, vma);
-		}
+		folio_remove_rmap_ptes(folio, subpage, nr_pages, vma);
 		if (vma->vm_flags & VM_LOCKED)
 			mlock_drain_local();
 		folio_put_refs(folio, nr_pages);

@@ -2398,7 +2411,8 @@ static int folio_not_mapped(struct folio *folio)
 void try_to_unmap(struct folio *folio, enum ttu_flags flags)
 {
 	struct rmap_walk_control rwc = {
-		.rmap_one = try_to_unmap_one,
+		.rmap_one = folio_test_hugetlb(folio) ?
+				try_to_unmap_hugetlb_one : try_to_unmap_one,
 		.arg = (void *)flags,
 		.done = folio_not_mapped,
 		.anon_lock = folio_lock_anon_vma_read,

Quoted 38 lines · Dev Jain

Simplify try_to_unmap_one() by separating the hugetlb parts into try_to_unmap_hugetlb_one().

To understand the correctness of the refactoring, the following points are noted:

1. try_to_unmap() is called for hugetlb folios only when they are hwpoisoned.

2. A hugetlb VMA cannot be mlocked.

3. page_vma_mapped_walk() returns at most one hugetlb mapping in a VMA, and that mapping points at the head PFN.

4. We won't ever process a softleaf entry that encodes a hugetlb folio; hugetlb folios are never swapped out, migration entries will be skipped (PVMW_MIGRATION not passed), and device-exclusive does not work for hugetlb.

5. The hwpoison entry is constructed from the poisoned folio, just as in the pre-refactor code. Any previous uffd-wp state is deliberately not preserved for the hwpoison entry.

6. TTU_HWPOISON is always present; for it to not be present, either the folio has to be in swapcache, or mapping_can_writeback() is true (see unmap_poisoned_folio), none of which is true for hugetlb folios.

7. Hugetlb uses separate counters from normal rss counters, therefore update_highwater_rss() need not be called.

While at it:

- Change VM_BUG_* to VM_WARN_*.

- Do not declare variables which are only used once.

- Constify some variables.

Quoted 8 lines

- Add some more VM_WARN_* to assert some invariants.

Except the above 4 points, no functional change intended.

Suggested-by David Hildenbrand (Arm) <[email protected]>

Acked-by David Hildenbrand (Arm) <[email protected]>

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

---

Andrew, could you squash the following diff:

mm/rmap.c +0 -0

index cf02b315a2ead..b61051b879c4c 100644

Quoted 59 lines · Dev Jain

Quoted 38 lines · Dev Jain

Simplify try_to_unmap_one() by separating the hugetlb parts into try_to_unmap_hugetlb_one().

To understand the correctness of the refactoring, the following points are noted:

1. try_to_unmap() is called for hugetlb folios only when they are hwpoisoned.

2. A hugetlb VMA cannot be mlocked.

3. page_vma_mapped_walk() returns at most one hugetlb mapping in a VMA, and that mapping points at the head PFN.

4. We won't ever process a softleaf entry that encodes a hugetlb folio; hugetlb folios are never swapped out, migration entries will be skipped (PVMW_MIGRATION not passed), and device-exclusive does not work for hugetlb.

5. The hwpoison entry is constructed from the poisoned folio, just as in the pre-refactor code. Any previous uffd-wp state is deliberately not preserved for the hwpoison entry.

6. TTU_HWPOISON is always present; for it to not be present, either the folio has to be in swapcache, or mapping_can_writeback() is true (see unmap_poisoned_folio), none of which is true for hugetlb folios.

7. Hugetlb uses separate counters from normal rss counters, therefore update_highwater_rss() need not be called.

While at it:

- Change VM_BUG_* to VM_WARN_*.

- Do not declare variables which are only used once.

- Constify some variables.

Quoted 8 lines

- Add some more VM_WARN_* to assert some invariants.

Except the above 4 points, no functional change intended.

Suggested-by David Hildenbrand (Arm) <[email protected]>

Acked-by David Hildenbrand (Arm) <[email protected]>

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

---

Andrew, could you squash the following diff:

mm/rmap.c +0 -0

index cf02b315a2ead..b61051b879c4c 100644
@@ -2229,8 +2229,7 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
                update_hiwater_rss(mm);
 

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

Quoted 7 lines

mm/rmap.c +1 -2

                /*
-                * With TTU_HWPOISON, we only expect small folios or hugetlb
-                * folios here for now.
+                * With TTU_HWPOISON, we only expect small folios here.
                 */
                if (folio_test_hwpoison(folio) && (flags & TTU_HWPOISON)) {
                        pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));

Note that the comment likely fits into a single line now.

Add ttu_anon_folio() as the common entry point for anonymous folio unmapping. It dispatches to the lazyfree or swapbacked helper as appropriate.

Keep restoration of cleared PTEs in the caller, where the PTEs were removed.

No functional change intended.

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

---
 mm/rmap.c | 47 +++++++++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 22 deletions(-)

mm/rmap.c +1 -0

index 68459f23f5c0a..55aa3ff8bf216 100644
@@ -2172,6 +2172,28 @@ static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
 	return true;
 }
 
+static bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,

Quoted 23 lines · Dev Jain

Add ttu_anon_folio() as the common entry point for anonymous folio unmapping. It dispatches to the lazyfree or swapbacked helper as appropriate.

Keep restoration of cleared PTEs in the caller, where the PTEs were removed.

No functional change intended.

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

---
 mm/rmap.c | 47 +++++++++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 22 deletions(-)

mm/rmap.c +1 -0

index 68459f23f5c0a..55aa3ff8bf216 100644
@@ -2172,6 +2172,28 @@ static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
 	return true;
 }
 
+static bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,

mm/rmap.c static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma, · line 2185

Quoted 10 lines

mm/rmap.c +10 -0

+		struct page *page, unsigned long address, pte_t *ptep,
+		pte_t pteval, unsigned long nr_pages)
+{
+	/*
+	 * Store the swap location in the pte.
+	 * See handle_pte_fault() ...
+	 */
+	if (unlikely(folio_test_swapbacked(folio) !=
+			folio_test_swapcache(folio))) {
+		WARN_ON_ONCE(1);

Can't you directly do

if (WARN_ON_ONCE(...))

	return false;

That even implies the unlikely() :)

mm/rmap.c static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma, · line 2193

Quoted 8 lines

mm/rmap.c +8 -0

+		return false;
+	}
+
+	if (!folio_test_swapbacked(folio))
+		return ttu_anon_lazyfree_folio(vma, folio, nr_pages);
+
+	/* nr_pages > 1 not supported yet */
+	return ttu_anon_swapbacked_folio(vma, folio, page, address, ptep,

LGTM

Acked-by David Hildenbrand (Arm) <[email protected]>

Andrew, could you you squash this diff:

Quoted 1 line · On Fri, 10 Jul 2026 18:43:14 +0530 Dev Jain <[email protected]> wrote:

Andrew, could you you squash this diff:

I think I'll await v3, please. I'm getting a big reject in rmap.c and methinks we're all a bit firehosed at present.

Are we all agreed that we'll be "batching anonymous large folio unmapping to optimize it"? If not, is this series pointful?

Quoted 1 line · On Fri, 10 Jul 2026 18:43:14 +0530 Dev Jain <[email protected]> wrote:

Andrew, could you you squash this diff:

I think I'll await v3, please. I'm getting a big reject in rmap.c and methinks we're all a bit firehosed at present.

Are we all agreed that we'll be "batching anonymous large folio unmapping to optimize it"? If not, is this series pointful?