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

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

[Hui Su]

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);
+	nr = folio_nr_pages(folio);
+
+	if (WARN_ON_ONCE(nr > remaining)) {
+		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
-- 
2.55.0

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

[Balbir Singh]

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>

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

[Hui Su]

Hi Balbir,

Thanks for the review and the Ack.

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

Yes. I'll avoid setting MIGRATE_PFN_COMPOUND when
migrate_device_pfn_lock() fails.

I don't think we should return 1 immediately when *src_pfn is zero,
though. We should still consume and clear the slots corresponding to the
whole compound folio. Otherwise, the outer loop would advance by only one
entry and the next iteration could treat a tail page of the same folio as
a new source PFN.

So the plan is to keep the folio-sized slot accounting, but only encode
MIGRATE_PFN_COMPOUND when *src_pfn is non-zero, and only do the
unlock/put when the lock succeeded.

> Can we please change this to VM_WARN_ON_ONCE?

Yes, agreed. I'll switch it to VM_WARN_ON_ONCE() and keep the existing
defensive clear-and-stop behavior.

I'll send a v3 with both changes.

Thanks,
Hui