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_*.

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

---

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.

- Use huge_pte_dirty() instead of pte_dirty().

- Add 3 VM_WARN_ON_ONCE asserting that TTU_HWPOISON should be present, pte_present() must be true, and the pfn derived from the huge pte must be that of the head of the hugetlb folio.

Except the BUG->WARN change, no functional change intended.

Suggested-by David Hildenbrand <[email protected]>

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

---
 include/linux/hugetlb.h |   1 +
 mm/rmap.c               | 177 +++++++++++++++++++++-------------------
 2 files changed, 96 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 +0 -0

index dbb077f8443e3..934773dfa2f2a 100644
@@ -1978,6 +1978,95 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 				     FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY);

mm/rmap.c +11 -0

+	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);

Quoted 71 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.

- Use huge_pte_dirty() instead of pte_dirty().

- Add 3 VM_WARN_ON_ONCE asserting that TTU_HWPOISON should be present, pte_present() must be true, and the pfn derived from the huge pte must be that of the head of the hugetlb folio.

Except the BUG->WARN change, no functional change intended.

Suggested-by David Hildenbrand <[email protected]>

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

---
 include/linux/hugetlb.h |   1 +
 mm/rmap.c               | 177 +++++++++++++++++++++-------------------
 2 files changed, 96 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 +0 -0

index dbb077f8443e3..934773dfa2f2a 100644
@@ -1978,6 +1978,95 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 				     FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY);
 }
 

mm/rmap.c static inline unsigned int folio_unmap_pte_batch(struct folio *folio, · line 1988

Quoted 8 lines

mm/rmap.c +8 -0

+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);
+	struct mmu_notifier_range range;
+	enum ttu_flags flags = (enum ttu_flags)(long)arg;
+	struct mm_struct *mm = vma->vm_mm;
+	const unsigned long hsz = huge_page_size(hstate_vma(vma));

I would move const variables all the way up.

Quoted 11 lines

mm/rmap.c +11 -0

+	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);

mm/rmap.c static inline unsigned int folio_unmap_pte_batch(struct folio *folio, · line 2009

Quoted 10 lines

mm/rmap.c +10 -0

+	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;
+
+	address = pvmw.address;

Will address ever change due to the page_vma_mapped_walk()? I don't think so, right?

So this can instead become a VM_WARN_ON_ONCE, maybe?

Nothing else jumped at me, much clearer to me.

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

For lazyfree folio unmapping, after clearing the ptes we must abort the operation if the folio got dirtied or it has unexpected references.

Refactor this logic into a function which will return whether we need to abort or not.

If we abort, we restore the ptes and bail out of try_to_unmap_one. Otherwise adjust the rss stats of the mm and jump to a label.

Also rename that label from "discard" to "finish_unmap"; the former is appropriate in the lazyfree context, but the code following the label is executed for other successful unmap code paths too, so 'discard' does not sound correct for them.

No functional change intended.

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

---
 mm/rmap.c | 87 +++++++++++++++++++++++++++++++------------------------
 1 file changed, 49 insertions(+), 38 deletions(-)

mm/rmap.c +34 -0

index 934773dfa2f2a..00b571c2a1bab 100644
@@ -2067,6 +2067,52 @@ static bool try_to_unmap_hugetlb_one(struct folio *folio,
 	return ret;
 }
 
+static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma,
+		struct folio *folio)
+{
+	int ref_count, map_count;
+
+	/*
+	 * Synchronize with gup_pte_range():
+	 * - clear PTE; barrier; read refcount
+	 * - inc refcount; barrier; read PTE
+	 */
+	smp_mb();
+
+	ref_count = folio_ref_count(folio);
+	map_count = folio_mapcount(folio);
+
+	/*
+	 * Order reads for page refcount and dirty flag
+	 * (see comments in __remove_mapping()).
+	 */
+	smp_rmb();
+
+	if (folio_test_dirty(folio) && !(vma->vm_flags & VM_DROPPABLE)) {
+		/*
+		 * redirtied either using the page table or a previously
+		 * obtained GUP reference.
+		 */
+		folio_set_swapbacked(folio);
+		return false;
+	}
+
+	if (ref_count != 1 + map_count) {
+		/*
+		 * Additional reference. Could be a GUP reference or any
+		 * speculative reference. GUP users must mark the folio

Quoted 63 lines · Dev Jain

For lazyfree folio unmapping, after clearing the ptes we must abort the operation if the folio got dirtied or it has unexpected references.

Refactor this logic into a function which will return whether we need to abort or not.

If we abort, we restore the ptes and bail out of try_to_unmap_one. Otherwise adjust the rss stats of the mm and jump to a label.

Also rename that label from "discard" to "finish_unmap"; the former is appropriate in the lazyfree context, but the code following the label is executed for other successful unmap code paths too, so 'discard' does not sound correct for them.

No functional change intended.

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

---
 mm/rmap.c | 87 +++++++++++++++++++++++++++++++------------------------
 1 file changed, 49 insertions(+), 38 deletions(-)

mm/rmap.c +34 -0

index 934773dfa2f2a..00b571c2a1bab 100644
@@ -2067,6 +2067,52 @@ static bool try_to_unmap_hugetlb_one(struct folio *folio,
 	return ret;
 }
 
+static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma,
+		struct folio *folio)
+{
+	int ref_count, map_count;
+
+	/*
+	 * Synchronize with gup_pte_range():
+	 * - clear PTE; barrier; read refcount
+	 * - inc refcount; barrier; read PTE
+	 */
+	smp_mb();
+
+	ref_count = folio_ref_count(folio);
+	map_count = folio_mapcount(folio);
+
+	/*
+	 * Order reads for page refcount and dirty flag
+	 * (see comments in __remove_mapping()).
+	 */
+	smp_rmb();
+
+	if (folio_test_dirty(folio) && !(vma->vm_flags & VM_DROPPABLE)) {
+		/*
+		 * redirtied either using the page table or a previously
+		 * obtained GUP reference.
+		 */
+		folio_set_swapbacked(folio);
+		return false;
+	}
+
+	if (ref_count != 1 + map_count) {
+		/*
+		 * Additional reference. Could be a GUP reference or any
+		 * speculative reference. GUP users must mark the folio

mm/rmap.c static bool try_to_unmap_hugetlb_one(struct folio *folio, · line 2113

Quoted 10 lines

mm/rmap.c +10 -0

+		 * dirty if there was a modification. This folio cannot be
+		 * reclaimed right now either way, so act just like nothing
+		 * happened.
+		 * We'll come back here later and detect if the folio was
+		 * dirtied when the additional reference is gone.
+		 */
+		return false;
+	}
+
+	return true;

You could simply do

/*
* Additional references could be due to GUP or from a speculative
* lookup. ...

*/
return ref_count == 1 + map_count;

Nothing else jumped at me

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

Quoted 78 lines · David Hildenbrand (Arm)

Quoted 63 lines · Dev Jain

For lazyfree folio unmapping, after clearing the ptes we must abort the operation if the folio got dirtied or it has unexpected references.

Refactor this logic into a function which will return whether we need to abort or not.

If we abort, we restore the ptes and bail out of try_to_unmap_one. Otherwise adjust the rss stats of the mm and jump to a label.

Also rename that label from "discard" to "finish_unmap"; the former is appropriate in the lazyfree context, but the code following the label is executed for other successful unmap code paths too, so 'discard' does not sound correct for them.

No functional change intended.

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

---
 mm/rmap.c | 87 +++++++++++++++++++++++++++++++------------------------
 1 file changed, 49 insertions(+), 38 deletions(-)

mm/rmap.c +34 -0

index 934773dfa2f2a..00b571c2a1bab 100644
@@ -2067,6 +2067,52 @@ static bool try_to_unmap_hugetlb_one(struct folio *folio,
 	return ret;
 }
 
+static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma,
+		struct folio *folio)
+{
+	int ref_count, map_count;
+
+	/*
+	 * Synchronize with gup_pte_range():
+	 * - clear PTE; barrier; read refcount
+	 * - inc refcount; barrier; read PTE
+	 */
+	smp_mb();
+
+	ref_count = folio_ref_count(folio);
+	map_count = folio_mapcount(folio);
+
+	/*
+	 * Order reads for page refcount and dirty flag
+	 * (see comments in __remove_mapping()).
+	 */
+	smp_rmb();
+
+	if (folio_test_dirty(folio) && !(vma->vm_flags & VM_DROPPABLE)) {
+		/*
+		 * redirtied either using the page table or a previously
+		 * obtained GUP reference.
+		 */
+		folio_set_swapbacked(folio);
+		return false;
+	}
+
+	if (ref_count != 1 + map_count) {
+		/*
+		 * Additional reference. Could be a GUP reference or any
+		 * speculative reference. GUP users must mark the folio

mm/rmap.c static bool try_to_unmap_hugetlb_one(struct folio *folio, · line 2113

Quoted 10 lines

mm/rmap.c +10 -0

+		 * dirty if there was a modification. This folio cannot be
+		 * reclaimed right now either way, so act just like nothing
+		 * happened.
+		 * We'll come back here later and detect if the folio was
+		 * dirtied when the additional reference is gone.
+		 */
+		return false;
+	}
+
+	return true;

You could simply do

Quoted 3 lines

/*
* Additional references could be due to GUP or from a speculative
* lookup. ...

Are you also suggesting to shorten the comment?

Refactor anonymous swapbacked folio unmap to ttu_anon_swapbacked_folio().

No functional change intended.

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

---
 mm/rmap.c | 105 ++++++++++++++++++++++++++++++++----------------------
 1 file changed, 62 insertions(+), 43 deletions(-)

mm/rmap.c +0 -0

index 00b571c2a1bab..ade78df5be2bd 100644

mm/rmap.c +14 -0

+{
+	pte_t swp_pte = swp_entry_to_pte(entry);
+
+	if (anon_exclusive)
+		swp_pte = pte_swp_mkexclusive(swp_pte);
+
+	if (likely(pte_present(pteval))) {
+		if (pte_soft_dirty(pteval))
+			swp_pte = pte_swp_mksoft_dirty(swp_pte);
+		if (pte_uffd_wp(pteval))
+			swp_pte = pte_swp_mkuffd_wp(swp_pte);
+	} else {
+		/* Device-exclusive entry */
+		if (pte_swp_soft_dirty(pteval))

mm/rmap.c +9 -0

+	swp_entry_t entry = page_swap_entry(subpage);
+	struct mm_struct *mm = vma->vm_mm;
+
+	if (folio_dup_swap(folio, subpage) < 0)
+		return false;
+
+	/*
+	 * arch_unmap_one() is expected to be a NOP on
+	 * architectures where we could have PFN swap PTEs,

mm/rmap.c +10 -0

+		return false;
+	}
+
+	mm_prepare_for_swap_entries(mm);
+	dec_mm_counter(mm, MM_ANONPAGES);
+	inc_mm_counter(mm, MM_SWAPENTS);
+	set_swp_pte_at(mm, address, ptep, entry, pteval, anon_exclusive);
+	return true;
+}
+

Quoted 9 lines · Dev Jain

Refactor anonymous swapbacked folio unmap to ttu_anon_swapbacked_folio().

No functional change intended.

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

---
 mm/rmap.c | 105 ++++++++++++++++++++++++++++++++----------------------
 1 file changed, 62 insertions(+), 43 deletions(-)

mm/rmap.c static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma, · line 2115

Quoted 8 lines

mm/rmap.c +0 -0

index 00b571c2a1bab..ade78df5be2bd 100644
@@ -2113,6 +2113,64 @@ static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma,
 	return true;
 }
 

BTW, no need for the "inline" for most of these functions (applies to other patches as well).

mm/rmap.c static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma, · line 2117

Quoted 3 lines

mm/rmap.c +2 -0

+static inline void set_swp_pte_at(struct mm_struct *mm, unsigned long address,
+		pte_t *ptep, swp_entry_t entry, pte_t pteval, bool anon_exclusive)

It's confusing that we say "set_swp_pte_at", but the pte_t we was is actually the old pte.

Can we instead have a function that creates us a swap_pte from the other data, and then do the set_pte_at() in the caller?

static pte_t swp_pte_prepare(...)

{
	...
	return swp_pte;
}

Quoted 19 lines

mm/rmap.c +18 -0

+{
+	pte_t swp_pte = swp_entry_to_pte(entry);
+
+	if (anon_exclusive)
+		swp_pte = pte_swp_mkexclusive(swp_pte);
+
+	if (likely(pte_present(pteval))) {
+		if (pte_soft_dirty(pteval))
+			swp_pte = pte_swp_mksoft_dirty(swp_pte);
+		if (pte_uffd_wp(pteval))
+			swp_pte = pte_swp_mkuffd_wp(swp_pte);
+	} else {
+		/* Device-exclusive entry */
+		if (pte_swp_soft_dirty(pteval))
+			swp_pte = pte_swp_mksoft_dirty(swp_pte);
+		if (pte_swp_uffd_wp(pteval))
+			swp_pte = pte_swp_mkuffd_wp(swp_pte);
+	}

mm/rmap.c static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma, · line 2141

Quoted 6 lines

mm/rmap.c +6 -0

+
+	set_pte_at(mm, address, ptep, swp_pte);
+}
+
+static inline bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma, struct folio *folio,
+		struct page *subpage, unsigned long address, pte_t *ptep,

Subpages do not exist :)

mm/rmap.c static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma, · line 2144

Quoted 4 lines

mm/rmap.c +3 -0

+		pte_t pteval)
+{
+	bool anon_exclusive = folio_test_anon(folio) && PageAnonExclusive(subpage);

const?

Quoted 15 lines

mm/rmap.c +14 -0

+	swp_entry_t entry = page_swap_entry(subpage);
+	struct mm_struct *mm = vma->vm_mm;
+
+	if (folio_dup_swap(folio, subpage) < 0)
+		return false;
+
+	/*
+	 * arch_unmap_one() is expected to be a NOP on
+	 * architectures where we could have PFN swap PTEs,
+	 * so we'll not check/care.
+	 */
+	if (arch_unmap_one(mm, vma, address, pteval) < 0) {
+		folio_put_swap(folio, subpage);
+		return false;

mm/rmap.c static inline bool ttu_anon_lazyfree_folio(struct vm_area_struct *vma, · line 2163

Quoted 5 lines

mm/rmap.c +5 -0

+	}
+
+	/* See folio_try_share_anon_rmap(): clear PTE first. */
+	if (anon_exclusive && folio_try_share_anon_rmap_pte(folio, subpage)) {
+		folio_put_swap(folio, subpage);

I recall I stumbled over this before, but there is no way to undo the arch_unmap_one(), right?

Quoted 11 lines

mm/rmap.c +10 -0

+		return false;
+	}
+
+	mm_prepare_for_swap_entries(mm);
+	dec_mm_counter(mm, MM_ANONPAGES);
+	inc_mm_counter(mm, MM_SWAPENTS);
+	set_swp_pte_at(mm, address, ptep, entry, pteval, anon_exclusive);
+	return true;
+}
+

Add ttu_anon_folio() as the common entry point for anonymous folio unmapping. It dispatches to the lazyfree or swapbacked helper as appropriate, and centralizes restoration of cleared PTEs on failure.

No functional change intended.

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

---
 mm/rmap.c | 59 +++++++++++++++++++++++++++++++++----------------------
 1 file changed, 36 insertions(+), 23 deletions(-)

mm/rmap.c +12 -0

index ade78df5be2bd..f021ecd51a4a2 100644
@@ -2171,6 +2171,40 @@ static inline bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma, struct
 	return true;
 }
 
+static inline bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,
+		struct page *subpage, unsigned long address, pte_t *ptep,
+		pte_t pteval, unsigned long nr_pages)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	bool ret;
+
+	/*
+	 * Store the swap location in the pte.
+	 * See handle_pte_fault() ...
+	 */
+	if (unlikely(folio_test_swapbacked(folio) !=

mm/rmap.c +6 -16

+
+	return ret;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument
  */

@@ -2351,31 +2385,10 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			dec_mm_counter(mm, mm_counter(folio));
 		} else if (folio_test_anon(folio)) {
-			/*
-			 * 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);
+			if (!ttu_anon_folio(vma, folio, subpage, address,
+					    pvmw.pte, pteval, nr_pages))
 				goto walk_abort;
-			}
 
-			/* MADV_FREE page check */
-			if (!folio_test_swapbacked(folio)) {
-				if (!ttu_anon_lazyfree_folio(vma, folio)) {
-					set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
-					goto walk_abort;
-				}
-				add_mm_counter(mm, MM_ANONPAGES, -nr_pages);
-				goto finish_unmap;

Quoted 35 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, and centralizes restoration of cleared PTEs on failure.

No functional change intended.

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

---
 mm/rmap.c | 59 +++++++++++++++++++++++++++++++++----------------------
 1 file changed, 36 insertions(+), 23 deletions(-)

mm/rmap.c +16 -0

index ade78df5be2bd..f021ecd51a4a2 100644
@@ -2171,6 +2171,40 @@ static inline bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma, struct
 	return true;
 }
 
+static inline bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,
+		struct page *subpage, unsigned long address, pte_t *ptep,
+		pte_t pteval, unsigned long nr_pages)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	bool ret;
+
+	/*
+	 * 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);
+		return false;
+	}

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

Quoted 6 lines

mm/rmap.c +6 -0

+
+	/* MADV_FREE page check */
+	if (!folio_test_swapbacked(folio)) {
+		ret = ttu_anon_lazyfree_folio(vma, folio);
+		if (ret)
+			add_mm_counter(mm, MM_ANONPAGES, -nr_pages);

Can't the be handled in there as well?

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

Quoted 9 lines

mm/rmap.c +8 -0

+	} else {
+		/* nr_pages > 1 not supported yet */
+		ret = ttu_anon_swapbacked_folio(vma, folio, subpage, address,
+						ptep, pteval);
+	}
+
+	if (!ret)
+		set_ptes(mm, address, ptep, pteval, nr_pages);

It might be cleaner to do that in the caller, where we actually removed the ptes in the first place?

Quoted 32 lines

mm/rmap.c +6 -16

+
+	return ret;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument
  */

@@ -2351,31 +2385,10 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			dec_mm_counter(mm, mm_counter(folio));
 		} else if (folio_test_anon(folio)) {
-			/*
-			 * 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);
+			if (!ttu_anon_folio(vma, folio, subpage, address,
+					    pvmw.pte, pteval, nr_pages))
 				goto walk_abort;
-			}
 
-			/* MADV_FREE page check */
-			if (!folio_test_swapbacked(folio)) {
-				if (!ttu_anon_lazyfree_folio(vma, folio)) {
-					set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
-					goto walk_abort;
-				}
-				add_mm_counter(mm, MM_ANONPAGES, -nr_pages);
-				goto finish_unmap;

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

Quoted 10 lines

mm/rmap.c +0 -7

-			}
-
-			if (!ttu_anon_swapbacked_folio(vma, folio, subpage, address,
-						pvmw.pte, pteval)) {
-				set_pte_at(mm, address, pvmw.pte, pteval);
-				goto walk_abort;
-			}
 			goto finish_unmap;
 		} else {
 			/*

Very nice.