[PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[Wandun Chen]

From: Wandun Chen <chenwandun@lixiang.com>

Commit ab521b4142aa ("mm/rmap: improve mlock tracking for large folios")
mlocks a large folio when all its pages are added in one rmap operation.
However, when a PTE-mapped large folio is migrated, remove_migration_pte()
restores PTEs individually, passing nr_pages = 1 to the rmap helpers. As a
result, PG_mlocked is not set on the destination folio.

Reclaim can restore mlock in try_to_unmap(), but Mlocked in /proc/meminfo
remains undercounted until that scan occurs.

Restore mlock once the full mapping has been restored in a locked VMA
under one page table lock, rather than waiting for try_to_unmap() to
restore the mlock state.

Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
---
 mm/migrate.c | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/mm/migrate.c b/mm/migrate.c
index 94033c222e04..140e291d4587 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
 {
 	struct rmap_walk_arg *rmap_walk_arg = arg;
 	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
+	unsigned long nr_ptes = 0;
 
 	while (page_vma_mapped_walk(&pvmw)) {
 		rmap_t rmap_flags = RMAP_NONE;
@@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
 			else
 				folio_add_file_rmap_pte(folio, new, vma);
 			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
+
+			nr_ptes++;
 		}
-		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
+		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
+			if (folio_test_large(folio) &&
+			    nr_ptes == folio_nr_pages(folio) &&
+			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
+				mlock_vma_folio(folio, vma);
 			mlock_drain_local();
+		}
 
 		trace_remove_migration_pte(pvmw.address, pte_val(pte),
 					   compound_order(new));
-- 
2.43.0

Re: [PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[Wandun]

On 11/10/26 20:32, Wandun Chen wrote:
> From: Wandun Chen <chenwandun@lixiang.com>
> 
> Commit ab521b4142aa ("mm/rmap: improve mlock tracking for large folios")
> mlocks a large folio when all its pages are added in one rmap operation.
> However, when a PTE-mapped large folio is migrated, remove_migration_pte()
> restores PTEs individually, passing nr_pages = 1 to the rmap helpers. As a
> result, PG_mlocked is not set on the destination folio.
> 
> Reclaim can restore mlock in try_to_unmap(), but Mlocked in /proc/meminfo
> remains undercounted until that scan occurs.
> 
> Restore mlock once the full mapping has been restored in a locked VMA
> under one page table lock, rather than waiting for try_to_unmap() to
> restore the mlock state.

I forgot to include the v1-to-v2 changelog:

Changes in v2:
1. Reword the subject and commit message.
2. Drop the Fixes tag.

Best regards,
Wandun

> 
> Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
> ---
>  mm/migrate.c | 10 +++++++++-
>  1 file changed, 9 insertions(+), 1 deletion(-)
> 
> diff --git a/mm/migrate.c b/mm/migrate.c
> index 94033c222e04..140e291d4587 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
>  {
>  	struct rmap_walk_arg *rmap_walk_arg = arg;
>  	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
> +	unsigned long nr_ptes = 0;
>  
>  	while (page_vma_mapped_walk(&pvmw)) {
>  		rmap_t rmap_flags = RMAP_NONE;
> @@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
>  			else
>  				folio_add_file_rmap_pte(folio, new, vma);
>  			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
> +
> +			nr_ptes++;
>  		}
> -		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
> +		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
> +			if (folio_test_large(folio) &&
> +			    nr_ptes == folio_nr_pages(folio) &&
> +			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
> +				mlock_vma_folio(folio, vma);
>  			mlock_drain_local();
> +		}
>  
>  		trace_remove_migration_pte(pvmw.address, pte_val(pte),
>  					   compound_order(new));

Re: [PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[David Hildenbrand (Arm)]

On 11/10/26 13:32, Wandun Chen wrote:
> From: Wandun Chen <chenwandun@lixiang.com>
> 
> Commit ab521b4142aa ("mm/rmap: improve mlock tracking for large folios")
> mlocks a large folio when all its pages are added in one rmap operation.
> However, when a PTE-mapped large folio is migrated, remove_migration_pte()
> restores PTEs individually, passing nr_pages = 1 to the rmap helpers. As a
> result, PG_mlocked is not set on the destination folio.
> 
> Reclaim can restore mlock in try_to_unmap(), but Mlocked in /proc/meminfo
> remains undercounted until that scan occurs.
> 
> Restore mlock once the full mapping has been restored in a locked VMA
> under one page table lock, rather than waiting for try_to_unmap() to
> restore the mlock state.
> 
> Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
> ---
>  mm/migrate.c | 10 +++++++++-
>  1 file changed, 9 insertions(+), 1 deletion(-)
> 
> diff --git a/mm/migrate.c b/mm/migrate.c
> index 94033c222e04..140e291d4587 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
>  {
>  	struct rmap_walk_arg *rmap_walk_arg = arg;
>  	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
> +	unsigned long nr_ptes = 0;
>  
>  	while (page_vma_mapped_walk(&pvmw)) {
>  		rmap_t rmap_flags = RMAP_NONE;
> @@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
>  			else
>  				folio_add_file_rmap_pte(folio, new, vma);
>  			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
> +
> +			nr_ptes++;
>  		}
> -		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
> +		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
> +			if (folio_test_large(folio) &&
> +			    nr_ptes == folio_nr_pages(folio) &&
> +			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
> +				mlock_vma_folio(folio, vma);
>  			mlock_drain_local();
> +		}
>  
>  		trace_remove_migration_pte(pvmw.address, pte_val(pte),
>  					   compound_order(new));

I guess pte batching would be one way of handling it, but that requires more churn.

The spirit is similar as in try_to_unmap_one(), where we count the number of
ptes in one VMA. For consistency, I would call it "ptes".

@Kiryl, can you take a look?

BTW, why don't we care about that in try_to_migrate?

-- 
Cheers,

David

Re: [PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[Kiryl Shutsemau]

On Wed, Nov 11, 2026 at 09:00:10AM +0100, David Hildenbrand (Arm) wrote:
> On 11/10/26 13:32, Wandun Chen wrote:
> > From: Wandun Chen <chenwandun@lixiang.com>
> > 
> > Commit ab521b4142aa ("mm/rmap: improve mlock tracking for large folios")
> > mlocks a large folio when all its pages are added in one rmap operation.
> > However, when a PTE-mapped large folio is migrated, remove_migration_pte()
> > restores PTEs individually, passing nr_pages = 1 to the rmap helpers. As a
> > result, PG_mlocked is not set on the destination folio.
> > 
> > Reclaim can restore mlock in try_to_unmap(), but Mlocked in /proc/meminfo
> > remains undercounted until that scan occurs.
> > 
> > Restore mlock once the full mapping has been restored in a locked VMA
> > under one page table lock, rather than waiting for try_to_unmap() to
> > restore the mlock state.
> > 
> > Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
> > ---
> >  mm/migrate.c | 10 +++++++++-
> >  1 file changed, 9 insertions(+), 1 deletion(-)
> > 
> > diff --git a/mm/migrate.c b/mm/migrate.c
> > index 94033c222e04..140e291d4587 100644
> > --- a/mm/migrate.c
> > +++ b/mm/migrate.c
> > @@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
> >  {
> >  	struct rmap_walk_arg *rmap_walk_arg = arg;
> >  	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
> > +	unsigned long nr_ptes = 0;
> >  
> >  	while (page_vma_mapped_walk(&pvmw)) {
> >  		rmap_t rmap_flags = RMAP_NONE;
> > @@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
> >  			else
> >  				folio_add_file_rmap_pte(folio, new, vma);
> >  			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
> > +
> > +			nr_ptes++;
> >  		}
> > -		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
> > +		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
> > +			if (folio_test_large(folio) &&
> > +			    nr_ptes == folio_nr_pages(folio) &&
> > +			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
> > +				mlock_vma_folio(folio, vma);
> >  			mlock_drain_local();
> > +		}
> >  
> >  		trace_remove_migration_pte(pvmw.address, pte_val(pte),
> >  					   compound_order(new));
> 
> I guess pte batching would be one way of handling it, but that requires more churn.
> 
> The spirit is similar as in try_to_unmap_one(), where we count the number of
> ptes in one VMA. For consistency, I would call it "ptes".
> 
> @Kiryl, can you take a look?

Looks right to me. Agreed on "ptes".

With that adjusted:

Reviewed-by: Kiryl Shutsemau (Meta) <kas@kernel.org>

> 
> BTW, why don't we care about that in try_to_migrate?

try_to_migrate_one() unmaps the folio whether the VMA is VM_LOCKED or
not, so there is no mapping left to keep mlocked; removing the rmap
munlocks the source.

The mlock is rebuilt on the destination in remove_migration_ptes(),
which is the path this patch fixes.

-- 
  Kiryl Shutsemau / Kirill A. Shutemov

Re: [PATCH v2] mm/migrate: improve mlock tracking for PTE-mapped large folios

[Wandun]

On 11/12/26 03:01, Kiryl Shutsemau wrote:
> On Wed, Nov 11, 2026 at 09:00:10AM +0100, David Hildenbrand (Arm) wrote:
>> On 11/10/26 13:32, Wandun Chen wrote:
>>> From: Wandun Chen <chenwandun@lixiang.com>
>>>
>>> Commit ab521b4142aa ("mm/rmap: improve mlock tracking for large folios")
>>> mlocks a large folio when all its pages are added in one rmap operation.
>>> However, when a PTE-mapped large folio is migrated, remove_migration_pte()
>>> restores PTEs individually, passing nr_pages = 1 to the rmap helpers. As a
>>> result, PG_mlocked is not set on the destination folio.
>>>
>>> Reclaim can restore mlock in try_to_unmap(), but Mlocked in /proc/meminfo
>>> remains undercounted until that scan occurs.
>>>
>>> Restore mlock once the full mapping has been restored in a locked VMA
>>> under one page table lock, rather than waiting for try_to_unmap() to
>>> restore the mlock state.
>>>
>>> Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
>>> ---
>>>  mm/migrate.c | 10 +++++++++-
>>>  1 file changed, 9 insertions(+), 1 deletion(-)
>>>
>>> diff --git a/mm/migrate.c b/mm/migrate.c
>>> index 94033c222e04..140e291d4587 100644
>>> --- a/mm/migrate.c
>>> +++ b/mm/migrate.c
>>> @@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
>>>  {
>>>  	struct rmap_walk_arg *rmap_walk_arg = arg;
>>>  	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
>>> +	unsigned long nr_ptes = 0;
>>>  
>>>  	while (page_vma_mapped_walk(&pvmw)) {
>>>  		rmap_t rmap_flags = RMAP_NONE;
>>> @@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
>>>  			else
>>>  				folio_add_file_rmap_pte(folio, new, vma);
>>>  			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
>>> +
>>> +			nr_ptes++;
>>>  		}
>>> -		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
>>> +		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
>>> +			if (folio_test_large(folio) &&
>>> +			    nr_ptes == folio_nr_pages(folio) &&
>>> +			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
>>> +				mlock_vma_folio(folio, vma);
>>>  			mlock_drain_local();
>>> +		}
>>>  
>>>  		trace_remove_migration_pte(pvmw.address, pte_val(pte),
>>>  					   compound_order(new));
>>
>> I guess pte batching would be one way of handling it, but that requires more churn.
>>
>> The spirit is similar as in try_to_unmap_one(), where we count the number of
>> ptes in one VMA. For consistency, I would call it "ptes".
>>
>> @Kiryl, can you take a look?
> 
> Looks right to me. Agreed on "ptes".
> 
> With that adjusted:
> 
> Reviewed-by: Kiryl Shutsemau (Meta) <kas@kernel.org>

I'll rename it to ptes and send a v3 with Reviewed-by tag.

On pte batching: following the swap_pte_batch()/do_swap_page()
precedent, restore could scan consecutive migration entries with
matching flags and do one set_ptes() plus one folio_add_*_rmap_ptes();
a full-folio batch would even set PG_mlocked for free via the rmap
helpers.

However, try_to_migrate_one() creates migration entries per-PTE with
per-PTE flags, so restore batches will often be truncated by flag
changes. For such folios the ptes counting in this patch remains the
only way to restore mlock (it survives batching with a one-line
change: ptes += nr). So I'd like to land this first and do the pte
batching as a separate follow-up series.

Thanks,
Wandun


> 
>>
>> BTW, why don't we care about that in try_to_migrate?
> 
> try_to_migrate_one() unmaps the folio whether the VMA is VM_LOCKED or
> not, so there is no mapping left to keep mlocked; removing the rmap
> munlocks the source.
> 
> The mlock is rebuilt on the destination in remove_migration_ptes(),
> which is the path this patch fixes.
>