Re: [PATCH v2 3/7] mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()

[Karim Manaouil <kmanaouil.dev@gmail.com>]

Hi Shivank,

On Thu, Aug 13, 2026 at 04:23:14AM +0000, Shivank Garg wrote:
> remove_migration_pte() interleaves hugetlb handling with the regular
> folio path. Move it into a dedicated callback selected by
> remove_migration_ptes(), leaving the generic callback focused on regular
> folios ahead of PTE batching.
> 
> With hugetlb folios routed separately, simplify the PMD mapping check to
> a one-time warning.
> 
> Signed-off-by: Shivank Garg <shivankg@amd.com>
> ---
>  mm/migrate.c | 97 +++++++++++++++++++++++++++++++++++++++---------------------
>  1 file changed, 63 insertions(+), 34 deletions(-)
> 
> diff --git a/mm/migrate.c b/mm/migrate.c
> index a3362cc9ef66..ee1b8a55a2a4 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -375,6 +375,57 @@ static pte_t migration_entry_to_pte(struct folio *folio, struct page *new,
>  	return pte;
>  }
>  
> +/*
> + * Restore a potential migration pte to a working pte entry for hugetlb folios.
> + */
> +#ifdef CONFIG_HUGETLB_PAGE
> +static bool remove_migration_pte_hugetlb(struct folio *folio,
> +		struct vm_area_struct *vma, unsigned long addr, void *arg)
> +{
> +	struct rmap_walk_arg *rmap_walk_arg = arg;
> +	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
> +	struct hstate *h = hstate_vma(vma);
> +	unsigned int shift = huge_page_shift(h);
> +	unsigned long psize = huge_page_size(h);
> +	struct page *new = folio_page(folio, 0);
> +	rmap_t rmap_flags = RMAP_NONE;
> +	pte_t old_pte, pte;
> +	softleaf_t entry;
> +
> +	/* There is only a single mapping in a VMA. */

I don't think this comment is necessary. page_vma_mapped_walk() clearly
explains what happens with HugeTLB entries.

> +	if (!page_vma_mapped_walk(&pvmw))
> +		return true;
> +
> +	old_pte = huge_ptep_get(vma->vm_mm, pvmw.address, pvmw.pte);
> +	entry = softleaf_from_pte(old_pte);
> +	folio_get(folio);
> +	pte = migration_entry_to_pte(folio, new, entry, old_pte, vma, &rmap_flags);
> +	pte = arch_make_huge_pte(pte, shift, vma->vm_flags);
> +	if (folio_test_anon(folio))
> +		hugetlb_add_anon_rmap(folio, vma, pvmw.address, rmap_flags);
> +	else
> +		hugetlb_add_file_rmap(folio);
> +	set_huge_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte, psize);

I know nothing in HugeTLB, but can't they be mapped as 4KiB entries? In
which case we have to use

	while (page_vma_mapped_walk(&pvmw)) {
		...
	}

It seems like it's not the case from my little investigation, but
I'll keep it here just to be sure.

> +	if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
> +		mlock_drain_local();
> +
> +	trace_remove_migration_pte(pvmw.address, pte_val(pte), compound_order(new));
> +
> +	/* No need to invalidate - it was non-present before */
> +	update_mmu_cache(vma, pvmw.address, pvmw.pte);
> +	page_vma_mapped_walk_done(&pvmw);
> +
> +	return true;
> +}
> +#else
> +static bool remove_migration_pte_hugetlb(struct folio *folio,
> +		struct vm_area_struct *vma, unsigned long addr, void *arg)
> +{
> +	return false;
> +}
> +#endif /* CONFIG_HUGETLB_PAGE */
> +
>  /*
>   * Restore a potential migration pte to a working pte entry
>   */
> @@ -395,20 +446,14 @@ static bool remove_migration_pte(struct folio *folio,
>  #ifdef CONFIG_ARCH_HAS_PMD_SOFTLEAVES
>  		/* PMD-mapped THP migration entry */
>  		if (!pvmw.pte) {
> -			VM_BUG_ON_FOLIO(folio_test_hugetlb(folio) ||
> -					!folio_test_pmd_mappable(folio), folio);
> +			VM_WARN_ON_ONCE_FOLIO(!folio_test_pmd_mappable(folio), folio);
>  			remove_migration_pmd(&pvmw, folio);
>  			continue;
>  		}
>  #endif
> -		if (folio_test_hugetlb(folio))
> -			old_pte = huge_ptep_get(vma->vm_mm, pvmw.address,
> -						pvmw.pte);
> -		else
> -			old_pte = ptep_get(pvmw.pte);
> -
> +		old_pte = ptep_get(pvmw.pte);
>  		entry = softleaf_from_pte(old_pte);
> -		if (folio_test_large(folio) && !folio_test_hugetlb(folio))
> +		if (folio_test_large(folio))
>  			idx = softleaf_to_pfn(entry) - pvmw.pfn;
>  
>  		if (rmap_walk_arg->map_unused_to_zeropage &&
> @@ -434,30 +479,12 @@ static bool remove_migration_pte(struct folio *folio,
>  				pte = pte_swp_mkuffd(pte);
>  		}
>  
> -#ifdef CONFIG_HUGETLB_PAGE
> -		if (folio_test_hugetlb(folio)) {
> -			struct hstate *h = hstate_vma(vma);
> -			unsigned int shift = huge_page_shift(h);
> -			unsigned long psize = huge_page_size(h);
> -
> -			pte = arch_make_huge_pte(pte, shift, vma->vm_flags);
> -			if (folio_test_anon(folio))
> -				hugetlb_add_anon_rmap(folio, vma, pvmw.address,
> -						      rmap_flags);
> -			else
> -				hugetlb_add_file_rmap(folio);
> -			set_huge_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte,
> -					psize);
> -		} else
> -#endif
> -		{
> -			if (folio_test_anon(folio))
> -				folio_add_anon_rmap_pte(folio, new, vma,
> -							pvmw.address, rmap_flags);
> -			else
> -				folio_add_file_rmap_pte(folio, new, vma);
> -			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
> -		}
> +		if (folio_test_anon(folio))
> +			folio_add_anon_rmap_pte(folio, new, vma,
> +						pvmw.address, rmap_flags);
> +		else
> +			folio_add_file_rmap_pte(folio, new, vma);
> +		set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
>  		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
>  			mlock_drain_local();
>  
> @@ -484,7 +511,9 @@ void remove_migration_ptes(struct folio *src, struct folio *dst,
>  	};
>  
>  	struct rmap_walk_control rwc = {
> -		.rmap_one = remove_migration_pte,
> +		.rmap_one = folio_test_hugetlb(src) ?
> +				remove_migration_pte_hugetlb :
> +				remove_migration_pte,
>  		.arg = &rmap_walk_arg,
>  	};
>  
> 
> -- 
> 2.43.0
> 

-- 
~karim