Re: [RFC 03/11] mm/thp: zone_device awareness in THP handling code

[Balbir Singh <balbirs@nvidia.com>]

On 7/9/25 00:10, David Hildenbrand wrote:
> On 06.03.25 05:42, Balbir Singh wrote:
>> Make THP handling code in the mm subsystem for THP pages
>> aware of zone device pages. Although the code is
>> designed to be generic when it comes to handling splitting
>> of pages, the code is designed to work for THP page sizes
>> corresponding to HPAGE_PMD_NR.
>>
>> Modify page_vma_mapped_walk() to return true when a zone
>> device huge entry is present, enabling try_to_migrate()
>> and other code migration paths to appropriately process the
>> entry
>>
>> pmd_pfn() does not work well with zone device entries, use
>> pfn_pmd_entry_to_swap() for checking and comparison as for
>> zone device entries.
>>
>> try_to_map_to_unused_zeropage() does not apply to zone device
>> entries, zone device entries are ignored in the call.
>>
>> Signed-off-by: Balbir Singh <balbirs@nvidia.com>
>> ---
>>   mm/huge_memory.c     | 151 +++++++++++++++++++++++++++++++------------
>>   mm/migrate.c         |   2 +
>>   mm/page_vma_mapped.c |  10 +++
>>   mm/rmap.c            |  19 +++++-
>>   4 files changed, 138 insertions(+), 44 deletions(-)
>>
>> diff --git a/mm/huge_memory.c b/mm/huge_memory.c
>> index 826bfe907017..d8e018d1bdbd 100644
>> --- a/mm/huge_memory.c
>> +++ b/mm/huge_memory.c
>> @@ -2247,10 +2247,17 @@ int zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
>>           } else if (thp_migration_supported()) {
>>               swp_entry_t entry;
>>   -            VM_BUG_ON(!is_pmd_migration_entry(orig_pmd));
>>               entry = pmd_to_swp_entry(orig_pmd);
>>               folio = pfn_swap_entry_folio(entry);
>>               flush_needed = 0;
>> +
>> +            VM_BUG_ON(!is_pmd_migration_entry(*pmd) &&
>> +                    !folio_is_device_private(folio));
> 
> Convert that to a VM_WARN_ON_ONCE() while you are at it.
> 

Ack

> But really, check that the *pmd* is as expected (device_pritavte entry), and not the folio after the effects.
> 
> Also, hiding all that under the thp_migration_supported() looks wrong.
> 
> Likely you must clean that up first, to have something that expresses that we support PMD swap entries or sth like that. Not just "migration entries".
> 

The logic for the check is

if (pmd_present()) {
   ...
} else if (thp_migration_supported()) {
  ...

}

PMD swap is supported for migration entries (and zone device private after these changes)



> 
>> +
>> +            if (folio_is_device_private(folio)) {
>> +                folio_remove_rmap_pmd(folio, folio_page(folio, 0), vma);
>> +                WARN_ON_ONCE(folio_mapcount(folio) < 0);
>> +            }
> 
> 
> zap_nonpresent_ptes() does
> 
> if (is_device_private_entry(entry)) {
>     ...
> } else if (is_migration_entry(entry)) {
>     ....
> }
> 
> Can we adjust the same way of foing things? (yes, we might want a thp_migration_supported() check somewhere)

Are you suggesting refactoring of the code to add zap_nonpresent_pmd()? There really isn't much to be
done for specifically for migration entries in this context

> 
>>           } else
>>               WARN_ONCE(1, "Non present huge pmd without pmd migration enabled!");
>>   @@ -2264,6 +2271,15 @@ int zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
>>                          -HPAGE_PMD_NR);
>>           }
>>   +        /*
>> +         * Do a folio put on zone device private pages after
>> +         * changes to mm_counter, because the folio_put() will
>> +         * clean folio->mapping and the folio_test_anon() check
>> +         * will not be usable.
>> +         */
>> +        if (folio_is_device_private(folio))
>> +            folio_put(folio);
>> +
>>           spin_unlock(ptl);
>>           if (flush_needed)
>>               tlb_remove_page_size(tlb, &folio->page, HPAGE_PMD_SIZE);
>> @@ -2392,7 +2408,8 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
>>           struct folio *folio = pfn_swap_entry_folio(entry);
>>           pmd_t newpmd;
>>   -        VM_BUG_ON(!is_pmd_migration_entry(*pmd));
>> +        VM_BUG_ON(!is_pmd_migration_entry(*pmd) &&
>> +              !folio_is_device_private(folio));
>>           if (is_writable_migration_entry(entry)) {
>>               /*
>>                * A protection check is difficult so
>> @@ -2405,9 +2422,11 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
>>               newpmd = swp_entry_to_pmd(entry);
>>               if (pmd_swp_soft_dirty(*pmd))
>>                   newpmd = pmd_swp_mksoft_dirty(newpmd);
>> -        } else {
>> +        } else if (is_writable_device_private_entry(entry)) {
>> +            newpmd = swp_entry_to_pmd(entry);
>> +            entry = make_device_exclusive_entry(swp_offset(entry));
>> +        } else
>>               newpmd = *pmd;
>> -        }
>>             if (uffd_wp)
>>               newpmd = pmd_swp_mkuffd_wp(newpmd);
>> @@ -2860,11 +2879,12 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
>>       struct page *page;
>>       pgtable_t pgtable;
>>       pmd_t old_pmd, _pmd;
>> -    bool young, write, soft_dirty, pmd_migration = false, uffd_wp = false;
>> -    bool anon_exclusive = false, dirty = false;
>> +    bool young, write, soft_dirty, uffd_wp = false;
>> +    bool anon_exclusive = false, dirty = false, present = false;
>>       unsigned long addr;
>>       pte_t *pte;
>>       int i;
>> +    swp_entry_t swp_entry;
>>         VM_BUG_ON(haddr & ~HPAGE_PMD_MASK);
>>       VM_BUG_ON_VMA(vma->vm_start > haddr, vma);
>> @@ -2918,20 +2938,25 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
>>           return __split_huge_zero_page_pmd(vma, haddr, pmd);
>>       }
>>   -    pmd_migration = is_pmd_migration_entry(*pmd);
>> -    if (unlikely(pmd_migration)) {
>> -        swp_entry_t entry;
>>   +    present = pmd_present(*pmd);
>> +    if (unlikely(!present)) {
>> +        swp_entry = pmd_to_swp_entry(*pmd);
>>           old_pmd = *pmd;
>> -        entry = pmd_to_swp_entry(old_pmd);
>> -        page = pfn_swap_entry_to_page(entry);
>> -        write = is_writable_migration_entry(entry);
>> +
>> +        folio = pfn_swap_entry_folio(swp_entry);
>> +        VM_BUG_ON(!is_migration_entry(swp_entry) &&
>> +                !is_device_private_entry(swp_entry));
>> +        page = pfn_swap_entry_to_page(swp_entry);
>> +        write = is_writable_migration_entry(swp_entry);
>> +
>>           if (PageAnon(page))
>> -            anon_exclusive = is_readable_exclusive_migration_entry(entry);
>> -        young = is_migration_entry_young(entry);
>> -        dirty = is_migration_entry_dirty(entry);
>> +            anon_exclusive =
>> +                is_readable_exclusive_migration_entry(swp_entry);
>>           soft_dirty = pmd_swp_soft_dirty(old_pmd);
>>           uffd_wp = pmd_swp_uffd_wp(old_pmd);
>> +        young = is_migration_entry_young(swp_entry);
>> +        dirty = is_migration_entry_dirty(swp_entry);
>>       } else {
>>           /*
>>            * Up to this point the pmd is present and huge and userland has
>> @@ -3015,30 +3040,45 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
>>        * Note that NUMA hinting access restrictions are not transferred to
>>        * avoid any possibility of altering permissions across VMAs.
>>        */
>> -    if (freeze || pmd_migration) {
>> +    if (freeze || !present) {
>>           for (i = 0, addr = haddr; i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE) {
>>               pte_t entry;
>> -            swp_entry_t swp_entry;
>> -
>> -            if (write)
>> -                swp_entry = make_writable_migration_entry(
>> -                            page_to_pfn(page + i));
>> -            else if (anon_exclusive)
>> -                swp_entry = make_readable_exclusive_migration_entry(
>> -                            page_to_pfn(page + i));
>> -            else
>> -                swp_entry = make_readable_migration_entry(
>> -                            page_to_pfn(page + i));
>> -            if (young)
>> -                swp_entry = make_migration_entry_young(swp_entry);
>> -            if (dirty)
>> -                swp_entry = make_migration_entry_dirty(swp_entry);
>> -            entry = swp_entry_to_pte(swp_entry);
>> -            if (soft_dirty)
>> -                entry = pte_swp_mksoft_dirty(entry);
>> -            if (uffd_wp)
>> -                entry = pte_swp_mkuffd_wp(entry);
>> -
>> +            if (freeze || is_migration_entry(swp_entry)) {
>> +                if (write)
>> +                    swp_entry = make_writable_migration_entry(
>> +                                page_to_pfn(page + i));
>> +                else if (anon_exclusive)
>> +                    swp_entry = make_readable_exclusive_migration_entry(
>> +                                page_to_pfn(page + i));
>> +                else
>> +                    swp_entry = make_readable_migration_entry(
>> +                                page_to_pfn(page + i));
>> +                if (young)
>> +                    swp_entry = make_migration_entry_young(swp_entry);
>> +                if (dirty)
>> +                    swp_entry = make_migration_entry_dirty(swp_entry);
>> +                entry = swp_entry_to_pte(swp_entry);
>> +                if (soft_dirty)
>> +                    entry = pte_swp_mksoft_dirty(entry);
>> +                if (uffd_wp)
>> +                    entry = pte_swp_mkuffd_wp(entry);
>> +            } else {
>> +                VM_BUG_ON(!is_device_private_entry(swp_entry));
>> +                if (write)
>> +                    swp_entry = make_writable_device_private_entry(
>> +                                page_to_pfn(page + i));
>> +                else if (anon_exclusive)
>> +                    swp_entry = make_device_exclusive_entry(
>> +                                page_to_pfn(page + i));
> 
> I am pretty sure this is wrong. You cannot suddenly mix in device-exclusive entries.
> 
> And now I am confused again how device-private, anon and GUP interact.
> 

:)

Yep, this is wrong, we don't need to anything specific for anon_exclusive since
device private pages cannot be pinned. Other reviews have pointed towards fixes
needed as well.

>> +                else
>> +                    swp_entry = make_readable_device_private_entry(
>> +                                page_to_pfn(page + i));
>> +                entry = swp_entry_to_pte(swp_entry);
>> +                if (soft_dirty)
>> +                    entry = pte_swp_mksoft_dirty(entry);
>> +                if (uffd_wp)
>> +                    entry = pte_swp_mkuffd_wp(entry);
>> +            }
>>               VM_WARN_ON(!pte_none(ptep_get(pte + i)));
>>               set_pte_at(mm, addr, pte + i, entry);
>>           }
>> @@ -3065,7 +3105,7 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
>>       }
>>       pte_unmap(pte);
>>   -    if (!pmd_migration)
>> +    if (present)
>>           folio_remove_rmap_pmd(folio, page, vma);
>>       if (freeze)
>>           put_page(page);
>> @@ -3077,6 +3117,7 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
>>   void split_huge_pmd_locked(struct vm_area_struct *vma, unsigned long address,
>>                  pmd_t *pmd, bool freeze, struct folio *folio)
>>   {
>> +    struct folio *pmd_folio;
>>       VM_WARN_ON_ONCE(folio && !folio_test_pmd_mappable(folio));
>>       VM_WARN_ON_ONCE(!IS_ALIGNED(address, HPAGE_PMD_SIZE));
>>       VM_WARN_ON_ONCE(folio && !folio_test_locked(folio));
>> @@ -3089,7 +3130,14 @@ void split_huge_pmd_locked(struct vm_area_struct *vma, unsigned long address,
>>        */
>>       if (pmd_trans_huge(*pmd) || pmd_devmap(*pmd) ||
>>           is_pmd_migration_entry(*pmd)) {
>> -        if (folio && folio != pmd_folio(*pmd))
>> +        if (folio && !pmd_present(*pmd)) {
>> +            swp_entry_t swp_entry = pmd_to_swp_entry(*pmd);
>> +
>> +            pmd_folio = page_folio(pfn_swap_entry_to_page(swp_entry));
>> +        } else {
>> +            pmd_folio = pmd_folio(*pmd);
>> +        }
>> +        if (folio && folio != pmd_folio)
>>               return;
>>           __split_huge_pmd_locked(vma, pmd, address, freeze);
>>       }
>> @@ -3581,11 +3629,16 @@ static int __split_unmapped_folio(struct folio *folio, int new_order,
>>                        folio_test_swapcache(origin_folio)) ?
>>                            folio_nr_pages(release) : 0));
>>   +            if (folio_is_device_private(release))
>> +                percpu_ref_get_many(&release->pgmap->ref,
>> +                            (1 << new_order) - 1);
>> +
>>               if (release == origin_folio)
>>                   continue;
>>   -            lru_add_page_tail(origin_folio, &release->page,
>> -                        lruvec, list);
>> +            if (!folio_is_device_private(origin_folio))
>> +                lru_add_page_tail(origin_folio, &release->page,
>> +                            lruvec, list);
>>                 /* Some pages can be beyond EOF: drop them from page cache */
>>               if (release->index >= end) {
>> @@ -4625,7 +4678,10 @@ int set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
>>           return 0;
>>         flush_cache_range(vma, address, address + HPAGE_PMD_SIZE);
>> -    pmdval = pmdp_invalidate(vma, address, pvmw->pmd);
>> +    if (!folio_is_device_private(folio))
>> +        pmdval = pmdp_invalidate(vma, address, pvmw->pmd);
>> +    else
>> +        pmdval = pmdp_huge_clear_flush(vma, address, pvmw->pmd);
> 
> Please handle this like we handle the PTE case -- checking for pmd_present() instead.
> 
> Avoid placing these nasty folio_is_device_private() all over the place where avoidable.
> 

Ack, I'll try and use the presence of pmd to create the migration entries as appropriate

>>         /* See folio_try_share_anon_rmap_pmd(): invalidate PMD first. */
>>       anon_exclusive = folio_test_anon(folio) && PageAnonExclusive(page);
>> @@ -4675,6 +4731,17 @@ void remove_migration_pmd(struct page_vma_mapped_walk *pvmw, struct page *new)
>>       entry = pmd_to_swp_entry(*pvmw->pmd);
>>       folio_get(folio);
>>       pmde = mk_huge_pmd(new, READ_ONCE(vma->vm_page_prot));
>> +
>> +    if (unlikely(folio_is_device_private(folio))) {
>> +        if (pmd_write(pmde))
>> +            entry = make_writable_device_private_entry(
>> +                            page_to_pfn(new));
>> +        else
>> +            entry = make_readable_device_private_entry(
>> +                            page_to_pfn(new));
>> +        pmde = swp_entry_to_pmd(entry);
>> +    }
>> +
>>       if (pmd_swp_soft_dirty(*pvmw->pmd))
>>           pmde = pmd_mksoft_dirty(pmde);
>>       if (is_writable_migration_entry(entry))
>> diff --git a/mm/migrate.c b/mm/migrate.c
>> index 59e39aaa74e7..0aa1bdb711c3 100644
>> --- a/mm/migrate.c
>> +++ b/mm/migrate.c
>> @@ -200,6 +200,8 @@ static bool try_to_map_unused_to_zeropage(struct page_vma_mapped_walk *pvmw,
>>         if (PageCompound(page))
>>           return false;
>> +    if (folio_is_device_private(folio))
>> +        return false;
> 
> Why is that check required when you are adding THP handling and there is a PageCompound check right there?
> 

Fair point, we might not need the check here for THP handling.

>>       VM_BUG_ON_PAGE(!PageAnon(page), page);
>>       VM_BUG_ON_PAGE(!PageLocked(page), page);
>>       VM_BUG_ON_PAGE(pte_present(*pvmw->pte), page);
>> diff --git a/mm/page_vma_mapped.c b/mm/page_vma_mapped.c
>> index e463c3be934a..5dd2e51477d3 100644
>> --- a/mm/page_vma_mapped.c
>> +++ b/mm/page_vma_mapped.c
>> @@ -278,6 +278,16 @@ bool page_vma_mapped_walk(struct page_vma_mapped_walk *pvmw)
>>                * cannot return prematurely, while zap_huge_pmd() has
>>                * cleared *pmd but not decremented compound_mapcount().
>>                */
>> +            swp_entry_t entry;
>> +
>> +            if (!thp_migration_supported())
>> +                return not_found(pvmw);
> 
> This check looks misplaced. We should follow the same model as check_pte().
> 
> Checking for THP migration support when you are actually caring about device-private entries is weird.
> 

The thp migration check is common to the pmd code checks even above the patched code, the
code checks for thp_migration() and PVMW_MIGRATION. If thp migration is not supported
is there any point in returning true?

> That is, I would expect something like
> 
> } else if (is_swap_pmd(pmde)) {
>     swp_entry_t entry;
> 
>     entry = pmd_to_swp_entry(pmde);
>     if (!is_device_private_entry(entry))
>         return false;
> 

I don't think the code above is correct, you'll notice that there is a specific race
that the code handles for the !pmd_present() case and zap_huge_pmd() when PVMW_SYNC is set

I get the idea you're driving towards, I'll see how I can refactor it better.


>     ...
> }
> 
>> +            entry = pmd_to_swp_entry(pmde);
>> +            if (is_device_private_entry(entry)) {
>> +                pvmw->ptl = pmd_lock(mm, pvmw->pmd);
>> +                return true;
>> +            }
>> +
>>               if ((pvmw->flags & PVMW_SYNC) &&
>>                   thp_vma_suitable_order(vma, pvmw->address,
>>                              PMD_ORDER) &&
>> diff --git a/mm/rmap.c b/mm/rmap.c
>> index 67bb273dfb80..67e99dc5f2ef 100644
>> --- a/mm/rmap.c
>> +++ b/mm/rmap.c
>> @@ -2326,8 +2326,23 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>   #ifdef CONFIG_ARCH_ENABLE_THP_MIGRATION
>>           /* PMD-mapped THP migration entry */
>>           if (!pvmw.pte) {
>> -            subpage = folio_page(folio,
>> -                pmd_pfn(*pvmw.pmd) - folio_pfn(folio));
>> +            /*
>> +             * Zone device private folios do not work well with
>> +             * pmd_pfn() on some architectures due to pte
>> +             * inversion.
>> +             */
>> +            if (folio_is_device_private(folio)) {
>> +                swp_entry_t entry = pmd_to_swp_entry(*pvmw.pmd);
>> +                unsigned long pfn = swp_offset_pfn(entry);
>> +
>> +                subpage = folio_page(folio, pfn
>> +                            - folio_pfn(folio));
>> +            } else {
>> +                subpage = folio_page(folio,
>> +                            pmd_pfn(*pvmw.pmd)
>> +                            - folio_pfn(folio));
>> +            }
>> +
> 
> 
> Please follow the same model we use for PTEs.
> 
> /*
>  * Handle PFN swap PMDs, such as device-exclusive ones, that
>  * actually map pages.
>  */
> if (likely(pmd_present(...))) {
> 
> }
> 
> 

Will refactor to check for pmd_present first

Balbir