Re: [PATCH] userfaultfd: opportunistic TLB-flush batching for present pages in MOVE

[Barry Song <21cnbao@gmail.com>]

On Thu, Jul 31, 2025 at 6:47 PM Lokesh Gidra <lokeshgidra@google.com> wrote:
>
> MOVE ioctl's runtime is dominated by TLB-flush cost, which is required
> for moving present pages. Mitigate this cost by opportunistically
> batching present contiguous pages for TLB flushing.
>
> Without batching, in our testing on an arm64 Android device with UFFD GC,
> which uses MOVE ioctl for compaction, we observed that out of the total
> time spent in move_pages_pte(), over 40% is in ptep_clear_flush(), and
> ~20% in vm_normal_folio().
>
> With batching, the proportion of vm_normal_folio() increases to over
> 70% of move_pages_pte() without any changes to vm_normal_folio().
> Furthermore, time spent within move_pages_pte() is only ~20%, which
> includes TLB-flush overhead.
>
> Cc: Suren Baghdasaryan <surenb@google.com>
> Cc: Kalesh Singh <kaleshsingh@google.com>
> Cc: Barry Song <v-songbaohua@oppo.com>
> Cc: David Hildenbrand <david@redhat.com>
> Cc: Peter Xu <peterx@redhat.com>
> Signed-off-by: Lokesh Gidra <lokeshgidra@google.com>
> ---
>  mm/userfaultfd.c | 179 +++++++++++++++++++++++++++++++++--------------
>  1 file changed, 127 insertions(+), 52 deletions(-)
>
> diff --git a/mm/userfaultfd.c b/mm/userfaultfd.c
> index 8253978ee0fb..2465fb234671 100644
> --- a/mm/userfaultfd.c
> +++ b/mm/userfaultfd.c
> @@ -1026,18 +1026,62 @@ static inline bool is_pte_pages_stable(pte_t *dst_pte, pte_t *src_pte,
>                pmd_same(dst_pmdval, pmdp_get_lockless(dst_pmd));
>  }
>
> -static int move_present_pte(struct mm_struct *mm,
> -                           struct vm_area_struct *dst_vma,
> -                           struct vm_area_struct *src_vma,
> -                           unsigned long dst_addr, unsigned long src_addr,
> -                           pte_t *dst_pte, pte_t *src_pte,
> -                           pte_t orig_dst_pte, pte_t orig_src_pte,
> -                           pmd_t *dst_pmd, pmd_t dst_pmdval,
> -                           spinlock_t *dst_ptl, spinlock_t *src_ptl,
> -                           struct folio *src_folio)
> +/*
> + * Checks if the two ptes and the corresponding folio are eligible for batched
> + * move. If so, then returns pointer to the folio, after locking it. Otherwise,
> + * returns NULL.
> + */
> +static struct folio *check_ptes_for_batched_move(struct vm_area_struct *src_vma,
> +                                                unsigned long src_addr,
> +                                                pte_t *src_pte, pte_t *dst_pte)
> +{
> +       pte_t orig_dst_pte, orig_src_pte;
> +       struct folio *folio;
> +
> +       orig_dst_pte = ptep_get(dst_pte);
> +       if (!pte_none(orig_dst_pte))
> +               return NULL;
> +
> +       orig_src_pte = ptep_get(src_pte);
> +       if (pte_none(orig_src_pte))
> +               return NULL;
> +       if (!pte_present(orig_src_pte) || is_zero_pfn(pte_pfn(orig_src_pte)))
> +               return NULL;
> +
> +       folio = vm_normal_folio(src_vma, src_addr, orig_src_pte);
> +       if (!folio || !folio_trylock(folio))
> +               return NULL;
> +       if (!PageAnonExclusive(&folio->page) || folio_test_large(folio)) {
> +               folio_unlock(folio);
> +               return NULL;
> +       }
> +       return folio;
> +}
> +
> +static long move_present_ptes(struct mm_struct *mm,
> +                             struct vm_area_struct *dst_vma,
> +                             struct vm_area_struct *src_vma,
> +                             unsigned long dst_addr, unsigned long src_addr,
> +                             pte_t *dst_pte, pte_t *src_pte,
> +                             pte_t orig_dst_pte, pte_t orig_src_pte,
> +                             pmd_t *dst_pmd, pmd_t dst_pmdval,
> +                             spinlock_t *dst_ptl, spinlock_t *src_ptl,
> +                             struct folio *src_folio, unsigned long len)
>  {
>         int err = 0;
> +       unsigned long src_start = src_addr;
> +       unsigned long addr_end;
> +
> +       if (len > PAGE_SIZE) {
> +               addr_end = (dst_addr + PMD_SIZE) & PMD_MASK;
> +               if (dst_addr + len > addr_end)
> +                       len = addr_end - dst_addr;
>
> +               addr_end = (src_addr + PMD_SIZE) & PMD_MASK;
> +               if (src_addr + len > addr_end)
> +                       len = addr_end - src_addr;
> +       }
> +       flush_cache_range(src_vma, src_addr, src_addr + len);
>         double_pt_lock(dst_ptl, src_ptl);
>
>         if (!is_pte_pages_stable(dst_pte, src_pte, orig_dst_pte, orig_src_pte,
> @@ -1051,31 +1095,60 @@ static int move_present_pte(struct mm_struct *mm,
>                 err = -EBUSY;
>                 goto out;
>         }
> +       /* Avoid batching overhead for single page case */
> +       if (len > PAGE_SIZE) {
> +               flush_tlb_batched_pending(mm);

What’s confusing to me is that they track the unmapping of multiple
consecutive PTEs and defer TLB invalidation until later.
In contrast, you’re not tracking anything and instead call
flush_tlb_range() directly, which triggers the flush immediately.

It seems you might be combining two different batching approaches.
From what I can tell, you're essentially using flush_range
as a replacement for flushing each entry individually.

> +               arch_enter_lazy_mmu_mode();
> +               orig_src_pte = ptep_get_and_clear(mm, src_addr, src_pte);
> +       } else
> +               orig_src_pte = ptep_clear_flush(src_vma, src_addr, src_pte);
> +
> +       addr_end = src_start + len;
> +       do {
> +               /* Folio got pinned from under us. Put it back and fail the move. */
> +               if (folio_maybe_dma_pinned(src_folio)) {
> +                       set_pte_at(mm, src_addr, src_pte, orig_src_pte);
> +                       err = -EBUSY;
> +                       break;
> +               }
>
> -       orig_src_pte = ptep_clear_flush(src_vma, src_addr, src_pte);
> -       /* Folio got pinned from under us. Put it back and fail the move. */
> -       if (folio_maybe_dma_pinned(src_folio)) {
> -               set_pte_at(mm, src_addr, src_pte, orig_src_pte);
> -               err = -EBUSY;
> -               goto out;
> -       }
> -
> -       folio_move_anon_rmap(src_folio, dst_vma);
> -       src_folio->index = linear_page_index(dst_vma, dst_addr);
> +               folio_move_anon_rmap(src_folio, dst_vma);
> +               src_folio->index = linear_page_index(dst_vma, dst_addr);
>
> -       orig_dst_pte = folio_mk_pte(src_folio, dst_vma->vm_page_prot);
> -       /* Set soft dirty bit so userspace can notice the pte was moved */
> +               orig_dst_pte = folio_mk_pte(src_folio, dst_vma->vm_page_prot);
> +               /* Set soft dirty bit so userspace can notice the pte was moved */
>  #ifdef CONFIG_MEM_SOFT_DIRTY
> -       orig_dst_pte = pte_mksoft_dirty(orig_dst_pte);
> +               orig_dst_pte = pte_mksoft_dirty(orig_dst_pte);
>  #endif
> -       if (pte_dirty(orig_src_pte))
> -               orig_dst_pte = pte_mkdirty(orig_dst_pte);
> -       orig_dst_pte = pte_mkwrite(orig_dst_pte, dst_vma);
> +               if (pte_dirty(orig_src_pte))
> +                       orig_dst_pte = pte_mkdirty(orig_dst_pte);
> +               orig_dst_pte = pte_mkwrite(orig_dst_pte, dst_vma);
> +               set_pte_at(mm, dst_addr, dst_pte, orig_dst_pte);
> +
> +               src_addr += PAGE_SIZE;
> +               if (src_addr == addr_end)
> +                       break;
> +               src_pte++;
> +               dst_pte++;
>
> -       set_pte_at(mm, dst_addr, dst_pte, orig_dst_pte);
> +               folio_unlock(src_folio);
> +               src_folio = check_ptes_for_batched_move(src_vma, src_addr, src_pte, dst_pte);
> +               if (!src_folio)
> +                       break;
> +               orig_src_pte = ptep_get_and_clear(mm, src_addr, src_pte);
> +               dst_addr += PAGE_SIZE;
> +       } while (true);
> +
> +       if (len > PAGE_SIZE) {
> +               arch_leave_lazy_mmu_mode();
> +               if (src_addr > src_start)
> +                       flush_tlb_range(src_vma, src_start, src_addr);
> +       }

Can't we just remove the `if (len > PAGE_SIZE)` check and unify the
handling for both single-page and multi-page cases?


Thanks
Barry