Re: [PATCH v2] mm/migrate_device: consolidate compound folio handling

[Balbir Singh <balbirs@nvidia.com>]

On 9/12/26 1:44 PM, Hui Su wrote:
> Commit dc41e961a269 ("mm/migrate_device: avoid out-of-bounds writes for
> compound folios") added handling for compound folios that do not fit in
> the remaining PFN array.
> 
> migrate_device_range() and migrate_device_pfns() duplicate the logic for
> locking device PFNs, encoding compound folios, and handling this boundary
> condition.
> 
> A compound folio cannot be represented partially for migration. Warn when
> one does not fit in the remaining PFN array, while retaining the existing
> defensive handling: release any lock and reference acquired for the
> current folio, clear the remaining entries, and stop collecting.
> 
> Move the shared collection and encoding logic into a helper so both
> interfaces handle compound folios consistently. Also use memset() for
> the compound-folio tail entries instead of open-coding the clearing loop.
> 
> Document that an encountered compound folio must fit entirely in the
> remaining range or PFN array.
> 
> Link: https://lore.kernel.org/r/c99ca53a-73ef-4a0c-8738-eba1cc89bea2@kernel.org
> Suggested-by: David Hildenbrand <david@kernel.org>
> Signed-off-by: Hui Su <sh_def@163.com>
> ---
> 
> Notes:
>     Changes in v2:
>     - Fix the helper parameter indentation and keep the declaration to two
>       lines, as suggested by David Hildenbrand.
> 
>  mm/migrate_device.c | 89 ++++++++++++++++++++++++---------------------
>  1 file changed, 48 insertions(+), 41 deletions(-)
> 
> diff --git a/mm/migrate_device.c b/mm/migrate_device.c
> index 009bfa8b212d..ace543fd6946 100644
> --- a/mm/migrate_device.c
> +++ b/mm/migrate_device.c
> @@ -1392,6 +1392,38 @@ static unsigned long migrate_device_pfn_lock(unsigned long pfn)
>  	return migrate_pfn(pfn) | MIGRATE_PFN_MIGRATE;
>  }
>  
> +/*
> + * Collect a device folio into the page-granular PFN array.
> + *
> + * Return the number of entries consumed, or 0 if the folio does not fit in
> + * the remaining array.
> + */
> +static unsigned int migrate_device_collect_folio(unsigned long *src_pfn,
> +		unsigned long pfn, unsigned long remaining)
> +{
> +	struct folio *folio = page_folio(pfn_to_page(pfn));
> +	unsigned int nr;
> +
> +	*src_pfn = migrate_device_pfn_lock(pfn);

Don't we want to check for *src_pfn == 0?

> +	nr = folio_nr_pages(folio);
> +
> +	if (WARN_ON_ONCE(nr > remaining)) {

Can we please change this to VM_WARN_ON_ONCE?

> +		if (*src_pfn & MIGRATE_PFN_MIGRATE) {
> +			folio_unlock(folio);
> +			folio_put(folio);
> +		}
> +		memset(src_pfn, 0, remaining * sizeof(*src_pfn));
> +		return 0;
> +	}
> +
> +	if (nr > 1) {
> +		*src_pfn |= MIGRATE_PFN_COMPOUND;
> +		memset(src_pfn + 1, 0, (nr - 1) * sizeof(*src_pfn));
> +	}
> +
> +	return nr;
> +}
> +
>  /**
>   * migrate_device_range() - migrate device private pfns to normal memory.
>   * @src_pfns: array large enough to hold migrating source device private pfns.
> @@ -1410,35 +1442,22 @@ static unsigned long migrate_device_pfn_lock(unsigned long pfn)
>   * migrating pages that aren't free before unmapping them. Drivers may then
>   * allocate destination pages and start copying data from the device to CPU
>   * memory before calling migrate_device_pages().
> + *
> + * A compound folio must fit entirely in the remaining range.
>   */
>  int migrate_device_range(unsigned long *src_pfns, unsigned long start,
>  			unsigned long npages)
>  {
> -	unsigned long i, j, pfn;
> +	unsigned long i, pfn;
>  
>  	for (pfn = start, i = 0; i < npages; pfn++, i++) {
> -		struct page *page = pfn_to_page(pfn);
> -		struct folio *folio = page_folio(page);
> -		unsigned int nr = 1;
> +		unsigned int nr;
>  
> -		src_pfns[i] = migrate_device_pfn_lock(pfn);
> -		nr = folio_nr_pages(folio);
> -		if (nr > npages - i) {
> -			if (src_pfns[i] & MIGRATE_PFN_MIGRATE) {
> -				folio_unlock(folio);
> -				folio_put(folio);
> -			}
> -			memset(&src_pfns[i], 0,
> -			       (npages - i) * sizeof(*src_pfns));
> +		nr = migrate_device_collect_folio(&src_pfns[i], pfn, npages - i);
> +		if (!nr)
>  			break;
> -		}
> -		if (nr > 1) {
> -			src_pfns[i] |= MIGRATE_PFN_COMPOUND;
> -			for (j = 1; j < nr; j++)
> -				src_pfns[i+j] = 0;
> -			i += j - 1;
> -			pfn += j - 1;
> -		}
> +		i += nr - 1;
> +		pfn += nr - 1;
>  	}
>  
>  	migrate_device_unmap(src_pfns, npages, NULL);
> @@ -1454,33 +1473,21 @@ EXPORT_SYMBOL(migrate_device_range);
>   *
>   * Similar to migrate_device_range() but supports non-contiguous pre-populated
>   * array of device pages to migrate.
> + *
> + * A compound folio must fit entirely in the remaining PFN array.
>   */
>  int migrate_device_pfns(unsigned long *src_pfns, unsigned long npages)
>  {
> -	unsigned long i, j;
> +	unsigned long i;
>  
>  	for (i = 0; i < npages; i++) {
> -		struct page *page = pfn_to_page(src_pfns[i]);
> -		struct folio *folio = page_folio(page);
> -		unsigned int nr = 1;
> +		unsigned long pfn = src_pfns[i];
> +		unsigned int nr;
>  
> -		src_pfns[i] = migrate_device_pfn_lock(src_pfns[i]);
> -		nr = folio_nr_pages(folio);
> -		if (nr > npages - i) {
> -			if (src_pfns[i] & MIGRATE_PFN_MIGRATE) {
> -				folio_unlock(folio);
> -				folio_put(folio);
> -			}
> -			memset(&src_pfns[i], 0,
> -			       (npages - i) * sizeof(*src_pfns));
> +		nr = migrate_device_collect_folio(&src_pfns[i], pfn, npages - i);
> +		if (!nr)
>  			break;
> -		}
> -		if (nr > 1) {
> -			src_pfns[i] |= MIGRATE_PFN_COMPOUND;
> -			for (j = 1; j < nr; j++)
> -				src_pfns[i+j] = 0;
> -			i += j - 1;
> -		}
> +		i += nr - 1;
>  	}
>  
>  	migrate_device_unmap(src_pfns, npages, NULL);
> 
> base-commit: df2908090cda368b01ff43709f51890076c56157


Otherwise, looks good to me

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