[v3,01/15] mm: Batch-copy PTE ranges during fork()

  Convert copy_pte_range() to copy a set of ptes in a batch. A given batch
maps a physically contiguous block of memory, all belonging to the same
folio. This will likely improve performance by a tiny amount due to
batching the folio reference count management and calling set_ptes()
rather than making individual calls to set_pte_at().

However, the primary motivation for this change is to reduce the number
of tlb maintenance operations that the arm64 backend has to perform
during fork, as it is about to add transparent support for the
"contiguous bit" in its ptes. By write-protecting the parent using the
new ptep_set_wrprotects() (note the 's' at the end) function, the
backend can avoid having to unfold contig ranges of PTEs, which is
expensive, when all ptes in the range are being write-protected.
Similarly, by using set_ptes() rather than set_pte_at() to set up ptes
in the child, the backend does not need to fold a contiguous range once
they are all populated - they can be initially populated as a contiguous
range in the first place.

This change addresses the core-mm refactoring only, and introduces
ptep_set_wrprotects() with a default implementation that calls
ptep_set_wrprotect() for each pte in the range. A separate change will
implement ptep_set_wrprotects() in the arm64 backend to realize the
performance improvement as part of the work to enable contpte mappings.

Signed-off-by: Ryan Roberts <ryan.roberts@arm.com>
---
 include/linux/pgtable.h |  13 +++
 mm/memory.c             | 195 ++++++++++++++++++++++++++++++----------
 2 files changed, 162 insertions(+), 46 deletions(-)
  

On 04.12.23 11:54, Ryan Roberts wrote:
> Convert copy_pte_range() to copy a set of ptes in a batch. A given batch
> maps a physically contiguous block of memory, all belonging to the same
> folio. This will likely improve performance by a tiny amount due to
> batching the folio reference count management and calling set_ptes()
> rather than making individual calls to set_pte_at().
> 
> However, the primary motivation for this change is to reduce the number
> of tlb maintenance operations that the arm64 backend has to perform
> during fork, as it is about to add transparent support for the
> "contiguous bit" in its ptes. By write-protecting the parent using the
> new ptep_set_wrprotects() (note the 's' at the end) function, the
> backend can avoid having to unfold contig ranges of PTEs, which is
> expensive, when all ptes in the range are being write-protected.
> Similarly, by using set_ptes() rather than set_pte_at() to set up ptes
> in the child, the backend does not need to fold a contiguous range once
> they are all populated - they can be initially populated as a contiguous
> range in the first place.
> 
> This change addresses the core-mm refactoring only, and introduces
> ptep_set_wrprotects() with a default implementation that calls
> ptep_set_wrprotect() for each pte in the range. A separate change will
> implement ptep_set_wrprotects() in the arm64 backend to realize the
> performance improvement as part of the work to enable contpte mappings.
> 
> Signed-off-by: Ryan Roberts <ryan.roberts@arm.com>
> ---
>   include/linux/pgtable.h |  13 +++
>   mm/memory.c             | 195 ++++++++++++++++++++++++++++++----------
>   2 files changed, 162 insertions(+), 46 deletions(-)
> 
> diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
> index af7639c3b0a3..1c50f8a0fdde 100644
> --- a/include/linux/pgtable.h
> +++ b/include/linux/pgtable.h
> @@ -622,6 +622,19 @@ static inline void ptep_set_wrprotect(struct mm_struct *mm, unsigned long addres
>   }
>   #endif
>   
> +#ifndef ptep_set_wrprotects
> +struct mm_struct;
> +static inline void ptep_set_wrprotects(struct mm_struct *mm,
> +				unsigned long address, pte_t *ptep,
> +				unsigned int nr)
> +{
> +	unsigned int i;
> +
> +	for (i = 0; i < nr; i++, address += PAGE_SIZE, ptep++)
> +		ptep_set_wrprotect(mm, address, ptep);
> +}
> +#endif
> +
>   /*
>    * On some architectures hardware does not set page access bit when accessing
>    * memory page, it is responsibility of software setting this bit. It brings
> diff --git a/mm/memory.c b/mm/memory.c
> index 1f18ed4a5497..8a87a488950c 100644
> --- a/mm/memory.c
> +++ b/mm/memory.c
> @@ -924,68 +924,162 @@ copy_present_page(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma
>   	return 0;
>   }
>   
> +static int folio_nr_pages_cont_mapped(struct folio *folio,
> +				      struct page *page, pte_t *pte,
> +				      unsigned long addr, unsigned long end,
> +				      pte_t ptent, bool enforce_uffd_wp,
> +				      int *dirty_nr, int *writable_nr)
> +{
> +	int floops;
> +	int i;
> +	unsigned long pfn;
> +	bool prot_none;
> +	bool uffd_wp;
> +
> +	if (!folio_test_large(folio))
> +		return 1;
> +
> +	/*
> +	 * Loop either to `end` or to end of folio if its contiguously mapped,
> +	 * whichever is smaller.
> +	 */
> +	floops = (end - addr) >> PAGE_SHIFT;
> +	floops = min_t(int, floops,
> +		       folio_pfn(folio_next(folio)) - page_to_pfn(page));
> +
> +	pfn = page_to_pfn(page);
> +	prot_none = pte_protnone(ptent);
> +	uffd_wp = pte_uffd_wp(ptent);
> +
> +	*dirty_nr = !!pte_dirty(ptent);
> +	*writable_nr = !!pte_write(ptent);
> +
> +	pfn++;
> +	pte++;
> +
> +	for (i = 1; i < floops; i++) {
> +		ptent = ptep_get(pte);
> +
> +		if (!pte_present(ptent) || pte_pfn(ptent) != pfn ||
> +		    prot_none != pte_protnone(ptent) ||
> +		    (enforce_uffd_wp && uffd_wp != pte_uffd_wp(ptent)))
> +			break;
> +
> +		if (pte_dirty(ptent))
> +			(*dirty_nr)++;
> +		if (pte_write(ptent))
> +			(*writable_nr)++;
> +
> +		pfn++;
> +		pte++;
> +	}
> +
> +	return i;
> +}
> +
>   /*
> - * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
> - * is required to copy this pte.
> + * Copy set of contiguous ptes.  Returns number of ptes copied if succeeded
> + * (always gte 1), or -EAGAIN if one preallocated page is required to copy the
> + * first pte.
>    */
>   static inline int
> -copy_present_pte(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
> -		 pte_t *dst_pte, pte_t *src_pte, unsigned long addr, int *rss,
> -		 struct folio **prealloc)
> +copy_present_ptes(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
> +		  pte_t *dst_pte, pte_t *src_pte,
> +		  unsigned long addr, unsigned long end,
> +		  int *rss, struct folio **prealloc)
>   {
>   	struct mm_struct *src_mm = src_vma->vm_mm;
>   	unsigned long vm_flags = src_vma->vm_flags;
>   	pte_t pte = ptep_get(src_pte);
>   	struct page *page;
>   	struct folio *folio;
> +	int nr = 1;
> +	bool anon = false;
> +	bool enforce_uffd_wp = userfaultfd_wp(dst_vma);
> +	int nr_dirty = !!pte_dirty(pte);
> +	int nr_writable = !!pte_write(pte);
> +	int i, ret;
>   
>   	page = vm_normal_page(src_vma, addr, pte);
> -	if (page)
> +	if (page) {
>   		folio = page_folio(page);
> -	if (page && folio_test_anon(folio)) {
> -		/*
> -		 * If this page may have been pinned by the parent process,
> -		 * copy the page immediately for the child so that we'll always
> -		 * guarantee the pinned page won't be randomly replaced in the
> -		 * future.
> -		 */
> -		folio_get(folio);
> -		if (unlikely(page_try_dup_anon_rmap(page, false, src_vma))) {
> -			/* Page may be pinned, we have to copy. */
> -			folio_put(folio);
> -			return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
> -						 addr, rss, prealloc, page);
> +		anon = folio_test_anon(folio);
> +		nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
> +						pte, enforce_uffd_wp, &nr_dirty,
> +						&nr_writable);
> +		folio_ref_add(folio, nr);
> +
> +		for (i = 0; i < nr; i++, page++) {
> +			if (anon) {
> +				/*
> +				 * If this page may have been pinned by the
> +				 * parent process, copy the page immediately for
> +				 * the child so that we'll always guarantee the
> +				 * pinned page won't be randomly replaced in the
> +				 * future.
> +				 */
> +				if (unlikely(page_try_dup_anon_rmap(
> +						page, false, src_vma))) {
> +					if (i != 0)
> +						break;
> +					/* Page may be pinned, we have to copy. */
> +					folio_ref_sub(folio, nr);
> +					ret = copy_present_page(
> +						dst_vma, src_vma, dst_pte,
> +						src_pte, addr, rss, prealloc,
> +						page);
> +					return ret == 0 ? 1 : ret;
> +				}
> +				rss[MM_ANONPAGES]++;
> +				VM_BUG_ON(PageAnonExclusive(page));
> +			} else {
> +				page_dup_file_rmap(page, false);
> +				rss[mm_counter_file(page)]++;
> +			}
>   		}
> -		rss[MM_ANONPAGES]++;
> -	} else if (page) {
> -		folio_get(folio);
> -		page_dup_file_rmap(page, false);
> -		rss[mm_counter_file(page)]++;
> -	}

This likely looks a lot neater if you keep the existing structure.

For example, you can simply have on the !anon path

} else if (page) {
	folio = page_folio(page);
	nr = folio_nr_pages_cont_mapped ...
	folio_ref_add(folio, nr);
	for (i = 0; i < nr; i++, page++)
		page_dup_file_rmap(page, false);
	rss[mm_counter_file(&folio->page)] += nr;
}
  

On 04.12.23 16:47, David Hildenbrand wrote:
> On 04.12.23 11:54, Ryan Roberts wrote:
>> Convert copy_pte_range() to copy a set of ptes in a batch. A given batch
>> maps a physically contiguous block of memory, all belonging to the same
>> folio. This will likely improve performance by a tiny amount due to
>> batching the folio reference count management and calling set_ptes()
>> rather than making individual calls to set_pte_at().
>>
>> However, the primary motivation for this change is to reduce the number
>> of tlb maintenance operations that the arm64 backend has to perform
>> during fork, as it is about to add transparent support for the
>> "contiguous bit" in its ptes. By write-protecting the parent using the
>> new ptep_set_wrprotects() (note the 's' at the end) function, the
>> backend can avoid having to unfold contig ranges of PTEs, which is
>> expensive, when all ptes in the range are being write-protected.
>> Similarly, by using set_ptes() rather than set_pte_at() to set up ptes
>> in the child, the backend does not need to fold a contiguous range once
>> they are all populated - they can be initially populated as a contiguous
>> range in the first place.
>>
>> This change addresses the core-mm refactoring only, and introduces
>> ptep_set_wrprotects() with a default implementation that calls
>> ptep_set_wrprotect() for each pte in the range. A separate change will
>> implement ptep_set_wrprotects() in the arm64 backend to realize the
>> performance improvement as part of the work to enable contpte mappings.
>>
>> Signed-off-by: Ryan Roberts <ryan.roberts@arm.com>
>> ---
>>    include/linux/pgtable.h |  13 +++
>>    mm/memory.c             | 195 ++++++++++++++++++++++++++++++----------
>>    2 files changed, 162 insertions(+), 46 deletions(-)
>>
>> diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
>> index af7639c3b0a3..1c50f8a0fdde 100644
>> --- a/include/linux/pgtable.h
>> +++ b/include/linux/pgtable.h
>> @@ -622,6 +622,19 @@ static inline void ptep_set_wrprotect(struct mm_struct *mm, unsigned long addres
>>    }
>>    #endif
>>    
>> +#ifndef ptep_set_wrprotects
>> +struct mm_struct;
>> +static inline void ptep_set_wrprotects(struct mm_struct *mm,
>> +				unsigned long address, pte_t *ptep,
>> +				unsigned int nr)
>> +{
>> +	unsigned int i;
>> +
>> +	for (i = 0; i < nr; i++, address += PAGE_SIZE, ptep++)
>> +		ptep_set_wrprotect(mm, address, ptep);
>> +}
>> +#endif
>> +
>>    /*
>>     * On some architectures hardware does not set page access bit when accessing
>>     * memory page, it is responsibility of software setting this bit. It brings
>> diff --git a/mm/memory.c b/mm/memory.c
>> index 1f18ed4a5497..8a87a488950c 100644
>> --- a/mm/memory.c
>> +++ b/mm/memory.c
>> @@ -924,68 +924,162 @@ copy_present_page(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma
>>    	return 0;
>>    }
>>    
>> +static int folio_nr_pages_cont_mapped(struct folio *folio,
>> +				      struct page *page, pte_t *pte,
>> +				      unsigned long addr, unsigned long end,
>> +				      pte_t ptent, bool enforce_uffd_wp,
>> +				      int *dirty_nr, int *writable_nr)
>> +{
>> +	int floops;
>> +	int i;
>> +	unsigned long pfn;
>> +	bool prot_none;
>> +	bool uffd_wp;
>> +
>> +	if (!folio_test_large(folio))
>> +		return 1;
>> +
>> +	/*
>> +	 * Loop either to `end` or to end of folio if its contiguously mapped,
>> +	 * whichever is smaller.
>> +	 */
>> +	floops = (end - addr) >> PAGE_SHIFT;
>> +	floops = min_t(int, floops,
>> +		       folio_pfn(folio_next(folio)) - page_to_pfn(page));
>> +
>> +	pfn = page_to_pfn(page);
>> +	prot_none = pte_protnone(ptent);
>> +	uffd_wp = pte_uffd_wp(ptent);
>> +
>> +	*dirty_nr = !!pte_dirty(ptent);
>> +	*writable_nr = !!pte_write(ptent);
>> +
>> +	pfn++;
>> +	pte++;
>> +
>> +	for (i = 1; i < floops; i++) {
>> +		ptent = ptep_get(pte);
>> +
>> +		if (!pte_present(ptent) || pte_pfn(ptent) != pfn ||
>> +		    prot_none != pte_protnone(ptent) ||
>> +		    (enforce_uffd_wp && uffd_wp != pte_uffd_wp(ptent)))
>> +			break;
>> +
>> +		if (pte_dirty(ptent))
>> +			(*dirty_nr)++;
>> +		if (pte_write(ptent))
>> +			(*writable_nr)++;
>> +
>> +		pfn++;
>> +		pte++;
>> +	}
>> +
>> +	return i;
>> +}
>> +
>>    /*
>> - * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
>> - * is required to copy this pte.
>> + * Copy set of contiguous ptes.  Returns number of ptes copied if succeeded
>> + * (always gte 1), or -EAGAIN if one preallocated page is required to copy the
>> + * first pte.
>>     */
>>    static inline int
>> -copy_present_pte(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>> -		 pte_t *dst_pte, pte_t *src_pte, unsigned long addr, int *rss,
>> -		 struct folio **prealloc)
>> +copy_present_ptes(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>> +		  pte_t *dst_pte, pte_t *src_pte,
>> +		  unsigned long addr, unsigned long end,
>> +		  int *rss, struct folio **prealloc)
>>    {
>>    	struct mm_struct *src_mm = src_vma->vm_mm;
>>    	unsigned long vm_flags = src_vma->vm_flags;
>>    	pte_t pte = ptep_get(src_pte);
>>    	struct page *page;
>>    	struct folio *folio;
>> +	int nr = 1;
>> +	bool anon = false;
>> +	bool enforce_uffd_wp = userfaultfd_wp(dst_vma);
>> +	int nr_dirty = !!pte_dirty(pte);
>> +	int nr_writable = !!pte_write(pte);
>> +	int i, ret;
>>    
>>    	page = vm_normal_page(src_vma, addr, pte);
>> -	if (page)
>> +	if (page) {
>>    		folio = page_folio(page);
>> -	if (page && folio_test_anon(folio)) {
>> -		/*
>> -		 * If this page may have been pinned by the parent process,
>> -		 * copy the page immediately for the child so that we'll always
>> -		 * guarantee the pinned page won't be randomly replaced in the
>> -		 * future.
>> -		 */
>> -		folio_get(folio);
>> -		if (unlikely(page_try_dup_anon_rmap(page, false, src_vma))) {
>> -			/* Page may be pinned, we have to copy. */
>> -			folio_put(folio);
>> -			return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
>> -						 addr, rss, prealloc, page);
>> +		anon = folio_test_anon(folio);
>> +		nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>> +						pte, enforce_uffd_wp, &nr_dirty,
>> +						&nr_writable);
>> +		folio_ref_add(folio, nr);
>> +
>> +		for (i = 0; i < nr; i++, page++) {
>> +			if (anon) {
>> +				/*
>> +				 * If this page may have been pinned by the
>> +				 * parent process, copy the page immediately for
>> +				 * the child so that we'll always guarantee the
>> +				 * pinned page won't be randomly replaced in the
>> +				 * future.
>> +				 */
>> +				if (unlikely(page_try_dup_anon_rmap(
>> +						page, false, src_vma))) {
>> +					if (i != 0)
>> +						break;
>> +					/* Page may be pinned, we have to copy. */
>> +					folio_ref_sub(folio, nr);
>> +					ret = copy_present_page(
>> +						dst_vma, src_vma, dst_pte,
>> +						src_pte, addr, rss, prealloc,
>> +						page);
>> +					return ret == 0 ? 1 : ret;
>> +				}
>> +				rss[MM_ANONPAGES]++;
>> +				VM_BUG_ON(PageAnonExclusive(page));
>> +			} else {
>> +				page_dup_file_rmap(page, false);
>> +				rss[mm_counter_file(page)]++;
>> +			}
>>    		}
>> -		rss[MM_ANONPAGES]++;
>> -	} else if (page) {
>> -		folio_get(folio);
>> -		page_dup_file_rmap(page, false);
>> -		rss[mm_counter_file(page)]++;
>> -	}
> 
> This likely looks a lot neater if you keep the existing structure.
> 
> For example, you can simply have on the !anon path
> 
> } else if (page) {
> 	folio = page_folio(page);
> 	nr = folio_nr_pages_cont_mapped ...
> 	folio_ref_add(folio, nr);
> 	for (i = 0; i < nr; i++, page++)
> 		page_dup_file_rmap(page, false);
> 	rss[mm_counter_file(&folio->page)] += nr;
> }
> 

With rmap batching from [1] -- rebased+changed on top of that -- we could turn
that into an effective (untested):

         if (page && folio_test_anon(folio)) {
+               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
+                                               pte, enforce_uffd_wp, &nr_dirty,
+                                               &nr_writable);
                 /*
                  * If this page may have been pinned by the parent process,
                  * copy the page immediately for the child so that we'll always
                  * guarantee the pinned page won't be randomly replaced in the
                  * future.
                  */
-               folio_get(folio);
-               if (unlikely(folio_try_dup_anon_rmap_pte(folio, page, src_vma))) {
+               folio_ref_add(folio, nr);
+               if (unlikely(folio_try_dup_anon_rmap_ptes(folio, page, nr, src_vma))) {
                         /* Page may be pinned, we have to copy. */
-                       folio_put(folio);
-                       return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
-                                                addr, rss, prealloc, page);
+                       folio_ref_sub(folio, nr);
+                       ret = copy_present_page(dst_vma, src_vma, dst_pte,
+                                               src_pte, addr, rss, prealloc,
+                                               page);
+                       return ret == 0 ? 1 : ret;
                 }
-               rss[MM_ANONPAGES]++;
+               rss[MM_ANONPAGES] += nr;
         } else if (page) {
-               folio_get(folio);
-               folio_dup_file_rmap_pte(folio, page);
-               rss[mm_counter_file(page)]++;
+               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
+                                               pte, enforce_uffd_wp, &nr_dirty,
+                                               &nr_writable);
+               folio_ref_add(folio, nr);
+               folio_dup_file_rmap_ptes(folio, page, nr);
+               rss[mm_counter_file(page)] += nr;
         }


We'll have to test performance, but it could be that we want to specialize
more on !folio_test_large(). That code is very performance-sensitive.


[1] https://lkml.kernel.org/r/20231204142146.91437-1-david@redhat.com
  

> 
> With rmap batching from [1] -- rebased+changed on top of that -- we could turn
> that into an effective (untested):
> 
>           if (page && folio_test_anon(folio)) {
> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
> +                                               pte, enforce_uffd_wp, &nr_dirty,
> +                                               &nr_writable);
>                   /*
>                    * If this page may have been pinned by the parent process,
>                    * copy the page immediately for the child so that we'll always
>                    * guarantee the pinned page won't be randomly replaced in the
>                    * future.
>                    */
> -               folio_get(folio);
> -               if (unlikely(folio_try_dup_anon_rmap_pte(folio, page, src_vma))) {
> +               folio_ref_add(folio, nr);
> +               if (unlikely(folio_try_dup_anon_rmap_ptes(folio, page, nr, src_vma))) {
>                           /* Page may be pinned, we have to copy. */
> -                       folio_put(folio);
> -                       return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
> -                                                addr, rss, prealloc, page);
> +                       folio_ref_sub(folio, nr);
> +                       ret = copy_present_page(dst_vma, src_vma, dst_pte,
> +                                               src_pte, addr, rss, prealloc,
> +                                               page);
> +                       return ret == 0 ? 1 : ret;
>                   }
> -               rss[MM_ANONPAGES]++;
> +               rss[MM_ANONPAGES] += nr;
>           } else if (page) {
> -               folio_get(folio);
> -               folio_dup_file_rmap_pte(folio, page);
> -               rss[mm_counter_file(page)]++;
> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
> +                                               pte, enforce_uffd_wp, &nr_dirty,
> +                                               &nr_writable);
> +               folio_ref_add(folio, nr);
> +               folio_dup_file_rmap_ptes(folio, page, nr);
> +               rss[mm_counter_file(page)] += nr;
>           }
> 
> 
> We'll have to test performance, but it could be that we want to specialize
> more on !folio_test_large(). That code is very performance-sensitive.
> 
> 
> [1] https://lkml.kernel.org/r/20231204142146.91437-1-david@redhat.com

So, on top of [1] without rmap batching but with a slightly modified 
version of yours (that keeps the existing code structure as pointed out 
and e.g., updates counter updates), running my fork() microbenchmark 
with a 1 GiB of memory:

Compared to [1], with all order-0 pages it gets 13--14% _slower_ and 
with all PTE-mapped THP (order-9) it gets ~29--30% _faster_.

So looks like we really want to have a completely seprate code path for 
"!folio_test_large()" to keep that case as fast as possible. And 
"Likely" we want to use "likely(!folio_test_large()". ;)

Performing rmap batching on top of that code only slightly (another 1% 
or so) improves performance in the PTE-mapped THP (order-9) case right 
now, in contrast to other rmap batching. Reason is as all rmap code gets 
inlined here and we're only doing subpage mapcount updates + PAE handling.
  

On 04/12/2023 17:27, David Hildenbrand wrote:
>>
>> With rmap batching from [1] -- rebased+changed on top of that -- we could turn
>> that into an effective (untested):
>>
>>           if (page && folio_test_anon(folio)) {
>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>> +                                               pte, enforce_uffd_wp, &nr_dirty,
>> +                                               &nr_writable);
>>                   /*
>>                    * If this page may have been pinned by the parent process,
>>                    * copy the page immediately for the child so that we'll always
>>                    * guarantee the pinned page won't be randomly replaced in the
>>                    * future.
>>                    */
>> -               folio_get(folio);
>> -               if (unlikely(folio_try_dup_anon_rmap_pte(folio, page,
>> src_vma))) {
>> +               folio_ref_add(folio, nr);
>> +               if (unlikely(folio_try_dup_anon_rmap_ptes(folio, page, nr,
>> src_vma))) {
>>                           /* Page may be pinned, we have to copy. */
>> -                       folio_put(folio);
>> -                       return copy_present_page(dst_vma, src_vma, dst_pte,
>> src_pte,
>> -                                                addr, rss, prealloc, page);
>> +                       folio_ref_sub(folio, nr);
>> +                       ret = copy_present_page(dst_vma, src_vma, dst_pte,
>> +                                               src_pte, addr, rss, prealloc,
>> +                                               page);
>> +                       return ret == 0 ? 1 : ret;
>>                   }
>> -               rss[MM_ANONPAGES]++;
>> +               rss[MM_ANONPAGES] += nr;
>>           } else if (page) {
>> -               folio_get(folio);
>> -               folio_dup_file_rmap_pte(folio, page);
>> -               rss[mm_counter_file(page)]++;
>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>> +                                               pte, enforce_uffd_wp, &nr_dirty,
>> +                                               &nr_writable);
>> +               folio_ref_add(folio, nr);
>> +               folio_dup_file_rmap_ptes(folio, page, nr);
>> +               rss[mm_counter_file(page)] += nr;
>>           }
>>
>>
>> We'll have to test performance, but it could be that we want to specialize
>> more on !folio_test_large(). That code is very performance-sensitive.
>>
>>
>> [1] https://lkml.kernel.org/r/20231204142146.91437-1-david@redhat.com
> 
> So, on top of [1] without rmap batching but with a slightly modified version of

Can you clarify what you mean by "without rmap batching"? I thought [1]
implicitly adds rmap batching? (e.g. folio_dup_file_rmap_ptes(), which you've
added in the code snippet above).

> yours (that keeps the existing code structure as pointed out and e.g., updates
> counter updates), running my fork() microbenchmark with a 1 GiB of memory:
> 
> Compared to [1], with all order-0 pages it gets 13--14% _slower_ and with all
> PTE-mapped THP (order-9) it gets ~29--30% _faster_.

What test are you running - I'd like to reproduce if possible, since it sounds
like I've got some work to do to remove the order-0 regression.

> 
> So looks like we really want to have a completely seprate code path for
> "!folio_test_large()" to keep that case as fast as possible. And "Likely" we
> want to use "likely(!folio_test_large()". ;)

Yuk, but fair enough. If I can repro the perf numbers, I'll have a go a
reworking this.

I think you're also implicitly suggesting that this change needs to depend on
[1]? Which is a shame...

I guess I should also go through a similar exercise for patch 2 in this series.

> 
> Performing rmap batching on top of that code only slightly (another 1% or so)
> improves performance in the PTE-mapped THP (order-9) case right now, in contrast
> to other rmap batching. Reason is as all rmap code gets inlined here and we're
> only doing subpage mapcount updates + PAE handling.
>
  

On 05.12.23 12:30, Ryan Roberts wrote:
> On 04/12/2023 17:27, David Hildenbrand wrote:
>>>
>>> With rmap batching from [1] -- rebased+changed on top of that -- we could turn
>>> that into an effective (untested):
>>>
>>>            if (page && folio_test_anon(folio)) {
>>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>>> +                                               pte, enforce_uffd_wp, &nr_dirty,
>>> +                                               &nr_writable);
>>>                    /*
>>>                     * If this page may have been pinned by the parent process,
>>>                     * copy the page immediately for the child so that we'll always
>>>                     * guarantee the pinned page won't be randomly replaced in the
>>>                     * future.
>>>                     */
>>> -               folio_get(folio);
>>> -               if (unlikely(folio_try_dup_anon_rmap_pte(folio, page,
>>> src_vma))) {
>>> +               folio_ref_add(folio, nr);
>>> +               if (unlikely(folio_try_dup_anon_rmap_ptes(folio, page, nr,
>>> src_vma))) {
>>>                            /* Page may be pinned, we have to copy. */
>>> -                       folio_put(folio);
>>> -                       return copy_present_page(dst_vma, src_vma, dst_pte,
>>> src_pte,
>>> -                                                addr, rss, prealloc, page);
>>> +                       folio_ref_sub(folio, nr);
>>> +                       ret = copy_present_page(dst_vma, src_vma, dst_pte,
>>> +                                               src_pte, addr, rss, prealloc,
>>> +                                               page);
>>> +                       return ret == 0 ? 1 : ret;
>>>                    }
>>> -               rss[MM_ANONPAGES]++;
>>> +               rss[MM_ANONPAGES] += nr;
>>>            } else if (page) {
>>> -               folio_get(folio);
>>> -               folio_dup_file_rmap_pte(folio, page);
>>> -               rss[mm_counter_file(page)]++;
>>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>>> +                                               pte, enforce_uffd_wp, &nr_dirty,
>>> +                                               &nr_writable);
>>> +               folio_ref_add(folio, nr);
>>> +               folio_dup_file_rmap_ptes(folio, page, nr);
>>> +               rss[mm_counter_file(page)] += nr;
>>>            }
>>>
>>>
>>> We'll have to test performance, but it could be that we want to specialize
>>> more on !folio_test_large(). That code is very performance-sensitive.
>>>
>>>
>>> [1] https://lkml.kernel.org/r/20231204142146.91437-1-david@redhat.com
>>
>> So, on top of [1] without rmap batching but with a slightly modified version of
> 
> Can you clarify what you mean by "without rmap batching"? I thought [1]
> implicitly adds rmap batching? (e.g. folio_dup_file_rmap_ptes(), which you've
> added in the code snippet above).

Not calling the batched variants but essentially doing what your code 
does (with some minor improvements, like updating the rss counters only 
once).

The snipped above is only linked below. I had the performance numbers 
for [1] ready, so I gave it a test on top of that.

To keep it simple, you might just benchmark w and w/o your patches.

> 
>> yours (that keeps the existing code structure as pointed out and e.g., updates
>> counter updates), running my fork() microbenchmark with a 1 GiB of memory:
>>
>> Compared to [1], with all order-0 pages it gets 13--14% _slower_ and with all
>> PTE-mapped THP (order-9) it gets ~29--30% _faster_.
> 
> What test are you running - I'd like to reproduce if possible, since it sounds
> like I've got some work to do to remove the order-0 regression.

Essentially just allocating 1 GiB of memory an measuring how long it 
takes to call fork().

order-0 benchmarks:

https://gitlab.com/davidhildenbrand/scratchspace/-/raw/main/order-0-benchmarks.c?ref_type=heads

e.g.,: $ ./order-0-benchmarks fork 100


pte-mapped-thp benchmarks:

https://gitlab.com/davidhildenbrand/scratchspace/-/raw/main/pte-mapped-thp-benchmarks.c?ref_type=heads

e.g.,: $ ./pte-mapped-thp-benchmarks fork 100


Ideally, pin to one CPU and get stable performance numbers by disabling 
SMT+turbo etc.

> 
>>
>> So looks like we really want to have a completely seprate code path for
>> "!folio_test_large()" to keep that case as fast as possible. And "Likely" we
>> want to use "likely(!folio_test_large()". ;)
> 
> Yuk, but fair enough. If I can repro the perf numbers, I'll have a go a
> reworking this.
> 
> I think you're also implicitly suggesting that this change needs to depend on
> [1]? Which is a shame...

Not necessarily. It certainly cleans up the code, but we can do that in 
any order reasonable.

> 
> I guess I should also go through a similar exercise for patch 2 in this series.


Yes. There are "unmap" and "pte-dontneed" benchmarks contained in both 
files above.
  

On 05/12/2023 12:04, David Hildenbrand wrote:
> On 05.12.23 12:30, Ryan Roberts wrote:
>> On 04/12/2023 17:27, David Hildenbrand wrote:
>>>>
>>>> With rmap batching from [1] -- rebased+changed on top of that -- we could turn
>>>> that into an effective (untested):
>>>>
>>>>            if (page && folio_test_anon(folio)) {
>>>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr,
>>>> end,
>>>> +                                               pte, enforce_uffd_wp,
>>>> &nr_dirty,
>>>> +                                               &nr_writable);
>>>>                    /*
>>>>                     * If this page may have been pinned by the parent process,
>>>>                     * copy the page immediately for the child so that we'll
>>>> always
>>>>                     * guarantee the pinned page won't be randomly replaced
>>>> in the
>>>>                     * future.
>>>>                     */
>>>> -               folio_get(folio);
>>>> -               if (unlikely(folio_try_dup_anon_rmap_pte(folio, page,
>>>> src_vma))) {
>>>> +               folio_ref_add(folio, nr);
>>>> +               if (unlikely(folio_try_dup_anon_rmap_ptes(folio, page, nr,
>>>> src_vma))) {
>>>>                            /* Page may be pinned, we have to copy. */
>>>> -                       folio_put(folio);
>>>> -                       return copy_present_page(dst_vma, src_vma, dst_pte,
>>>> src_pte,
>>>> -                                                addr, rss, prealloc, page);
>>>> +                       folio_ref_sub(folio, nr);
>>>> +                       ret = copy_present_page(dst_vma, src_vma, dst_pte,
>>>> +                                               src_pte, addr, rss, prealloc,
>>>> +                                               page);
>>>> +                       return ret == 0 ? 1 : ret;
>>>>                    }
>>>> -               rss[MM_ANONPAGES]++;
>>>> +               rss[MM_ANONPAGES] += nr;
>>>>            } else if (page) {
>>>> -               folio_get(folio);
>>>> -               folio_dup_file_rmap_pte(folio, page);
>>>> -               rss[mm_counter_file(page)]++;
>>>> +               nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr,
>>>> end,
>>>> +                                               pte, enforce_uffd_wp,
>>>> &nr_dirty,
>>>> +                                               &nr_writable);
>>>> +               folio_ref_add(folio, nr);
>>>> +               folio_dup_file_rmap_ptes(folio, page, nr);
>>>> +               rss[mm_counter_file(page)] += nr;
>>>>            }
>>>>
>>>>
>>>> We'll have to test performance, but it could be that we want to specialize
>>>> more on !folio_test_large(). That code is very performance-sensitive.
>>>>
>>>>
>>>> [1] https://lkml.kernel.org/r/20231204142146.91437-1-david@redhat.com
>>>
>>> So, on top of [1] without rmap batching but with a slightly modified version of
>>
>> Can you clarify what you mean by "without rmap batching"? I thought [1]
>> implicitly adds rmap batching? (e.g. folio_dup_file_rmap_ptes(), which you've
>> added in the code snippet above).
> 
> Not calling the batched variants but essentially doing what your code does (with
> some minor improvements, like updating the rss counters only once).
> 
> The snipped above is only linked below. I had the performance numbers for [1]
> ready, so I gave it a test on top of that.
> 
> To keep it simple, you might just benchmark w and w/o your patches.
> 
>>
>>> yours (that keeps the existing code structure as pointed out and e.g., updates
>>> counter updates), running my fork() microbenchmark with a 1 GiB of memory:
>>>
>>> Compared to [1], with all order-0 pages it gets 13--14% _slower_ and with all
>>> PTE-mapped THP (order-9) it gets ~29--30% _faster_.
>>
>> What test are you running - I'd like to reproduce if possible, since it sounds
>> like I've got some work to do to remove the order-0 regression.
> 
> Essentially just allocating 1 GiB of memory an measuring how long it takes to
> call fork().
> 
> order-0 benchmarks:
> 
> https://gitlab.com/davidhildenbrand/scratchspace/-/raw/main/order-0-benchmarks.c?ref_type=heads
> 
> e.g.,: $ ./order-0-benchmarks fork 100
> 
> 
> pte-mapped-thp benchmarks:
> 
> https://gitlab.com/davidhildenbrand/scratchspace/-/raw/main/pte-mapped-thp-benchmarks.c?ref_type=heads
> 
> e.g.,: $ ./pte-mapped-thp-benchmarks fork 100
> 
> 
> Ideally, pin to one CPU and get stable performance numbers by disabling
> SMT+turbo etc.

This is great - thanks! I'll get to work...

> 
>>
>>>
>>> So looks like we really want to have a completely seprate code path for
>>> "!folio_test_large()" to keep that case as fast as possible. And "Likely" we
>>> want to use "likely(!folio_test_large()". ;)
>>
>> Yuk, but fair enough. If I can repro the perf numbers, I'll have a go a
>> reworking this.
>>
>> I think you're also implicitly suggesting that this change needs to depend on
>> [1]? Which is a shame...
> 
> Not necessarily. It certainly cleans up the code, but we can do that in any
> order reasonable.
> 
>>
>> I guess I should also go through a similar exercise for patch 2 in this series.
> 
> 
> Yes. There are "unmap" and "pte-dontneed" benchmarks contained in both files above.
>
  

Ryan Roberts <ryan.roberts@arm.com> writes:

<snip>

>  /*
>   * On some architectures hardware does not set page access bit when accessing
>   * memory page, it is responsibility of software setting this bit. It brings
> diff --git a/mm/memory.c b/mm/memory.c
> index 1f18ed4a5497..8a87a488950c 100644
> --- a/mm/memory.c
> +++ b/mm/memory.c
> @@ -924,68 +924,162 @@ copy_present_page(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma
>  	return 0;
>  }
>  
> +static int folio_nr_pages_cont_mapped(struct folio *folio,
> +				      struct page *page, pte_t *pte,
> +				      unsigned long addr, unsigned long end,
> +				      pte_t ptent, bool enforce_uffd_wp,
> +				      int *dirty_nr, int *writable_nr)
> +{
> +	int floops;
> +	int i;
> +	unsigned long pfn;
> +	bool prot_none;
> +	bool uffd_wp;
> +
> +	if (!folio_test_large(folio))
> +		return 1;
> +
> +	/*
> +	 * Loop either to `end` or to end of folio if its contiguously mapped,
> +	 * whichever is smaller.
> +	 */
> +	floops = (end - addr) >> PAGE_SHIFT;
> +	floops = min_t(int, floops,
> +		       folio_pfn(folio_next(folio)) - page_to_pfn(page));

Much better, thanks for addressing my comments here.

> +
> +	pfn = page_to_pfn(page);
> +	prot_none = pte_protnone(ptent);
> +	uffd_wp = pte_uffd_wp(ptent);
> +
> +	*dirty_nr = !!pte_dirty(ptent);
> +	*writable_nr = !!pte_write(ptent);
> +
> +	pfn++;
> +	pte++;
> +
> +	for (i = 1; i < floops; i++) {
> +		ptent = ptep_get(pte);
> +
> +		if (!pte_present(ptent) || pte_pfn(ptent) != pfn ||
> +		    prot_none != pte_protnone(ptent) ||
> +		    (enforce_uffd_wp && uffd_wp != pte_uffd_wp(ptent)))
> +			break;
> +
> +		if (pte_dirty(ptent))
> +			(*dirty_nr)++;
> +		if (pte_write(ptent))
> +			(*writable_nr)++;
> +
> +		pfn++;
> +		pte++;
> +	}
> +
> +	return i;
> +}
> +
>  /*
> - * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
> - * is required to copy this pte.
> + * Copy set of contiguous ptes.  Returns number of ptes copied if succeeded
> + * (always gte 1), or -EAGAIN if one preallocated page is required to copy the
> + * first pte.
>   */
>  static inline int
> -copy_present_pte(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
> -		 pte_t *dst_pte, pte_t *src_pte, unsigned long addr, int *rss,
> -		 struct folio **prealloc)
> +copy_present_ptes(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
> +		  pte_t *dst_pte, pte_t *src_pte,
> +		  unsigned long addr, unsigned long end,
> +		  int *rss, struct folio **prealloc)
>  {
>  	struct mm_struct *src_mm = src_vma->vm_mm;
>  	unsigned long vm_flags = src_vma->vm_flags;
>  	pte_t pte = ptep_get(src_pte);
>  	struct page *page;
>  	struct folio *folio;
> +	int nr = 1;
> +	bool anon = false;
> +	bool enforce_uffd_wp = userfaultfd_wp(dst_vma);
> +	int nr_dirty = !!pte_dirty(pte);
> +	int nr_writable = !!pte_write(pte);
> +	int i, ret;
>  
>  	page = vm_normal_page(src_vma, addr, pte);
> -	if (page)
> +	if (page) {
>  		folio = page_folio(page);
> -	if (page && folio_test_anon(folio)) {
> -		/*
> -		 * If this page may have been pinned by the parent process,
> -		 * copy the page immediately for the child so that we'll always
> -		 * guarantee the pinned page won't be randomly replaced in the
> -		 * future.
> -		 */
> -		folio_get(folio);
> -		if (unlikely(page_try_dup_anon_rmap(page, false, src_vma))) {
> -			/* Page may be pinned, we have to copy. */
> -			folio_put(folio);
> -			return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
> -						 addr, rss, prealloc, page);
> +		anon = folio_test_anon(folio);
> +		nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
> +						pte, enforce_uffd_wp, &nr_dirty,
> +						&nr_writable);
> +		folio_ref_add(folio, nr);
> +
> +		for (i = 0; i < nr; i++, page++) {
> +			if (anon) {
> +				/*
> +				 * If this page may have been pinned by the
> +				 * parent process, copy the page immediately for
> +				 * the child so that we'll always guarantee the
> +				 * pinned page won't be randomly replaced in the
> +				 * future.
> +				 */
> +				if (unlikely(page_try_dup_anon_rmap(
> +						page, false, src_vma))) {
> +					if (i != 0)
> +						break;
> +					/* Page may be pinned, we have to copy. */
> +					folio_ref_sub(folio, nr);
> +					ret = copy_present_page(
> +						dst_vma, src_vma, dst_pte,
> +						src_pte, addr, rss, prealloc,
> +						page);
> +					return ret == 0 ? 1 : ret;
> +				}
> +				rss[MM_ANONPAGES]++;
> +				VM_BUG_ON(PageAnonExclusive(page));
> +			} else {
> +				page_dup_file_rmap(page, false);
> +				rss[mm_counter_file(page)]++;
> +			}
>  		}
> -		rss[MM_ANONPAGES]++;
> -	} else if (page) {
> -		folio_get(folio);
> -		page_dup_file_rmap(page, false);
> -		rss[mm_counter_file(page)]++;
> -	}
>  
> -	/*
> -	 * If it's a COW mapping, write protect it both
> -	 * in the parent and the child
> -	 */
> -	if (is_cow_mapping(vm_flags) && pte_write(pte)) {
> -		ptep_set_wrprotect(src_mm, addr, src_pte);
> -		pte = pte_wrprotect(pte);
> +		if (i < nr) {
> +			folio_ref_sub(folio, nr - i);
> +			nr = i;
> +		}
>  	}
> -	VM_BUG_ON(page && folio_test_anon(folio) && PageAnonExclusive(page));
>  
>  	/*
> -	 * If it's a shared mapping, mark it clean in
> -	 * the child
> +	 * If it's a shared mapping, mark it clean and write protected in the
> +	 * child, and rely on a write fault to fix up the permissions. This
> +	 * allows determining batch size without having to consider RO/RW
> +	 * permissions. As an optimization, skip wrprotect if all ptes in the
> +	 * batch have the same permissions.
> +	 *
> +	 * If its a private (CoW) mapping, mark it dirty in the child if _any_
> +	 * of the parent mappings in the block were marked dirty. The contiguous
> +	 * block of mappings are all backed by the same folio, so if any are
> +	 * dirty then the whole folio is dirty. This allows determining batch
> +	 * size without having to consider the dirty bit. Further, write protect
> +	 * it both in the parent and the child so that a future write will cause
> +	 * a CoW. As as an optimization, skip the wrprotect if all the ptes in
> +	 * the batch are already readonly.
>  	 */
> -	if (vm_flags & VM_SHARED)
> +	if (vm_flags & VM_SHARED) {
>  		pte = pte_mkclean(pte);
> -	pte = pte_mkold(pte);
> +		if (nr_writable > 0 && nr_writable < nr)
> +			pte = pte_wrprotect(pte);
> +	} else {
> +		if (nr_dirty)
> +			pte = pte_mkdirty(pte);
> +		if (nr_writable) {
> +			ptep_set_wrprotects(src_mm, addr, src_pte, nr);
> +			pte = pte_wrprotect(pte);
> +		}
> +	}
>  
> -	if (!userfaultfd_wp(dst_vma))
> +	pte = pte_mkold(pte);
> +	pte = pte_clear_soft_dirty(pte);
> +	if (!enforce_uffd_wp)
>  		pte = pte_clear_uffd_wp(pte);
>  
> -	set_pte_at(dst_vma->vm_mm, addr, dst_pte, pte);
> -	return 0;
> +	set_ptes(dst_vma->vm_mm, addr, dst_pte, pte, nr);
> +	return nr;

I don't have any further comments and you have addressed my previous
ones so feel free to add:

Reviewed-by: Alistair Popple <apopple@nvidia.com>

However whilst I think the above CoW sequence looks correct it would be
nice if someone else could take a look as well.

>  }
>  
>  static inline struct folio *page_copy_prealloc(struct mm_struct *src_mm,
> @@ -1021,6 +1115,7 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>  	int rss[NR_MM_COUNTERS];
>  	swp_entry_t entry = (swp_entry_t){0};
>  	struct folio *prealloc = NULL;
> +	int nr_ptes;
>  
>  again:
>  	progress = 0;
> @@ -1051,6 +1146,8 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>  	arch_enter_lazy_mmu_mode();
>  
>  	do {
> +		nr_ptes = 1;
> +
>  		/*
>  		 * We are holding two locks at this point - either of them
>  		 * could generate latencies in another task on another CPU.
> @@ -1086,16 +1183,21 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>  			 * the now present pte.
>  			 */
>  			WARN_ON_ONCE(ret != -ENOENT);
> +			ret = 0;
>  		}
> -		/* copy_present_pte() will clear `*prealloc' if consumed */
> -		ret = copy_present_pte(dst_vma, src_vma, dst_pte, src_pte,
> -				       addr, rss, &prealloc);
> +		/* copy_present_ptes() will clear `*prealloc' if consumed */
> +		nr_ptes = copy_present_ptes(dst_vma, src_vma, dst_pte, src_pte,
> +					    addr, end, rss, &prealloc);
> +
>  		/*
>  		 * If we need a pre-allocated page for this pte, drop the
>  		 * locks, allocate, and try again.
>  		 */
> -		if (unlikely(ret == -EAGAIN))
> +		if (unlikely(nr_ptes == -EAGAIN)) {
> +			ret = -EAGAIN;
>  			break;
> +		}
> +
>  		if (unlikely(prealloc)) {
>  			/*
>  			 * pre-alloc page cannot be reused by next time so as
> @@ -1106,8 +1208,9 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>  			folio_put(prealloc);
>  			prealloc = NULL;
>  		}
> -		progress += 8;
> -	} while (dst_pte++, src_pte++, addr += PAGE_SIZE, addr != end);
> +		progress += 8 * nr_ptes;
> +	} while (dst_pte += nr_ptes, src_pte += nr_ptes,
> +		 addr += PAGE_SIZE * nr_ptes, addr != end);
>  
>  	arch_leave_lazy_mmu_mode();
>  	pte_unmap_unlock(orig_src_pte, src_ptl);
  

On 08/12/2023 00:32, Alistair Popple wrote:
> 
> Ryan Roberts <ryan.roberts@arm.com> writes:
> 
> <snip>
> 
>>  /*
>>   * On some architectures hardware does not set page access bit when accessing
>>   * memory page, it is responsibility of software setting this bit. It brings
>> diff --git a/mm/memory.c b/mm/memory.c
>> index 1f18ed4a5497..8a87a488950c 100644
>> --- a/mm/memory.c
>> +++ b/mm/memory.c
>> @@ -924,68 +924,162 @@ copy_present_page(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma
>>  	return 0;
>>  }
>>  
>> +static int folio_nr_pages_cont_mapped(struct folio *folio,
>> +				      struct page *page, pte_t *pte,
>> +				      unsigned long addr, unsigned long end,
>> +				      pte_t ptent, bool enforce_uffd_wp,
>> +				      int *dirty_nr, int *writable_nr)
>> +{
>> +	int floops;
>> +	int i;
>> +	unsigned long pfn;
>> +	bool prot_none;
>> +	bool uffd_wp;
>> +
>> +	if (!folio_test_large(folio))
>> +		return 1;
>> +
>> +	/*
>> +	 * Loop either to `end` or to end of folio if its contiguously mapped,
>> +	 * whichever is smaller.
>> +	 */
>> +	floops = (end - addr) >> PAGE_SHIFT;
>> +	floops = min_t(int, floops,
>> +		       folio_pfn(folio_next(folio)) - page_to_pfn(page));
> 
> Much better, thanks for addressing my comments here.
> 
>> +
>> +	pfn = page_to_pfn(page);
>> +	prot_none = pte_protnone(ptent);
>> +	uffd_wp = pte_uffd_wp(ptent);
>> +
>> +	*dirty_nr = !!pte_dirty(ptent);
>> +	*writable_nr = !!pte_write(ptent);
>> +
>> +	pfn++;
>> +	pte++;
>> +
>> +	for (i = 1; i < floops; i++) {
>> +		ptent = ptep_get(pte);
>> +
>> +		if (!pte_present(ptent) || pte_pfn(ptent) != pfn ||
>> +		    prot_none != pte_protnone(ptent) ||
>> +		    (enforce_uffd_wp && uffd_wp != pte_uffd_wp(ptent)))
>> +			break;
>> +
>> +		if (pte_dirty(ptent))
>> +			(*dirty_nr)++;
>> +		if (pte_write(ptent))
>> +			(*writable_nr)++;
>> +
>> +		pfn++;
>> +		pte++;
>> +	}
>> +
>> +	return i;
>> +}
>> +
>>  /*
>> - * Copy one pte.  Returns 0 if succeeded, or -EAGAIN if one preallocated page
>> - * is required to copy this pte.
>> + * Copy set of contiguous ptes.  Returns number of ptes copied if succeeded
>> + * (always gte 1), or -EAGAIN if one preallocated page is required to copy the
>> + * first pte.
>>   */
>>  static inline int
>> -copy_present_pte(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>> -		 pte_t *dst_pte, pte_t *src_pte, unsigned long addr, int *rss,
>> -		 struct folio **prealloc)
>> +copy_present_ptes(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>> +		  pte_t *dst_pte, pte_t *src_pte,
>> +		  unsigned long addr, unsigned long end,
>> +		  int *rss, struct folio **prealloc)
>>  {
>>  	struct mm_struct *src_mm = src_vma->vm_mm;
>>  	unsigned long vm_flags = src_vma->vm_flags;
>>  	pte_t pte = ptep_get(src_pte);
>>  	struct page *page;
>>  	struct folio *folio;
>> +	int nr = 1;
>> +	bool anon = false;
>> +	bool enforce_uffd_wp = userfaultfd_wp(dst_vma);
>> +	int nr_dirty = !!pte_dirty(pte);
>> +	int nr_writable = !!pte_write(pte);
>> +	int i, ret;
>>  
>>  	page = vm_normal_page(src_vma, addr, pte);
>> -	if (page)
>> +	if (page) {
>>  		folio = page_folio(page);
>> -	if (page && folio_test_anon(folio)) {
>> -		/*
>> -		 * If this page may have been pinned by the parent process,
>> -		 * copy the page immediately for the child so that we'll always
>> -		 * guarantee the pinned page won't be randomly replaced in the
>> -		 * future.
>> -		 */
>> -		folio_get(folio);
>> -		if (unlikely(page_try_dup_anon_rmap(page, false, src_vma))) {
>> -			/* Page may be pinned, we have to copy. */
>> -			folio_put(folio);
>> -			return copy_present_page(dst_vma, src_vma, dst_pte, src_pte,
>> -						 addr, rss, prealloc, page);
>> +		anon = folio_test_anon(folio);
>> +		nr = folio_nr_pages_cont_mapped(folio, page, src_pte, addr, end,
>> +						pte, enforce_uffd_wp, &nr_dirty,
>> +						&nr_writable);
>> +		folio_ref_add(folio, nr);
>> +
>> +		for (i = 0; i < nr; i++, page++) {
>> +			if (anon) {
>> +				/*
>> +				 * If this page may have been pinned by the
>> +				 * parent process, copy the page immediately for
>> +				 * the child so that we'll always guarantee the
>> +				 * pinned page won't be randomly replaced in the
>> +				 * future.
>> +				 */
>> +				if (unlikely(page_try_dup_anon_rmap(
>> +						page, false, src_vma))) {
>> +					if (i != 0)
>> +						break;
>> +					/* Page may be pinned, we have to copy. */
>> +					folio_ref_sub(folio, nr);
>> +					ret = copy_present_page(
>> +						dst_vma, src_vma, dst_pte,
>> +						src_pte, addr, rss, prealloc,
>> +						page);
>> +					return ret == 0 ? 1 : ret;
>> +				}
>> +				rss[MM_ANONPAGES]++;
>> +				VM_BUG_ON(PageAnonExclusive(page));
>> +			} else {
>> +				page_dup_file_rmap(page, false);
>> +				rss[mm_counter_file(page)]++;
>> +			}
>>  		}
>> -		rss[MM_ANONPAGES]++;
>> -	} else if (page) {
>> -		folio_get(folio);
>> -		page_dup_file_rmap(page, false);
>> -		rss[mm_counter_file(page)]++;
>> -	}
>>  
>> -	/*
>> -	 * If it's a COW mapping, write protect it both
>> -	 * in the parent and the child
>> -	 */
>> -	if (is_cow_mapping(vm_flags) && pte_write(pte)) {
>> -		ptep_set_wrprotect(src_mm, addr, src_pte);
>> -		pte = pte_wrprotect(pte);
>> +		if (i < nr) {
>> +			folio_ref_sub(folio, nr - i);
>> +			nr = i;
>> +		}
>>  	}
>> -	VM_BUG_ON(page && folio_test_anon(folio) && PageAnonExclusive(page));
>>  
>>  	/*
>> -	 * If it's a shared mapping, mark it clean in
>> -	 * the child
>> +	 * If it's a shared mapping, mark it clean and write protected in the
>> +	 * child, and rely on a write fault to fix up the permissions. This
>> +	 * allows determining batch size without having to consider RO/RW
>> +	 * permissions. As an optimization, skip wrprotect if all ptes in the
>> +	 * batch have the same permissions.
>> +	 *
>> +	 * If its a private (CoW) mapping, mark it dirty in the child if _any_
>> +	 * of the parent mappings in the block were marked dirty. The contiguous
>> +	 * block of mappings are all backed by the same folio, so if any are
>> +	 * dirty then the whole folio is dirty. This allows determining batch
>> +	 * size without having to consider the dirty bit. Further, write protect
>> +	 * it both in the parent and the child so that a future write will cause
>> +	 * a CoW. As as an optimization, skip the wrprotect if all the ptes in
>> +	 * the batch are already readonly.
>>  	 */
>> -	if (vm_flags & VM_SHARED)
>> +	if (vm_flags & VM_SHARED) {
>>  		pte = pte_mkclean(pte);
>> -	pte = pte_mkold(pte);
>> +		if (nr_writable > 0 && nr_writable < nr)
>> +			pte = pte_wrprotect(pte);
>> +	} else {
>> +		if (nr_dirty)
>> +			pte = pte_mkdirty(pte);
>> +		if (nr_writable) {
>> +			ptep_set_wrprotects(src_mm, addr, src_pte, nr);
>> +			pte = pte_wrprotect(pte);
>> +		}
>> +	}
>>  
>> -	if (!userfaultfd_wp(dst_vma))
>> +	pte = pte_mkold(pte);
>> +	pte = pte_clear_soft_dirty(pte);
>> +	if (!enforce_uffd_wp)
>>  		pte = pte_clear_uffd_wp(pte);
>>  
>> -	set_pte_at(dst_vma->vm_mm, addr, dst_pte, pte);
>> -	return 0;
>> +	set_ptes(dst_vma->vm_mm, addr, dst_pte, pte, nr);
>> +	return nr;
> 
> I don't have any further comments and you have addressed my previous
> ones so feel free to add:
> 
> Reviewed-by: Alistair Popple <apopple@nvidia.com>
> 
> However whilst I think the above CoW sequence looks correct it would be
> nice if someone else could take a look as well.

Thanks for the RB! David has taken a look at the CoW part and helped develop the
logic, so I'm pretty confident in it.

However, David also sent me some microbenchmarks for fork, DONTNEED, munmap, etc
for order-0 and PTE-mapped THP (2M). I'm seeing a ferw performance regressions
with those, which I'm currently trying to resolve. At the moment it's looking
like I'll have to expose some function to allow the core code to skip forward a
number of ptes so that in the contpte-mapped case, the core code only does
ptep_get() once per contpte block. As a result there will be some churn here.

I'm currently working out some bugs and hope to post an updated series with perf
numbers for those microbenchmarks by the end of the week, all being well.

> 
>>  }
>>  
>>  static inline struct folio *page_copy_prealloc(struct mm_struct *src_mm,
>> @@ -1021,6 +1115,7 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>>  	int rss[NR_MM_COUNTERS];
>>  	swp_entry_t entry = (swp_entry_t){0};
>>  	struct folio *prealloc = NULL;
>> +	int nr_ptes;
>>  
>>  again:
>>  	progress = 0;
>> @@ -1051,6 +1146,8 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>>  	arch_enter_lazy_mmu_mode();
>>  
>>  	do {
>> +		nr_ptes = 1;
>> +
>>  		/*
>>  		 * We are holding two locks at this point - either of them
>>  		 * could generate latencies in another task on another CPU.
>> @@ -1086,16 +1183,21 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>>  			 * the now present pte.
>>  			 */
>>  			WARN_ON_ONCE(ret != -ENOENT);
>> +			ret = 0;
>>  		}
>> -		/* copy_present_pte() will clear `*prealloc' if consumed */
>> -		ret = copy_present_pte(dst_vma, src_vma, dst_pte, src_pte,
>> -				       addr, rss, &prealloc);
>> +		/* copy_present_ptes() will clear `*prealloc' if consumed */
>> +		nr_ptes = copy_present_ptes(dst_vma, src_vma, dst_pte, src_pte,
>> +					    addr, end, rss, &prealloc);
>> +
>>  		/*
>>  		 * If we need a pre-allocated page for this pte, drop the
>>  		 * locks, allocate, and try again.
>>  		 */
>> -		if (unlikely(ret == -EAGAIN))
>> +		if (unlikely(nr_ptes == -EAGAIN)) {
>> +			ret = -EAGAIN;
>>  			break;
>> +		}
>> +
>>  		if (unlikely(prealloc)) {
>>  			/*
>>  			 * pre-alloc page cannot be reused by next time so as
>> @@ -1106,8 +1208,9 @@ copy_pte_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
>>  			folio_put(prealloc);
>>  			prealloc = NULL;
>>  		}
>> -		progress += 8;
>> -	} while (dst_pte++, src_pte++, addr += PAGE_SIZE, addr != end);
>> +		progress += 8 * nr_ptes;
>> +	} while (dst_pte += nr_ptes, src_pte += nr_ptes,
>> +		 addr += PAGE_SIZE * nr_ptes, addr != end);
>>  
>>  	arch_leave_lazy_mmu_mode();
>>  	pte_unmap_unlock(orig_src_pte, src_ptl);
>