Optimize anonymous swapbacked large folio unmapping

In preparation for optimizing large folio unmapping, we need to reuse the page_anon_exclusive_batch helper in rmap.c and rmap.h and obey the existing use in mprotect.c .

Therefore, move it from mprotect.c to mm.h. Gate with CONFIG_MMU since both rmap and mprotect users are built only for CONFIG_MMU.

While at it, change start_idx and max_len to unsigned long type for future proofing against THP support at >= PUD level. Also shorten expected_anon_exclusive -> anon_exclusive.

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

---
 include/linux/mm.h | 19 +++++++++++++++++++
 mm/mprotect.c      | 17 -----------------
 2 files changed, 19 insertions(+), 17 deletions(-)

include/linux/mm.h +9 -0

index dd09c438fa23e..63da8813bc5df 100644
@@ -244,6 +244,25 @@ static inline unsigned long folio_page_idx(const struct folio *folio,
        return page - &folio->page;
 }
 
+#ifdef CONFIG_MMU
+/*
+ * Get max length of consecutive PTEs pointing to PageAnonExclusive() pages or
+ * !PageAnonExclusive() pages, starting from start_idx. Caller must enforce
+ * that the PTEs point to consecutive pages of the same anon large folio.
+ */
+static __always_inline int page_anon_exclusive_batch(unsigned long start_idx,
+               unsigned long max_len, struct page *first_page, bool anon_exclusive)
+{

Quoted 35 lines · Dev Jain

In preparation for optimizing large folio unmapping, we need to reuse the page_anon_exclusive_batch helper in rmap.c and rmap.h and obey the existing use in mprotect.c .

Therefore, move it from mprotect.c to mm.h. Gate with CONFIG_MMU since both rmap and mprotect users are built only for CONFIG_MMU.

While at it, change start_idx and max_len to unsigned long type for future proofing against THP support at >= PUD level. Also shorten expected_anon_exclusive -> anon_exclusive.

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

---
 include/linux/mm.h | 19 +++++++++++++++++++
 mm/mprotect.c      | 17 -----------------
 2 files changed, 19 insertions(+), 17 deletions(-)

include/linux/mm.h +9 -0

index dd09c438fa23e..63da8813bc5df 100644
@@ -244,6 +244,25 @@ static inline unsigned long folio_page_idx(const struct folio *folio,
        return page - &folio->page;
 }
 
+#ifdef CONFIG_MMU
+/*
+ * Get max length of consecutive PTEs pointing to PageAnonExclusive() pages or
+ * !PageAnonExclusive() pages, starting from start_idx. Caller must enforce
+ * that the PTEs point to consecutive pages of the same anon large folio.
+ */
+static __always_inline int page_anon_exclusive_batch(unsigned long start_idx,
+               unsigned long max_len, struct page *first_page, bool anon_exclusive)
+{

include/linux/mm.h static inline unsigned long folio_page_idx(const struct folio *folio, · line 265

Quoted 10 lines

include/linux/mm.h +10 -0

+       int idx;
+
+       for (idx = start_idx + 1; idx < start_idx + max_len; ++idx) {
+               if (anon_exclusive != PageAnonExclusive(first_page + idx))
+                       break;
+       }
+       return idx - start_idx;
+}
+#endif
+

Hi Dev,

I really don't want everything to become a top-level header file. If these are only mm-internal things, could we move them to mm/internal.h or another more appropriate header?

Best Regards
Barry

Quoted 6 lines

Hi Dev,

I really don't want everything to become a top-level header file. If these are only mm-internal things, could we move them to mm/internal.h or another more appropriate header?

What do you suggest? I had it in mm/internal.h initially, but because I had to use the helper this time in patch 4, I couldn't do that. The helper will now be used by mprotect.c, rmap.c and rmap.h.

I can keep it in rmap.h and make mprotect.c include rmap.h too.

In preparation for optimizing large folio unmapping, we need to reuse the page_anon_exclusive_batch helper in rmap.c and rmap.h and obey the existing use in mprotect.c .

Therefore, move it from mprotect.c to mm.h. Gate with CONFIG_MMU since both rmap and mprotect users are built only for CONFIG_MMU.

While at it, change start_idx and max_len to unsigned long type for future proofing against THP support at >= PUD level. Also shorten expected_anon_exclusive -> anon_exclusive.

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

---
 include/linux/mm.h | 19 +++++++++++++++++++
 mm/mprotect.c      | 17 -----------------
 2 files changed, 19 insertions(+), 17 deletions(-)

include/linux/mm.h +0 -0

index dd09c438fa23e..63da8813bc5df 100644

include/linux/mm.h +12 -0

+		unsigned long max_len, struct page *first_page, bool anon_exclusive)
+{
+	int idx;
+
+	for (idx = start_idx + 1; idx < start_idx + max_len; ++idx) {
+		if (anon_exclusive != PageAnonExclusive(first_page + idx))
+			break;
+	}
+	return idx - start_idx;
+}
+#endif
+
 static inline struct folio *lru_to_folio(struct list_head *head)
 {
 	return list_entry((head)->prev, struct folio, lru);

mm/mprotect.c +0 -17

index 2888ee638d872..75c2b1a00eb36 100644
@@ -138,23 +138,6 @@ static __always_inline void prot_commit_flush_ptes(struct vm_area_struct *vma,
 		tlb_flush_pte_range(tlb, addr, nr_ptes * PAGE_SIZE);
 }
 
-/*
- * Get max length of consecutive ptes pointing to PageAnonExclusive() pages or
- * !PageAnonExclusive() pages, starting from start_idx. Caller must enforce
- * that the ptes point to consecutive pages of the same anon large folio.
- */
-static __always_inline int page_anon_exclusive_batch(int start_idx, int max_len,
-		struct page *first_page, bool expected_anon_exclusive)
-{
-	int idx;
-
-	for (idx = start_idx + 1; idx < start_idx + max_len; ++idx) {
-		if (expected_anon_exclusive != PageAnonExclusive(first_page + idx))
-			break;
-	}
-	return idx - start_idx;
-}
-
 /*
  * This function is a result of trying our very best to retain the
  * "avoid the write-fault handler" optimization. In can_change_pte_writable(),

To enable batched unmapping of anonymous folios, we need to handle the sharing of exclusive pages. Hence, a batched version of folio_try_share_anon_rmap_pte is required.

Currently, the sole purpose of nr_pages in __folio_try_share_anon_rmap is to do some rmap sanity checks. Now, clear the PageAnonExclusive bit on a batch of nr_pages. Refactor the function such that the clearing of the bit can be done at one place without duplication.

Note that __folio_try_share_anon_rmap can receive nr_pages == HPAGE_PMD_NR from the PMD path, but currently we only clear the bit on the head page. Retain this behaviour by setting nr_pages = 1 in case the caller is folio_try_share_anon_rmap_pmd.

While at it, convert nr_pages to unsigned long to future-proof from overflow in case P4D-huge mappings etc get supported down the road. I haven't made such a change in each function receiving nr_pages in try_to_unmap_one - perhaps this can be done incrementally.

Add two WARN's: check that the batch is entirely exclusive (for PMD callers, need to check only head page), and that there are only PTE/PMD paths converging into __folio_try_share_anon_rmap.

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

---
 include/linux/rmap.h | 56 ++++++++++++++++++++++++++++++--------------
 1 file changed, 39 insertions(+), 17 deletions(-)

include/linux/rmap.h +4 -1

index 0b332770abeed..320f9f14f6020 100644
@@ -706,17 +706,23 @@ static inline int folio_try_dup_anon_rmap_pmd(struct folio *folio,
 }
 
 static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
-               struct page *page, int nr_pages, enum pgtable_level level)
+               struct page *page, unsigned long nr_pages, enum pgtable_level level)
 {
+       /* device private folios cannot get pinned via GUP. */
+       const bool pinnable = !folio_is_device_private(folio);
+

include/linux/rmap.h +6 -0

+
+       /* We only clear anon-exclusive from head page of PMD folio. */
+       if (level == PGTABLE_LEVEL_PMD)
+               nr_pages = 1;
+
+       VM_WARN_ON_FOLIO(page_anon_exclusive_batch(0, nr_pages, page, true) != nr_pages, folio);
 
        /*
         * We have to make sure that when we clear PageAnonExclusive, that

include/linux/rmap.h +0 -3

-       if (unlikely(folio_maybe_dma_pinned(folio)))
-               return -EBUSY;
-       ClearPageAnonExclusive(page);

Quoted 43 lines · Dev Jain

To enable batched unmapping of anonymous folios, we need to handle the sharing of exclusive pages. Hence, a batched version of folio_try_share_anon_rmap_pte is required.

Currently, the sole purpose of nr_pages in __folio_try_share_anon_rmap is to do some rmap sanity checks. Now, clear the PageAnonExclusive bit on a batch of nr_pages. Refactor the function such that the clearing of the bit can be done at one place without duplication.

Note that __folio_try_share_anon_rmap can receive nr_pages == HPAGE_PMD_NR from the PMD path, but currently we only clear the bit on the head page. Retain this behaviour by setting nr_pages = 1 in case the caller is folio_try_share_anon_rmap_pmd.

While at it, convert nr_pages to unsigned long to future-proof from overflow in case P4D-huge mappings etc get supported down the road. I haven't made such a change in each function receiving nr_pages in try_to_unmap_one - perhaps this can be done incrementally.

Add two WARN's: check that the batch is entirely exclusive (for PMD callers, need to check only head page), and that there are only PTE/PMD paths converging into __folio_try_share_anon_rmap.

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

---
 include/linux/rmap.h | 56 ++++++++++++++++++++++++++++++--------------
 1 file changed, 39 insertions(+), 17 deletions(-)

include/linux/rmap.h +4 -1

index 0b332770abeed..320f9f14f6020 100644
@@ -706,17 +706,23 @@ static inline int folio_try_dup_anon_rmap_pmd(struct folio *folio,
 }
 
 static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
-               struct page *page, int nr_pages, enum pgtable_level level)
+               struct page *page, unsigned long nr_pages, enum pgtable_level level)
 {
+       /* device private folios cannot get pinned via GUP. */
+       const bool pinnable = !folio_is_device_private(folio);
+

include/linux/rmap.h static inline int folio_try_dup_anon_rmap_pmd(struct folio *folio, · line 719

Quoted 10 lines

include/linux/rmap.h +1 -5

        VM_WARN_ON_FOLIO(!folio_test_anon(folio), folio);
        VM_WARN_ON_FOLIO(!PageAnonExclusive(page), folio);
+
        __folio_rmap_sanity_checks(folio, page, nr_pages, level);
 
-       /* device private folios cannot get pinned via GUP. */
-       if (unlikely(folio_is_device_private(folio))) {
-               ClearPageAnonExclusive(page);
-               return 0;
-       }

Somehow, I feel the early return for folio_is_device_private(folio) is more readable. Can we keep it? Then we can avoid many if (pinnable) checks later.

include/linux/rmap.h static inline int folio_try_dup_anon_rmap_pmd(struct folio *folio, · line 719

Quoted 1 line

include/linux/rmap.h +1 -0

+       VM_WARN_ON_ONCE(level > PGTABLE_LEVEL_PMD);

Maybe the below would be better, as it avoids depending on the exact value of PGTABLE_LEVEL_PMD and above.

VM_WARN_ON_ONCE(level != PGTABLE_LEVEL_PTE && level != PGTABLE_LEVEL_PMD);

Quoted 14 lines

include/linux/rmap.h +6 -0

+
+       /* We only clear anon-exclusive from head page of PMD folio. */
+       if (level == PGTABLE_LEVEL_PMD)
+               nr_pages = 1;
+
+       VM_WARN_ON_FOLIO(page_anon_exclusive_batch(0, nr_pages, page, true) != nr_pages, folio);
 
        /*
         * We have to make sure that when we clear PageAnonExclusive, that

@@ -760,29 +766,38 @@ static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
         * so we use explicit ones here.
         */
 

include/linux/rmap.h static __always_inline int __folio_try_share_anon_rmap(struct folio *folio, · line 772

Quoted 7 lines

include/linux/rmap.h +4 -3

-       /* Paired with the memory barrier in try_grab_folio(). */
-       if (IS_ENABLED(CONFIG_HAVE_GUP_FAST))
-               smp_mb();
+       if (likely(pinnable)) {
+               /* Paired with the memory barrier in try_grab_folio(). */
+               if (IS_ENABLED(CONFIG_HAVE_GUP_FAST))
+                       smp_mb();

If we return early for !pinnable, shouldn't we be able to avoid this? Is the reason you don't do the early return that you want to batch the folio_is_device_private(folio) case as well? If so, that seems sensible.

Is this a real use case that you're supporting with your patchset?

Quoted 4 lines

include/linux/rmap.h +0 -3

-       if (unlikely(folio_maybe_dma_pinned(folio)))
-               return -EBUSY;
-       ClearPageAnonExclusive(page);

include/linux/rmap.h static __always_inline int __folio_try_share_anon_rmap(struct folio *folio, · line 783

Quoted 10 lines

include/linux/rmap.h +10 -0

+               if (unlikely(folio_maybe_dma_pinned(folio)))
+                       return -EBUSY;
+       }
+
+       for (;;) {
+               ClearPageAnonExclusive(page);
+               if (--nr_pages == 0)
+                       break;
+               page++;
+       }

Maybe ?

     while (nr_pages--)
         ClearPageAnonExclusive(page++);

Best Regards
Barry

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

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

---
 mm/internal.h | 29 +++++++++++++++--------------
 mm/memory.c   |  4 ++--
 2 files changed, 17 insertions(+), 16 deletions(-)

Quoted 10 lines · Dev Jain

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

Quoted 3 lines

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

I don't quite understand why you mention this. Is anyone calling pte_swp_exclusive() on a migration entry? Shouldn't it only be called when softleaf_is_swap() is true?

Quoted 10 lines · Dev Jain

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

Quoted 3 lines

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

I don't quite understand why you mention this. Is anyone calling pte_swp_exclusive() on a migration entry? Shouldn't it only be called when softleaf_is_swap() is true?

Quoted 3 lines

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

Reviewed-by Barry Song <[email protected]>

Quoted 6 lines

---
 mm/internal.h | 29 +++++++++++++++--------------
 mm/memory.c   |  4 ++--
 2 files changed, 17 insertions(+), 16 deletions(-)

[...]

Quoted 1 line

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 526

Quoted 5 lines

/**

mm/internal.h +0 -0

@@ -523,7 +524,7 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {

We might find a user for this in the future, in which case we might want to rename swap_pte_batch() to swap_softleaf_batch()?

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 527

Quoted 3 lines

mm/internal.h +1 -1

-       pte_t expected_pte = pte_next_swp_offset(pte);
+       pte_t expected_pte = pte_next_softleaf_offset(pte);

Best Regards
Barry

Quoted 15 lines · Barry Song

Quoted 13 lines · Dev Jain

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

Quoted 3 lines

I don't quite understand why you mention this. Is anyone calling pte_swp_exclusive() on a migration entry? Shouldn't it only be called when softleaf_is_swap() is true?

You are right. I just wanted to emphasize the second paragraph - that the pte_move_swp_offset will also work for softleaf entries. But I think the third para confuses more, I'll drop it.

Quoted 6 lines

Quoted 4 lines

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

Reviewed-by Barry Song <[email protected]>

Thanks.

Quoted 8 lines

Quoted 6 lines

---
 mm/internal.h | 29 +++++++++++++++--------------
 mm/memory.c   |  4 ++--
 2 files changed, 17 insertions(+), 16 deletions(-)

[...]

Quoted 1 line

Quoted 8 lines

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 526

Quoted 5 lines

/**

mm/internal.h +0 -0

@@ -523,7 +524,7 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {

We might find a user for this in the future, in which case we might want to rename swap_pte_batch() to swap_softleaf_batch()?

That is what is being done here:
https://lore.kernel.org/all/[email protected]/

I don't have a strong opinion, I can also generalize this right now.

Quoted 6 lines

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 527

Quoted 5 lines

mm/internal.h +1 -1

-       pte_t expected_pte = pte_next_swp_offset(pte);
+       pte_t expected_pte = pte_next_softleaf_offset(pte);

Quoted 23 lines · Dev Jain

Quoted 18 lines · Barry Song

Quoted 5 lines · Dev Jain

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Quoted 8 lines

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

I don't quite understand why you mention this. Is anyone calling pte_swp_exclusive() on a migration entry? Shouldn't it only be called when softleaf_is_swap() is true?

Quoted 3 lines

You are right. I just wanted to emphasize the second paragraph - that the pte_move_swp_offset will also work for softleaf entries. But I think the third para confuses more, I'll drop it.

I would rather interpret this as meaning that pte_move_softleaf_offset will also work for swap softleafs, since pte_move_softleaf_offset will call some pure-swap functions？

Would it be possible to use if (softleaf_is_swap()) for these cases to make the intent clearer? or we can keep both pte_move_softleaf_offset and pte_move_swap_offset?

As pte_move_softleaf_offset() is getting a bit weird now, a generic softleaf function has a lot of swap-specific code in it:

static inline pte_t pte_move_softleaf_offset(pte_t pte, long delta)

{
        const softleaf_t entry = softleaf_from_pte(pte);
        pte_t new = __swp_entry_to_pte(__swp_entry(swp_type(entry),
                                                   (swp_offset(entry)

+ delta)));

        if (pte_swp_soft_dirty(pte))
                new = pte_swp_mksoft_dirty(new);
        if (pte_swp_exclusive(pte))
                new = pte_swp_mkexclusive(new);
        if (pte_swp_uffd(pte))
                new = pte_swp_mkuffd(new);

        return new;
}

Am I missing something here?

Quoted 27 lines

Quoted 6 lines

Quoted 4 lines

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

Reviewed-by Barry Song <[email protected]>

Thanks.

Quoted 16 lines

Quoted 6 lines

---
 mm/internal.h | 29 +++++++++++++++--------------
 mm/memory.c   |  4 ++--
 2 files changed, 17 insertions(+), 16 deletions(-)

[...]

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 526

Quoted 6 lines

/**

mm/internal.h +0 -0

@@ -523,7 +524,7 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {

We might find a user for this in the future, in which case we might want to rename swap_pte_batch() to swap_softleaf_batch()?

Quoted 4 lines

That is what is being done here:
https://lore.kernel.org/all/[email protected]/

I don't have a strong opinion, I can also generalize this right now.

I don't know how you and Shivank are collaborating, since you're both changing the same thing at the same time :-)

Hopefully, we can get one of yours into mm-new first, and the other one can find a way to resolve the conflicts.

Quoted 7 lines

Quoted 7 lines

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 527

Quoted 6 lines

mm/internal.h +1 -1

-       pte_t expected_pte = pte_next_swp_offset(pte);
+       pte_t expected_pte = pte_next_softleaf_offset(pte);

Best Regards
Barry

Quoted 24 lines · Barry Song

Quoted 23 lines · Dev Jain

Quoted 18 lines · Barry Song

Quoted 5 lines · Dev Jain

In preparation for adding a helper to set softleaf ptes in one go, generalize the swap entry helpers shifting the swap offset by delta, for softleaves.

Quoted 8 lines

Note that the soft-dirty bit, exclusive bit and uffd bit preservation will still work for non-swap softleaves, since a softleaf entry is constructed out of a type and offset, and those bits are ahead of the soft-dirty, exclusive and uffd bits.

For example, for a migration entry, pte_swp_exclusive() will return false, as the exclusivity is encoded in the type itself (SOFTLEAF_MIGRATION_READ_EXCLUSIVE).

I don't quite understand why you mention this. Is anyone calling pte_swp_exclusive() on a migration entry? Shouldn't it only be called when softleaf_is_swap() is true?

Quoted 8 lines

Quoted 3 lines

You are right. I just wanted to emphasize the second paragraph - that the pte_move_swp_offset will also work for softleaf entries. But I think the third para confuses more, I'll drop it.

I would rather interpret this as meaning that pte_move_softleaf_offset will also work for swap softleafs, since pte_move_softleaf_offset will call some pure-swap functions？

Yes, so if you see remove_migration_pte, there pte_swp_uffd and pte_swp_soft_dirty are being used. Same with restore_exclusive_pte (for device-exclusive stuff).

Which means that pte_swp_exclusive is the only one being used *only* for swap entries.

Quoted 4 lines

Would it be possible to use if (softleaf_is_swap()) for these cases to make the intent clearer? or we can keep both pte_move_softleaf_offset and pte_move_swap_offset?

So in my patch I can do

if (softleaf_is_swap(entry) && pte_swp_exclusive(pte))

That would make the intent clear.

Quoted 53 lines

As pte_move_softleaf_offset() is getting a bit weird now, a generic softleaf function has a lot of swap-specific code in it:

static inline pte_t pte_move_softleaf_offset(pte_t pte, long delta)

{
        const softleaf_t entry = softleaf_from_pte(pte);
        pte_t new = __swp_entry_to_pte(__swp_entry(swp_type(entry),
                                                   (swp_offset(entry)

+ delta)));

        if (pte_swp_soft_dirty(pte))
                new = pte_swp_mksoft_dirty(new);
        if (pte_swp_exclusive(pte))
                new = pte_swp_mkexclusive(new);
        if (pte_swp_uffd(pte))
                new = pte_swp_mkuffd(new);

        return new;
}

Am I missing something here?

Quoted 30 lines

Quoted 6 lines

Quoted 4 lines

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

Reviewed-by Barry Song <[email protected]>

Thanks.

Quoted 8 lines

Quoted 6 lines

---
 mm/internal.h | 29 +++++++++++++++--------------
 mm/memory.c   |  4 ++--
 2 files changed, 17 insertions(+), 16 deletions(-)

[...]

Quoted 1 line

Quoted 8 lines

mm/internal.h static inline pte_t pte_next_swp_offset(pte_t pte) · line 526

Quoted 5 lines

/**

mm/internal.h +0 -0

@@ -523,7 +524,7 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {

We might find a user for this in the future, in which case we might want to rename swap_pte_batch() to swap_softleaf_batch()?

That is what is being done here:
https://lore.kernel.org/all/[email protected]/

Quoted 4 lines

Quoted 1 line

I don't have a strong opinion, I can also generalize this right now.

I don't know how you and Shivank are collaborating, since you're both changing the same thing at the same time :-)

We are not collaborating really : ) he should have ideally waited for my patchset to go in first since I was doing big changes in rmap.c.

Quoted 3 lines

Hopefully, we can get one of yours into mm-new first, and the other one can find a way to resolve the conflicts.

Yes.

Currently we have a helper called set_ptes() which is used to set consecutive present ptes in the pgtables. To do the same operation but to set "consecutive" nonpresent (softleaf) ptes, add set_softleaf_ptes().

The softleaves which have a notion of consecutivity is swap softleaf, and those grouped by softeaf_has_pfn(). The latter is trivial; future code can convert present ptes pointing to the same large folio to swap softleaves with consecutive offsets using this helper.

The other case is softleaf markers. They do not have a notion of a swap offset or PFN, so future code can use set_softleaf_ptes() to store multiple (same) markers on the ptes.

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

---

Use set_softleaf_ptes() to store multiple uffd-wp marker entries instead of open coding the batched setting.

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

---

Enable batch clearing of ptes, and batch swap setting of ptes for anon swap-backed folio unmapping.

Processing all ptes of a large folio in one go helps us batch across atomics (add_mm_counter etc), barriers (in the function __folio_try_share_anon_rmap), repeated calls to page_vma_mapped_walk(), to name a few. In general, batching helps us to execute similar code together, making the execution of the program more memory and CPU friendly.

On arm64-contpte, batching also helps us avoid redundant ptep_get() calls and TLB flushes while breaking the contpte mapping.

The handling of anon-exclusivity is very similar to commit cac1db8c3aad ("mm: optimize mprotect() by PTE batching"). Since folio_unmap_pte_batch() won't look at the bits of the underlying page, we need to process sub-batches of ptes pointing to pages which are same w.r.t exclusivity, and batch set only those ptes to swap ptes in one go.

arch_unmap_one() is only defined for sparc64; I am not comfortable regarding the nuances between retrieving the pfn from pte_pfn() or from (paddr = pte_val(oldpte) & _PAGE_PADDR_4V).

(And, pte_next_pfn() can't even be called from arch_unmap_one() because that file does not include pgtable.h) So just disable the "sparc64-anon-swapbacked" case for now.

We need to take care of rmap accounting (folio_remove_rmap_ptes) and reference accounting (folio_put_refs) when anon folio unmap succeeds. In case we partially batch the large folio and fail, we need to correctly do the accounting for pages which were successfully unmapped. So, put this accounting code (which is finish_folio_unmap()) in __ttu_anon_swapbacked_folio() itself, instead of doing some horrible goto jumping at the callsite of ttu_anon_folio().

Similarly, do the finish_folio_unmap() in ttu_anon_folio itself for the non-swapbacked (lazyfree) case.

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 swap pte, so check_pte returns false and we keep skipping until we hit a present pte, which is where we want to start unmapping from next.

Remove the label finish_unmap since no goto callers are left now.

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

---
 mm/rmap.c | 110 ++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 81 insertions(+), 29 deletions(-)

mm/rmap.c +0 -0

index 1b9f07d4d1be9..68e0201ffd003 100644
@@ -1964,12 +1964,14 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
        end_addr = pmd_addr_end(addr, vma->vm_end);

mm/rmap.c +26 -11

+{
+       folio_remove_rmap_ptes(folio, page, nr_pages, vma);
+       if (vma->vm_flags & VM_LOCKED)
+               mlock_drain_local();
+       folio_put_refs(folio, nr_pages);
+}
+
+static bool __ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
                struct folio *folio, struct page *page, unsigned long address,
-               pte_t *ptep, pte_t pteval)
+               pte_t *ptep, pte_t pteval, unsigned long nr_pages,
+               bool anon_exclusive)
 {
-       const bool anon_exclusive = folio_test_anon(folio) &&
-                                   PageAnonExclusive(page);
        swp_entry_t entry = page_swap_entry(page);
        struct mm_struct *mm = vma->vm_mm;
+       pte_t swp_pte;
 
-       if (folio_dup_swap_pages(folio, page, 1) < 0)
+       if (folio_dup_swap_pages(folio, page, nr_pages) < 0)
                return false;
 
        /*

@@ -2157,21 +2168,57 @@ static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
         * so we'll not check/care.
         */
        if (arch_unmap_one(mm, vma, address, pteval) < 0) {
-               folio_put_swap_pages(folio, page, 1);
+               VM_WARN_ON(nr_pages != 1);
+               folio_put_swap_pages(folio, page, nr_pages);
                return false;
        }
 
        /* See folio_try_share_anon_rmap(): clear PTE first. */
-       if (anon_exclusive && folio_try_share_anon_rmap_pte(folio, page)) {
-               folio_put_swap_pages(folio, page, 1);
+       if (anon_exclusive &&
+           folio_try_share_anon_rmap_ptes(folio, page, nr_pages)) {
+               folio_put_swap_pages(folio, page, nr_pages);
                return false;
        }
 
        mm_prepare_for_swap_entries(mm);
-       dec_mm_counter(mm, MM_ANONPAGES);
-       inc_mm_counter(mm, MM_SWAPENTS);
-       set_pte_at(mm, address, ptep,
-                  swp_pte_prepare(entry, pteval, anon_exclusive));
+       add_mm_counter(mm, MM_ANONPAGES, -nr_pages);
+       add_mm_counter(mm, MM_SWAPENTS, nr_pages);
+       swp_pte = swp_pte_prepare(entry, pteval, anon_exclusive);
+       set_softleaf_ptes(mm, address, ptep, swp_pte, nr_pages);
+       finish_folio_unmap(vma, folio, page, nr_pages);
+       return true;
+}
+
+static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,

mm/rmap.c +11 -0

+
+               if (!__ttu_anon_swapbacked_folio(vma, folio,
+                               first_page + batch_idx, address, ptep, pteval,
+                               len, anon_exclusive)) {
+                       /* Restore the remaining PTEs that were cleared. */
+                       set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
+                       return false;
+               }
+
+               nr_pages -= len;
+               if (!nr_pages)

Quoted 61 lines · Dev Jain

Enable batch clearing of ptes, and batch swap setting of ptes for anon swap-backed folio unmapping.

Processing all ptes of a large folio in one go helps us batch across atomics (add_mm_counter etc), barriers (in the function __folio_try_share_anon_rmap), repeated calls to page_vma_mapped_walk(), to name a few. In general, batching helps us to execute similar code together, making the execution of the program more memory and CPU friendly.

On arm64-contpte, batching also helps us avoid redundant ptep_get() calls and TLB flushes while breaking the contpte mapping.

The handling of anon-exclusivity is very similar to commit cac1db8c3aad ("mm: optimize mprotect() by PTE batching"). Since folio_unmap_pte_batch() won't look at the bits of the underlying page, we need to process sub-batches of ptes pointing to pages which are same w.r.t exclusivity, and batch set only those ptes to swap ptes in one go.

arch_unmap_one() is only defined for sparc64; I am not comfortable regarding the nuances between retrieving the pfn from pte_pfn() or from (paddr = pte_val(oldpte) & _PAGE_PADDR_4V).

(And, pte_next_pfn() can't even be called from arch_unmap_one() because that file does not include pgtable.h) So just disable the "sparc64-anon-swapbacked" case for now.

We need to take care of rmap accounting (folio_remove_rmap_ptes) and reference accounting (folio_put_refs) when anon folio unmap succeeds. In case we partially batch the large folio and fail, we need to correctly do the accounting for pages which were successfully unmapped. So, put this accounting code (which is finish_folio_unmap()) in __ttu_anon_swapbacked_folio() itself, instead of doing some horrible goto jumping at the callsite of ttu_anon_folio().

Similarly, do the finish_folio_unmap() in ttu_anon_folio itself for the non-swapbacked (lazyfree) case.

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 swap pte, so check_pte returns false and we keep skipping until we hit a present pte, which is where we want to start unmapping from next.

Remove the label finish_unmap since no goto callers are left now.

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

---
 mm/rmap.c | 110 ++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 81 insertions(+), 29 deletions(-)

mm/rmap.c +0 -0

index 1b9f07d4d1be9..68e0201ffd003 100644
@@ -1964,12 +1964,14 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
        end_addr = pmd_addr_end(addr, vma->vm_end);

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

Quoted 10 lines

mm/rmap.c +3 -3

        max_nr = (end_addr - addr) >> PAGE_SHIFT;
 
-       /* We only support lazyfree or file folios batching for now ... */
-       if (folio_test_anon(folio) && folio_test_swapbacked(folio))
+       if (pte_unused(pte))
                return 1;
 
-       if (pte_unused(pte))
+#ifdef __HAVE_ARCH_UNMAP_ONE
+       /* Add batching support to arch_unmap_one() to remove this */

I'd like to make this clearer. For example, could we say that sparc has arch_unmap_one(), which doesn't support batching?

BTW, it shouldn't be too hard to save nr_pages tags, looking at the code:

static inline int arch_unmap_one(struct mm_struct *mm,
struct vm_area_struct *vma, unsigned long addr, pte_t oldpte)

{
        if (adi_state.enabled && (pte_val(oldpte) & _PAGE_MCD_4V))
                return adi_save_tags(mm, vma, addr, oldpte);
        return 0;
}

Maybe the sparc folks can handle this.

Quoted 8 lines

mm/rmap.c +2 -0

+       if (folio_test_anon(folio) && folio_test_swapbacked(folio))
                return 1;
+#endif
 
        /*
         * If unmap fails, we need to restore the ptes. To avoid accidentally

mm/rmap.c static pte_t swp_pte_prepare(swp_entry_t entry, pte_t old_pte, · line 2145

Quoted 6 lines

mm/rmap.c +2 -1

@@ -2139,16 +2141,25 @@ static pte_t swp_pte_prepare(swp_entry_t entry, pte_t old_pte,
        return swp_pte;
 }
 
-static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
+static void finish_folio_unmap(struct vm_area_struct *vma,
+               struct folio *folio, struct page *page, unsigned long nr_pages)

We are not necessarily finishing the whole folio here, right? The name is a bit misleading to me, as it sounds like we're finishing the whole folio.

Maybe finish_folio_unmap_batch()?

Quoted 63 lines

mm/rmap.c +31 -11

+{
+       folio_remove_rmap_ptes(folio, page, nr_pages, vma);
+       if (vma->vm_flags & VM_LOCKED)
+               mlock_drain_local();
+       folio_put_refs(folio, nr_pages);
+}
+
+static bool __ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
                struct folio *folio, struct page *page, unsigned long address,
-               pte_t *ptep, pte_t pteval)
+               pte_t *ptep, pte_t pteval, unsigned long nr_pages,
+               bool anon_exclusive)
 {
-       const bool anon_exclusive = folio_test_anon(folio) &&
-                                   PageAnonExclusive(page);
        swp_entry_t entry = page_swap_entry(page);
        struct mm_struct *mm = vma->vm_mm;
+       pte_t swp_pte;
 
-       if (folio_dup_swap_pages(folio, page, 1) < 0)
+       if (folio_dup_swap_pages(folio, page, nr_pages) < 0)
                return false;
 
        /*

@@ -2157,21 +2168,57 @@ static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
         * so we'll not check/care.
         */
        if (arch_unmap_one(mm, vma, address, pteval) < 0) {
-               folio_put_swap_pages(folio, page, 1);
+               VM_WARN_ON(nr_pages != 1);
+               folio_put_swap_pages(folio, page, nr_pages);
                return false;
        }
 
        /* See folio_try_share_anon_rmap(): clear PTE first. */
-       if (anon_exclusive && folio_try_share_anon_rmap_pte(folio, page)) {
-               folio_put_swap_pages(folio, page, 1);
+       if (anon_exclusive &&
+           folio_try_share_anon_rmap_ptes(folio, page, nr_pages)) {
+               folio_put_swap_pages(folio, page, nr_pages);
                return false;
        }
 
        mm_prepare_for_swap_entries(mm);
-       dec_mm_counter(mm, MM_ANONPAGES);
-       inc_mm_counter(mm, MM_SWAPENTS);
-       set_pte_at(mm, address, ptep,
-                  swp_pte_prepare(entry, pteval, anon_exclusive));
+       add_mm_counter(mm, MM_ANONPAGES, -nr_pages);
+       add_mm_counter(mm, MM_SWAPENTS, nr_pages);
+       swp_pte = swp_pte_prepare(entry, pteval, anon_exclusive);
+       set_softleaf_ptes(mm, address, ptep, swp_pte, nr_pages);
+       finish_folio_unmap(vma, folio, page, nr_pages);
+       return true;
+}
+
+static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
+               struct folio *folio, struct page *first_page,
+               unsigned long address, pte_t *ptep, pte_t pteval,
+               unsigned long nr_pages)
+{
+       unsigned long batch_idx = 0;

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

Quoted 5 lines

mm/rmap.c +5 -0

+
+       while (nr_pages) {
+               bool anon_exclusive = PageAnonExclusive(first_page + batch_idx);
+               unsigned long len = page_anon_exclusive_batch(batch_idx,
+                               nr_pages, first_page, anon_exclusive);

len is really a bad name, as len usually describes a size. Maybe batch_pages?

Quoted 12 lines

mm/rmap.c +11 -0

+
+               if (!__ttu_anon_swapbacked_folio(vma, folio,
+                               first_page + batch_idx, address, ptep, pteval,
+                               len, anon_exclusive)) {
+                       /* Restore the remaining PTEs that were cleared. */
+                       set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
+                       return false;
+               }
+
+               nr_pages -= len;
+               if (!nr_pages)

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

Quoted 10 lines

mm/rmap.c +8 -0

+                       break;
+
+               pteval = pte_advance_pfn(pteval, len);
+               address += len * PAGE_SIZE;
+               batch_idx += len;
+               ptep += len;
+       }
+
        return true;
 }

Best Regards
Barry

Quoted 6 lines

Quoted 6 lines

mm/rmap.c +6 -0

+static bool ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
+               struct folio *folio, struct page *first_page,
+               unsigned long address, pte_t *ptep, pte_t pteval,
+               unsigned long nr_pages)
+{
+       unsigned long batch_idx = 0;

Quoted 8 lines

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

Quoted 5 lines

mm/rmap.c +5 -0

+
+       while (nr_pages) {
+               bool anon_exclusive = PageAnonExclusive(first_page + batch_idx);
+               unsigned long len = page_anon_exclusive_batch(batch_idx,
+                               nr_pages, first_page, anon_exclusive);

len is really a bad name, as len usually describes a size. Maybe batch_pages?

I disagree here : ) I don't think someone should mistake len with size. len is ... "length". So in this case it is the length of pages in the array, starting from batch_idx, upto nr_pages, which are all exclusive or not. Also I would prefer short variable names.