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

[Usama Arif]

On 06/07/2026 18:56, Kairui Song wrote:
> 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.

Hi Kairui,

Thanks! and Thanks for the review comments!

> 
>> ---
>>  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).

Yes, I’ll add this to the commit message. Thanks!

Yes readahead won't be useful for PMD as you said.

> 
> 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.

The main issue over here is zswap.
zswap still stores the range as order-0 objects. If we call PMD-order
swapin_sync() when the cache is empty but any slot is in zswap, zswap_load()
fails the large-folio load and leaves a non-uptodate PMD-sized folio in
swapcache. Splitting the PMD after that is not enough unless we also delete that
failed large folio, because the subsequent PTE fault can find the non-uptodate
large folio and SIGBUS instead of loading the individual zswap entries.


> 
>> +               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?

Agreed that PMD-sized IO can keep the folio locked longer, especially on
slow swap devices. (Which hopefully are dying away :))
The current behavior matches do_swap_page(): on the normal retryable fault
path it drops the fault lock and waits for the folio unlock before returning
VM_FAULT_RETRY, so it should not be a tight retry loop.
The wait can be longer for 2MB IO, but the retry usually comes back to
either an uptodate folio or a PMD that another fault already  installed.
I think this is the right baseline behavior, and any improvement here
belongs with the broader swap readahead / async swapin work.