Re: [PATCH v4 3/4] mm: Optimize mprotect() by PTE-batching

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

On 01/07/25 1:33 pm, Ryan Roberts wrote:
> On 30/06/2025 13:52, Lorenzo Stoakes wrote:
>> On Sat, Jun 28, 2025 at 05:04:34PM +0530, Dev Jain wrote:
>>> Use folio_pte_batch to batch process a large folio. Reuse the folio from
>>> prot_numa case if possible.
>>>
>>> For all cases other than the PageAnonExclusive case, if the case holds true
>>> for one pte in the batch, one can confirm that that case will hold true for
>>> other ptes in the batch too; for pte_needs_soft_dirty_wp(), we do not pass
>>> FPB_IGNORE_SOFT_DIRTY. modify_prot_start_ptes() collects the dirty
>>> and access bits across the batch, therefore batching across
>>> pte_dirty(): this is correct since the dirty bit on the PTE really is
>>> just an indication that the folio got written to, so even if the PTE is
>>> not actually dirty (but one of the PTEs in the batch is), the wp-fault
>>> optimization can be made.
>>>
>>> The crux now is how to batch around the PageAnonExclusive case; we must
>>> check the corresponding condition for every single page. Therefore, from
>>> the large folio batch, we process sub batches of ptes mapping pages with
>>> the same PageAnonExclusive condition, and process that sub batch, then
>>> determine and process the next sub batch, and so on. Note that this does
>>> not cause any extra overhead; if suppose the size of the folio batch
>>> is 512, then the sub batch processing in total will take 512 iterations,
>>> which is the same as what we would have done before.
>>>
>>> Signed-off-by: Dev Jain <dev.jain@arm.com>
>>> ---
>>>   mm/mprotect.c | 143 +++++++++++++++++++++++++++++++++++++++++---------
>>>   1 file changed, 117 insertions(+), 26 deletions(-)
>>>
>>> diff --git a/mm/mprotect.c b/mm/mprotect.c
>>> index 627b0d67cc4a..28c7ce7728ff 100644
>>> --- a/mm/mprotect.c
>>> +++ b/mm/mprotect.c
>>> @@ -40,35 +40,47 @@
>>>
>>>   #include "internal.h"
>>>
>>> -bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>>> -			     pte_t pte)
>>> -{
>>> -	struct page *page;
>>> +enum tristate {
>>> +	TRI_FALSE = 0,
>>> +	TRI_TRUE = 1,
>>> +	TRI_MAYBE = -1,
>>> +};
>> Yeah no, absolutely not, this is horrible, I don't want to see an arbitrary type
>> like this added, to a random file, and I absolutely think this adds confusion
>> and does not in any way help clarify things.
>>
>>> +/*
>>> + * Returns enum tristate indicating whether the pte can be changed to writable.
>>> + * If TRI_MAYBE is returned, then the folio is anonymous and the user must
>>> + * additionally check PageAnonExclusive() for every page in the desired range.
>>> + */
>>> +static int maybe_change_pte_writable(struct vm_area_struct *vma,
>>> +				     unsigned long addr, pte_t pte,
>>> +				     struct folio *folio)
>>> +{
>>>   	if (WARN_ON_ONCE(!(vma->vm_flags & VM_WRITE)))
>>> -		return false;
>>> +		return TRI_FALSE;
>>>
>>>   	/* Don't touch entries that are not even readable. */
>>>   	if (pte_protnone(pte))
>>> -		return false;
>>> +		return TRI_FALSE;
>>>
>>>   	/* Do we need write faults for softdirty tracking? */
>>>   	if (pte_needs_soft_dirty_wp(vma, pte))
>>> -		return false;
>>> +		return TRI_FALSE;
>>>
>>>   	/* Do we need write faults for uffd-wp tracking? */
>>>   	if (userfaultfd_pte_wp(vma, pte))
>>> -		return false;
>>> +		return TRI_FALSE;
>>>
>>>   	if (!(vma->vm_flags & VM_SHARED)) {
>>>   		/*
>>>   		 * Writable MAP_PRIVATE mapping: We can only special-case on
>>>   		 * exclusive anonymous pages, because we know that our
>>>   		 * write-fault handler similarly would map them writable without
>>> -		 * any additional checks while holding the PT lock.
>>> +		 * any additional checks while holding the PT lock. So if the
>>> +		 * folio is not anonymous, we know we cannot change pte to
>>> +		 * writable. If it is anonymous then the caller must further
>>> +		 * check that the page is AnonExclusive().
>>>   		 */
>>> -		page = vm_normal_page(vma, addr, pte);
>>> -		return page && PageAnon(page) && PageAnonExclusive(page);
>>> +		return (!folio || folio_test_anon(folio)) ? TRI_MAYBE : TRI_FALSE;
>>>   	}
>>>
>>>   	VM_WARN_ON_ONCE(is_zero_pfn(pte_pfn(pte)) && pte_dirty(pte));
>>> @@ -80,15 +92,61 @@ bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>>>   	 * FS was already notified and we can simply mark the PTE writable
>>>   	 * just like the write-fault handler would do.
>>>   	 */
>>> -	return pte_dirty(pte);
>>> +	return pte_dirty(pte) ? TRI_TRUE : TRI_FALSE;
>>> +}
>> Yeah not a fan of this at all.
>>
>> This is squashing all the logic into one place when we don't really need to.
>>
>> We can separate out the shared logic and just do something like:
>>
>> ////// Lorenzo's suggestion //////
>>
>> -bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>> -			     pte_t pte)
>> +static bool maybe_change_pte_writable(struct vm_area_struct *vma,
>> +		pte_t pte)
>>   {
>> -	struct page *page;
>> -
>>   	if (WARN_ON_ONCE(!(vma->vm_flags & VM_WRITE)))
>>   		return false;
>>
>> @@ -60,16 +58,14 @@ bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>>   	if (userfaultfd_pte_wp(vma, pte))
>>   		return false;
>>
>> -	if (!(vma->vm_flags & VM_SHARED)) {
>> -		/*
>> -		 * Writable MAP_PRIVATE mapping: We can only special-case on
>> -		 * exclusive anonymous pages, because we know that our
>> -		 * write-fault handler similarly would map them writable without
>> -		 * any additional checks while holding the PT lock.
>> -		 */
>> -		page = vm_normal_page(vma, addr, pte);
>> -		return page && PageAnon(page) && PageAnonExclusive(page);
>> -	}
>> +	return true;
>> +}
>> +
>> +static bool can_change_shared_pte_writable(struct vm_area_struct *vma,
>> +		pte_t pte)
>> +{
>> +	if (!maybe_change_pte_writable(vma, pte))
>> +		return false;
>>
>>   	VM_WARN_ON_ONCE(is_zero_pfn(pte_pfn(pte)) && pte_dirty(pte));
>>
>> @@ -83,6 +79,33 @@ bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>>   	return pte_dirty(pte);
>>   }
>>
>> +static bool can_change_private_pte_writable(struct vm_area_struct *vma,
>> +		unsigned long addr, pte_t pte)
>> +{
>> +	struct page *page;
>> +
>> +	if (!maybe_change_pte_writable(vma, pte))
>> +		return false;
>> +
>> +	/*
>> +	 * Writable MAP_PRIVATE mapping: We can only special-case on
>> +	 * exclusive anonymous pages, because we know that our
>> +	 * write-fault handler similarly would map them writable without
>> +	 * any additional checks while holding the PT lock.
>> +	 */
>> +	page = vm_normal_page(vma, addr, pte);
>> +	return page && PageAnon(page) && PageAnonExclusive(page);
>> +}
>> +
>> +bool can_change_pte_writable(struct vm_area_struct *vma,
>> +		unsigned long addr, pte_t pte)
>> +{
>> +	if (vma->vm_flags & VM_SHARED)
>> +		return can_change_shared_pte_writable(vma, pte);
>> +
>> +	return can_change_private_pte_writable(vma, addr, pte);
>> +}
>> +
>>
>> ////// end of Lorenzo's suggestion //////
>>
>> You can obviously modify this to change other stuff like whether you feed back
>> the PAE or not in private case for use in your code.
> This sugestion for this part of the problem looks much cleaner!
>
> Sorry; this whole struct tristate thing was my idea. I never really liked it but
> I was more focussed on trying to illustrate the big picture flow that I thought
> would work well with a batch and sub-batches (which it seems below that you
> hate... but let's talk about that down there).
>
>>> +
>>> +/*
>>> + * Returns the number of pages within the folio, starting from the page
>>> + * indicated by pgidx and up to pgidx + max_nr, that have the same value of
>>> + * PageAnonExclusive(). Must only be called for anonymous folios. Value of
>>> + * PageAnonExclusive() is returned in *exclusive.
>>> + */
>>> +static int anon_exclusive_batch(struct folio *folio, int pgidx, int max_nr,
>>> +				bool *exclusive)
>> Let's generalise it to something like count_folio_fungible_pages()
>>
>> or maybe count_folio_batchable_pages()?
>>
>> Yes naming is hard... :P but right now it reads like this is returning a batch
>> or doing something with a batch.
>>
>>> +{
>>> +	struct page *page;
>>> +	int nr = 1;
>>> +
>>> +	if (!folio) {
>>> +		*exclusive = false;
>>> +		return nr;
>>> +	}
>>> +
>>> +	page = folio_page(folio, pgidx++);
>>> +	*exclusive = PageAnonExclusive(page);
>>> +	while (nr < max_nr) {
>> The C programming language asks why you don't like using for :)
>>
>>> +		page = folio_page(folio, pgidx++);
>>> +		if ((*exclusive) != PageAnonExclusive(page))
>>> +			break;
>>> +		nr++;
>> This *exclusive stuff makes me want to cry :)
>>
>> Just set a local variable and hand it back at the end.
>>
>>> +	}
>>> +
>>> +	return nr;
>>> +}
>>> +
>>> +bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
>>> +			     pte_t pte)
>>> +{
>>> +	struct page *page;
>>> +	int ret;
>>> +
>>> +	ret = maybe_change_pte_writable(vma, addr, pte, NULL);
>>> +	if (ret == TRI_MAYBE) {
>>> +		page = vm_normal_page(vma, addr, pte);
>>> +		ret = page && PageAnon(page) && PageAnonExclusive(page);
>>> +	}
>>> +
>>> +	return ret;
>>>   }
>> See above comments on this stuff.
>>
>>>   static int mprotect_folio_pte_batch(struct folio *folio, unsigned long addr,
>>> -		pte_t *ptep, pte_t pte, int max_nr_ptes)
>>> +		pte_t *ptep, pte_t pte, int max_nr_ptes, fpb_t switch_off_flags)
>> This last parameter is pretty horrible. It's a negative mask so now you're
>> passing 'ignore soft dirty' to the function meaning 'don't ignore it'. This is
>> just planting land mines.
>>
>> Obviously David's flag changes will also alter all this.
>>
>> Just add a boolean re: soft dirty.
> Dev had a boolean for this in the last round. I've seen various functions expand
> over time with increasing numbers of bool flags. So I asked to convert to a
> flags parameter and just pass in the flags we need. Then it's a bit more future
> proof and self documenting. (For the record I dislike the "switch_off_flags"
> approach taken here).
>
>>>   {
>>> -	const fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>>> +	fpb_t flags = FPB_IGNORE_DIRTY | FPB_IGNORE_SOFT_DIRTY;
>>> +
>>> +	flags &= ~switch_off_flags;
>>>
>>> -	if (!folio || !folio_test_large(folio) || (max_nr_ptes == 1))
>>> +	if (!folio || !folio_test_large(folio))
>>>   		return 1;
>> Why remove this last check?
>>
>>>   	return folio_pte_batch(folio, addr, ptep, pte, max_nr_ptes, flags,
>>> @@ -154,7 +212,8 @@ static int prot_numa_skip_ptes(struct folio **foliop, struct vm_area_struct *vma
>>>   	}
>>>
>>>   skip_batch:
>>> -	nr_ptes = mprotect_folio_pte_batch(folio, addr, pte, oldpte, max_nr_ptes);
>>> +	nr_ptes = mprotect_folio_pte_batch(folio, addr, pte, oldpte,
>>> +					   max_nr_ptes, 0);
>> See above about flag param. If you change to boolean, please prefix this with
>> e.g. /* set_soft_dirty= */ true or whatever the flag ends up being :)
>>
>>>   out:
>>>   	*foliop = folio;
>>>   	return nr_ptes;
>>> @@ -191,7 +250,10 @@ static long change_pte_range(struct mmu_gather *tlb,
>>>   		if (pte_present(oldpte)) {
>>>   			int max_nr_ptes = (end - addr) >> PAGE_SHIFT;
>>>   			struct folio *folio = NULL;
>>> -			pte_t ptent;
>>> +			int sub_nr_ptes, pgidx = 0;
>>> +			pte_t ptent, newpte;
>>> +			bool sub_set_write;
>>> +			int set_write;
>>>
>>>   			/*
>>>   			 * Avoid trapping faults against the zero or KSM
>>> @@ -206,6 +268,11 @@ static long change_pte_range(struct mmu_gather *tlb,
>>>   					continue;
>>>   			}
>>>
>>> +			if (!folio)
>>> +				folio = vm_normal_folio(vma, addr, oldpte);
>>> +
>>> +			nr_ptes = mprotect_folio_pte_batch(folio, addr, pte, oldpte,
>>> +							   max_nr_ptes, FPB_IGNORE_SOFT_DIRTY);
>> Don't we only care about S/D if pte_needs_soft_dirty_wp()?
>>
>>>   			oldpte = modify_prot_start_ptes(vma, addr, pte, nr_ptes);
>>>   			ptent = pte_modify(oldpte, newprot);
>>>
>>> @@ -227,15 +294,39 @@ static long change_pte_range(struct mmu_gather *tlb,
>>>   			 * example, if a PTE is already dirty and no other
>>>   			 * COW or special handling is required.
>>>   			 */
>>> -			if ((cp_flags & MM_CP_TRY_CHANGE_WRITABLE) &&
>>> -			    !pte_write(ptent) &&
>>> -			    can_change_pte_writable(vma, addr, ptent))
>>> -				ptent = pte_mkwrite(ptent, vma);
>>> -
>>> -			modify_prot_commit_ptes(vma, addr, pte, oldpte, ptent, nr_ptes);
>>> -			if (pte_needs_flush(oldpte, ptent))
>>> -				tlb_flush_pte_range(tlb, addr, PAGE_SIZE);
>>> -			pages++;
>>> +			set_write = (cp_flags & MM_CP_TRY_CHANGE_WRITABLE) &&
>>> +				    !pte_write(ptent);
>>> +			if (set_write)
>>> +				set_write = maybe_change_pte_writable(vma, addr, ptent, folio);
>>> +
>>> +			while (nr_ptes) {
>>> +				if (set_write == TRI_MAYBE) {
>>> +					sub_nr_ptes = anon_exclusive_batch(folio,
>>> +						pgidx, nr_ptes, &sub_set_write);
>>> +				} else {
>>> +					sub_nr_ptes = nr_ptes;
>>> +					sub_set_write = (set_write == TRI_TRUE);
>>> +				}
>>> +
>>> +				if (sub_set_write)
>>> +					newpte = pte_mkwrite(ptent, vma);
>>> +				else
>>> +					newpte = ptent;
>>> +
>>> +				modify_prot_commit_ptes(vma, addr, pte, oldpte,
>>> +							newpte, sub_nr_ptes);
>>> +				if (pte_needs_flush(oldpte, newpte))
>>> +					tlb_flush_pte_range(tlb, addr,
>>> +						sub_nr_ptes * PAGE_SIZE);
>>> +
>>> +				addr += sub_nr_ptes * PAGE_SIZE;
>>> +				pte += sub_nr_ptes;
>>> +				oldpte = pte_advance_pfn(oldpte, sub_nr_ptes);
>>> +				ptent = pte_advance_pfn(ptent, sub_nr_ptes);
>>> +				nr_ptes -= sub_nr_ptes;
>>> +				pages += sub_nr_ptes;
>>> +				pgidx += sub_nr_ptes;
>>> +			}
>> I hate hate hate having this loop here, let's abstract this please.
>>
>> I mean I think we can just use mprotect_folio_pte_batch() no? It's not
>> abstracting much here, and we can just do all this handling there. Maybe have to
>> pass in a bunch more params, but it saves us having to do all this.
> In an ideal world we would flatten and just have mprotect_folio_pte_batch()
> return the batch size considering all the relevant PTE bits AND the
> AnonExclusive bit on the pages. IIRC one of Dev's earlier versions modified the
> core folio_pte_batch() function to also look at the AnonExclusive bit, but I
> really disliked changing that core function (I think others did too?).

That patch was in our private exchange, not on the list.

>
> So barring that approach, we are really only left with the batch and sub-batch
> approach - although, yes, it could be abstracted more. We could maintain a
> context struct that persists across all calls to mprotect_folio_pte_batch() and
> it can use that to keep it's state to remember if we are in the middle of a
> sub-batch and decide either to call folio_pte_batch() to get a new batch, or
> call anon_exclusive_batch() to get the next sub-batch within the current batch.
> But that started to feel overly abstracted to me.
>
> This loop approach, as written, felt more straightforward for the reader to
> understand (i.e. the least-worst option). Is the context approach what you are
> suggesting or do you have something else in mind?
>
>> Alternatively, we could add a new wrapper function, but yeah definitely not
>> this.
>>
>> Also the C programming language asks... etc etc. ;)
>>
>> Overall since you always end up processing folio_nr_pages(folio) you can just
>> have the batch function or a wrapper return this and do updates as necessary
>> here on that basis, and leave the 'sub' batching to that function.
> Sorry I don't understand this statement - could you clarify? Especially the bit
> about "always ... processing folio_nr_pages(folio)"; I don't think we do. In
> various corner cases the size of the folio has no relationship to the way the
> PTEs are mapped.
>
> Thanks,
> Ryan
>
>>
>>>   		} else if (is_swap_pte(oldpte)) {
>>>   			swp_entry_t entry = pte_to_swp_entry(oldpte);
>>>   			pte_t newpte;
>>> --
>>> 2.30.2
>>>