Re: [PATCH v3 09/11] mm: handle PMD swap entry faults on swap-in

[Kairui Song]

On Sat, Jul 4, 2026 at 1:41 AM Usama Arif <usama.arif@xxxxxxxxx> wrote:
>
> Add do_huge_pmd_swap_page() and dispatch to it from __handle_mm_fault()
> when vmf->orig_pmd encodes a swap entry.  The handler resolves the
> entire 2 MB mapping in one shot, mirroring do_swap_page() (PTE path)
> at PMD granularity:
>
>   - Look up the folio in the swap cache; on a miss, allocate a
>     PMD-order folio via swap_cache_alloc_folio() and read from swap.
>
>   - After locking, re-validate that the folio still corresponds to our
>     entry and is still PMD-sized.  Between the unlocked cache lookup
>     and the lock, a racing swap-in on the same entry may have removed
>     it from the cache via folio_free_swap(), or reclaim / memory_failure
>     / deferred-split may have split the folio into smaller folios.
>
>   - Restore soft_dirty and uffd_wp from the swap PMD.  Map writable
>     only when the entry was exclusive, the VMA permits writes, and
>     uffd-wp is not armed.  Drop the exclusive marker when the cached
>     folio is under writeback to an SWP_STABLE_WRITES backend (zram,
>     encrypted) so the PMD is mapped read-only; a later write COWs
>     into a fresh folio rather than corrupting the in-flight writeback.
>     Mirrors do_swap_page().
>
>   - When the resulting PMD is read-only but the fault was a write,
>     update vmf->orig_pmd and call wp_huge_pmd() in the same handler
>     to COW immediately rather than forcing a second fault.  Mask
>     VM_FAULT_FALLBACK from its return: a PMD-COW that splits to
>     PTE-level is normal, but the bit is part of VM_FAULT_ERROR and
>     arch fault handlers BUG() on it without SIGBUS/HWPOISON/SIGSEGV.
>     Requires exposing wp_huge_pmd() via mm/internal.h.
>
>   - Free the swap slot via should_try_to_free_swap() (hoisted from
>     mm/memory.c into mm/internal.h so PTE- and PMD-level swap-in
>     share the heuristic).
>
> When PMD-order resources are unavailable (folio allocation fails,
> the cached folio was split, memcg charge fails, or swapin_folio()
> races) split the PMD swap entry into 512 PTE swap entries via
> __split_huge_pmd() and return 0.  The fault retries and do_swap_page()
> takes over per-PTE.  This avoids returning VM_FAULT_OOM for transient
> PMD-order allocation failures.
>
> Signed-off-by: Usama Arif <usama.arif@xxxxxxxxx>

Hello Usama

Amazing work.

> ---
>  include/linux/huge_mm.h |   9 ++
>  mm/huge_memory.c        | 216 ++++++++++++++++++++++++++++++++++++++++
>  mm/internal.h           |  36 +++++++
>  mm/memory.c             |  40 +-------
>  4 files changed, 265 insertions(+), 36 deletions(-)
>
> diff --git a/include/linux/huge_mm.h b/include/linux/huge_mm.h
> index 1487bf4af1a7..9ec475ccfc91 100644
> --- a/include/linux/huge_mm.h
> +++ b/include/linux/huge_mm.h
> @@ -531,6 +531,15 @@ vm_fault_t do_huge_pmd_numa_page(struct vm_fault *vmf);
>
>  vm_fault_t do_huge_pmd_device_private(struct vm_fault *vmf);
>
> +#ifdef CONFIG_THP_SWAP
> +vm_fault_t do_huge_pmd_swap_page(struct vm_fault *vmf);
> +#else
> +static inline vm_fault_t do_huge_pmd_swap_page(struct vm_fault *vmf)
> +{
> +       return 0;
> +}
> +#endif
> +
>  extern struct folio *huge_zero_folio;
>  extern unsigned long huge_zero_pfn;
>
> diff --git a/mm/huge_memory.c b/mm/huge_memory.c
> index fdc1a503c609..5fa60324a2f0 100644
> --- a/mm/huge_memory.c
> +++ b/mm/huge_memory.c
> @@ -41,6 +41,7 @@
>  #include <linux/pgalloc.h>
>  #include <linux/pgalloc_tag.h>
>  #include <linux/pagewalk.h>
> +#include <linux/zswap.h>
>
>  #include <asm/tlb.h>
>  #include "internal.h"
> @@ -2312,6 +2313,221 @@ vm_fault_t do_huge_pmd_numa_page(struct vm_fault *vmf)
>         return 0;
>  }
>
> +#ifdef CONFIG_THP_SWAP
> +/**
> + * do_huge_pmd_swap_page() - Handle a fault on a PMD-level swap entry.
> + * @vmf: Fault context. vmf->orig_pmd contains the swap PMD.
> + *
> + * A PMD swap entry is a compact encoding for HPAGE_PMD_NR consecutive swap
> + * slots. If the swap cache still has one PMD-sized folio covering the range,
> + * map it directly at PMD level. If the range has been split into per-page
> + * cache state, or zswap may have per-page state for it, split the PMD swap
> + * entry and retry at PTE granularity.
> + *
> + * Return: VM_FAULT_* flags.
> + */
> +vm_fault_t do_huge_pmd_swap_page(struct vm_fault *vmf)
> +{
> +       struct vm_area_struct *vma = vmf->vma;
> +       struct mm_struct *mm = vma->vm_mm;
> +       struct folio *folio;
> +       struct page *page;
> +       struct swap_info_struct *si;
> +       unsigned long haddr = vmf->address & HPAGE_PMD_MASK;
> +       softleaf_t entry;
> +       swp_entry_t swp_entry;
> +       pmd_t pmd;
> +       vm_fault_t ret = 0;
> +       bool exclusive;
> +       rmap_t rmap_flags = RMAP_NONE;
> +       enum swap_pmd_cache cache_state;
> +
> +       entry = softleaf_from_pmd(vmf->orig_pmd);
> +       if (unlikely(!softleaf_is_swap(entry)))
> +               return 0;
> +
> +       swp_entry = entry;
> +
> +       /* Prevent swapoff from happening to us. */
> +       si = get_swap_device(swp_entry);
> +       if (unlikely(!si))
> +               return 0;
> +
> +       cache_state = swap_pmd_cache_lookup(swp_entry, &folio);
> +       if (cache_state == SWAP_PMD_CACHE_SPLIT)
> +               goto split_fallback;
> +       if (!folio) {
> +               /*
> +                * PMD swap entries encode ordinary per-page swap slots. If any
> +                * slot is in zswap, split and let the PTE swap path load the
> +                * range per page. Otherwise the range is all on disk and can be
> +                * read back as one PMD-sized folio.
> +                */
> +               if (zswap_range_has_entry(swp_entry, HPAGE_PMD_NR))
> +                       goto split_fallback;
> +
> +               folio = swapin_sync(swp_entry, GFP_HIGHUSER_MOVABLE,
> +                                   BIT(HPAGE_PMD_ORDER), vmf, NULL, 0);

Hmm, so here we always do a sync swapin which means there is no
readahead, kind of make sense since PMD is huge, readahead might be
not helpful and a order 0 readahead breaks following PMD. But maybe
worth mention it in the commit message at least so it's easiler to
follow. (partly because the name swapin_sync is a bit confusing at the
moment, we could change that name as we improve swap readahead later).

And just an idea that will it be suffient if we just use the return
value of swapin_sync directly? If a smaller folio covers the PMD start
entry, the smaller folio is returned and you catched that just fine.
If a smaller folio covers a follow-up sub-entry or a sub-entry is
freed, -EBUSY is returned here and you just do the split. And it
doesn't do the allocation for the first lookup check, so there should
be no thrashing issue.

> +               if (IS_ERR_OR_NULL(folio))
> +                       goto split_fallback;
> +
> +               /* Had to read from swap area: Major fault */
> +               ret = VM_FAULT_MAJOR;
> +               count_vm_event(PGMAJFAULT);
> +               count_memcg_event_mm(mm, PGMAJFAULT);
> +       }
> +
> +       ret |= folio_lock_or_retry(folio, vmf);

With PMD sized folio on a slower device, this try lock could fail much
more frequently, leading to repeated faults. That's not a new issue
and I remember Barry talked about it at the LSFMM this year, just I'm
not sure if that will get much worse for PMD?