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

[Dev Jain <dev.jain@arm.com>]

On 18/06/25 11:20 pm, Lorenzo Stoakes wrote:
> This series has a lot of duplication in it esp vs. your other series [0], but
> perhaps something we can tackle in a follow up.
>
> It'd be nice if we could find a way to de-duplicate some of the near-identical
> code though.
>
> But it's a 'maybe' on that because hey, the code in this file is hideous anyway
> and needs a mega rework in any case...
>
> [0]: https://lore.kernel.org/all/20250618102607.10551-1-dev.jain@arm.com/
>
> On Wed, Jun 18, 2025 at 09:26:08PM +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>
>> ---
>>
>> This is rebased on:
>> https://lore.kernel.org/all/20250618102607.10551-1-dev.jain@arm.com/
>> If there will be a v2 of either version I'll send them together.
> Hmmm I say again - slow down a bit :) there's no need to shoot out multiple
> patches in a single day and you'd maybe avoid some of this kind of thing.
>
> It's really preferable to avoid possible conflicts like this or at least reduce
> the chance by having review on one thing done first.
>
> I mean, why not just put both of these in a series for the respin? Just a
> thought ;) in fact this is probably an ideal use of a series for that as you can
> ensure you deal with both if any conflicts arise.

Sorry for this. I found these two patches independently with a couple of
hours difference, and I posted this one yesterday because stupid me thought
someone will, after seeing my first patch, point out that the optimization
can be made at one more place. So I will send this and the other patch as
a series anyway for v2.

>
>>   mm/khugepaged.c | 38 +++++++++++++++++++++++++-------------
>>   1 file changed, 25 insertions(+), 13 deletions(-)
>>
>> diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>> index 649ccb2670f8..7d37058eda5b 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;
> NIT: I don't know why you're moving this, and while y'know it's kind of the fun
> of subjective stuff I'd rather the assigned values and unassigned values be on
> different lines (yes I know this codebase violates this with the pml, ptl below
> but hey :P)

To maintain a reverse Xmas fashion. Now I know that the declarations are already
not in an Xmas fashion, but I wanted to make sure the code I change maintains
that for the part I am changing :)

>
>>   	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);
>> @@ -1620,12 +1621,17 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   	if (unlikely(!pmd_same(pgt_pmd, pmdp_get_lockless(pmd))))
>>   		goto abort;
>>
>> +	i = 0, addr = haddr, pte = start_pte;
> This is horrid, no absolutely not. This is not how we do assignment in arbitrary
> C code.
>
> I don't know why we need a do/while here in general, I think the for loop should
> still work ok no?

I don't recall now and I cannot even find it, but I have been following this
pattern for some time, by god I cannot remember which pattern I am copying
from. Anyhow I also hate the do/while thingy so I will change this to a
for loop.

>
>>   	/* 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++) {
>> +	do {
>> +		const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>> +		int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
>> +		struct folio *this_folio;
> Hate this name. We are not C#... ;)
>
> Just call it folio no? The 'this_' is redundant.

There is already a struct folio *folio which we retrieve from filemap_lock_folio.
So wanted to differentiate, I didn't think hard about the name. How about mapped_folio?

>
>
>>   		struct page *page;
>>   		pte_t ptent = ptep_get(pte);
>>
>> +		nr_batch_ptes = 1;
>> +
>>   		if (pte_none(ptent))
>>   			continue;
>>   		/*
>> @@ -1639,6 +1645,11 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   			goto abort;
>>   		}
>>   		page = vm_normal_page(vma, addr, ptent);
>> +		this_folio = page_folio(page);
>> +		if (folio_test_large(this_folio) && max_nr_batch_ptes != 1)
>> +			nr_batch_ptes = folio_pte_batch(this_folio, addr, pte, ptent,
>> +					max_nr_batch_ptes, flags, NULL, NULL, NULL);
>> +
>>   		if (folio_page(folio, i) != page)
>>   			goto abort;
>>
>> @@ -1647,18 +1658,19 @@ int collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
>>   		 * 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, false);
>> +		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
>> +		nr_mapped_ptes += nr_batch_ptes;
>> +	} while (i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
>> +		 pte += nr_batch_ptes, i < HPAGE_PMD_NR);
>>
>>   	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 +1703,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
>>
> Logic looks generally sane though... :)