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

Dev Jain posted 3 patches 3 months, 2 weeks ago
There is a newer version of this series
[PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching
Posted by Dev Jain 3 months, 2 weeks ago
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
Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching
Posted by Lorenzo Stoakes 3 months, 2 weeks ago
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 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching
Posted by Dev Jain 3 months, 2 weeks ago
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 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching
Posted by Lorenzo Stoakes 3 months, 2 weeks ago
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
Posted by Dev Jain 2 months, 3 weeks ago
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
Posted by Lorenzo Stoakes 2 months, 3 weeks ago
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
Posted by David Hildenbrand 2 months, 3 weeks ago
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
Posted by Lorenzo Stoakes 2 months, 3 weeks ago
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
Posted by Dev Jain 2 months, 3 weeks ago
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
Posted by Lorenzo Stoakes 2 months, 3 weeks ago
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
Re: [PATCH v2 2/3] khugepaged: Optimize collapse_pte_mapped_thp() for large folios by PTE batching
Posted by Dev Jain 3 months, 2 weeks ago
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