[PATCH 2/2] mm/memory: reuse 16 PTEs of an exclusive large folio on a write fault

[David Hildenbrand (Arm)]

Signed-off-by: David Hildenbrand (Arm) <david@xxxxxxxxxx>
---
 mm/memory.c | 120 +++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 115 insertions(+), 5 deletions(-)

diff --git a/mm/memory.c b/mm/memory.c
index 67fcf67bc64fd..dab2a274783bf 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -4414,10 +4414,78 @@ static bool wp_can_reuse_anon_folio(struct folio *folio,
 	return true;
 }
 
+#define WP_REUSE_MAX_NR_PTES	16
+
+static unsigned int wp_anon_folio_pte_batch(struct vm_fault *vmf,
+		struct folio *folio, unsigned long *addr, struct page **page,
+		pte_t *pte, pte_t **ptep)
+{
+	/* modify_prot_start_ptes() needs most PTE bits to match. */
+	const fpb_t flags = FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY;
+	struct vm_area_struct *vma = vmf->vma;
+	unsigned long batch_start_addr, batch_size, folio_idx, nr_before;
+	unsigned int batch_nr_pages;
+	pte_t *batch_start_ptep;
+	pte_t batch_start_pte, expected_pte;
+
+	if (!IS_ENABLED(CONFIG_TRANSPARENT_HUGEPAGE) ||
+	    !folio_test_large(folio))
+		return 1;
+
+	/*
+	 * We'll try batching in a naturally aligned block surrounding our
+	 * faulting PTE.
+	 */
+	batch_nr_pages = min(folio_large_nr_pages(folio), WP_REUSE_MAX_NR_PTES);
+	batch_size = batch_nr_pages << PAGE_SHIFT;
+	batch_start_addr = ALIGN_DOWN(*addr, batch_size);
+	folio_idx = folio_page_idx(folio, *page);
+	nr_before = (*addr - batch_start_addr) >> PAGE_SHIFT;
+
+	/* Stay within the folio. */
+	if (nr_before > folio_idx ||
+	    folio_idx - nr_before + batch_nr_pages > folio_large_nr_pages(folio))
+		return 1;
+
+	/* Stay within the VMA. */
+	if (batch_start_addr < vma->vm_start ||
+	    batch_start_addr + batch_size > vma->vm_end)
+		return 1;
+
+	batch_start_ptep = *ptep - nr_before;
+	batch_start_pte = ptep_get(batch_start_ptep);
+
+	expected_pte = pte_advance_pfn(batch_start_pte, nr_before);
+	if (!pte_same(__pte_batch_clear_ignored(expected_pte, flags),
+		      __pte_batch_clear_ignored(*pte, flags)))
+		return 1;
+
+	if (folio_pte_batch_flags(folio, NULL, batch_start_ptep,
+				  &batch_start_pte, batch_nr_pages,
+				  flags) != batch_nr_pages)
+		return 1;
+
+	/*
+	 * Our faulting PTE is guaranteed to be part of the batch, and all
+	 * PTE bits are compatible.
+	 */
+	*addr = batch_start_addr;
+	*page = *page - nr_before;
+	*pte = batch_start_pte;
+	*ptep = batch_start_ptep;
+	return batch_nr_pages;
+}
+
 static bool wp_try_reuse_anon_page(struct vm_fault *vmf, struct folio *folio)
 	__cond_releases(true, vmf->ptl)
 {
 	const bool unshare = vmf->flags & FAULT_FLAG_UNSHARE;
+	struct vm_area_struct *vma = vmf->vma;
+	unsigned long addr = vmf->address;
+	struct page *page = vmf->page;
+	pte_t new_pte, pte = vmf->orig_pte;
+	pte_t *ptep = vmf->pte;
+	unsigned int i, nr = 1;
 
 	VM_WARN_ON_ONCE(!folio_test_anon(folio));
 
@@ -4426,14 +4494,56 @@ static bool wp_try_reuse_anon_page(struct vm_fault *vmf, struct folio *folio)
 	 * is impossible. We might miss VM_WRITE for FOLL_FORCE handling.
 	 *
 	 * If we encounter a page that is marked exclusive, we must reuse
-	 * the page without further checks.
+	 * the page without further checks. Don't process more than a single
+	 * PTE in that case.
 	 */
-	if (!PageAnonExclusive(vmf->page)) {
-		if (!wp_can_reuse_anon_folio(folio, vmf->vma))
-			return false;
-		SetPageAnonExclusive(vmf->page);
+	if (PageAnonExclusive(page))
+		goto reuse_single_page;
+
+	if (!wp_can_reuse_anon_folio(folio, vma))
+		return false;
+
+	/* We can use any folio page that is mapped in this page table. */
+	nr = wp_anon_folio_pte_batch(vmf, folio, &addr, &page, &pte, &ptep);
+	if (nr == 1) {
+		SetPageAnonExclusive(page);
+		goto reuse_single_page;
 	}
 
+	for (i = 0; i < nr; i++)
+		if (!PageAnonExclusive(page + i))
+			SetPageAnonExclusive(page + i);
+
+	/* Careful: don't mark unrelated PTEs soft-dirty by batching. */
+	if (unlikely(pte_needs_soft_dirty_wp(vma, pte)))
+		goto reuse_single_page;
+
+	if (unlikely(unshare)) {
+		pte_unmap_unlock(vmf->pte, vmf->ptl);
+		return true;
+	}
+
+	/* See wp_page_reuse() */
+	folio_xchg_last_cpupid(folio, (1 << LAST_CPUPID_SHIFT) - 1);
+
+	for (i = 0; i < nr; i++)
+		flush_cache_page(vma, addr + (i << PAGE_SHIFT), pte_pfn(pte) + i);
+
+	pte = modify_prot_start_ptes(vma, addr, ptep, nr);
+	new_pte = maybe_mkwrite(pte_mkdirty(pte_mkyoung(pte)), vmf->vma);
+	modify_prot_commit_ptes(vma, addr, ptep, pte, new_pte, nr);
+
+	/* Remove stale read-only TLB entry for the faulting PTE only. */
+	flush_tlb_fix_spurious_fault(vma, vmf->address, vmf->pte);
+
+	/* But update the MMU cache of all changed PTEs. */
+	update_mmu_cache_range(vmf, vma, addr, ptep, nr);
+
+	pte_unmap_unlock(vmf->pte, vmf->ptl);
+	count_vm_event(PGREUSE);
+	return true;
+
+reuse_single_page:
 	if (unlikely(unshare)) {
 		pte_unmap_unlock(vmf->pte, vmf->ptl);
 		return true;
-- 
2.43.0



-- 
Cheers,

David