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

[Yuan-Hao Hsu <aa9736195201@gmail.com>]

fork() maps the anonymous pages of the parent read-only and clears
PageAnonExclusive on them.  Once the child has exec'ed or exited, the
parent's write fault takes the reuse path of do_wp_page():
wp_can_reuse_anon_folio() finds that all references to the folio come
from this MM, the page is marked exclusive again and its PTE is made
writable.

For a large folio that check is about the folio and holds for every
page of it, but only the page that faulted is marked exclusive and
made writable.  Each other page takes a write fault of its own and
takes the large mapcount lock to find out the same thing again: 16
faults for a 64K folio.  The same THP mapped by a PMD is reused by one
fault in do_huge_pmd_wp_page(), and do_swap_page() maps all PTEs of an
exclusive large folio writable at once.

Barry proposed reusing the whole mTHP from one fault in 2024 [1].  The
reservations then were the latency of the individual write fault and
how far to go around the faulting PTE: a contpte-sized block was fine,
anything bigger not yet convincing (David's replies, linked below).
Commit 1da190f4d0a6 ("mm: Copy-on-Write (COW) reuse support for
PTE-mapped THP") then added the per-folio check and left faulting
around for later.

This is the contpte-sized part.  Once the folio is known to be
exclusive, walk the aligned block of 16 PTEs around the fault, within
the folio and the VMA.  PTEs that still map the folio read-only are
batched with folio_pte_batch_flags(), their pages are marked
exclusive, and where can_change_pte_writable() agrees the batch is
made writable with modify_prot_start_ptes()/modify_prot_commit_ptes(),
as mprotect() does it.  That leaves NUMA hinting and uffd-wp PTEs
alone, keeps soft-dirty tracking exact, and on arm64 writes a contpte
block back as a whole where ptep_set_access_flags() on one PTE has to
unfold it.  Pages that cannot be made writable are still marked
exclusive, so their own fault skips the folio check.  The PTE that
faulted is in one of the batches and is completed by wp_page_reuse()
as before.  Small folios, PMD-mapped THPs, unsharing faults and the
copy path are unchanged.

The cost is bounded by the block: at most 16 PTEs read once by
folio_pte_batch_flags(), the scan fork() and mprotect() already run
over them, 16 pages marked exclusive and 16 PTEs written.  On x86-64
(i7-12700KF; medians of 15 runs, two boots of each kernel taken
alternately) the fault that does that for a 64K folio takes 750 ns,
against 420-440 ns for a reuse fault today and 4,600-5,000 ns for the
fault that allocated the folio.

256 MiB of 64K folios after fork() and the child's exit, faults and
time of the pass:

                            v7.3-rc3+             patched
  one byte per page, seq    65,601  30.8/30.5 ms   4,164   5.7/ 5.7 ms
  one byte per page, random 65,601  39.5/41.6 ms   4,161   7.4/ 7.6 ms
  memset()                  65,541  57.4/59.5 ms   4,102  39.5/39.8 ms
  8 threads, random order   65,541   6.1/ 6.2 ms   4,155   0.9/ 0.9 ms
  one store per 64K folio    4,101   1.9/ 2.0 ms   4,101   3.5/ 3.3 ms

Redis 7.0.15, 1.3 M keys of 512 bytes on 64K mTHP, BGSAVE and then
1,000,000 SETs: the faults of redis-server during the SETs go from
262,100 to 18,550 and its CPU time from 1.48-1.53 s to 1.31-1.41 s;
the requests per second stay within the boot-to-boot spread.  With
THP off nothing changes.

arm64, under QEMU for the counters: the first pass after fork()
unfolds every contpte block of the 64K folios today (512
contpte_convert() calls for 512 blocks, and nothing folds them again);
with this patch none is unfolded.

[1] https://lore.kernel.org/r/20240831092339.66085-1-21cnbao@gmail.com

Link: https://lore.kernel.org/all/20240831092339.66085-1-21cnbao@gmail.com/
Link: https://lore.kernel.org/all/b7853f0f-7044-4c49-931c-c61900229b19@redhat.com/
Link: https://lore.kernel.org/all/36933711-ae0f-468c-93bd-d6a67d974c9d@redhat.com/
Assisted-by: LLM sparse
Signed-off-by: Yuan-Hao Hsu <aa9736195201@gmail.com>
---
 mm/memory.c | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 73 insertions(+), 2 deletions(-)

diff --git a/mm/memory.c b/mm/memory.c
index 8b0c2c735d3d..73e5691b3ed8 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -4355,6 +4355,71 @@ static bool wp_can_reuse_anon_folio(struct folio *folio,
 	return true;
 }
 
+/*
+ * The pages of the folio around the one that faulted are handled in aligned
+ * blocks of this many PTEs: a 64K folio with 4K pages, and the size of a
+ * contpte block on arm64.
+ */
+#define WP_REUSE_NR_PTES	16
+
+/*
+ * wp_can_reuse_anon_folio() found a large folio to be exclusive to this MM.
+ * That holds for all of its pages and not only for the one that faulted: mark
+ * the ones in the same block exclusive as well and map them writable, like
+ * mprotect() would. Each of them would otherwise take a write fault of its own
+ * that repeats the check on the very same folio.
+ *
+ * The PTE that faulted is among them; wp_page_reuse() completes it.
+ */
+static void wp_reuse_large_anon_folio(struct vm_fault *vmf,
+		struct folio *folio)
+{
+	const fpb_t flags = FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY;
+	const unsigned long idx = folio_page_idx(folio, vmf->page);
+	struct vm_area_struct *vma = vmf->vma;
+	unsigned long addr = vmf->address;
+	unsigned long block = ALIGN_DOWN(addr, WP_REUSE_NR_PTES * PAGE_SIZE);
+	unsigned long nr_before, nr_after, end;
+	struct page *page;
+	unsigned int nr, i;
+	pte_t *ptep, pte;
+
+	/* Stay within the folio, the VMA and the block. */
+	nr_before = min3(idx, (addr - block) >> PAGE_SHIFT,
+			 (addr - vma->vm_start) >> PAGE_SHIFT);
+	nr_after = min3(folio_nr_pages(folio) - idx,
+			(block + WP_REUSE_NR_PTES * PAGE_SIZE - addr) >> PAGE_SHIFT,
+			(vma->vm_end - addr) >> PAGE_SHIFT);
+	end = addr + (nr_after << PAGE_SHIFT);
+	addr -= nr_before << PAGE_SHIFT;
+	ptep = vmf->pte - nr_before;
+	page = vmf->page - nr_before;
+
+	for (; addr != end; addr += nr * PAGE_SIZE, ptep += nr, page += nr) {
+		pte = ptep_get(ptep);
+		nr = 1;
+
+		/* Unmapped or replaced since, or writable already. */
+		if (!pte_present(pte) || pte_pfn(pte) != page_to_pfn(page) ||
+		    pte_write(pte))
+			continue;
+
+		nr = folio_pte_batch_flags(folio, NULL, ptep, &pte,
+					   (end - addr) >> PAGE_SHIFT, flags);
+		for (i = 0; i < nr; i++)
+			if (!PageAnonExclusive(page + i))
+				SetPageAnonExclusive(page + i);
+
+		/* The PTEs of a batch agree on everything this looks at. */
+		if (!can_change_pte_writable(vma, addr, pte))
+			continue;
+
+		pte = modify_prot_start_ptes(vma, addr, ptep, nr);
+		modify_prot_commit_ptes(vma, addr, ptep, pte,
+					pte_mkwrite(pte, vma), nr);
+	}
+}
+
 /*
  * This routine handles present pages, when
  * * users try to write to a shared page (FAULT_FLAG_WRITE)
@@ -4449,8 +4514,14 @@ static vm_fault_t do_wp_page(struct vm_fault *vmf)
 	 */
 	if (folio && folio_test_anon(folio) &&
 	    (PageAnonExclusive(vmf->page) || wp_can_reuse_anon_folio(folio, vma))) {
-		if (!PageAnonExclusive(vmf->page))
-			SetPageAnonExclusive(vmf->page);
+		if (!PageAnonExclusive(vmf->page)) {
+			if (IS_ENABLED(CONFIG_TRANSPARENT_HUGEPAGE) &&
+			    folio_test_large(folio) && likely(!unshare) &&
+			    likely(vma->vm_flags & VM_WRITE))
+				wp_reuse_large_anon_folio(vmf, folio);
+			else
+				SetPageAnonExclusive(vmf->page);
+		}
 		if (unlikely(unshare)) {
 			pte_unmap_unlock(vmf->pte, vmf->ptl);
 			return 0;
-- 
2.43.0