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

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

[Zi Yan]

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;
 

---
base-commit: 302dfbff3e73be2ba5723237c6303244ec27cebb
change-id: 20260707-fix-array-overflow-in-migrate_vma_collect_pmd-c504b92613de

Best regards,
-- 
Yan, Zi

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

[Balbir Singh]

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>

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

[David Hildenbrand (Arm)]

On 7/8/26 03:50, 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;

It's hard to express which feelings reading migrate_vma_collect_pmd() gives me,
haha :)


In case we split ... couldn't we just undo what we already did, before doing the
"goto again" ?

-- 
Cheers,

David

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

[Zi Yan]

On Wed Jul 8, 2026 at 4:43 AM EDT, David Hildenbrand (Arm) wrote:
> On 7/8/26 03:50, 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;
>
> It's hard to express which feelings reading migrate_vma_collect_pmd() gives me,
> haha :)

I guess I have the same feelings.
>
>
> In case we split ... couldn't we just undo what we already did, before doing the
> "goto again" ?

You mean we reset migrate->cpages and migrate->npages and restart from
the beginning? But it is not only that, since the code below also
changes PTEs into migration entries. We will need to revert them as
well.

Hmm, migrate_vma_collect() documents itself as "update the src array and
"takes a reference on the page" without mentioning changing PTEs. I
wonder why changing PTEs is necessary, since later migrate_vma_unmap()
also changes page table entries to migration entries, although
migrate_vma_unmap() changes all entries to a folio, whereas
migrate_vma_collect() only changes PTEs from the specified VMA.

-- 
Best Regards,
Yan, Zi

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

[Zi Yan]

On Wed Jul 8, 2026 at 10:35 AM EDT, Zi Yan wrote:
> On Wed Jul 8, 2026 at 4:43 AM EDT, David Hildenbrand (Arm) wrote:
>> On 7/8/26 03:50, 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;
>>
>> It's hard to express which feelings reading migrate_vma_collect_pmd() gives me,
>> haha :)
>
> I guess I have the same feelings.
>>
>>
>> In case we split ... couldn't we just undo what we already did, before doing the
>> "goto again" ?
>
> You mean we reset migrate->cpages and migrate->npages and restart from
> the beginning? But it is not only that, since the code below also
> changes PTEs into migration entries. We will need to revert them as
> well.
>
> Hmm, migrate_vma_collect() documents itself as "update the src array and
> "takes a reference on the page" without mentioning changing PTEs. I
> wonder why changing PTEs is necessary, since later migrate_vma_unmap()
> also changes page table entries to migration entries, although
> migrate_vma_unmap() changes all entries to a folio, whereas
> migrate_vma_collect() only changes PTEs from the specified VMA.

OK, it is an optimization[1] when migrate_vma*() was introduced. If
there is only one mapping, migrate_vma_collect() will set migration
entry immediately without waiting until migrate_vma_unmap(). Fun.

[1] Commit 8c3328f1f36a5 ("mm/migrate: migrate_vma() unmap page from vma while collecting pages")
-- 
Best Regards,
Yan, Zi

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

[Zi Yan]

On Tue Jul 7, 2026 at 9:50 PM EDT, 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(-)
>

Sashiko has a new concern:

https://sashiko.dev/#/patchset/20260707-fix-array-overflow-in-migrate_vma_collect_pmd-v1-1-ce3ff4627653@nvidia.com

Q: migrate_vma_collect_pmd() drops the page table lock without flushing
the TLB leave stale entries active? When an order-0 page is collected and
its present PTE becomes a migration entry without flushing TLBs, later
CPUs with the stale TLB can still write to the page, potentially leading
to data corruption.

Answer: No, migrate_vma_collect_pmd() adds a reference to the collected
pages, so these pages are not going away. In addition, unmapped is
incremented after a page collection and warrants a TLB flush before the
code exits migrate_vma_collect_pmd():

1. pte_offset_map_lock() fails, this patch adds a flush.
2. split fails, a flush is already in the code.
3. the whole range is processed, a flush is at the end of the function.


-- 
Best Regards,
Yan, Zi

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

[Zi Yan]

On Wed Jul 8, 2026 at 10:43 AM EDT, Zi Yan wrote:
> On Wed Jul 8, 2026 at 10:35 AM EDT, Zi Yan wrote:
>> On Wed Jul 8, 2026 at 4:43 AM EDT, David Hildenbrand (Arm) wrote:
>>> On 7/8/26 03:50, 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;
>>>
>>> It's hard to express which feelings reading migrate_vma_collect_pmd() gives me,
>>> haha :)
>>
>> I guess I have the same feelings.
>>>
>>>
>>> In case we split ... couldn't we just undo what we already did, before doing the
>>> "goto again" ?
>>
>> You mean we reset migrate->cpages and migrate->npages and restart from
>> the beginning? But it is not only that, since the code below also
>> changes PTEs into migration entries. We will need to revert them as
>> well.
>>
>> Hmm, migrate_vma_collect() documents itself as "update the src array and
>> "takes a reference on the page" without mentioning changing PTEs. I
>> wonder why changing PTEs is necessary, since later migrate_vma_unmap()
>> also changes page table entries to migration entries, although
>> migrate_vma_unmap() changes all entries to a folio, whereas
>> migrate_vma_collect() only changes PTEs from the specified VMA.
>
> OK, it is an optimization[1] when migrate_vma*() was introduced. If
> there is only one mapping, migrate_vma_collect() will set migration
> entry immediately without waiting until migrate_vma_unmap(). Fun.
>
> [1] Commit 8c3328f1f36a5 ("mm/migrate: migrate_vma() unmap page from vma while collecting pages")

Hi David,

Do you think this patch is good to get in? I would like to get ack from
you.

Thanks.

-- 
Best Regards,
Yan, Zi

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

[David Hildenbrand (Arm)]

On 7/10/26 13:30, Zi Yan wrote:
> On Wed Jul 8, 2026 at 10:43 AM EDT, Zi Yan wrote:
>> On Wed Jul 8, 2026 at 10:35 AM EDT, Zi Yan wrote:
>>>
>>> I guess I have the same feelings.
>>>
>>> You mean we reset migrate->cpages and migrate->npages and restart from
>>> the beginning? But it is not only that, since the code below also
>>> changes PTEs into migration entries. We will need to revert them as
>>> well.
>>>
>>> Hmm, migrate_vma_collect() documents itself as "update the src array and
>>> "takes a reference on the page" without mentioning changing PTEs. I
>>> wonder why changing PTEs is necessary, since later migrate_vma_unmap()
>>> also changes page table entries to migration entries, although
>>> migrate_vma_unmap() changes all entries to a folio, whereas
>>> migrate_vma_collect() only changes PTEs from the specified VMA.
>>
>> OK, it is an optimization[1] when migrate_vma*() was introduced. If
>> there is only one mapping, migrate_vma_collect() will set migration
>> entry immediately without waiting until migrate_vma_unmap(). Fun.
>>
>> [1] Commit 8c3328f1f36a5 ("mm/migrate: migrate_vma() unmap page from vma while collecting pages")
> 
> Hi David,
> 

Hi,

> Do you think this patch is good to get in? I would like to get ack from
> you.

I dislike that we fail collecting even though there is something (different)
mapped now.

In case we stumble over a PMD, there sure are no migration entries on the PTE
level anymore?

-- 
Cheers,

David

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

[Zi Yan]

On Fri Jul 10, 2026 at 7:56 AM EDT, David Hildenbrand (Arm) wrote:
> On 7/10/26 13:30, Zi Yan wrote:
>> On Wed Jul 8, 2026 at 10:43 AM EDT, Zi Yan wrote:
>>> On Wed Jul 8, 2026 at 10:35 AM EDT, Zi Yan wrote:
>>>>
>>>> I guess I have the same feelings.
>>>>
>>>> You mean we reset migrate->cpages and migrate->npages and restart from
>>>> the beginning? But it is not only that, since the code below also
>>>> changes PTEs into migration entries. We will need to revert them as
>>>> well.
>>>>
>>>> Hmm, migrate_vma_collect() documents itself as "update the src array and
>>>> "takes a reference on the page" without mentioning changing PTEs. I
>>>> wonder why changing PTEs is necessary, since later migrate_vma_unmap()
>>>> also changes page table entries to migration entries, although
>>>> migrate_vma_unmap() changes all entries to a folio, whereas
>>>> migrate_vma_collect() only changes PTEs from the specified VMA.
>>>
>>> OK, it is an optimization[1] when migrate_vma*() was introduced. If
>>> there is only one mapping, migrate_vma_collect() will set migration
>>> entry immediately without waiting until migrate_vma_unmap(). Fun.
>>>
>>> [1] Commit 8c3328f1f36a5 ("mm/migrate: migrate_vma() unmap page from vma while collecting pages")
>> 
>> Hi David,
>> 
>
> Hi,
>
>> Do you think this patch is good to get in? I would like to get ack from
>> you.
>
> I dislike that we fail collecting even though there is something (different)
> mapped now.

Not all the time. After a folio split and PTE lock is dropped, there are
different cases:

1. no page table change, the collection grabs the PTE lock and continues;

2. some PTE(s) is changed, the same as 1 and the collection continues;

3. the PMD pointing to the PTE page table is changed to either a leaf
   PMD or an invalid PMD, the collection cannot grab the PTE lock to
   work on the remaming PTEs, since they are gone. For the collected
   PFNs (they are mapped more than once) and folios with elevated
   refcount (they are mapped once and unmapped here as an optimization),
   they will be processed later in migrate_vma_unmap(). Since
   migrate_vma_collect() is best effort, there is no need to revert and
   try to recollect from the beginning (to get that possible large
   folio or skip).

BTW, recollection will be more feasible if migrate_vma_collect() does
not do unmap singly-mapped optimization, since in the case, no PTE is
changed, we just need to reset migrate->->cpages and migrate->npages and
restart from the beginning.

>
> In case we stumble over a PMD, there sure are no migration entries on the PTE
> level anymore?

Right. The PTE page table should be gone and the original PMD, pointing
to the PTE page table, becomes a leaf PMD or an invalid PMD, since
someone changed it when the PTE lock is dropped.

-- 
Best Regards,
Yan, Zi

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

[David Hildenbrand (Arm)]

On 7/10/26 17:35, Zi Yan wrote:
> On Fri Jul 10, 2026 at 7:56 AM EDT, David Hildenbrand (Arm) wrote:
>> On 7/10/26 13:30, Zi Yan wrote:
>>>
>>> Hi David,
>>>
>>
>> Hi,
>>
>>> Do you think this patch is good to get in? I would like to get ack from
>>> you.
>>
>> I dislike that we fail collecting even though there is something (different)
>> mapped now.
> 
> Not all the time. After a folio split and PTE lock is dropped, there are
> different cases:
> 
> 1. no page table change, the collection grabs the PTE lock and continues;

Yes.

> 
> 2. some PTE(s) is changed, the same as 1 and the collection continues;

Yes.

> 
> 3. the PMD pointing to the PTE page table is changed to either a leaf
>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
>    work on the remaming PTEs, since they are gone. For the collected
>    PFNs (they are mapped more than once) and folios with elevated
>    refcount (they are mapped once and unmapped here as an optimization),
>    they will be processed later in migrate_vma_unmap(). Since
>    migrate_vma_collect() is best effort, there is no need to revert and
>    try to recollect from the beginning (to get that possible large
>    folio or skip).

Yes.

> 
> BTW, recollection will be more feasible if migrate_vma_collect() does
> not do unmap singly-mapped optimization, since in the case, no PTE is
> changed, we just need to reset migrate->->cpages and migrate->npages and
> restart from the beginning.
> 
>>
>> In case we stumble over a PMD, there sure are no migration entries on the PTE
>> level anymore?
> 
> Right. The PTE page table should be gone and the original PMD, pointing
> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
> someone changed it when the PTE lock is dropped.

Right, so I guess we'd have to remember how far to rollback, and rollback would
mean resetting migrate->->cpages and migrate->npages.

And for the entries we rollback, we have to decide whether to folio_put() and
whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
that the code is rather "interesting".

-- 
Cheers,

David

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

[Zi Yan]

On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:

> On 7/10/26 17:35, Zi Yan wrote:
>> On Fri Jul 10, 2026 at 7:56 AM EDT, David Hildenbrand (Arm) wrote:
>>> On 7/10/26 13:30, Zi Yan wrote:
>>>>
>>>> Hi David,
>>>>
>>>
>>> Hi,
>>>
>>>> Do you think this patch is good to get in? I would like to get ack from
>>>> you.
>>>
>>> I dislike that we fail collecting even though there is something (different)
>>> mapped now.
>>
>> Not all the time. After a folio split and PTE lock is dropped, there are
>> different cases:
>>
>> 1. no page table change, the collection grabs the PTE lock and continues;
>
> Yes.
>
>>
>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
>
> Yes.
>
>>
>> 3. the PMD pointing to the PTE page table is changed to either a leaf
>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
>>    work on the remaming PTEs, since they are gone. For the collected
>>    PFNs (they are mapped more than once) and folios with elevated
>>    refcount (they are mapped once and unmapped here as an optimization),
>>    they will be processed later in migrate_vma_unmap(). Since
>>    migrate_vma_collect() is best effort, there is no need to revert and
>>    try to recollect from the beginning (to get that possible large
>>    folio or skip).
>
> Yes.
>
>>
>> BTW, recollection will be more feasible if migrate_vma_collect() does
>> not do unmap singly-mapped optimization, since in the case, no PTE is
>> changed, we just need to reset migrate->->cpages and migrate->npages and
>> restart from the beginning.
>>
>>>
>>> In case we stumble over a PMD, there sure are no migration entries on the PTE
>>> level anymore?
>>
>> Right. The PTE page table should be gone and the original PMD, pointing
>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
>> someone changed it when the PTE lock is dropped.
>
> Right, so I guess we'd have to remember how far to rollback, and rollback would
> mean resetting migrate->->cpages and migrate->npages.

For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
we just roll back to the start address and set migrate->cpages and migrate->npages
both to 0.

>
> And for the entries we rollback, we have to decide whether to folio_put() and
> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
> that the code is rather "interesting".

IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
only mapped once” optimization, so there is no need to roll back at all.
migrate_vma_collect() is best effort, so there is no hard requirement about
a folio has to be collected. But I am happy to be proven wrong if Balbir,
Alistair, and Matthew give any counterexample.


Best Regards,
Yan, Zi

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

[David Hildenbrand (Arm)]

On 7/13/26 17:20, Zi Yan wrote:
> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
> 
>> On 7/10/26 17:35, Zi Yan wrote:
>>>
>>> Not all the time. After a folio split and PTE lock is dropped, there are
>>> different cases:
>>>
>>> 1. no page table change, the collection grabs the PTE lock and continues;
>>
>> Yes.
>>
>>>
>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
>>
>> Yes.
>>
>>>
>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
>>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
>>>    work on the remaming PTEs, since they are gone. For the collected
>>>    PFNs (they are mapped more than once) and folios with elevated
>>>    refcount (they are mapped once and unmapped here as an optimization),
>>>    they will be processed later in migrate_vma_unmap(). Since
>>>    migrate_vma_collect() is best effort, there is no need to revert and
>>>    try to recollect from the beginning (to get that possible large
>>>    folio or skip).
>>
>> Yes.
>>
>>>
>>> BTW, recollection will be more feasible if migrate_vma_collect() does
>>> not do unmap singly-mapped optimization, since in the case, no PTE is
>>> changed, we just need to reset migrate->->cpages and migrate->npages and
>>> restart from the beginning.
>>>
>>>
>>> Right. The PTE page table should be gone and the original PMD, pointing
>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
>>> someone changed it when the PTE lock is dropped.
>>
>> Right, so I guess we'd have to remember how far to rollback, and rollback would
>> mean resetting migrate->->cpages and migrate->npages.
> 
> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
> we just roll back to the start address and set migrate->cpages and migrate->npages
> both to 0.

Can't we walk multiple PMDs and have data already in there?

>>
>> And for the entries we rollback, we have to decide whether to folio_put() and
>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
>> that the code is rather "interesting".
> 
> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
> only mapped once” optimization, so there is no need to roll back at all.

I guess most folios are only mapped once. Do we have any numbers on that?

-- 
Cheers,

David

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

[Zi Yan]

On 13 Jul 2026, at 11:25, David Hildenbrand (Arm) wrote:

> On 7/13/26 17:20, Zi Yan wrote:
>> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
>>
>>> On 7/10/26 17:35, Zi Yan wrote:
>>>>
>>>> Not all the time. After a folio split and PTE lock is dropped, there are
>>>> different cases:
>>>>
>>>> 1. no page table change, the collection grabs the PTE lock and continues;
>>>
>>> Yes.
>>>
>>>>
>>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
>>>
>>> Yes.
>>>
>>>>
>>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
>>>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
>>>>    work on the remaming PTEs, since they are gone. For the collected
>>>>    PFNs (they are mapped more than once) and folios with elevated
>>>>    refcount (they are mapped once and unmapped here as an optimization),
>>>>    they will be processed later in migrate_vma_unmap(). Since
>>>>    migrate_vma_collect() is best effort, there is no need to revert and
>>>>    try to recollect from the beginning (to get that possible large
>>>>    folio or skip).
>>>
>>> Yes.
>>>
>>>>
>>>> BTW, recollection will be more feasible if migrate_vma_collect() does
>>>> not do unmap singly-mapped optimization, since in the case, no PTE is
>>>> changed, we just need to reset migrate->->cpages and migrate->npages and
>>>> restart from the beginning.
>>>>
>>>>
>>>> Right. The PTE page table should be gone and the original PMD, pointing
>>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
>>>> someone changed it when the PTE lock is dropped.
>>>
>>> Right, so I guess we'd have to remember how far to rollback, and rollback would
>>> mean resetting migrate->->cpages and migrate->npages.
>>
>> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
>> we just roll back to the start address and set migrate->cpages and migrate->npages
>> both to 0.
>
> Can't we walk multiple PMDs and have data already in there?

migrate_vma_collect() is the pmd_entry of struct mm_walk_ops, so it only
walks a single PMD at a time.

>
>>>
>>> And for the entries we rollback, we have to decide whether to folio_put() and
>>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
>>> that the code is rather "interesting".
>>
>> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
>> only mapped once” optimization, so there is no need to roll back at all.
>
> I guess most folios are only mapped once. Do we have any numbers on that?

I will defer the question to Balbir, Alistair, and Matthew.

BTW, for this patch, since there is no bug report on it, we would like to
fix it properly without worrying about backport, right?

Best Regards,
Yan, Zi

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

[Alistair Popple]

On 2026-07-14 at 01:39 +1000, Zi Yan <ziy@nvidia.com> wrote...
> On 13 Jul 2026, at 11:25, David Hildenbrand (Arm) wrote:
> 
> > On 7/13/26 17:20, Zi Yan wrote:
> >> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
> >>
> >>> On 7/10/26 17:35, Zi Yan wrote:
> >>>>
> >>>> Not all the time. After a folio split and PTE lock is dropped, there are
> >>>> different cases:
> >>>>
> >>>> 1. no page table change, the collection grabs the PTE lock and continues;
> >>>
> >>> Yes.
> >>>
> >>>>
> >>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
> >>>
> >>> Yes.
> >>>
> >>>>
> >>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
> >>>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
> >>>>    work on the remaming PTEs, since they are gone. For the collected
> >>>>    PFNs (they are mapped more than once) and folios with elevated
> >>>>    refcount (they are mapped once and unmapped here as an optimization),
> >>>>    they will be processed later in migrate_vma_unmap(). Since
> >>>>    migrate_vma_collect() is best effort, there is no need to revert and
> >>>>    try to recollect from the beginning (to get that possible large
> >>>>    folio or skip).
> >>>
> >>> Yes.
> >>>
> >>>>
> >>>> BTW, recollection will be more feasible if migrate_vma_collect() does
> >>>> not do unmap singly-mapped optimization, since in the case, no PTE is
> >>>> changed, we just need to reset migrate->->cpages and migrate->npages and
> >>>> restart from the beginning.
> >>>>
> >>>>
> >>>> Right. The PTE page table should be gone and the original PMD, pointing
> >>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
> >>>> someone changed it when the PTE lock is dropped.
> >>>
> >>> Right, so I guess we'd have to remember how far to rollback, and rollback would
> >>> mean resetting migrate->->cpages and migrate->npages.
> >>
> >> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
> >> we just roll back to the start address and set migrate->cpages and migrate->npages
> >> both to 0.
> >
> > Can't we walk multiple PMDs and have data already in there?
> 
> migrate_vma_collect() is the pmd_entry of struct mm_walk_ops, so it only
> walks a single PMD at a time.
> 
> >
> >>>
> >>> And for the entries we rollback, we have to decide whether to folio_put() and
> >>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
> >>> that the code is rather "interesting".
> >>
> >> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
> >> only mapped once” optimization, so there is no need to roll back at all.
> >
> > I guess most folios are only mapped once. Do we have any numbers on that?
> 
> I will defer the question to Balbir, Alistair, and Matthew.

Unfortunately previous tests that I've done showed the optimisation is critical
to performance in the common case (ie. singly mapped folios).

I say unfortunately because I have similar feelings to others here about this
code :) I tested this a while ago because I was trying to clean up the collect
step by mostly removing it and replacing it with a hmm_range_fault() type thing,
but the optimisation proved too important due to unmap having to do another
rmap/page table walk.

I don't have the numbers on hand atm, but will recreate them just in case my
recollection is faulty.

 - Alistair

> BTW, for this patch, since there is no bug report on it, we would like to
> fix it properly without worrying about backport, right?
> 
> Best Regards,
> Yan, Zi

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

[Balbir Singh]

On 7/14/26 8:24 AM, Alistair Popple wrote:
> On 2026-07-14 at 01:39 +1000, Zi Yan <ziy@nvidia.com> wrote...
>> On 13 Jul 2026, at 11:25, David Hildenbrand (Arm) wrote:
>>
>>> On 7/13/26 17:20, Zi Yan wrote:
>>>> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
>>>>
>>>>> On 7/10/26 17:35, Zi Yan wrote:
>>>>>>
>>>>>> Not all the time. After a folio split and PTE lock is dropped, there are
>>>>>> different cases:
>>>>>>
>>>>>> 1. no page table change, the collection grabs the PTE lock and continues;
>>>>>
>>>>> Yes.
>>>>>
>>>>>>
>>>>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
>>>>>
>>>>> Yes.
>>>>>
>>>>>>
>>>>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
>>>>>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
>>>>>>    work on the remaming PTEs, since they are gone. For the collected
>>>>>>    PFNs (they are mapped more than once) and folios with elevated
>>>>>>    refcount (they are mapped once and unmapped here as an optimization),
>>>>>>    they will be processed later in migrate_vma_unmap(). Since
>>>>>>    migrate_vma_collect() is best effort, there is no need to revert and
>>>>>>    try to recollect from the beginning (to get that possible large
>>>>>>    folio or skip).
>>>>>
>>>>> Yes.
>>>>>
>>>>>>
>>>>>> BTW, recollection will be more feasible if migrate_vma_collect() does
>>>>>> not do unmap singly-mapped optimization, since in the case, no PTE is
>>>>>> changed, we just need to reset migrate->->cpages and migrate->npages and
>>>>>> restart from the beginning.
>>>>>>
>>>>>>
>>>>>> Right. The PTE page table should be gone and the original PMD, pointing
>>>>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
>>>>>> someone changed it when the PTE lock is dropped.
>>>>>
>>>>> Right, so I guess we'd have to remember how far to rollback, and rollback would
>>>>> mean resetting migrate->->cpages and migrate->npages.
>>>>
>>>> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
>>>> we just roll back to the start address and set migrate->cpages and migrate->npages
>>>> both to 0.
>>>
>>> Can't we walk multiple PMDs and have data already in there?
>>
>> migrate_vma_collect() is the pmd_entry of struct mm_walk_ops, so it only
>> walks a single PMD at a time.
>>
>>>
>>>>>
>>>>> And for the entries we rollback, we have to decide whether to folio_put() and
>>>>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
>>>>> that the code is rather "interesting".
>>>>
>>>> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
>>>> only mapped once” optimization, so there is no need to roll back at all.
>>>
>>> I guess most folios are only mapped once. Do we have any numbers on that?
>>
>> I will defer the question to Balbir, Alistair, and Matthew.
> 
> Unfortunately previous tests that I've done showed the optimisation is critical
> to performance in the common case (ie. singly mapped folios).
> 
> I say unfortunately because I have similar feelings to others here about this
> code :) I tested this a while ago because I was trying to clean up the collect
> step by mostly removing it and replacing it with a hmm_range_fault() type thing,
> but the optimisation proved too important due to unmap having to do another
> rmap/page table walk.
> 
> I don't have the numbers on hand atm, but will recreate them just in case my
> recollection is faulty.
> 

I had a patch to add tracepoints to the code and some of that instrumentation
included this data. A combination of the trace points + hmm_tests will help us
get some data to start with.

Balbir

>> BTW, for this patch, since there is no bug report on it, we would like to
>> fix it properly without worrying about backport, right?
>>

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

[Alistair Popple]

On 2026-07-14 at 10:10 +1000, Balbir Singh <balbirs@nvidia.com> wrote...
> On 7/14/26 8:24 AM, Alistair Popple wrote:
> > On 2026-07-14 at 01:39 +1000, Zi Yan <ziy@nvidia.com> wrote...
> >> On 13 Jul 2026, at 11:25, David Hildenbrand (Arm) wrote:
> >>
> >>> On 7/13/26 17:20, Zi Yan wrote:
> >>>> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
> >>>>
> >>>>> On 7/10/26 17:35, Zi Yan wrote:
> >>>>>>
> >>>>>> Not all the time. After a folio split and PTE lock is dropped, there are
> >>>>>> different cases:
> >>>>>>
> >>>>>> 1. no page table change, the collection grabs the PTE lock and continues;
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
> >>>>>>    PMD or an invalid PMD, the collection cannot grab the PTE lock to
> >>>>>>    work on the remaming PTEs, since they are gone. For the collected
> >>>>>>    PFNs (they are mapped more than once) and folios with elevated
> >>>>>>    refcount (they are mapped once and unmapped here as an optimization),
> >>>>>>    they will be processed later in migrate_vma_unmap(). Since
> >>>>>>    migrate_vma_collect() is best effort, there is no need to revert and
> >>>>>>    try to recollect from the beginning (to get that possible large
> >>>>>>    folio or skip).
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> BTW, recollection will be more feasible if migrate_vma_collect() does
> >>>>>> not do unmap singly-mapped optimization, since in the case, no PTE is
> >>>>>> changed, we just need to reset migrate->->cpages and migrate->npages and
> >>>>>> restart from the beginning.
> >>>>>>
> >>>>>>
> >>>>>> Right. The PTE page table should be gone and the original PMD, pointing
> >>>>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
> >>>>>> someone changed it when the PTE lock is dropped.
> >>>>>
> >>>>> Right, so I guess we'd have to remember how far to rollback, and rollback would
> >>>>> mean resetting migrate->->cpages and migrate->npages.
> >>>>
> >>>> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
> >>>> we just roll back to the start address and set migrate->cpages and migrate->npages
> >>>> both to 0.
> >>>
> >>> Can't we walk multiple PMDs and have data already in there?
> >>
> >> migrate_vma_collect() is the pmd_entry of struct mm_walk_ops, so it only
> >> walks a single PMD at a time.
> >>
> >>>
> >>>>>
> >>>>> And for the entries we rollback, we have to decide whether to folio_put() and
> >>>>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
> >>>>> that the code is rather "interesting".
> >>>>
> >>>> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
> >>>> only mapped once” optimization, so there is no need to roll back at all.
> >>>
> >>> I guess most folios are only mapped once. Do we have any numbers on that?
> >>
> >> I will defer the question to Balbir, Alistair, and Matthew.
> > 
> > Unfortunately previous tests that I've done showed the optimisation is critical
> > to performance in the common case (ie. singly mapped folios).
> > 
> > I say unfortunately because I have similar feelings to others here about this
> > code :) I tested this a while ago because I was trying to clean up the collect
> > step by mostly removing it and replacing it with a hmm_range_fault() type thing,
> > but the optimisation proved too important due to unmap having to do another
> > rmap/page table walk.
> > 
> > I don't have the numbers on hand atm, but will recreate them just in case my
> > recollection is faulty.
> > 
> 
> I had a patch to add tracepoints to the code and some of that instrumentation
> included this data. A combination of the trace points + hmm_tests will help us
> get some data to start with.

But hmm-test already has a nice bandwidth benchmark :) We just need to force
the slow path. That used to be easier before I simplified the trylock handling
but with the below patch applied I get quite a lot better perf from the
optimisation. Note the patch isn't entirely correct as it assumes any locked
page was locked by migrate_vma_collect_pmd but it's good enough for the purposes
here.

Optimised results:

HMM THP Migration Benchmark
---------------------------
System page size: 4096 bytes

=== Small Buffer (0.5 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.097 ms        | 0.120 ms        | 19.5%
Dev->Sys Migration   | 0.061 ms        | 0.062 ms        | 0.4%
S->D Throughput      | 5.05 GB/s      | 4.07 GB/s      | 24.2%
D->S Throughput      | 7.96 GB/s      | 7.92 GB/s      | 0.4%

=== Half THP Size (1.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.217 ms        | 0.212 ms        | -2.1%
Dev->Sys Migration   | 0.136 ms        | 0.138 ms        | 1.2%
S->D Throughput      | 4.51 GB/s      | 4.60 GB/s      | -2.1%
D->S Throughput      | 7.17 GB/s      | 7.09 GB/s      | 1.2%

=== Single THP Size (2.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.290 ms        | 0.506 ms        | 42.6%
Dev->Sys Migration   | 0.095 ms        | 0.283 ms        | 66.3%
S->D Throughput      | 6.72 GB/s      | 3.86 GB/s      | 74.3%
D->S Throughput      | 20.52 GB/s      | 6.91 GB/s      | 197.0%

=== Two THP Size (4.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.634 ms        | 1.026 ms        | 38.2%
Dev->Sys Migration   | 0.195 ms        | 0.567 ms        | 65.6%
S->D Throughput      | 6.16 GB/s      | 3.81 GB/s      | 61.8%
D->S Throughput      | 20.05 GB/s      | 6.89 GB/s      | 191.1%

=== Four THP Size (8.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 1.301 ms        | 2.593 ms        | 49.8%
Dev->Sys Migration   | 0.541 ms        | 1.444 ms        | 62.5%
S->D Throughput      | 6.01 GB/s      | 3.01 GB/s      | 99.3%
D->S Throughput      | 14.45 GB/s      | 5.41 GB/s      | 167.0%

=== Eight THP Size (16.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 3.853 ms        | 6.119 ms        | 37.0%
Dev->Sys Migration   | 1.731 ms        | 3.354 ms        | 48.4%
S->D Throughput      | 4.06 GB/s      | 2.55 GB/s      | 58.8%
D->S Throughput      | 9.03 GB/s      | 4.66 GB/s      | 93.8%

=== One twenty eight THP Size (256.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 70.022 ms        | 105.921 ms        | 33.9%
Dev->Sys Migration   | 31.175 ms        | 55.393 ms        | 43.7%
S->D Throughput      | 3.57 GB/s      | 2.36 GB/s      | 51.3%
D->S Throughput      | 8.02 GB/s      | 4.51 GB/s      | 77.7%


Non-optmisied results (ignore the THP column, I wasn't able to test the slow
path so easily there so just skipped that bit of the test):

=== Small Buffer (0.5 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 0.321 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 0.276 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.52 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.77 GB/s      | -100.0%

=== Half THP Size (1.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 0.660 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 0.571 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.48 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.71 GB/s      | -100.0%

=== Single THP Size (2.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 1.385 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 1.152 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.41 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.69 GB/s      | -100.0%

=== Two THP Size (4.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 2.863 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 2.323 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.36 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.68 GB/s      | -100.0%

=== Four THP Size (8.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 6.199 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 4.904 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.26 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.59 GB/s      | -100.0%

=== Eight THP Size (16.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 14.260 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 10.457 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.10 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.49 GB/s      | -100.0%

=== One twenty eight THP Size (256.0 MB) ===
                     | With THP        | Without THP     | Improvement
---------------------------------------------------------------------
Sys->Dev Migration   | 0.000 ms        | 234.215 ms        | 100.0%
Dev->Sys Migration   | 0.000 ms        | 167.061 ms        | 100.0%
S->D Throughput      | 0.00 GB/s      | 1.07 GB/s      | -100.0%
D->S Throughput      | 0.00 GB/s      | 1.50 GB/s      | -100.0%

---

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 554754eb26ff..9f5563af8cd9 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -397,7 +397,7 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
                 * optimisation to avoid walking the rmap later with
                 * try_to_migrate().
                 */
-               if (fault_folio == folio || folio_trylock(folio)) {
+               if (fault_folio == folio) {
                        bool anon_exclusive;
                        pte_t swp_pte;
 
@@ -466,8 +466,16 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
                        if (pte_present(pte))
                                unmapped++;
                } else {
-                       folio_put(folio);
-                       mpfn = 0;
+                       /*
+                        * Slow path: couldn't lock the folio. Leave the PTE
+                        * intact and collect it anyway; migrate_device_unmap()
+                        * will lock it and use try_to_migrate() to install the
+                        * migration entry via rmap walk.
+                        *
+                        * Keep the reference taken above for
+                        * migrate_device_unmap() to drop.
+                        */
+                       migrate->cpages++;
                }
 
 next:
@@ -605,6 +613,16 @@ static unsigned long migrate_device_unmap(unsigned long *src_pfns,
                        folio_put(folio);
                }
 
+               /*
+                * Fast-path pages (folio_trylock succeeded in collect) are
+                * already locked. Slow-path pages were not locked during
+                * collect; lock them now so try_to_migrate() and the finalize
+                * path can rely on the folio being locked. This applies to
+                * both regular and zone-device folios.
+                */
+               if (!folio_test_locked(folio))
+                       folio_lock(folio);
+
                if (folio_mapped(folio))
                        try_to_migrate(folio, 0);
 
diff --git a/tools/testing/selftests/mm/hmm-tests.c b/tools/testing/selftests/mm/hmm-tests.c
index e4c49699f3f7..3a673181b9ad 100644
--- a/tools/testing/selftests/mm/hmm-tests.c
+++ b/tools/testing/selftests/mm/hmm-tests.c
@@ -2936,8 +2936,8 @@ TEST_F_TIMEOUT(hmm, benchmark_thp_migration, 120)
                        break;
 
                /* Test with THP */
-               ASSERT_EQ(run_migration_benchmark(self->fd, 1, test_sizes[i],
-                                       iterations, &thp_results), 0);
+               // ASSERT_EQ(run_migration_benchmark(self->fd, 1, test_sizes[i],
+               //                      iterations, &thp_results), 0);
 
                /* Test without THP */
                ASSERT_EQ(run_migration_benchmark(self->fd, 0, test_sizes[i],