[PATCH v2 0/3] Optimizations for khugepaged

[PATCH v2 0/3] Optimizations for khugepaged

[Dev Jain]

This patch series uses PTE batching to optimize two functions, and
also addresses a race condition.

---
@David I did not make the clear_full_ptes clean up since that will require
some more work than I originally thought :)

v1->v2:
 - Use for loop instead of do-while loop (Lorenzo)
 - Remove folio_test_large check since the subpage-check condition
   will imply that (Baolin)
 - Combine patch 1 and 2 into this series, add new patch 3

Dev Jain (3):
  khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE
    batching
  khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE
    batching
  khugepaged: Reduce race probability between migration and khugepaged

 mm/khugepaged.c | 77 ++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 57 insertions(+), 20 deletions(-)

-- 
2.30.2

[PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[Dev Jain]

Use PTE batching to optimize __collapse_huge_page_copy_succeeded().

On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
Then, calling ptep_clear() for every pte will cause a TLB flush for every
contpte block. Instead, clear_full_ptes() does a
contpte_try_unfold_partial() which will flush the TLB only for the (if any)
starting and ending contpte block, if they partially overlap with the range
khugepaged is looking at.

For all arches, there should be a benefit due to batching atomic operations
on mapcounts due to folio_remove_rmap_ptes().

No issues were observed with mm-selftests.

Signed-off-by: Dev Jain <dev.jain@arm.com>
---
 mm/khugepaged.c | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index d45d08b521f6..3944b112d452 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -700,12 +700,15 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
 						spinlock_t *ptl,
 						struct list_head *compound_pagelist)
 {
+	unsigned long end = address + HPAGE_PMD_SIZE;
 	struct folio *src, *tmp;
-	pte_t *_pte;
 	pte_t pteval;
+	pte_t *_pte;
+	int nr_ptes;
 
-	for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
-	     _pte++, address += PAGE_SIZE) {
+	for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
+	     address += nr_ptes * PAGE_SIZE) {
+		nr_ptes = 1;
 		pteval = ptep_get(_pte);
 		if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
 			add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
@@ -719,21 +722,33 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
 				ksm_might_unmap_zero_page(vma->vm_mm, pteval);
 			}
 		} else {
+			const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
+			int max_nr_ptes;
+
 			struct page *src_page = pte_page(pteval);
 
 			src = page_folio(src_page);
 			if (!folio_test_large(src))
 				release_pte_folio(src);
+
+			max_nr_ptes = (end - address) >> PAGE_SHIFT;
+			if (folio_test_large(src))
+				nr_ptes = folio_pte_batch(src, address, _pte,
+							  pteval, max_nr_ptes,
+							  flags, NULL, NULL, NULL);
+
 			/*
 			 * ptl mostly unnecessary, but preempt has to
 			 * be disabled to update the per-cpu stats
 			 * inside folio_remove_rmap_pte().
 			 */
 			spin_lock(ptl);
-			ptep_clear(vma->vm_mm, address, _pte);
-			folio_remove_rmap_pte(src, src_page, vma);
+			clear_full_ptes(vma->vm_mm, address, _pte, nr_ptes,
+					/* full = */ false);
+			folio_remove_rmap_ptes(src, src_page, nr_ptes, vma);
 			spin_unlock(ptl);
-			free_folio_and_swap_cache(src);
+			free_swap_cache(src);
+			folio_put_refs(src, nr_ptes);
 		}
 	}
 
-- 
2.30.2

[PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Dev Jain]

Use PTE batching to optimize collapse_pte_mapped_thp().

On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
Then, calling ptep_clear() for every pte will cause a TLB flush for every
contpte block. Instead, clear_full_ptes() does a
contpte_try_unfold_partial() which will flush the TLB only for the (if any)
starting and ending contpte block, if they partially overlap with the range
khugepaged is looking at.

For all arches, there should be a benefit due to batching atomic operations
on mapcounts due to folio_remove_rmap_ptes().

Note that we do not need to make a change to the check
"if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
pages of the folio will be equal to the corresponding pages of our
batch mapping consecutive pages.

No issues were observed with mm-selftests.

Signed-off-by: Dev Jain <dev.jain@arm.com>
---
 mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 3944b112d452..4c8d33abfbd8 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
 int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
 			    bool install_pmd)
 {
+	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
 	struct mmu_notifier_range range;
 	bool notified = false;
 	unsigned long haddr = addr & HPAGE_PMD_MASK;
+	unsigned long end = haddr + HPAGE_PMD_SIZE;
 	struct vm_area_struct *vma = vma_lookup(mm, haddr);
 	struct folio *folio;
 	pte_t *start_pte, *pte;
 	pmd_t *pmd, pgt_pmd;
 	spinlock_t *pml = NULL, *ptl;
-	int nr_ptes = 0, result = SCAN_FAIL;
 	int i;
 
 	mmap_assert_locked(mm);
@@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
 		goto abort;
 
 	/* step 2: clear page table and adjust rmap */
-	for (i = 0, addr = haddr, pte = start_pte;
-	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
+	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
+	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
+	     pte += nr_batch_ptes) {
+		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
+		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
+		struct folio *mapped_folio;
 		struct page *page;
 		pte_t ptent = ptep_get(pte);
 
+		nr_batch_ptes = 1;
+
 		if (pte_none(ptent))
 			continue;
 		/*
@@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
 			goto abort;
 		}
 		page = vm_normal_page(vma, addr, ptent);
+		mapped_folio = page_folio(page);
+
 		if (folio_page(folio, i) != page)
 			goto abort;
 
+		mapped_folio = page_folio(page);
+		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
+						max_nr_batch_ptes, flags,
+						NULL, NULL, NULL);
+
 		/*
 		 * Must clear entry, or a racing truncate may re-remove it.
 		 * TLB flush can be left until pmdp_collapse_flush() does it.
 		 * PTE dirty? Shmem page is already dirty; file is read-only.
 		 */
-		ptep_clear(mm, addr, pte);
-		folio_remove_rmap_pte(folio, page, vma);
-		nr_ptes++;
+		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
+		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
+		nr_mapped_ptes += nr_batch_ptes;
 	}
 
 	if (!pml)
 		spin_unlock(ptl);
 
 	/* step 3: set proper refcount and mm_counters. */
-	if (nr_ptes) {
-		folio_ref_sub(folio, nr_ptes);
-		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
+	if (nr_mapped_ptes) {
+		folio_ref_sub(folio, nr_mapped_ptes);
+		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
 	}
 
 	/* step 4: remove empty page table */
@@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
 			: SCAN_SUCCEED;
 	goto drop_folio;
 abort:
-	if (nr_ptes) {
+	if (nr_mapped_ptes) {
 		flush_tlb_mm(mm);
-		folio_ref_sub(folio, nr_ptes);
-		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
+		folio_ref_sub(folio, nr_mapped_ptes);
+		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
 	}
 unlock:
 	if (start_pte)
-- 
2.30.2

[PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Dev Jain]

Suppose a folio is under migration, and khugepaged is also trying to
collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
page cache via filemap_lock_folio(), thus taking a reference on the folio
and sleeping on the folio lock, since the lock is held by the migration
path. Migration will then fail in
__folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
such a race happening (leading to migration failure) by bailing out
if we detect a PMD is marked with a migration entry.

This fixes the migration-shared-anon-thp testcase failure on Apple M3.

Signed-off-by: Dev Jain <dev.jain@arm.com>
---
 mm/khugepaged.c | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 4c8d33abfbd8..bc8774f62e86 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -31,6 +31,7 @@ enum scan_result {
 	SCAN_FAIL,
 	SCAN_SUCCEED,
 	SCAN_PMD_NULL,
+	SCAN_PMD_MIGRATION,
 	SCAN_PMD_NONE,
 	SCAN_PMD_MAPPED,
 	SCAN_EXCEED_NONE_PTE,
@@ -956,6 +957,8 @@ static inline int check_pmd_state(pmd_t *pmd)
 
 	if (pmd_none(pmde))
 		return SCAN_PMD_NONE;
+	if (is_pmd_migration_entry(pmde))
+		return SCAN_PMD_MIGRATION;
 	if (!pmd_present(pmde))
 		return SCAN_PMD_NULL;
 	if (pmd_trans_huge(pmde))
@@ -1518,9 +1521,12 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
 	    !range_in_vma(vma, haddr, haddr + HPAGE_PMD_SIZE))
 		return SCAN_VMA_CHECK;
 
-	/* Fast check before locking page if already PMD-mapped */
+	/*
+	 * Fast check before locking folio if already PMD-mapped, or if the
+	 * folio is under migration
+	 */
 	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
-	if (result == SCAN_PMD_MAPPED)
+	if (result == SCAN_PMD_MAPPED || result == SCAN_PMD_MIGRATION)
 		return result;
 
 	/*
@@ -2745,6 +2751,7 @@ static int madvise_collapse_errno(enum scan_result r)
 	case SCAN_PAGE_LRU:
 	case SCAN_DEL_PAGE_LRU:
 	case SCAN_PAGE_FILLED:
+	case SCAN_PMD_MIGRATION:
 		return -EAGAIN;
 	/*
 	 * Other: Trying again likely not to succeed / error intrinsic to
@@ -2834,6 +2841,7 @@ int madvise_collapse(struct vm_area_struct *vma, struct vm_area_struct **prev,
 			goto handle_result;
 		/* Whitelisted set of results where continuing OK */
 		case SCAN_PMD_NULL:
+		case SCAN_PMD_MIGRATION:
 		case SCAN_PTE_NON_PRESENT:
 		case SCAN_PTE_UFFD_WP:
 		case SCAN_PAGE_RO:
-- 
2.30.2

Re: [PATCH v2 0/3] Optimizations for khugepaged

[Lorenzo Stoakes]

On Wed, Jun 25, 2025 at 11:28:03AM +0530, Dev Jain wrote:
> This patch series uses PTE batching to optimize two functions, and
> also addresses a race condition.

Kinda succinct, be nice to expand a little. I know it can end up being a
little duplicative, but I think more information is always better than less
(as long as not truly overwrought)

>
> ---
> @David I did not make the clear_full_ptes clean up since that will require
> some more work than I originally thought :)
>
> v1->v2:
>  - Use for loop instead of do-while loop (Lorenzo)
>  - Remove folio_test_large check since the subpage-check condition
>    will imply that (Baolin)
>  - Combine patch 1 and 2 into this series, add new patch 3

One small nit, it'd be really handy to include links to previous series to
help us find them :P

As this one combines a couple earlier (which is nice!) but handy just for
comparison's sake :>)

>
> Dev Jain (3):
>   khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE
>     batching
>   khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE
>     batching
>   khugepaged: Reduce race probability between migration and khugepaged
>
>  mm/khugepaged.c | 77 ++++++++++++++++++++++++++++++++++++-------------
>  1 file changed, 57 insertions(+), 20 deletions(-)
>
> --
> 2.30.2
>

Re: [PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[David Hildenbrand]

On 25.06.25 07:58, Dev Jain wrote:
> Use PTE batching to optimize __collapse_huge_page_copy_succeeded().
> 
> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
> Then, calling ptep_clear() for every pte will cause a TLB flush for every
> contpte block. Instead, clear_full_ptes() does a
> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
> starting and ending contpte block, if they partially overlap with the range
> khugepaged is looking at.
> 
> For all arches, there should be a benefit due to batching atomic operations
> on mapcounts due to folio_remove_rmap_ptes().
> 
> No issues were observed with mm-selftests.
> 
> Signed-off-by: Dev Jain <dev.jain@arm.com>
> ---
>   mm/khugepaged.c | 27 +++++++++++++++++++++------
>   1 file changed, 21 insertions(+), 6 deletions(-)
> 
> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> index d45d08b521f6..3944b112d452 100644
> --- a/mm/khugepaged.c
> +++ b/mm/khugepaged.c
> @@ -700,12 +700,15 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>   						spinlock_t *ptl,
>   						struct list_head *compound_pagelist)
>   {
> +	unsigned long end = address + HPAGE_PMD_SIZE;
>   	struct folio *src, *tmp;
> -	pte_t *_pte;
>   	pte_t pteval;
> +	pte_t *_pte;
> +	int nr_ptes;
>   
> -	for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
> -	     _pte++, address += PAGE_SIZE) {
> +	for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
> +	     address += nr_ptes * PAGE_SIZE) {
> +		nr_ptes = 1;
>   		pteval = ptep_get(_pte);
>   		if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
>   			add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
> @@ -719,21 +722,33 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>   				ksm_might_unmap_zero_page(vma->vm_mm, pteval);
>   			}
>   		} else {
> +			const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
> +			int max_nr_ptes;
> +
>   			struct page *src_page = pte_page(pteval);
>   
>   			src = page_folio(src_page);
>   			if (!folio_test_large(src))
>   				release_pte_folio(src);
> +
> +			max_nr_ptes = (end - address) >> PAGE_SHIFT;
> +			if (folio_test_large(src))
> +				nr_ptes = folio_pte_batch(src, address, _pte,
> +							  pteval, max_nr_ptes,
> +							  flags, NULL, NULL, NULL);
> +
>   			/*
>   			 * ptl mostly unnecessary, but preempt has to
>   			 * be disabled to update the per-cpu stats
>   			 * inside folio_remove_rmap_pte().
>   			 */
>   			spin_lock(ptl);
> -			ptep_clear(vma->vm_mm, address, _pte);
> -			folio_remove_rmap_pte(src, src_page, vma);
> +			clear_full_ptes(vma->vm_mm, address, _pte, nr_ptes,
> +					/* full = */ false);

Can you add this patch to your series if nobody objects and use clear_ptes()
instead?

 From 95e20ab0ff62bbbdcd89898c9d76fdc1ea961257 Mon Sep 17 00:00:00 2001
From: David Hildenbrand <david@redhat.com>
Date: Wed, 25 Jun 2025 12:55:20 +0200
Subject: [PATCH] mm: add get_and_clear_ptes() and clear_ptes()

Let's add variants to be used where "full" does not apply -- which will
be the majority of cases in the future. "full" really only applies if
we are about to tear down a full MM.

Use get_and_clear_ptes() in existing code, clear_ptes() users will
be added next.

Should we make these inline functions instead and add separate docs?
Probably not worth it for now.

Signed-off-by: David Hildenbrand <david@redhat.com>
---
  include/linux/pgtable.h | 6 ++++++
  mm/mremap.c             | 2 +-
  mm/rmap.c               | 2 +-
  3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index cf1515c163e26..28679254b4f65 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -736,6 +736,9 @@ static inline pte_t get_and_clear_full_ptes(struct mm_struct *mm,
  }
  #endif
  
+#define get_and_clear_ptes(_mm, _addr, _ptep, _nr) \
+	get_and_clear_full_ptes(_mm, _addr, _ptep, _nr, 0)
+
  #ifndef clear_full_ptes
  /**
   * clear_full_ptes - Clear present PTEs that map consecutive pages of the same
@@ -768,6 +771,9 @@ static inline void clear_full_ptes(struct mm_struct *mm, unsigned long addr,
  }
  #endif
  
+#define clear_ptes(_mm, _addr, _ptep, _nr) \
+	clear_full_ptes(_mm, _addr, _ptep, _nr, 0)
+
  /*
   * If two threads concurrently fault at the same page, the thread that
   * won the race updates the PTE and its local TLB/Cache. The other thread
diff --git a/mm/mremap.c b/mm/mremap.c
index b31740f77b840..92890f8367574 100644
--- a/mm/mremap.c
+++ b/mm/mremap.c
@@ -322,7 +322,7 @@ static int move_ptes(struct pagetable_move_control *pmc,
  							 old_pte, max_nr_ptes);
  			force_flush = true;
  		}
-		pte = get_and_clear_full_ptes(mm, old_addr, old_ptep, nr_ptes, 0);
+		pte = get_and_clear_ptes(mm, old_addr, old_ptep, nr_ptes);
  		pte = move_pte(pte, old_addr, new_addr);
  		pte = move_soft_dirty_pte(pte);
  
diff --git a/mm/rmap.c b/mm/rmap.c
index 3b74bb19c11dd..8200d705fe4ac 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2031,7 +2031,7 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
  			flush_cache_range(vma, address, end_addr);
  
  			/* Nuke the page table entry. */
-			pteval = get_and_clear_full_ptes(mm, address, pvmw.pte, nr_pages, 0);
+			pteval = get_and_clear_ptes(mm, address, pvmw.pte, nr_pages);
  			/*
  			 * We clear the PTE but do not flush so potentially
  			 * a remote CPU could still be writing to the folio.
-- 
2.49.0


-- 
Cheers,

David / dhildenb

Re: [PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[Dev Jain]

On 25/06/25 4:44 pm, David Hildenbrand wrote:
> On 25.06.25 07:58, Dev Jain wrote:
>> Use PTE batching to optimize __collapse_huge_page_copy_succeeded().
>>
>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for 
>> collapse.
>> Then, calling ptep_clear() for every pte will cause a TLB flush for 
>> every
>> contpte block. Instead, clear_full_ptes() does a
>> contpte_try_unfold_partial() which will flush the TLB only for the 
>> (if any)
>> starting and ending contpte block, if they partially overlap with the 
>> range
>> khugepaged is looking at.
>>
>> For all arches, there should be a benefit due to batching atomic 
>> operations
>> on mapcounts due to folio_remove_rmap_ptes().
>>
>> No issues were observed with mm-selftests.
>>
>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>> ---
>>   mm/khugepaged.c | 27 +++++++++++++++++++++------
>>   1 file changed, 21 insertions(+), 6 deletions(-)
>>
>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>> index d45d08b521f6..3944b112d452 100644
>> --- a/mm/khugepaged.c
>> +++ b/mm/khugepaged.c
>> @@ -700,12 +700,15 @@ static void 
>> __collapse_huge_page_copy_succeeded(pte_t *pte,
>>                           spinlock_t *ptl,
>>                           struct list_head *compound_pagelist)
>>   {
>> +    unsigned long end = address + HPAGE_PMD_SIZE;
>>       struct folio *src, *tmp;
>> -    pte_t *_pte;
>>       pte_t pteval;
>> +    pte_t *_pte;
>> +    int nr_ptes;
>>   -    for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
>> -         _pte++, address += PAGE_SIZE) {
>> +    for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
>> +         address += nr_ptes * PAGE_SIZE) {
>> +        nr_ptes = 1;
>>           pteval = ptep_get(_pte);
>>           if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
>>               add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
>> @@ -719,21 +722,33 @@ static void 
>> __collapse_huge_page_copy_succeeded(pte_t *pte,
>>                   ksm_might_unmap_zero_page(vma->vm_mm, pteval);
>>               }
>>           } else {
>> +            const fpb_t flags = FPB_IGNORE_DIRTY | 
>> FPB_IGNORE_SOFT_DIRTY;
>> +            int max_nr_ptes;
>> +
>>               struct page *src_page = pte_page(pteval);
>>                 src = page_folio(src_page);
>>               if (!folio_test_large(src))
>>                   release_pte_folio(src);
>> +
>> +            max_nr_ptes = (end - address) >> PAGE_SHIFT;
>> +            if (folio_test_large(src))
>> +                nr_ptes = folio_pte_batch(src, address, _pte,
>> +                              pteval, max_nr_ptes,
>> +                              flags, NULL, NULL, NULL);
>> +
>>               /*
>>                * ptl mostly unnecessary, but preempt has to
>>                * be disabled to update the per-cpu stats
>>                * inside folio_remove_rmap_pte().
>>                */
>>               spin_lock(ptl);
>> -            ptep_clear(vma->vm_mm, address, _pte);
>> -            folio_remove_rmap_pte(src, src_page, vma);
>> +            clear_full_ptes(vma->vm_mm, address, _pte, nr_ptes,
>> +                    /* full = */ false);
>
> Can you add this patch to your series if nobody objects and use 
> clear_ptes()
> instead?

Thanks! Stupid me thought, because of arm64-specific impl of these, it

will be churn to do this, I am indeed lazy :)


>
> From 95e20ab0ff62bbbdcd89898c9d76fdc1ea961257 Mon Sep 17 00:00:00 2001
> From: David Hildenbrand <david@redhat.com>
> Date: Wed, 25 Jun 2025 12:55:20 +0200
> Subject: [PATCH] mm: add get_and_clear_ptes() and clear_ptes()
>
> Let's add variants to be used where "full" does not apply -- which will
> be the majority of cases in the future. "full" really only applies if
> we are about to tear down a full MM.
>
> Use get_and_clear_ptes() in existing code, clear_ptes() users will
> be added next.
>
> Should we make these inline functions instead and add separate docs?
> Probably not worth it for now.
>
> Signed-off-by: David Hildenbrand <david@redhat.com>
> ---
>  include/linux/pgtable.h | 6 ++++++
>  mm/mremap.c             | 2 +-
>  mm/rmap.c               | 2 +-
>  3 files changed, 8 insertions(+), 2 deletions(-)
>
> diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
> index cf1515c163e26..28679254b4f65 100644
> --- a/include/linux/pgtable.h
> +++ b/include/linux/pgtable.h
> @@ -736,6 +736,9 @@ static inline pte_t get_and_clear_full_ptes(struct 
> mm_struct *mm,
>  }
>  #endif
>
> +#define get_and_clear_ptes(_mm, _addr, _ptep, _nr) \
> +    get_and_clear_full_ptes(_mm, _addr, _ptep, _nr, 0)
> +
>  #ifndef clear_full_ptes
>  /**
>   * clear_full_ptes - Clear present PTEs that map consecutive pages of 
> the same
> @@ -768,6 +771,9 @@ static inline void clear_full_ptes(struct 
> mm_struct *mm, unsigned long addr,
>  }
>  #endif
>
> +#define clear_ptes(_mm, _addr, _ptep, _nr) \
> +    clear_full_ptes(_mm, _addr, _ptep, _nr, 0)
> +
>  /*
>   * If two threads concurrently fault at the same page, the thread that
>   * won the race updates the PTE and its local TLB/Cache. The other 
> thread
> diff --git a/mm/mremap.c b/mm/mremap.c
> index b31740f77b840..92890f8367574 100644
> --- a/mm/mremap.c
> +++ b/mm/mremap.c
> @@ -322,7 +322,7 @@ static int move_ptes(struct pagetable_move_control 
> *pmc,
>                               old_pte, max_nr_ptes);
>              force_flush = true;
>          }
> -        pte = get_and_clear_full_ptes(mm, old_addr, old_ptep, 
> nr_ptes, 0);
> +        pte = get_and_clear_ptes(mm, old_addr, old_ptep, nr_ptes);
>          pte = move_pte(pte, old_addr, new_addr);
>          pte = move_soft_dirty_pte(pte);
>
> diff --git a/mm/rmap.c b/mm/rmap.c
> index 3b74bb19c11dd..8200d705fe4ac 100644
> --- a/mm/rmap.c
> +++ b/mm/rmap.c
> @@ -2031,7 +2031,7 @@ static bool try_to_unmap_one(struct folio 
> *folio, struct vm_area_struct *vma,
>              flush_cache_range(vma, address, end_addr);
>
>              /* Nuke the page table entry. */
> -            pteval = get_and_clear_full_ptes(mm, address, pvmw.pte, 
> nr_pages, 0);
> +            pteval = get_and_clear_ptes(mm, address, pvmw.pte, 
> nr_pages);
>              /*
>               * We clear the PTE but do not flush so potentially
>               * a remote CPU could still be writing to the folio.

Re: [PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[Lorenzo Stoakes]

You forgot the v2 here :) this breaks b4 shazam...

I managed to do this on the cover letter (but not patches) of a series
before. So you're in good company... ;)

On Wed, Jun 25, 2025 at 11:28:04AM +0530, Dev Jain wrote:
> Use PTE batching to optimize __collapse_huge_page_copy_succeeded().
>
> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
> Then, calling ptep_clear() for every pte will cause a TLB flush for every
> contpte block. Instead, clear_full_ptes() does a
> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
> starting and ending contpte block, if they partially overlap with the range
> khugepaged is looking at.
>
> For all arches, there should be a benefit due to batching atomic operations
> on mapcounts due to folio_remove_rmap_ptes().
>
> No issues were observed with mm-selftests.
>
> Signed-off-by: Dev Jain <dev.jain@arm.com>

Overall looking way way better! Just some nits below.

> ---
>  mm/khugepaged.c | 27 +++++++++++++++++++++------
>  1 file changed, 21 insertions(+), 6 deletions(-)
>
> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> index d45d08b521f6..3944b112d452 100644
> --- a/mm/khugepaged.c
> +++ b/mm/khugepaged.c
> @@ -700,12 +700,15 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>  						spinlock_t *ptl,
>  						struct list_head *compound_pagelist)
>  {
> +	unsigned long end = address + HPAGE_PMD_SIZE;
>  	struct folio *src, *tmp;
> -	pte_t *_pte;
>  	pte_t pteval;
> +	pte_t *_pte;
> +	int nr_ptes;
>
> -	for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
> -	     _pte++, address += PAGE_SIZE) {
> +	for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
> +	     address += nr_ptes * PAGE_SIZE) {

Thanks this is much better.

> +		nr_ptes = 1;
>  		pteval = ptep_get(_pte);
>  		if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
>  			add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
> @@ -719,21 +722,33 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>  				ksm_might_unmap_zero_page(vma->vm_mm, pteval);
>  			}
>  		} else {
> +			const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
> +			int max_nr_ptes;
> +
>  			struct page *src_page = pte_page(pteval);
>
>  			src = page_folio(src_page);
>  			if (!folio_test_large(src))
>  				release_pte_folio(src);
> +
> +			max_nr_ptes = (end - address) >> PAGE_SHIFT;
> +			if (folio_test_large(src))
> +				nr_ptes = folio_pte_batch(src, address, _pte,
> +							  pteval, max_nr_ptes,
> +							  flags, NULL, NULL, NULL);

Nit, but we only use max_nr_ptes here so could declare and set here, e.g.:

			if (folio_test_large(src)) {
				int max_nr_ptes = (end - address) >> PAGE_SHIFT;

				nr_ptes = folio_pte_batch(src, address, _pte,
							  pteval, max_nr_ptes,
							  flags, NULL, NULL, NULL);
			}

BTW I think David raised it, but is there a way to wrap folio_pte_batch() to not
have to NULL, NULL, NULL here? :)


oh and if we do this, we can also combine this line with above so:

			if (folio_test_large(src)) {
				int max_nr_ptes = (end - address) >> PAGE_SHIFT;

				nr_ptes = folio_pte_batch(src, address, _pte,
							  pteval, max_nr_ptes,
							  flags, NULL, NULL, NULL);
			} else {
  				release_pte_folio(src);
			}

Which is neater.

> +
>  			/*
>  			 * ptl mostly unnecessary, but preempt has to
>  			 * be disabled to update the per-cpu stats
>  			 * inside folio_remove_rmap_pte().
>  			 */
>  			spin_lock(ptl);
> -			ptep_clear(vma->vm_mm, address, _pte);
> -			folio_remove_rmap_pte(src, src_page, vma);
> +			clear_full_ptes(vma->vm_mm, address, _pte, nr_ptes,
> +					/* full = */ false);
> +			folio_remove_rmap_ptes(src, src_page, nr_ptes, vma);
>  			spin_unlock(ptl);
> -			free_folio_and_swap_cache(src);
> +			free_swap_cache(src);
> +			folio_put_refs(src, nr_ptes);
>  		}
>  	}
>
> --
> 2.30.2
>

Re: [PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[David Hildenbrand]

> 			if (folio_test_large(src)) {
> 				int max_nr_ptes = (end - address) >> PAGE_SHIFT;
> 
> 				nr_ptes = folio_pte_batch(src, address, _pte,
> 							  pteval, max_nr_ptes,
> 							  flags, NULL, NULL, NULL);
> 			}
> 
> BTW I think David raised it, but is there a way to wrap folio_pte_batch() to not
> have to NULL, NULL, NULL here? :)

I'm working on a cleanup for all that. So stay tuned :)

-- 
Cheers,

David / dhildenb

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Lorenzo Stoakes]

On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
> Use PTE batching to optimize collapse_pte_mapped_thp().
>
> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
> Then, calling ptep_clear() for every pte will cause a TLB flush for every
> contpte block. Instead, clear_full_ptes() does a
> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
> starting and ending contpte block, if they partially overlap with the range
> khugepaged is looking at.
>
> For all arches, there should be a benefit due to batching atomic operations
> on mapcounts due to folio_remove_rmap_ptes().
>
> Note that we do not need to make a change to the check
> "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
> to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
> pages of the folio will be equal to the corresponding pages of our
> batch mapping consecutive pages.
>
> No issues were observed with mm-selftests.
>
> Signed-off-by: Dev Jain <dev.jain@arm.com>
> ---
>  mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
>  1 file changed, 26 insertions(+), 12 deletions(-)
>
> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> index 3944b112d452..4c8d33abfbd8 100644
> --- a/mm/khugepaged.c
> +++ b/mm/khugepaged.c
> @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
>  int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>  			    bool install_pmd)
>  {
> +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
>  	struct mmu_notifier_range range;
>  	bool notified = false;
>  	unsigned long haddr = addr & HPAGE_PMD_MASK;
> +	unsigned long end = haddr + HPAGE_PMD_SIZE;
>  	struct vm_area_struct *vma = vma_lookup(mm, haddr);
>  	struct folio *folio;
>  	pte_t *start_pte, *pte;
>  	pmd_t *pmd, pgt_pmd;
>  	spinlock_t *pml = NULL, *ptl;
> -	int nr_ptes = 0, result = SCAN_FAIL;
>  	int i;
>
>  	mmap_assert_locked(mm);
> @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>  		goto abort;
>
>  	/* step 2: clear page table and adjust rmap */
> -	for (i = 0, addr = haddr, pte = start_pte;
> -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
> +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
> +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
> +	     pte += nr_batch_ptes) {
> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
> +		struct folio *mapped_folio;
>  		struct page *page;
>  		pte_t ptent = ptep_get(pte);
>
> +		nr_batch_ptes = 1;
> +
>  		if (pte_none(ptent))
>  			continue;
>  		/*
> @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>  			goto abort;
>  		}
>  		page = vm_normal_page(vma, addr, ptent);
> +		mapped_folio = page_folio(page);
> +
>  		if (folio_page(folio, i) != page)
>  			goto abort;

Isn't this asserting that folio == mapped_folio here? We're saying page is the
ith page of folio, so why do we need to look up mapped_folio?

>
> +		mapped_folio = page_folio(page);

You're assigning this twice.

> +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
> +						max_nr_batch_ptes, flags,
> +						NULL, NULL, NULL);
> +
>  		/*
>  		 * Must clear entry, or a racing truncate may re-remove it.
>  		 * TLB flush can be left until pmdp_collapse_flush() does it.
>  		 * PTE dirty? Shmem page is already dirty; file is read-only.
>  		 */
> -		ptep_clear(mm, addr, pte);
> -		folio_remove_rmap_pte(folio, page, vma);
> -		nr_ptes++;
> +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
> +		nr_mapped_ptes += nr_batch_ptes;
>  	}
>
>  	if (!pml)
>  		spin_unlock(ptl);
>
>  	/* step 3: set proper refcount and mm_counters. */
> -	if (nr_ptes) {
> -		folio_ref_sub(folio, nr_ptes);
> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> +	if (nr_mapped_ptes) {
> +		folio_ref_sub(folio, nr_mapped_ptes);
> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>  	}
>
>  	/* step 4: remove empty page table */
> @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>  			: SCAN_SUCCEED;
>  	goto drop_folio;
>  abort:
> -	if (nr_ptes) {
> +	if (nr_mapped_ptes) {

I know it's ironic coming from me :P but I'm not sure why we need to churn this
up by renaming?

>  		flush_tlb_mm(mm);
> -		folio_ref_sub(folio, nr_ptes);
> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> +		folio_ref_sub(folio, nr_mapped_ptes);
> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>  	}
>  unlock:
>  	if (start_pte)
> --
> 2.30.2
>
V

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Lorenzo Stoakes]

On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
> Suppose a folio is under migration, and khugepaged is also trying to
> collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
> page cache via filemap_lock_folio(), thus taking a reference on the folio
> and sleeping on the folio lock, since the lock is held by the migration
> path. Migration will then fail in
> __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
> such a race happening (leading to migration failure) by bailing out
> if we detect a PMD is marked with a migration entry.
>
> This fixes the migration-shared-anon-thp testcase failure on Apple M3.

Hm is this related to the series at all? Seems somewhat unrelated?

Is there a Fixes, Closes, etc.? Do we need something in stable?

>
> Signed-off-by: Dev Jain <dev.jain@arm.com>
> ---
>  mm/khugepaged.c | 12 ++++++++++--
>  1 file changed, 10 insertions(+), 2 deletions(-)
>
> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> index 4c8d33abfbd8..bc8774f62e86 100644
> --- a/mm/khugepaged.c
> +++ b/mm/khugepaged.c
> @@ -31,6 +31,7 @@ enum scan_result {
>  	SCAN_FAIL,
>  	SCAN_SUCCEED,
>  	SCAN_PMD_NULL,
> +	SCAN_PMD_MIGRATION,
>  	SCAN_PMD_NONE,
>  	SCAN_PMD_MAPPED,
>  	SCAN_EXCEED_NONE_PTE,
> @@ -956,6 +957,8 @@ static inline int check_pmd_state(pmd_t *pmd)
>
>  	if (pmd_none(pmde))
>  		return SCAN_PMD_NONE;
> +	if (is_pmd_migration_entry(pmde))
> +		return SCAN_PMD_MIGRATION;
>  	if (!pmd_present(pmde))
>  		return SCAN_PMD_NULL;
>  	if (pmd_trans_huge(pmde))
> @@ -1518,9 +1521,12 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>  	    !range_in_vma(vma, haddr, haddr + HPAGE_PMD_SIZE))
>  		return SCAN_VMA_CHECK;
>
> -	/* Fast check before locking page if already PMD-mapped */
> +	/*
> +	 * Fast check before locking folio if already PMD-mapped, or if the
> +	 * folio is under migration
> +	 */
>  	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
> -	if (result == SCAN_PMD_MAPPED)
> +	if (result == SCAN_PMD_MAPPED || result == SCAN_PMD_MIGRATION)
>  		return result;
>
>  	/*
> @@ -2745,6 +2751,7 @@ static int madvise_collapse_errno(enum scan_result r)
>  	case SCAN_PAGE_LRU:
>  	case SCAN_DEL_PAGE_LRU:
>  	case SCAN_PAGE_FILLED:
> +	case SCAN_PMD_MIGRATION:
>  		return -EAGAIN;
>  	/*
>  	 * Other: Trying again likely not to succeed / error intrinsic to
> @@ -2834,6 +2841,7 @@ int madvise_collapse(struct vm_area_struct *vma, struct vm_area_struct **prev,
>  			goto handle_result;
>  		/* Whitelisted set of results where continuing OK */
>  		case SCAN_PMD_NULL:
> +		case SCAN_PMD_MIGRATION:
>  		case SCAN_PTE_NON_PRESENT:
>  		case SCAN_PTE_UFFD_WP:
>  		case SCAN_PAGE_RO:
> --
> 2.30.2
>

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Dev Jain]

On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
> On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
>> Use PTE batching to optimize collapse_pte_mapped_thp().
>>
>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
>> Then, calling ptep_clear() for every pte will cause a TLB flush for every
>> contpte block. Instead, clear_full_ptes() does a
>> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
>> starting and ending contpte block, if they partially overlap with the range
>> khugepaged is looking at.
>>
>> For all arches, there should be a benefit due to batching atomic operations
>> on mapcounts due to folio_remove_rmap_ptes().
>>
>> Note that we do not need to make a change to the check
>> "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
>> to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
>> pages of the folio will be equal to the corresponding pages of our
>> batch mapping consecutive pages.
>>
>> No issues were observed with mm-selftests.
>>
>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>> ---
>>   mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
>>   1 file changed, 26 insertions(+), 12 deletions(-)
>>
>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>> index 3944b112d452..4c8d33abfbd8 100644
>> --- a/mm/khugepaged.c
>> +++ b/mm/khugepaged.c
>> @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
>>   int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   			    bool install_pmd)
>>   {
>> +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
>>   	struct mmu_notifier_range range;
>>   	bool notified = false;
>>   	unsigned long haddr = addr & HPAGE_PMD_MASK;
>> +	unsigned long end = haddr + HPAGE_PMD_SIZE;
>>   	struct vm_area_struct *vma = vma_lookup(mm, haddr);
>>   	struct folio *folio;
>>   	pte_t *start_pte, *pte;
>>   	pmd_t *pmd, pgt_pmd;
>>   	spinlock_t *pml = NULL, *ptl;
>> -	int nr_ptes = 0, result = SCAN_FAIL;
>>   	int i;
>>
>>   	mmap_assert_locked(mm);
>> @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   		goto abort;
>>
>>   	/* step 2: clear page table and adjust rmap */
>> -	for (i = 0, addr = haddr, pte = start_pte;
>> -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
>> +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
>> +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
>> +	     pte += nr_batch_ptes) {
>> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
>> +		struct folio *mapped_folio;
>>   		struct page *page;
>>   		pte_t ptent = ptep_get(pte);
>>
>> +		nr_batch_ptes = 1;
>> +
>>   		if (pte_none(ptent))
>>   			continue;
>>   		/*
>> @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   			goto abort;
>>   		}
>>   		page = vm_normal_page(vma, addr, ptent);
>> +		mapped_folio = page_folio(page);
>> +
>>   		if (folio_page(folio, i) != page)
>>   			goto abort;
> Isn't this asserting that folio == mapped_folio here? We're saying page is the
> ith page of folio, so why do we need to look up mapped_folio?

We need to check for all PTEs whether they map the right page or not. This may
get disturbed due to mremap and stuff.

>
>> +		mapped_folio = page_folio(page);
> You're assigning this twice.

Forgot to remove, thanks.

>
>> +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
>> +						max_nr_batch_ptes, flags,
>> +						NULL, NULL, NULL);
>> +
>>   		/*
>>   		 * Must clear entry, or a racing truncate may re-remove it.
>>   		 * TLB flush can be left until pmdp_collapse_flush() does it.
>>   		 * PTE dirty? Shmem page is already dirty; file is read-only.
>>   		 */
>> -		ptep_clear(mm, addr, pte);
>> -		folio_remove_rmap_pte(folio, page, vma);
>> -		nr_ptes++;
>> +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
>> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
>> +		nr_mapped_ptes += nr_batch_ptes;
>>   	}
>>
>>   	if (!pml)
>>   		spin_unlock(ptl);
>>
>>   	/* step 3: set proper refcount and mm_counters. */
>> -	if (nr_ptes) {
>> -		folio_ref_sub(folio, nr_ptes);
>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>> +	if (nr_mapped_ptes) {
>> +		folio_ref_sub(folio, nr_mapped_ptes);
>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>   	}
>>
>>   	/* step 4: remove empty page table */
>> @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   			: SCAN_SUCCEED;
>>   	goto drop_folio;
>>   abort:
>> -	if (nr_ptes) {
>> +	if (nr_mapped_ptes) {
> I know it's ironic coming from me :P but I'm not sure why we need to churn this
> up by renaming?

Because nr_ptes is an existing variable and I need a new variable to make
the jump at the end of the PTE batch.

>
>>   		flush_tlb_mm(mm);
>> -		folio_ref_sub(folio, nr_ptes);
>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>> +		folio_ref_sub(folio, nr_mapped_ptes);
>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>   	}
>>   unlock:
>>   	if (start_pte)
>> --
>> 2.30.2
>>
> V

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Dev Jain]

On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
> On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
>> Suppose a folio is under migration, and khugepaged is also trying to
>> collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
>> page cache via filemap_lock_folio(), thus taking a reference on the folio
>> and sleeping on the folio lock, since the lock is held by the migration
>> path. Migration will then fail in
>> __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
>> such a race happening (leading to migration failure) by bailing out
>> if we detect a PMD is marked with a migration entry.
>>
>> This fixes the migration-shared-anon-thp testcase failure on Apple M3.
> Hm is this related to the series at all? Seems somewhat unrelated?

Not related.

>
> Is there a Fixes, Closes, etc.? Do we need something in stable?

We don't need anything. This is an "expected race" in the sense that
both migration and khugepaged collapse are best effort algorithms.
I am just seeing a test failure on my system because my system hits
the race more often. So this patch reduces the window for the race.

For some previous context, you may look at:

https://lore.kernel.org/all/20240801081657.1386743-1-dev.jain@arm.com/
and
https://lore.kernel.org/all/20240830051609.4037834-1-dev.jain@arm.com/

>
>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>> ---
>>   mm/khugepaged.c | 12 ++++++++++--
>>   1 file changed, 10 insertions(+), 2 deletions(-)
>>
>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>> index 4c8d33abfbd8..bc8774f62e86 100644
>> --- a/mm/khugepaged.c
>> +++ b/mm/khugepaged.c
>> @@ -31,6 +31,7 @@ enum scan_result {
>>   	SCAN_FAIL,
>>   	SCAN_SUCCEED,
>>   	SCAN_PMD_NULL,
>> +	SCAN_PMD_MIGRATION,
>>   	SCAN_PMD_NONE,
>>   	SCAN_PMD_MAPPED,
>>   	SCAN_EXCEED_NONE_PTE,
>> @@ -956,6 +957,8 @@ static inline int check_pmd_state(pmd_t *pmd)
>>
>>   	if (pmd_none(pmde))
>>   		return SCAN_PMD_NONE;
>> +	if (is_pmd_migration_entry(pmde))
>> +		return SCAN_PMD_MIGRATION;
>>   	if (!pmd_present(pmde))
>>   		return SCAN_PMD_NULL;
>>   	if (pmd_trans_huge(pmde))
>> @@ -1518,9 +1521,12 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   	    !range_in_vma(vma, haddr, haddr + HPAGE_PMD_SIZE))
>>   		return SCAN_VMA_CHECK;
>>
>> -	/* Fast check before locking page if already PMD-mapped */
>> +	/*
>> +	 * Fast check before locking folio if already PMD-mapped, or if the
>> +	 * folio is under migration
>> +	 */
>>   	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
>> -	if (result == SCAN_PMD_MAPPED)
>> +	if (result == SCAN_PMD_MAPPED || result == SCAN_PMD_MIGRATION)
>>   		return result;
>>
>>   	/*
>> @@ -2745,6 +2751,7 @@ static int madvise_collapse_errno(enum scan_result r)
>>   	case SCAN_PAGE_LRU:
>>   	case SCAN_DEL_PAGE_LRU:
>>   	case SCAN_PAGE_FILLED:
>> +	case SCAN_PMD_MIGRATION:
>>   		return -EAGAIN;
>>   	/*
>>   	 * Other: Trying again likely not to succeed / error intrinsic to
>> @@ -2834,6 +2841,7 @@ int madvise_collapse(struct vm_area_struct *vma, struct vm_area_struct **prev,
>>   			goto handle_result;
>>   		/* Whitelisted set of results where continuing OK */
>>   		case SCAN_PMD_NULL:
>> +		case SCAN_PMD_MIGRATION:
>>   		case SCAN_PTE_NON_PRESENT:
>>   		case SCAN_PTE_UFFD_WP:
>>   		case SCAN_PAGE_RO:
>> --
>> 2.30.2
>>

Re: [PATCH 1/3] khugepaged: Optimize __collapse_huge_page_copy_succeeded() by PTE batching

[Dev Jain]

On 25/06/25 5:44 pm, Lorenzo Stoakes wrote:
> You forgot the v2 here :) this breaks b4 shazam...
>
> I managed to do this on the cover letter (but not patches) of a series
> before. So you're in good company... ;)
>
> On Wed, Jun 25, 2025 at 11:28:04AM +0530, Dev Jain wrote:
>> Use PTE batching to optimize __collapse_huge_page_copy_succeeded().
>>
>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
>> Then, calling ptep_clear() for every pte will cause a TLB flush for every
>> contpte block. Instead, clear_full_ptes() does a
>> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
>> starting and ending contpte block, if they partially overlap with the range
>> khugepaged is looking at.
>>
>> For all arches, there should be a benefit due to batching atomic operations
>> on mapcounts due to folio_remove_rmap_ptes().
>>
>> No issues were observed with mm-selftests.
>>
>> Signed-off-by: Dev Jain <dev.jain@arm.com>
> Overall looking way way better! Just some nits below.
>
>> ---
>>   mm/khugepaged.c | 27 +++++++++++++++++++++------
>>   1 file changed, 21 insertions(+), 6 deletions(-)
>>
>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>> index d45d08b521f6..3944b112d452 100644
>> --- a/mm/khugepaged.c
>> +++ b/mm/khugepaged.c
>> @@ -700,12 +700,15 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>>   						spinlock_t *ptl,
>>   						struct list_head *compound_pagelist)
>>   {
>> +	unsigned long end = address + HPAGE_PMD_SIZE;
>>   	struct folio *src, *tmp;
>> -	pte_t *_pte;
>>   	pte_t pteval;
>> +	pte_t *_pte;
>> +	int nr_ptes;
>>
>> -	for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
>> -	     _pte++, address += PAGE_SIZE) {
>> +	for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
>> +	     address += nr_ptes * PAGE_SIZE) {
> Thanks this is much better.
>
>> +		nr_ptes = 1;
>>   		pteval = ptep_get(_pte);
>>   		if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
>>   			add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
>> @@ -719,21 +722,33 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
>>   				ksm_might_unmap_zero_page(vma->vm_mm, pteval);
>>   			}
>>   		} else {
>> +			const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>> +			int max_nr_ptes;
>> +
>>   			struct page *src_page = pte_page(pteval);
>>
>>   			src = page_folio(src_page);
>>   			if (!folio_test_large(src))
>>   				release_pte_folio(src);
>> +
>> +			max_nr_ptes = (end - address) >> PAGE_SHIFT;
>> +			if (folio_test_large(src))
>> +				nr_ptes = folio_pte_batch(src, address, _pte,
>> +							  pteval, max_nr_ptes,
>> +							  flags, NULL, NULL, NULL);
> Nit, but we only use max_nr_ptes here so could declare and set here, e.g.:
>
> 			if (folio_test_large(src)) {
> 				int max_nr_ptes = (end - address) >> PAGE_SHIFT;
>
> 				nr_ptes = folio_pte_batch(src, address, _pte,
> 							  pteval, max_nr_ptes,
> 							  flags, NULL, NULL, NULL);
> 			}
>
> BTW I think David raised it, but is there a way to wrap folio_pte_batch() to not
> have to NULL, NULL, NULL here? :)
>
>
> oh and if we do this, we can also combine this line with above so:
>
> 			if (folio_test_large(src)) {
> 				int max_nr_ptes = (end - address) >> PAGE_SHIFT;
>
> 				nr_ptes = folio_pte_batch(src, address, _pte,
> 							  pteval, max_nr_ptes,
> 							  flags, NULL, NULL, NULL);
> 			} else {
>    				release_pte_folio(src);
> 			}
>
> Which is neater.

Okay.

>
>> +
>>   			/*
>>   			 * ptl mostly unnecessary, but preempt has to
>>   			 * be disabled to update the per-cpu stats
>>   			 * inside folio_remove_rmap_pte().
>>   			 */
>>   			spin_lock(ptl);
>> -			ptep_clear(vma->vm_mm, address, _pte);
>> -			folio_remove_rmap_pte(src, src_page, vma);
>> +			clear_full_ptes(vma->vm_mm, address, _pte, nr_ptes,
>> +					/* full = */ false);
>> +			folio_remove_rmap_ptes(src, src_page, nr_ptes, vma);
>>   			spin_unlock(ptl);
>> -			free_folio_and_swap_cache(src);
>> +			free_swap_cache(src);
>> +			folio_put_refs(src, nr_ptes);
>>   		}
>>   	}
>>
>> --
>> 2.30.2
>>

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Lorenzo Stoakes]

On Thu, Jun 26, 2025 at 09:18:47AM +0530, Dev Jain wrote:
>
> On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
> > On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
> > > Use PTE batching to optimize collapse_pte_mapped_thp().
> > >
> > > On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
> > > Then, calling ptep_clear() for every pte will cause a TLB flush for every
> > > contpte block. Instead, clear_full_ptes() does a
> > > contpte_try_unfold_partial() which will flush the TLB only for the (if any)
> > > starting and ending contpte block, if they partially overlap with the range
> > > khugepaged is looking at.
> > >
> > > For all arches, there should be a benefit due to batching atomic operations
> > > on mapcounts due to folio_remove_rmap_ptes().
> > >
> > > Note that we do not need to make a change to the check
> > > "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
> > > to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
> > > pages of the folio will be equal to the corresponding pages of our
> > > batch mapping consecutive pages.
> > >
> > > No issues were observed with mm-selftests.
> > >
> > > Signed-off-by: Dev Jain <dev.jain@arm.com>
> > > ---
> > >   mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
> > >   1 file changed, 26 insertions(+), 12 deletions(-)
> > >
> > > diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> > > index 3944b112d452..4c8d33abfbd8 100644
> > > --- a/mm/khugepaged.c
> > > +++ b/mm/khugepaged.c
> > > @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
> > >   int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > >   			    bool install_pmd)
> > >   {
> > > +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
> > >   	struct mmu_notifier_range range;
> > >   	bool notified = false;
> > >   	unsigned long haddr = addr & HPAGE_PMD_MASK;
> > > +	unsigned long end = haddr + HPAGE_PMD_SIZE;
> > >   	struct vm_area_struct *vma = vma_lookup(mm, haddr);
> > >   	struct folio *folio;
> > >   	pte_t *start_pte, *pte;
> > >   	pmd_t *pmd, pgt_pmd;
> > >   	spinlock_t *pml = NULL, *ptl;
> > > -	int nr_ptes = 0, result = SCAN_FAIL;
> > >   	int i;
> > >
> > >   	mmap_assert_locked(mm);
> > > @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > >   		goto abort;
> > >
> > >   	/* step 2: clear page table and adjust rmap */
> > > -	for (i = 0, addr = haddr, pte = start_pte;
> > > -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
> > > +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
> > > +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
> > > +	     pte += nr_batch_ptes) {
> > > +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
> > > +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
> > > +		struct folio *mapped_folio;
> > >   		struct page *page;
> > >   		pte_t ptent = ptep_get(pte);
> > >
> > > +		nr_batch_ptes = 1;
> > > +
> > >   		if (pte_none(ptent))
> > >   			continue;
> > >   		/*
> > > @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > >   			goto abort;
> > >   		}
> > >   		page = vm_normal_page(vma, addr, ptent);
> > > +		mapped_folio = page_folio(page);
> > > +
> > >   		if (folio_page(folio, i) != page)
> > >   			goto abort;
> > Isn't this asserting that folio == mapped_folio here? We're saying page is the
> > ith page of folio, so why do we need to look up mapped_folio?
>
> We need to check for all PTEs whether they map the right page or not. This may
> get disturbed due to mremap and stuff.

Right but I'm saying mapped_folio == folio right? You're literally asserting it
here? So there's no need to assign mapped_folio at all, just reference folio no?

>
> >
> > > +		mapped_folio = page_folio(page);
> > You're assigning this twice.
>
> Forgot to remove, thanks.
>
> >
> > > +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
> > > +						max_nr_batch_ptes, flags,
> > > +						NULL, NULL, NULL);
> > > +
> > >   		/*
> > >   		 * Must clear entry, or a racing truncate may re-remove it.
> > >   		 * TLB flush can be left until pmdp_collapse_flush() does it.
> > >   		 * PTE dirty? Shmem page is already dirty; file is read-only.
> > >   		 */
> > > -		ptep_clear(mm, addr, pte);
> > > -		folio_remove_rmap_pte(folio, page, vma);
> > > -		nr_ptes++;
> > > +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
> > > +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
> > > +		nr_mapped_ptes += nr_batch_ptes;
> > >   	}
> > >
> > >   	if (!pml)
> > >   		spin_unlock(ptl);
> > >
> > >   	/* step 3: set proper refcount and mm_counters. */
> > > -	if (nr_ptes) {
> > > -		folio_ref_sub(folio, nr_ptes);
> > > -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> > > +	if (nr_mapped_ptes) {
> > > +		folio_ref_sub(folio, nr_mapped_ptes);
> > > +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
> > >   	}
> > >
> > >   	/* step 4: remove empty page table */
> > > @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > >   			: SCAN_SUCCEED;
> > >   	goto drop_folio;
> > >   abort:
> > > -	if (nr_ptes) {
> > > +	if (nr_mapped_ptes) {
> > I know it's ironic coming from me :P but I'm not sure why we need to churn this
> > up by renaming?
>
> Because nr_ptes is an existing variable and I need a new variable to make
> the jump at the end of the PTE batch.

I thought you eliminated nr_ptes as a variable here? Where else is it used?

Oh how this code needs refactoring...

>
> >
> > >   		flush_tlb_mm(mm);
> > > -		folio_ref_sub(folio, nr_ptes);
> > > -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> > > +		folio_ref_sub(folio, nr_mapped_ptes);
> > > +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
> > >   	}
> > >   unlock:
> > >   	if (start_pte)
> > > --
> > > 2.30.2
> > >
> > V

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Dev Jain]

On 26/06/25 10:17 am, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 09:18:47AM +0530, Dev Jain wrote:
>> On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
>>> On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
>>>> Use PTE batching to optimize collapse_pte_mapped_thp().
>>>>
>>>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
>>>> Then, calling ptep_clear() for every pte will cause a TLB flush for every
>>>> contpte block. Instead, clear_full_ptes() does a
>>>> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
>>>> starting and ending contpte block, if they partially overlap with the range
>>>> khugepaged is looking at.
>>>>
>>>> For all arches, there should be a benefit due to batching atomic operations
>>>> on mapcounts due to folio_remove_rmap_ptes().
>>>>
>>>> Note that we do not need to make a change to the check
>>>> "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
>>>> to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
>>>> pages of the folio will be equal to the corresponding pages of our
>>>> batch mapping consecutive pages.
>>>>
>>>> No issues were observed with mm-selftests.
>>>>
>>>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>>>> ---
>>>>    mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
>>>>    1 file changed, 26 insertions(+), 12 deletions(-)
>>>>
>>>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>>>> index 3944b112d452..4c8d33abfbd8 100644
>>>> --- a/mm/khugepaged.c
>>>> +++ b/mm/khugepaged.c
>>>> @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
>>>>    int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			    bool install_pmd)
>>>>    {
>>>> +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
>>>>    	struct mmu_notifier_range range;
>>>>    	bool notified = false;
>>>>    	unsigned long haddr = addr & HPAGE_PMD_MASK;
>>>> +	unsigned long end = haddr + HPAGE_PMD_SIZE;
>>>>    	struct vm_area_struct *vma = vma_lookup(mm, haddr);
>>>>    	struct folio *folio;
>>>>    	pte_t *start_pte, *pte;
>>>>    	pmd_t *pmd, pgt_pmd;
>>>>    	spinlock_t *pml = NULL, *ptl;
>>>> -	int nr_ptes = 0, result = SCAN_FAIL;
>>>>    	int i;
>>>>
>>>>    	mmap_assert_locked(mm);
>>>> @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    		goto abort;
>>>>
>>>>    	/* step 2: clear page table and adjust rmap */
>>>> -	for (i = 0, addr = haddr, pte = start_pte;
>>>> -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
>>>> +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
>>>> +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
>>>> +	     pte += nr_batch_ptes) {
>>>> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>>>> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
>>>> +		struct folio *mapped_folio;
>>>>    		struct page *page;
>>>>    		pte_t ptent = ptep_get(pte);
>>>>
>>>> +		nr_batch_ptes = 1;
>>>> +
>>>>    		if (pte_none(ptent))
>>>>    			continue;
>>>>    		/*
>>>> @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			goto abort;
>>>>    		}
>>>>    		page = vm_normal_page(vma, addr, ptent);
>>>> +		mapped_folio = page_folio(page);
>>>> +
>>>>    		if (folio_page(folio, i) != page)
>>>>    			goto abort;
>>> Isn't this asserting that folio == mapped_folio here? We're saying page is the
>>> ith page of folio, so why do we need to look up mapped_folio?
>> We need to check for all PTEs whether they map the right page or not. This may
>> get disturbed due to mremap and stuff.
> Right but I'm saying mapped_folio == folio right? You're literally asserting it
> here? So there's no need to assign mapped_folio at all, just reference folio no?

Ah I get you now, we can do the PTE batching directly on the folio, thanks.

>
>>>> +		mapped_folio = page_folio(page);
>>> You're assigning this twice.
>> Forgot to remove, thanks.
>>
>>>> +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
>>>> +						max_nr_batch_ptes, flags,
>>>> +						NULL, NULL, NULL);
>>>> +
>>>>    		/*
>>>>    		 * Must clear entry, or a racing truncate may re-remove it.
>>>>    		 * TLB flush can be left until pmdp_collapse_flush() does it.
>>>>    		 * PTE dirty? Shmem page is already dirty; file is read-only.
>>>>    		 */
>>>> -		ptep_clear(mm, addr, pte);
>>>> -		folio_remove_rmap_pte(folio, page, vma);
>>>> -		nr_ptes++;
>>>> +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
>>>> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
>>>> +		nr_mapped_ptes += nr_batch_ptes;
>>>>    	}
>>>>
>>>>    	if (!pml)
>>>>    		spin_unlock(ptl);
>>>>
>>>>    	/* step 3: set proper refcount and mm_counters. */
>>>> -	if (nr_ptes) {
>>>> -		folio_ref_sub(folio, nr_ptes);
>>>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>>>> +	if (nr_mapped_ptes) {
>>>> +		folio_ref_sub(folio, nr_mapped_ptes);
>>>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>>>    	}
>>>>
>>>>    	/* step 4: remove empty page table */
>>>> @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			: SCAN_SUCCEED;
>>>>    	goto drop_folio;
>>>>    abort:
>>>> -	if (nr_ptes) {
>>>> +	if (nr_mapped_ptes) {
>>> I know it's ironic coming from me :P but I'm not sure why we need to churn this
>>> up by renaming?
>> Because nr_ptes is an existing variable and I need a new variable to make
>> the jump at the end of the PTE batch.
> I thought you eliminated nr_ptes as a variable here? Where else is it used?
>
> Oh how this code needs refactoring...
>
>>>>    		flush_tlb_mm(mm);
>>>> -		folio_ref_sub(folio, nr_ptes);
>>>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>>>> +		folio_ref_sub(folio, nr_mapped_ptes);
>>>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>>>    	}
>>>>    unlock:
>>>>    	if (start_pte)
>>>> --
>>>> 2.30.2
>>>>
>>> V

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Lorenzo Stoakes]

On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
>
> On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
> > On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
> > > Suppose a folio is under migration, and khugepaged is also trying to
> > > collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
> > > page cache via filemap_lock_folio(), thus taking a reference on the folio
> > > and sleeping on the folio lock, since the lock is held by the migration
> > > path. Migration will then fail in
> > > __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
> > > such a race happening (leading to migration failure) by bailing out
> > > if we detect a PMD is marked with a migration entry.
> > >
> > > This fixes the migration-shared-anon-thp testcase failure on Apple M3.
> > Hm is this related to the series at all? Seems somewhat unrelated?
>
> Not related.
>
> >
> > Is there a Fixes, Closes, etc.? Do we need something in stable?
>
> We don't need anything. This is an "expected race" in the sense that
> both migration and khugepaged collapse are best effort algorithms.
> I am just seeing a test failure on my system because my system hits
> the race more often. So this patch reduces the window for the race.

Does it rely on previous patches? If not probably best to send this one
separately :)

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Dev Jain]

On 26/06/25 10:27 am, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
>> On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
>>> On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
>>>> Suppose a folio is under migration, and khugepaged is also trying to
>>>> collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
>>>> page cache via filemap_lock_folio(), thus taking a reference on the folio
>>>> and sleeping on the folio lock, since the lock is held by the migration
>>>> path. Migration will then fail in
>>>> __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
>>>> such a race happening (leading to migration failure) by bailing out
>>>> if we detect a PMD is marked with a migration entry.
>>>>
>>>> This fixes the migration-shared-anon-thp testcase failure on Apple M3.
>>> Hm is this related to the series at all? Seems somewhat unrelated?
>> Not related.
>>
>>> Is there a Fixes, Closes, etc.? Do we need something in stable?
>> We don't need anything. This is an "expected race" in the sense that
>> both migration and khugepaged collapse are best effort algorithms.
>> I am just seeing a test failure on my system because my system hits
>> the race more often. So this patch reduces the window for the race.
> Does it rely on previous patches? If not probably best to send this one
> separately :)

To prevent rebasing headaches for others (if any) I thought to send all together.
I'll send it separately if still that is the preference.

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Lorenzo Stoakes]

On Thu, Jun 26, 2025 at 10:29:22AM +0530, Dev Jain wrote:
>
> On 26/06/25 10:27 am, Lorenzo Stoakes wrote:
> > On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
> > > On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
> > > > On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
> > > > > Suppose a folio is under migration, and khugepaged is also trying to
> > > > > collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
> > > > > page cache via filemap_lock_folio(), thus taking a reference on the folio
> > > > > and sleeping on the folio lock, since the lock is held by the migration
> > > > > path. Migration will then fail in
> > > > > __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
> > > > > such a race happening (leading to migration failure) by bailing out
> > > > > if we detect a PMD is marked with a migration entry.
> > > > >
> > > > > This fixes the migration-shared-anon-thp testcase failure on Apple M3.
> > > > Hm is this related to the series at all? Seems somewhat unrelated?
> > > Not related.
> > >
> > > > Is there a Fixes, Closes, etc.? Do we need something in stable?
> > > We don't need anything. This is an "expected race" in the sense that
> > > both migration and khugepaged collapse are best effort algorithms.
> > > I am just seeing a test failure on my system because my system hits
> > > the race more often. So this patch reduces the window for the race.
> > Does it rely on previous patches? If not probably best to send this one
> > separately :)
>
> To prevent rebasing headaches for others (if any) I thought to send all together.
> I'll send it separately if still that is the preference.
>
>

Oh actually would it be a pain to rebase given the previous 2 patches? Maybe
leave it then. And I can actually finally review it... :)

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Dev Jain]

On 26/06/25 10:32 am, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 10:29:22AM +0530, Dev Jain wrote:
>> On 26/06/25 10:27 am, Lorenzo Stoakes wrote:
>>> On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
>>>> On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
>>>>> On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
>>>>>> Suppose a folio is under migration, and khugepaged is also trying to
>>>>>> collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
>>>>>> page cache via filemap_lock_folio(), thus taking a reference on the folio
>>>>>> and sleeping on the folio lock, since the lock is held by the migration
>>>>>> path. Migration will then fail in
>>>>>> __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
>>>>>> such a race happening (leading to migration failure) by bailing out
>>>>>> if we detect a PMD is marked with a migration entry.
>>>>>>
>>>>>> This fixes the migration-shared-anon-thp testcase failure on Apple M3.
>>>>> Hm is this related to the series at all? Seems somewhat unrelated?
>>>> Not related.
>>>>
>>>>> Is there a Fixes, Closes, etc.? Do we need something in stable?
>>>> We don't need anything. This is an "expected race" in the sense that
>>>> both migration and khugepaged collapse are best effort algorithms.
>>>> I am just seeing a test failure on my system because my system hits
>>>> the race more often. So this patch reduces the window for the race.
>>> Does it rely on previous patches? If not probably best to send this one
>>> separately :)
>> To prevent rebasing headaches for others (if any) I thought to send all together.
>> I'll send it separately if still that is the preference.
>>
>>
> Oh actually would it be a pain to rebase given the previous 2 patches? Maybe

Didn't take the time to actually check that, more of a "it *may* be a pain
so let's send this together".

> leave it then. And I can actually finally review it... :)

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Lorenzo Stoakes]

On Thu, Jun 26, 2025 at 06:02:25AM +0100, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 10:29:22AM +0530, Dev Jain wrote:
> >
> > On 26/06/25 10:27 am, Lorenzo Stoakes wrote:
> > > On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
> > > > On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
> > > > > On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
> > > > > > Suppose a folio is under migration, and khugepaged is also trying to
> > > > > > collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
> > > > > > page cache via filemap_lock_folio(), thus taking a reference on the folio
> > > > > > and sleeping on the folio lock, since the lock is held by the migration
> > > > > > path. Migration will then fail in
> > > > > > __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
> > > > > > such a race happening (leading to migration failure) by bailing out
> > > > > > if we detect a PMD is marked with a migration entry.
> > > > > >
> > > > > > This fixes the migration-shared-anon-thp testcase failure on Apple M3.
> > > > > Hm is this related to the series at all? Seems somewhat unrelated?
> > > > Not related.
> > > >
> > > > > Is there a Fixes, Closes, etc.? Do we need something in stable?
> > > > We don't need anything. This is an "expected race" in the sense that
> > > > both migration and khugepaged collapse are best effort algorithms.
> > > > I am just seeing a test failure on my system because my system hits
> > > > the race more often. So this patch reduces the window for the race.
> > > Does it rely on previous patches? If not probably best to send this one
> > > separately :)
> >
> > To prevent rebasing headaches for others (if any) I thought to send all together.
> > I'll send it separately if still that is the preference.
> >
> >
>
> Oh actually would it be a pain to rebase given the previous 2 patches? Maybe
> leave it then. And I can actually finally review it... :)

At the same time you'd have to update cover letter to mention and it's really
not good practice to include unrelated patches to a series.

So probably best overall to send separately but _wait_ for this series to be
merged first :>) then you can based it on mm-new without problems.

Re: [PATCH v2 3/3] khugepaged: Reduce race probability between migration and khugepaged

[Dev Jain]

On 26/06/25 10:36 am, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 06:02:25AM +0100, Lorenzo Stoakes wrote:
>> On Thu, Jun 26, 2025 at 10:29:22AM +0530, Dev Jain wrote:
>>> On 26/06/25 10:27 am, Lorenzo Stoakes wrote:
>>>> On Thu, Jun 26, 2025 at 09:22:28AM +0530, Dev Jain wrote:
>>>>> On 25/06/25 6:58 pm, Lorenzo Stoakes wrote:
>>>>>> On Wed, Jun 25, 2025 at 11:28:06AM +0530, Dev Jain wrote:
>>>>>>> Suppose a folio is under migration, and khugepaged is also trying to
>>>>>>> collapse it. collapse_pte_mapped_thp() will retrieve the folio from the
>>>>>>> page cache via filemap_lock_folio(), thus taking a reference on the folio
>>>>>>> and sleeping on the folio lock, since the lock is held by the migration
>>>>>>> path. Migration will then fail in
>>>>>>> __folio_migrate_mapping -> folio_ref_freeze. Reduce the probability of
>>>>>>> such a race happening (leading to migration failure) by bailing out
>>>>>>> if we detect a PMD is marked with a migration entry.
>>>>>>>
>>>>>>> This fixes the migration-shared-anon-thp testcase failure on Apple M3.
>>>>>> Hm is this related to the series at all? Seems somewhat unrelated?
>>>>> Not related.
>>>>>
>>>>>> Is there a Fixes, Closes, etc.? Do we need something in stable?
>>>>> We don't need anything. This is an "expected race" in the sense that
>>>>> both migration and khugepaged collapse are best effort algorithms.
>>>>> I am just seeing a test failure on my system because my system hits
>>>>> the race more often. So this patch reduces the window for the race.
>>>> Does it rely on previous patches? If not probably best to send this one
>>>> separately :)
>>> To prevent rebasing headaches for others (if any) I thought to send all together.
>>> I'll send it separately if still that is the preference.
>>>
>>>
>> Oh actually would it be a pain to rebase given the previous 2 patches? Maybe
>> leave it then. And I can actually finally review it... :)
> At the same time you'd have to update cover letter to mention and it's really
> not good practice to include unrelated patches to a series.
>
> So probably best overall to send separately but _wait_ for this series to be
> merged first :>) then you can based it on mm-new without problems.

Sure.

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Dev Jain]

On 26/06/25 10:17 am, Lorenzo Stoakes wrote:
> On Thu, Jun 26, 2025 at 09:18:47AM +0530, Dev Jain wrote:
>> On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
>>> On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
>>>> Use PTE batching to optimize collapse_pte_mapped_thp().
>>>>
>>>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
>>>> Then, calling ptep_clear() for every pte will cause a TLB flush for every
>>>> contpte block. Instead, clear_full_ptes() does a
>>>> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
>>>> starting and ending contpte block, if they partially overlap with the range
>>>> khugepaged is looking at.
>>>>
>>>> For all arches, there should be a benefit due to batching atomic operations
>>>> on mapcounts due to folio_remove_rmap_ptes().
>>>>
>>>> Note that we do not need to make a change to the check
>>>> "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
>>>> to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
>>>> pages of the folio will be equal to the corresponding pages of our
>>>> batch mapping consecutive pages.
>>>>
>>>> No issues were observed with mm-selftests.
>>>>
>>>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>>>> ---
>>>>    mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
>>>>    1 file changed, 26 insertions(+), 12 deletions(-)
>>>>
>>>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>>>> index 3944b112d452..4c8d33abfbd8 100644
>>>> --- a/mm/khugepaged.c
>>>> +++ b/mm/khugepaged.c
>>>> @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
>>>>    int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			    bool install_pmd)
>>>>    {
>>>> +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
>>>>    	struct mmu_notifier_range range;
>>>>    	bool notified = false;
>>>>    	unsigned long haddr = addr & HPAGE_PMD_MASK;
>>>> +	unsigned long end = haddr + HPAGE_PMD_SIZE;
>>>>    	struct vm_area_struct *vma = vma_lookup(mm, haddr);
>>>>    	struct folio *folio;
>>>>    	pte_t *start_pte, *pte;
>>>>    	pmd_t *pmd, pgt_pmd;
>>>>    	spinlock_t *pml = NULL, *ptl;
>>>> -	int nr_ptes = 0, result = SCAN_FAIL;
>>>>    	int i;
>>>>
>>>>    	mmap_assert_locked(mm);
>>>> @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    		goto abort;
>>>>
>>>>    	/* step 2: clear page table and adjust rmap */
>>>> -	for (i = 0, addr = haddr, pte = start_pte;
>>>> -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
>>>> +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
>>>> +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
>>>> +	     pte += nr_batch_ptes) {
>>>> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>>>> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
>>>> +		struct folio *mapped_folio;
>>>>    		struct page *page;
>>>>    		pte_t ptent = ptep_get(pte);
>>>>
>>>> +		nr_batch_ptes = 1;
>>>> +
>>>>    		if (pte_none(ptent))
>>>>    			continue;
>>>>    		/*
>>>> @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			goto abort;
>>>>    		}
>>>>    		page = vm_normal_page(vma, addr, ptent);
>>>> +		mapped_folio = page_folio(page);
>>>> +
>>>>    		if (folio_page(folio, i) != page)
>>>>    			goto abort;
>>> Isn't this asserting that folio == mapped_folio here? We're saying page is the
>>> ith page of folio, so why do we need to look up mapped_folio?
>> We need to check for all PTEs whether they map the right page or not. This may
>> get disturbed due to mremap and stuff.
> Right but I'm saying mapped_folio == folio right? You're literally asserting it
> here? So there's no need to assign mapped_folio at all, just reference folio no?
>
>>>> +		mapped_folio = page_folio(page);
>>> You're assigning this twice.
>> Forgot to remove, thanks.
>>
>>>> +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
>>>> +						max_nr_batch_ptes, flags,
>>>> +						NULL, NULL, NULL);
>>>> +
>>>>    		/*
>>>>    		 * Must clear entry, or a racing truncate may re-remove it.
>>>>    		 * TLB flush can be left until pmdp_collapse_flush() does it.
>>>>    		 * PTE dirty? Shmem page is already dirty; file is read-only.
>>>>    		 */
>>>> -		ptep_clear(mm, addr, pte);
>>>> -		folio_remove_rmap_pte(folio, page, vma);
>>>> -		nr_ptes++;
>>>> +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
>>>> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
>>>> +		nr_mapped_ptes += nr_batch_ptes;
>>>>    	}
>>>>
>>>>    	if (!pml)
>>>>    		spin_unlock(ptl);
>>>>
>>>>    	/* step 3: set proper refcount and mm_counters. */
>>>> -	if (nr_ptes) {
>>>> -		folio_ref_sub(folio, nr_ptes);
>>>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>>>> +	if (nr_mapped_ptes) {
>>>> +		folio_ref_sub(folio, nr_mapped_ptes);
>>>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>>>    	}
>>>>
>>>>    	/* step 4: remove empty page table */
>>>> @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>    			: SCAN_SUCCEED;
>>>>    	goto drop_folio;
>>>>    abort:
>>>> -	if (nr_ptes) {
>>>> +	if (nr_mapped_ptes) {
>>> I know it's ironic coming from me :P but I'm not sure why we need to churn this
>>> up by renaming?
>> Because nr_ptes is an existing variable and I need a new variable to make
>> the jump at the end of the PTE batch.
> I thought you eliminated nr_ptes as a variable here? Where else is it used?
>
> Oh how this code needs refactoring...

If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
which is confusing since the former is plain and the latter has a _mapped_ thingy
in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.

>
>>>>    		flush_tlb_mm(mm);
>>>> -		folio_ref_sub(folio, nr_ptes);
>>>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>>>> +		folio_ref_sub(folio, nr_mapped_ptes);
>>>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>>>    	}
>>>>    unlock:
>>>>    	if (start_pte)
>>>> --
>>>> 2.30.2
>>>>
>>> V

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Lorenzo Stoakes]

On Tue, Jul 15, 2025 at 12:04:56PM +0530, Dev Jain wrote:
>
> On 26/06/25 10:17 am, Lorenzo Stoakes wrote:
> > On Thu, Jun 26, 2025 at 09:18:47AM +0530, Dev Jain wrote:
> > > On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
> > > > On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
> > > > > Use PTE batching to optimize collapse_pte_mapped_thp().
> > > > >
> > > > > On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
> > > > > Then, calling ptep_clear() for every pte will cause a TLB flush for every
> > > > > contpte block. Instead, clear_full_ptes() does a
> > > > > contpte_try_unfold_partial() which will flush the TLB only for the (if any)
> > > > > starting and ending contpte block, if they partially overlap with the range
> > > > > khugepaged is looking at.
> > > > >
> > > > > For all arches, there should be a benefit due to batching atomic operations
> > > > > on mapcounts due to folio_remove_rmap_ptes().
> > > > >
> > > > > Note that we do not need to make a change to the check
> > > > > "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
> > > > > to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
> > > > > pages of the folio will be equal to the corresponding pages of our
> > > > > batch mapping consecutive pages.
> > > > >
> > > > > No issues were observed with mm-selftests.
> > > > >
> > > > > Signed-off-by: Dev Jain <dev.jain@arm.com>
> > > > > ---
> > > > >    mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
> > > > >    1 file changed, 26 insertions(+), 12 deletions(-)
> > > > >
> > > > > diff --git a/mm/khugepaged.c b/mm/khugepaged.c
> > > > > index 3944b112d452..4c8d33abfbd8 100644
> > > > > --- a/mm/khugepaged.c
> > > > > +++ b/mm/khugepaged.c
> > > > > @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
> > > > >    int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > > > >    			    bool install_pmd)
> > > > >    {
> > > > > +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
> > > > >    	struct mmu_notifier_range range;
> > > > >    	bool notified = false;
> > > > >    	unsigned long haddr = addr & HPAGE_PMD_MASK;
> > > > > +	unsigned long end = haddr + HPAGE_PMD_SIZE;
> > > > >    	struct vm_area_struct *vma = vma_lookup(mm, haddr);
> > > > >    	struct folio *folio;
> > > > >    	pte_t *start_pte, *pte;
> > > > >    	pmd_t *pmd, pgt_pmd;
> > > > >    	spinlock_t *pml = NULL, *ptl;
> > > > > -	int nr_ptes = 0, result = SCAN_FAIL;
> > > > >    	int i;
> > > > >
> > > > >    	mmap_assert_locked(mm);
> > > > > @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > > > >    		goto abort;
> > > > >
> > > > >    	/* step 2: clear page table and adjust rmap */
> > > > > -	for (i = 0, addr = haddr, pte = start_pte;
> > > > > -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
> > > > > +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
> > > > > +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
> > > > > +	     pte += nr_batch_ptes) {
> > > > > +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
> > > > > +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
> > > > > +		struct folio *mapped_folio;
> > > > >    		struct page *page;
> > > > >    		pte_t ptent = ptep_get(pte);
> > > > >
> > > > > +		nr_batch_ptes = 1;
> > > > > +
> > > > >    		if (pte_none(ptent))
> > > > >    			continue;
> > > > >    		/*
> > > > > @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > > > >    			goto abort;
> > > > >    		}
> > > > >    		page = vm_normal_page(vma, addr, ptent);
> > > > > +		mapped_folio = page_folio(page);
> > > > > +
> > > > >    		if (folio_page(folio, i) != page)
> > > > >    			goto abort;
> > > > Isn't this asserting that folio == mapped_folio here? We're saying page is the
> > > > ith page of folio, so why do we need to look up mapped_folio?
> > > We need to check for all PTEs whether they map the right page or not. This may
> > > get disturbed due to mremap and stuff.
> > Right but I'm saying mapped_folio == folio right? You're literally asserting it
> > here? So there's no need to assign mapped_folio at all, just reference folio no?
> >
> > > > > +		mapped_folio = page_folio(page);
> > > > You're assigning this twice.
> > > Forgot to remove, thanks.
> > >
> > > > > +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
> > > > > +						max_nr_batch_ptes, flags,
> > > > > +						NULL, NULL, NULL);
> > > > > +
> > > > >    		/*
> > > > >    		 * Must clear entry, or a racing truncate may re-remove it.
> > > > >    		 * TLB flush can be left until pmdp_collapse_flush() does it.
> > > > >    		 * PTE dirty? Shmem page is already dirty; file is read-only.
> > > > >    		 */
> > > > > -		ptep_clear(mm, addr, pte);
> > > > > -		folio_remove_rmap_pte(folio, page, vma);
> > > > > -		nr_ptes++;
> > > > > +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
> > > > > +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
> > > > > +		nr_mapped_ptes += nr_batch_ptes;
> > > > >    	}
> > > > >
> > > > >    	if (!pml)
> > > > >    		spin_unlock(ptl);
> > > > >
> > > > >    	/* step 3: set proper refcount and mm_counters. */
> > > > > -	if (nr_ptes) {
> > > > > -		folio_ref_sub(folio, nr_ptes);
> > > > > -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> > > > > +	if (nr_mapped_ptes) {
> > > > > +		folio_ref_sub(folio, nr_mapped_ptes);
> > > > > +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
> > > > >    	}
> > > > >
> > > > >    	/* step 4: remove empty page table */
> > > > > @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
> > > > >    			: SCAN_SUCCEED;
> > > > >    	goto drop_folio;
> > > > >    abort:
> > > > > -	if (nr_ptes) {
> > > > > +	if (nr_mapped_ptes) {
> > > > I know it's ironic coming from me :P but I'm not sure why we need to churn this
> > > > up by renaming?
> > > Because nr_ptes is an existing variable and I need a new variable to make
> > > the jump at the end of the PTE batch.
> > I thought you eliminated nr_ptes as a variable here? Where else is it used?
> >
> > Oh how this code needs refactoring...
>
> If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
> which is confusing since the former is plain and the latter has a _mapped_ thingy
> in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.
>

Sigh, this is still awful. But probably just existing awfulness. This whole
thing needs a tent thrown over it and fumigation... but again not your fault :)

I mean fine, this is fine then.

> >
> > > > >    		flush_tlb_mm(mm);
> > > > > -		folio_ref_sub(folio, nr_ptes);
> > > > > -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
> > > > > +		folio_ref_sub(folio, nr_mapped_ptes);
> > > > > +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
> > > > >    	}
> > > > >    unlock:
> > > > >    	if (start_pte)
> > > > > --
> > > > > 2.30.2
> > > > >
> > > > V

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[David Hildenbrand]

On 15.07.25 11:43, Lorenzo Stoakes wrote:
> On Tue, Jul 15, 2025 at 12:04:56PM +0530, Dev Jain wrote:
>>
>> On 26/06/25 10:17 am, Lorenzo Stoakes wrote:
>>> On Thu, Jun 26, 2025 at 09:18:47AM +0530, Dev Jain wrote:
>>>> On 25/06/25 6:41 pm, Lorenzo Stoakes wrote:
>>>>> On Wed, Jun 25, 2025 at 11:28:05AM +0530, Dev Jain wrote:
>>>>>> Use PTE batching to optimize collapse_pte_mapped_thp().
>>>>>>
>>>>>> On arm64, suppose khugepaged is scanning a pte-mapped 2MB THP for collapse.
>>>>>> Then, calling ptep_clear() for every pte will cause a TLB flush for every
>>>>>> contpte block. Instead, clear_full_ptes() does a
>>>>>> contpte_try_unfold_partial() which will flush the TLB only for the (if any)
>>>>>> starting and ending contpte block, if they partially overlap with the range
>>>>>> khugepaged is looking at.
>>>>>>
>>>>>> For all arches, there should be a benefit due to batching atomic operations
>>>>>> on mapcounts due to folio_remove_rmap_ptes().
>>>>>>
>>>>>> Note that we do not need to make a change to the check
>>>>>> "if (folio_page(folio, i) != page)"; if i'th page of the folio is equal
>>>>>> to the first page of our batch, then i + 1, .... i + nr_batch_ptes - 1
>>>>>> pages of the folio will be equal to the corresponding pages of our
>>>>>> batch mapping consecutive pages.
>>>>>>
>>>>>> No issues were observed with mm-selftests.
>>>>>>
>>>>>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>>>>>> ---
>>>>>>     mm/khugepaged.c | 38 ++++++++++++++++++++++++++------------
>>>>>>     1 file changed, 26 insertions(+), 12 deletions(-)
>>>>>>
>>>>>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>>>>>> index 3944b112d452..4c8d33abfbd8 100644
>>>>>> --- a/mm/khugepaged.c
>>>>>> +++ b/mm/khugepaged.c
>>>>>> @@ -1499,15 +1499,16 @@ static int set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
>>>>>>     int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>>>     			    bool install_pmd)
>>>>>>     {
>>>>>> +	int nr_mapped_ptes = 0, nr_batch_ptes, result = SCAN_FAIL;
>>>>>>     	struct mmu_notifier_range range;
>>>>>>     	bool notified = false;
>>>>>>     	unsigned long haddr = addr & HPAGE_PMD_MASK;
>>>>>> +	unsigned long end = haddr + HPAGE_PMD_SIZE;
>>>>>>     	struct vm_area_struct *vma = vma_lookup(mm, haddr);
>>>>>>     	struct folio *folio;
>>>>>>     	pte_t *start_pte, *pte;
>>>>>>     	pmd_t *pmd, pgt_pmd;
>>>>>>     	spinlock_t *pml = NULL, *ptl;
>>>>>> -	int nr_ptes = 0, result = SCAN_FAIL;
>>>>>>     	int i;
>>>>>>
>>>>>>     	mmap_assert_locked(mm);
>>>>>> @@ -1621,11 +1622,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>>>     		goto abort;
>>>>>>
>>>>>>     	/* step 2: clear page table and adjust rmap */
>>>>>> -	for (i = 0, addr = haddr, pte = start_pte;
>>>>>> -	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
>>>>>> +	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
>>>>>> +	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
>>>>>> +	     pte += nr_batch_ptes) {
>>>>>> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>>>>>> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
>>>>>> +		struct folio *mapped_folio;
>>>>>>     		struct page *page;
>>>>>>     		pte_t ptent = ptep_get(pte);
>>>>>>
>>>>>> +		nr_batch_ptes = 1;
>>>>>> +
>>>>>>     		if (pte_none(ptent))
>>>>>>     			continue;
>>>>>>     		/*
>>>>>> @@ -1639,26 +1646,33 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>>>     			goto abort;
>>>>>>     		}
>>>>>>     		page = vm_normal_page(vma, addr, ptent);
>>>>>> +		mapped_folio = page_folio(page);
>>>>>> +
>>>>>>     		if (folio_page(folio, i) != page)
>>>>>>     			goto abort;
>>>>> Isn't this asserting that folio == mapped_folio here? We're saying page is the
>>>>> ith page of folio, so why do we need to look up mapped_folio?
>>>> We need to check for all PTEs whether they map the right page or not. This may
>>>> get disturbed due to mremap and stuff.
>>> Right but I'm saying mapped_folio == folio right? You're literally asserting it
>>> here? So there's no need to assign mapped_folio at all, just reference folio no?
>>>
>>>>>> +		mapped_folio = page_folio(page);
>>>>> You're assigning this twice.
>>>> Forgot to remove, thanks.
>>>>
>>>>>> +		nr_batch_ptes = folio_pte_batch(mapped_folio, addr, pte, ptent,
>>>>>> +						max_nr_batch_ptes, flags,
>>>>>> +						NULL, NULL, NULL);
>>>>>> +
>>>>>>     		/*
>>>>>>     		 * Must clear entry, or a racing truncate may re-remove it.
>>>>>>     		 * TLB flush can be left until pmdp_collapse_flush() does it.
>>>>>>     		 * PTE dirty? Shmem page is already dirty; file is read-only.
>>>>>>     		 */
>>>>>> -		ptep_clear(mm, addr, pte);
>>>>>> -		folio_remove_rmap_pte(folio, page, vma);
>>>>>> -		nr_ptes++;
>>>>>> +		clear_full_ptes(mm, addr, pte, nr_batch_ptes, /* full = */ false);
>>>>>> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
>>>>>> +		nr_mapped_ptes += nr_batch_ptes;
>>>>>>     	}
>>>>>>
>>>>>>     	if (!pml)
>>>>>>     		spin_unlock(ptl);
>>>>>>
>>>>>>     	/* step 3: set proper refcount and mm_counters. */
>>>>>> -	if (nr_ptes) {
>>>>>> -		folio_ref_sub(folio, nr_ptes);
>>>>>> -		add_mm_counter(mm, mm_counter_file(folio), -nr_ptes);
>>>>>> +	if (nr_mapped_ptes) {
>>>>>> +		folio_ref_sub(folio, nr_mapped_ptes);
>>>>>> +		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
>>>>>>     	}
>>>>>>
>>>>>>     	/* step 4: remove empty page table */
>>>>>> @@ -1691,10 +1705,10 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>>>>>     			: SCAN_SUCCEED;
>>>>>>     	goto drop_folio;
>>>>>>     abort:
>>>>>> -	if (nr_ptes) {
>>>>>> +	if (nr_mapped_ptes) {
>>>>> I know it's ironic coming from me :P but I'm not sure why we need to churn this
>>>>> up by renaming?
>>>> Because nr_ptes is an existing variable and I need a new variable to make
>>>> the jump at the end of the PTE batch.
>>> I thought you eliminated nr_ptes as a variable here? Where else is it used?
>>>
>>> Oh how this code needs refactoring...
>>
>> If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
>> which is confusing since the former is plain and the latter has a _mapped_ thingy
>> in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.
>>
> 
> Sigh, this is still awful. But probably just existing awfulness. This whole
> thing needs a tent thrown over it and fumigation... but again not your fault :)
> 
> I mean fine, this is fine then.

Probably best to be pragmatic here: as long as the educated reader 
understands the code, all good. I hope I'll never have to explain it to 
my daughters.

I know, I'm a perfectionist myself ;)

-- 
Cheers,

David / dhildenb

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Lorenzo Stoakes]

On Tue, Jul 15, 2025 at 11:56:26AM +0200, David Hildenbrand wrote:
> > > > > > > -	if (nr_ptes) {
> > > > > > > +	if (nr_mapped_ptes) {
> > > > > > I know it's ironic coming from me :P but I'm not sure why we need to churn this
> > > > > > up by renaming?
> > > > > Because nr_ptes is an existing variable and I need a new variable to make
> > > > > the jump at the end of the PTE batch.
> > > > I thought you eliminated nr_ptes as a variable here? Where else is it used?
> > > >
> > > > Oh how this code needs refactoring...
> > >
> > > If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
> > > which is confusing since the former is plain and the latter has a _mapped_ thingy
> > > in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.
> > >
> >
> > Sigh, this is still awful. But probably just existing awfulness. This whole
> > thing needs a tent thrown over it and fumigation... but again not your fault :)
> >
> > I mean fine, this is fine then.
>
> Probably best to be pragmatic here: as long as the educated reader
> understands the code, all good. I hope I'll never have to explain it to my
> daughters.

:)

To be crystal clear here - Dev - I am not asking you to refactor the existing
mess, and obviously I concede your point on the issue raised here.

Did you need to respin other things here? If not I can relook and give tags.

>
> I know, I'm a perfectionist myself ;)

I mean if I wanted perfection I'd have run screaming from the kernel right at
the start :P

But in general on the refactoring point - I think this code is hard to
understand even for the informed, the endless nesting and nesting and levels of
'you just have to know what's happening' make it really easy to miss details
here.

What I'm saying is - we can improve on this, and really all I'm doing is
todo++. Possibly even dave_todo++?? ;) Or could even be $<name here>_todo++ :P

I find it impossible to not mention this when it seems apparent to me... force
of habit :)

Cheers, Lorenzo

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Dev Jain]

On 15/07/25 3:32 pm, Lorenzo Stoakes wrote:
> On Tue, Jul 15, 2025 at 11:56:26AM +0200, David Hildenbrand wrote:
>>>>>>>> -	if (nr_ptes) {
>>>>>>>> +	if (nr_mapped_ptes) {
>>>>>>> I know it's ironic coming from me :P but I'm not sure why we need to churn this
>>>>>>> up by renaming?
>>>>>> Because nr_ptes is an existing variable and I need a new variable to make
>>>>>> the jump at the end of the PTE batch.
>>>>> I thought you eliminated nr_ptes as a variable here? Where else is it used?
>>>>>
>>>>> Oh how this code needs refactoring...
>>>> If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
>>>> which is confusing since the former is plain and the latter has a _mapped_ thingy
>>>> in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.
>>>>
>>> Sigh, this is still awful. But probably just existing awfulness. This whole
>>> thing needs a tent thrown over it and fumigation... but again not your fault :)
>>>
>>> I mean fine, this is fine then.
>> Probably best to be pragmatic here: as long as the educated reader
>> understands the code, all good. I hope I'll never have to explain it to my
>> daughters.
> :)
>
> To be crystal clear here - Dev - I am not asking you to refactor the existing
> mess, and obviously I concede your point on the issue raised here.
>
> Did you need to respin other things here? If not I can relook and give tags.


Yes, I will have to respin because of folio_pte_batch() changes from David and
have to add a simplification patch to this series authored by David, I was just
doing that and was reading the email thread to see if I didn't miss any review
comment and then I noticed I forgot to reply to the above point.


>
>> I know, I'm a perfectionist myself ;)
> I mean if I wanted perfection I'd have run screaming from the kernel right at
> the start :P
>
> But in general on the refactoring point - I think this code is hard to
> understand even for the informed, the endless nesting and nesting and levels of
> 'you just have to know what's happening' make it really easy to miss details
> here.
>
> What I'm saying is - we can improve on this, and really all I'm doing is
> todo++. Possibly even dave_todo++?? ;) Or could even be $<name here>_todo++ :P
>
> I find it impossible to not mention this when it seems apparent to me... force
> of habit :)

I personally suck at refactoring code - it is something I have been trying to get
better at. Now that I understand khugepaged a code a lot better, I have been
thinking for a long time to clean this file up, maybe someday I will gather
the courage to do that : )
  

>
> Cheers, Lorenzo

Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching

[Lorenzo Stoakes]

On Tue, Jul 15, 2025 at 04:10:11PM +0530, Dev Jain wrote:
>
> On 15/07/25 3:32 pm, Lorenzo Stoakes wrote:
> > On Tue, Jul 15, 2025 at 11:56:26AM +0200, David Hildenbrand wrote:
> > > > > > > > > -	if (nr_ptes) {
> > > > > > > > > +	if (nr_mapped_ptes) {
> > > > > > > > I know it's ironic coming from me :P but I'm not sure why we need to churn this
> > > > > > > > up by renaming?
> > > > > > > Because nr_ptes is an existing variable and I need a new variable to make
> > > > > > > the jump at the end of the PTE batch.
> > > > > > I thought you eliminated nr_ptes as a variable here? Where else is it used?
> > > > > >
> > > > > > Oh how this code needs refactoring...
> > > > > If we retain nr_ptes, then the two variables will be nr_ptes and nr_mapped_ptes,
> > > > > which is confusing since the former is plain and the latter has a _mapped_ thingy
> > > > > in it, so instead now we call them nr_batch_ptes and nr_mapped_ptes.
> > > > >
> > > > Sigh, this is still awful. But probably just existing awfulness. This whole
> > > > thing needs a tent thrown over it and fumigation... but again not your fault :)
> > > >
> > > > I mean fine, this is fine then.
> > > Probably best to be pragmatic here: as long as the educated reader
> > > understands the code, all good. I hope I'll never have to explain it to my
> > > daughters.
> > :)
> >
> > To be crystal clear here - Dev - I am not asking you to refactor the existing
> > mess, and obviously I concede your point on the issue raised here.
> >
> > Did you need to respin other things here? If not I can relook and give tags.
>
>
> Yes, I will have to respin because of folio_pte_batch() changes from David and
> have to add a simplification patch to this series authored by David, I was just
> doing that and was reading the email thread to see if I didn't miss any review
> comment and then I noticed I forgot to reply to the above point.

Cool thanks :)

Will have a look on respin!

>
>
> >
> > > I know, I'm a perfectionist myself ;)
> > I mean if I wanted perfection I'd have run screaming from the kernel right at
> > the start :P
> >
> > But in general on the refactoring point - I think this code is hard to
> > understand even for the informed, the endless nesting and nesting and levels of
> > 'you just have to know what's happening' make it really easy to miss details
> > here.
> >
> > What I'm saying is - we can improve on this, and really all I'm doing is
> > todo++. Possibly even dave_todo++?? ;) Or could even be $<name here>_todo++ :P
> >
> > I find it impossible to not mention this when it seems apparent to me... force
> > of habit :)
>
> I personally suck at refactoring code - it is something I have been trying to get
> better at. Now that I understand khugepaged a code a lot better, I have been
> thinking for a long time to clean this file up, maybe someday I will gather
> the courage to do that : )

:))

>
> >
> > Cheers, Lorenzo