Re: [PATCH] mm/migrate_device: avoid overflowing migrate_vma collection arrays

[Balbir Singh <balbirs@nvidia.com>]

On Tue, Jul 07, 2026 at 09:50:20PM -0400, Zi Yan wrote:
> migrate_vma_collect_pmd() can drop pte lock to split a large folio and
> restart. But the code does not handle restart properly when pmd becomes
> huge or cleared. It can overflow migrate->dst and migrate->src arrays
> during the hole or skip collection. Fix it by:
> 1. avoiding migrate_vma_collect_huge_pmd() if some collection is done,
> 2. skipping the rest of the range if pmd no longer points to a pte page
>    table.
>
> Fixes: a30b48bf1b244 ("mm/migrate_device: implement THP migration of zone device pages")
> Reported-by: Sashiko <sashiko-bot@kernel.org>
> Closes: https://sashiko.dev/#/patchset/20260706111958.3649651-1-wangkefeng.wang@huawei.com
> Assisted-by: Claude:claude-opus-4-8
> Signed-off-by: Zi Yan <ziy@nvidia.com>
> ---
> The issue is spot by Sashiko during a patch[1] review as a pre-existing one.
> This patch has the minimal change. An alternative is to reset
> migrate->cpages and migrate->npages and restart the whole range from the
> beginning, but that also requires a restoration of no-longer-present PTEs.
> 
> Link: https://lore.kernel.org/all/20260706111958.3649651-1-wangkefeng.wang@huawei.com/ [1]
> ---
>  mm/migrate_device.c | 19 +++++++++++++++++--
>  1 file changed, 17 insertions(+), 2 deletions(-)
> 
> diff --git a/mm/migrate_device.c b/mm/migrate_device.c
> index 2fffeb1f99694..6ceb47ec1da24 100644
> --- a/mm/migrate_device.c
> +++ b/mm/migrate_device.c
> @@ -257,7 +257,12 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
>  	pte_t *ptep;
>  
>  again:
> -	if (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp)) {
> +	/*
> +	 * Only check pmd when addr is at start, namely no pte is collected.
> +	 * It avoids collecting the same address range [start, addr) twice
> +	 * and overflowing the collection arrays.
> +	 */
> +	if (addr == start && (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp))) {
>  		int ret = migrate_vma_collect_huge_pmd(pmdp, start, end, walk, fault_folio);
>  
>  		if (ret == -EAGAIN)
> @@ -267,8 +272,18 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
>  	}
>  
>  	ptep = pte_offset_map_lock(mm, pmdp, start, &ptl);
> -	if (!ptep)
> +	if (!ptep) {
> +		/*
> +		 * Skip the rest if pmd becomes huge or cleared. Flush if any
> +		 * pte is modified
> +		 */
> +		if (addr != start) {
> +			if (unmapped)
> +				flush_tlb_range(walk->vma, start, end);
> +			return migrate_vma_collect_skip(addr, end, walk);
> +		}
>  		goto again;
> +	}
>  	lazy_mmu_mode_enable();
>  	ptep += (addr - start) / PAGE_SIZE;
>  
> 

Looks good, Thanks!

Reviewed-by: Balbir Singh <balbirs@nvidia.com>