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

[Dev Jain]

On 25/09/26 11:45 am, Barry Song wrote:
> 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?

I am not sure what is the "right place" you are referring to : ) More
than that, this comments feels very out of date - it is well-know that
the "swap location" is stored in the pte.

In fact, this comment starts on top of the wrong block, which is a
debug assertion for "when we are here, swapbacked if and only if swapcache".
So I'd rather just remove it.

> 
>>
>> -       /* 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?

The core problem is that we may succeed on a partial range.

ttu_anon_swapbacked_folio can succeed on x < nr_pages. So for x, we
should call finish_folio_unmap_batch. Either we do it my way, or,
ttu_anon_folio returns the number of pages successfully unmapped = x,
and then even in case of "failure" (we should now jump to
walk_abort since we failed to unmap 'nr_pages - x' pages) we call
finish_folio_unmap_batch for x pages.

IMHO that is uglier. I wanted to avoid this label jumping. I would
be interested in hearing a better alternate : )
> 
> That would also avoid the duplication.
> 
> Best Regards
> Barry