Re: [PATCH v3 8/9] mm/rmap: batch unmap anonymous swap-backed large folios

[Barry Song]

On Thu, Sep 24, 2026 at 9:12 PM Dev Jain <dev.jain@xxxxxxx> wrote:
>
> 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 __folio_try_share_anon_rmap(),
> and repeated calls to page_vma_mapped_walk(). In general, batching helps
> execute similar code together, making the path more memory and CPU
> friendly.
>
> On arm64-contpte, batching also helps 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() does not look at the bits of the underlying page,
> process sub-batches of PTEs pointing to pages with the same exclusivity
> state, and batch set only those PTEs to swap PTEs in one go.
>
> Disable batching for sparc (because of arch_unmap_one), we will batch this
> later.
>
> Rmap accounting and reference accounting must happen when anon folio unmap
> succeeds. If a large folio is only partially batched or a later sub-batch
> fails, account only the pages that were actually unmapped. Put that
> accounting in __ttu_anon_swapbacked_folio() itself instead of using goto
> jumps at the try_to_unmap_one() callsite.
>
> Similarly, do the finish_folio_unmap_batch() 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, skip
> over this batch. page_vma_mapped_walk() handles this: check_pte() returns
> true only if any of [pvmw->pfn, pvmw->pfn + nr_pages) is mapped by the
> PTE. Swap PTEs have no underlying PFN, so check_pte() returns false until
> the walk reaches the next present PTE to unmap.
>
> Remove the finish_unmap label since no goto callers are left now.
>
> Signed-off-by: Dev Jain <dev.jain@xxxxxxx>
> ---
>  mm/rmap.c | 109 +++++++++++++++++++++++++++++++++++++++---------------
>  1 file changed, 80 insertions(+), 29 deletions(-)
>
> diff --git a/mm/rmap.c b/mm/rmap.c
> index fa9fc8374fc26..cdda9efd80bc3 100644
> --- a/mm/rmap.c
> +++ b/mm/rmap.c
> @@ -1966,12 +1966,13 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
>         end_addr = pmd_addr_end(addr, vma->vm_end);
>         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
> +       if (folio_test_anon(folio) && folio_test_swapbacked(folio))
>                 return 1;
> +#endif
>
>         /*
>          * If unmap fails, we need to restore the ptes. To avoid accidentally
> @@ -2141,16 +2142,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_batch(struct vm_area_struct *vma,
> +               struct folio *folio, struct page *page, unsigned long nr_pages)
> +{
> +       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 = folio_page_swap_entry(folio, 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;
>
>         /*
> @@ -2159,21 +2169,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_batch(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;
> +
> +       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);
> +
> +               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)
> +                       break;
> +
> +               pteval = pte_advance_pfn(pteval, len);
> +               address += len * PAGE_SIZE;
> +               batch_idx += len;
> +               ptep += len;
> +       }
> +
>         return true;
>  }
>
> @@ -2186,15 +2232,22 @@ static bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,
>          * See handle_pte_fault() ...
>          */
>         if (WARN_ON_ONCE(folio_test_swapbacked(folio) !=
> -                        folio_test_swapcache(folio)))
> +                        folio_test_swapcache(folio))) {
> +               set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
>                 return false;
> +       }
>
> -       if (!folio_test_swapbacked(folio))
> -               return ttu_anon_lazyfree_folio(vma, folio, nr_pages);
> +       if (!folio_test_swapbacked(folio)) {
> +               if (!ttu_anon_lazyfree_folio(vma, folio, nr_pages)) {
> +                       set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
> +                       return false;
> +               }
> +               finish_folio_unmap_batch(vma, folio, page, nr_pages);
> +               return true;
> +       }

Hi Dev,

if we look at the whole `ttu_anon_folio()` code:

static bool ttu_anon_folio(struct vm_area_struct *vma, struct folio *folio,
                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 (WARN_ON_ONCE(folio_test_swapbacked(folio) !=
                         folio_test_swapcache(folio))) {
                set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
                return false;
        }

        if (!folio_test_swapbacked(folio)) {
                if (!ttu_anon_lazyfree_folio(vma, folio, nr_pages)) {
                        set_ptes(vma->vm_mm, address, ptep, pteval, nr_pages);
                        return false;
                }
                finish_folio_unmap_batch(vma, folio, page, nr_pages);
                return true;
        }

        return ttu_anon_swapbacked_folio(vma, folio, page, address, ptep,
                                         pteval, nr_pages);
}

The comment "Store the swap location in the pte" is really out of
place. Since we are already here, can we move it to the right place?

>
> -       /* nr_pages > 1 not supported yet */
>         return ttu_anon_swapbacked_folio(vma, folio, page, address, ptep,
> -                                        pteval);
> +                                        pteval, nr_pages);
>  }
>
>  /*
> @@ -2374,13 +2427,14 @@ 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)) {
> +                       /* finish_folio_unmap_batch handled internally */


let me paste the final code here:

                } else if (folio_test_anon(folio)) {
                        /* finish_folio_unmap_batch handled internally */
                        if (!ttu_anon_folio(vma, folio, page, address,
                                            pvmw.pte, pteval, nr_pages))
                                goto walk_abort;

                        if (nr_pages == folio_nr_pages(folio))
                                goto walk_done;
                        continue;
                } else {
                        ...
                }

                finish_folio_unmap_batch(vma, folio, page, nr_pages);

                /*
                 * If we are sure that we batched the entire folio and cleared
                 * all PTEs, we can just optimize and stop right here.
                 */
                if (nr_pages == folio_nr_pages(folio))
                        goto walk_done;
                continue;

I am really struggling to read the code here. For anon folios, we
handle `finish_folio_unmap_batch()` internally, while for other folios
we still do it outside. This feels a bit inconsistent.

Also, we are duplicating the code below twice:

                if (nr_pages == folio_nr_pages(folio))
                        goto walk_done;

Can't we just do `finish_folio_unmap_batch()` like the other cases?

That would also avoid the duplication.

Best Regards
Barry