[RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Yes, I know, this is a lot of changes. But I'm happy with the overall state
of the patchset and the only reason I tag it as RFC is that it is tricky
to get 57 patches upstream.

I wanted to give a view of the end state first. I will suggest a possible
way to split it below.

I would appreciate any feedback.

TL;DR
=====

This replaces khugepaged's anonymous collapse with an engine that
can collapse sub-PMD ranges. It is built around migration entries and
frozen folios instead of heavy locking and isolation, aiming for better
scalability and less disruption to the workload being collapsed.

Why
===

mTHP collapse landed in khugepaged in 7.2 and I was glad to see it.  We
at Meta run arm64 with 64K base pages, where a PMD is 512M: PMD-order THP
is of limited use at that size, and mTHP is exactly what we want.

It turned out not to help us.

khugepaged only ever looks at PMD-aligned windows, and it is not an easy
limitation to lift.

Fixing the alignment is a one-line change, but what it feeds assumes the
PMD everywhere that matters: collapse_huge_page() clears and flushes the
whole PMD whatever order it is collapsing, installs a PMD leaf because
that is the only thing it can produce, and keeps everyone out with
mmap_write_lock, anon_vma_lock_write() and an IPI broadcast while it
does.

Which is why hugepage_vma_revalidate() demands that the VMA span the
whole PMD even for an mTHP order -- "we'd need to lock all VMAs in the
PMD range to support this", as the comment there puts it.  A PMD-granular
operation is only safe when one VMA owns the PMD, and that is exactly the
restriction in the way.  The alignment is the symptom; the PMD is the
design.

So both roots have to go.

Design
======

The old mechanism holds the address space still because it has nothing
else stopping the sources from moving under the copy.  The new engine
makes the sources themselves inert instead, with the two barriers
migration already uses, raised in that order:

  1. migration entries replace the source PTEs.  Faults and GUP-slow
     now wait on the source folio's lock, which is taken before the
     first entry becomes visible.
  2. the source folio's refcount is frozen to its expected value.
     GUP-fast, pfn walkers, reclaim, compaction and memory-failure all
     fail folio_try_get() and back off.

Between the two, nothing can reach a source, so the copy runs with no
lock held at all -- and the address space is left alone while it does.

What that removes from every collapse path:

  mmap_write_lock              -> mmap_read
  anon_vma_lock_write()        -> nothing: an rmap walk needs the folio
                                  locked, and the engine holds that lock
                                  from freeze to putback
  tlb_remove_table_sync_one()  -> nothing: one ranged flush per round
  LRU isolation                -> nothing: sources are inert in place

Working in windows rather than whole PMDs takes care of the other root.
A sub-PMD window is collapsed under the page table lock, so a collapse
disturbs only the window it collapses, and each candidate is validated
at its own order -- a window need only fit its own VMA.  A PMD-order
candidate still has to own the whole PMD, which is the old rule kept
where it is still needed.

Candidates are carried through the passes a batch at a time rather than
one window at a time, so a round pays for its flush and its lock
acquisitions once.

With the barriers holding the sources still, which read lock the engine
takes stops being part of the design.  A round works inside a single
VMA, so patches 43-49 switch it from mmap_read to per-VMA locking: an
mmap_write elsewhere in the mm then stops waiting for a collapse that
has nothing to do with it.  That block is the only part of the series
that needs per-VMA locking to be unconditional, and it is a separate
dependency (see below); everything before it runs under mmap_read and
does not care.

Patch 7 sketches the engine as a comment naming every pass, what lock it
takes and what it may sleep on; the details are there rather than here.

What falls out beyond the lock diet:

 - mTHP collapse in VMAs smaller than a PMD, which is the arm64 case
   above: a 2M VMA on an arm64/64K machine collapses nothing today at
   any order, and collapses to mTHP here.
 - Hole and zeropage population at every order, so partially populated
   windows collapse to mTHP under the same max_ptes_none policy as PMD.
 - Sources come in spans -- any stretch of consecutive PTEs mapping
   consecutive pages of one folio -- so partially mapped and scrambled
   compound sources (the PTE-mapped-THP re-collapse class) work at
   every order.
 - A table that cannot become one huge page still yields the largest
   windows inside it, where before a single disqualified PTE gave up
   the whole table.

Reading the series
==================

57 patches is a lot to land on a list.  They go in blocks:

  1-6    helpers and shared state: pte_folio(), pte_none_or_zero(),
         mm/collapse.h, and the policy that replaces asking whether
         khugepaged started a collapse
  7-8    the engine's shape: entry points, a call-tree comment naming
         every pass, and the scan filled in
  9-23   the collapse half, top down: the round frame, then each pass
         in turn, then selection and the retry store
  24     per-candidate tracing, before the switch takes the old
         tracepoints away
  25-28  the switch: point the anon path at the engine, widen coverage
         to sub-PMD VMAs, delete the mechanism it replaces
  29-35  move what is left of collapse out of khugepaged.c, and
         MADV_COLLAPSE into madvise.c
  36-42  tracing: the engine's own events and trace header
  43-49  per-VMA locking, and the mm reference that makes it safe
  50-56  selftests for what the engine can now do
  57     MAINTAINERS

The two patches worth reading first if you read nothing else are 7 (the
design, as a comment naming the whole call tree) and 16 (the freeze,
which is where the safety argument lives).

A possible split, if that helps:

  1-2    two mm helpers, pte_folio() and pte_none_or_zero().  Both
         convert callers outside collapse and are useful on their own
  3-27   the engine and the switch-over.  This is the smallest unit
         that does anything: stop earlier and the tree carries an
         engine nothing calls
  28     remove the mechanism the engine replaces
  29-42  moving what is left of collapse out of khugepaged.c, and the
         engine's own tracepoints
  43-49  per-VMA locking
  50-57  selftests and MAINTAINERS

Keeping the removal separate leaves both engines in the tree with only
the new one reachable, so the switch can be reverted on its own if
something turns up.  The old mechanism is already carried that way for
three patches inside the series, so this costs nothing but 975 lines of
unreferenced code until 28 lands.  That safety net only lasts until the
blocks after it land, though: once collapse has moved out of
khugepaged.c and the locking has changed, reverting the switch no longer
gives back a working old engine.

Base and dependencies
=====================

This applies on the selftests series, not on plain mm-new:

  [PATCH v4 00/19] selftests/mm: improve khugepaged coverage
  https://lore.kernel.org/all/20260815015901.1236937-1-kirill@shutemov.name/

which is on mm-new 33f61b12d297.

Patches 43-49 depend on Suren's unconditional per-VMA locks:

  [PATCH v6 0/5] mm: Unconditional per-VMA locks and cleanups
  https://lore.kernel.org/all/20260813193433.3318288-1-surenb@google.com/

That series is not in mm-new yet, and with patch 46 applied SMP=n does
not build without it: lock_next_vma() is behind CONFIG_PER_VMA_LOCK in
mmap_lock.h.  Everything up to patch 42 builds and runs on mm-new as it
stands.  There is no fallback path by choice -- adding one would mean
carrying two locking models through every pass.

Both branches are available at

  git://git.kernel.org/pub/scm/linux/kernel/git/kas/linux.git collapse/rfc-v1

and the benchmark used for the numbers below, which is unposted and not a
dependency, at

  git://git.kernel.org/pub/scm/linux/kernel/git/kas/linux.git perf/bench-usemem

Performance
===========

Measuring khugepaged is awkward.  It is a background daemon, so what
matters is what a workload feels while it runs, not what the daemon
reports about itself -- and the usual coverage instrument is no help
below the PMD: smaps AnonHugePages only counts PMD-order folios, so it
reads zero however much mTHP has been collapsed.

So I wrote "perf bench mem usemem" for this.  It touches a region while
khugepaged works on it and reports the workload's own latency
percentiles and throughput, against per-size counters that can see
sub-PMD folios.  The branch is above; it is unposted and not a
dependency.

x86-64, production configs (no KASAN, no lockdep, no DEBUG_VM, no
PAGE_TABLE_CHECK), interleaved rounds on an idle host, equal work on
every arm.

I measured three kernels, so the two halves of the series can be told
apart in the numbers below:

  A   the base
  B   the new engine, still under mmap_read
  C   B plus per-VMA locking -- what this series ends up with

The engine: a sub-PMD collapse stops blanking the surrounding 2M
-----------------------------------------------------------------

base routes sub-PMD collapse through collapse_huge_page(), whose
pmdp_collapse_flush() and tlb_remove_table_sync_one() are not gated on
order: to collapse an order-4 window of 16 pages it clears and flushes
the whole 512-page PMD and IPIs, then repopulates.  The engine does the
window under the PTL.

A thread reading and writing a 32G region while it is collapsed at
order-4, 16384 collapses on every arm:

                         A        B        C
  read p99 (ns)       3071     1023     1023   -66.7%
  write p99 (ns)      3071      927      927   -69.8%

and the workload's read rate rises by 68% on both engine kernels.

B == C, so this is the engine, not the locking.

At PMD order the same workload is flat, and that is expected rather
than disappointing: it is the one configuration where both mechanisms
disturb exactly the same 2M.  Read it as no regression at PMD order.

Per-VMA locking: address-space operations stop waiting on the scan
-------------------------------------------------------------------

MADV_HUGEPAGE/MADV_NOHUGEPAGE toggling against a scanning mm, which is
what jemalloc does with its arenas.  4096 collapses on every arm:

                         A        B        C
  ops/sec           340656   340820   603305   +77%
  p99 (ns)           77823    86015     3327   -96%
  p99.9 (ns)         86015    94207     9215   -89%

B is about 11% worse than base at p99 here, consistently across runs:
the engine alone slightly worsens hint-toggle latency, and per-VMA
locking is what turns it into a win.  Both halves are in this series, so
C is what a reviewer gets, but the middle column is the honest one.

The trade is real in the other direction too.  On settled memory with
nothing to collapse and scan_sleep_millisecs=0, per-VMA locking costs
about 47% of scan throughput against one mmap_read for the whole walk.
That is a synthetic worst case -- the daemon wraps 8000 times a second
there, where production defaults to 10s between passes -- and it buys
mmap/munmap p99 of 56us against 1.4us.

Collapse itself is not slower
-----------------------------

One complete pass over a 32G region, 16384 collapses, khugepaged CPU
from /proc/<pid>/stat, 7 repeats:

  A base    median 5760 ms   spread 12.7%
  B engine  median 4990 ms   spread  3.8%
  C pervma  median 5020 ms   spread 12.4%

The base arm is bimodal, so its median moves with sampling and the
percentage is soft.  The distribution-free statement is better: every
engine run used less CPU than every base run.

The engine also allocates one destination per folio installed, where
base allocates 5.25 and frees the rest again: nothing is allocated until
the sources are frozen and the collapse can no longer be refused.

A measurement note, since an earlier version of this series quoted worse
figures.  khugepaged CPU has to be measured per collapse or per
completed pass, never over a fixed window with scan_sleep_millisecs=0:
the daemon never sleeps, so whichever kernel finishes the work sooner
spends the rest of the window scanning settled memory and is charged for
it.  Measured that way the engine appeared to cost 10% more CPU;
measured per unit of work it costs less.

Costs
=====

At PMD order the engine issues two TLB flushes per collapse where the
old mechanism issues one: the freeze's ranged flush plus the terminal
layer's pmdp_collapse_flush().  A PMD candidate is alone in its round,
so nothing amortizes the first.  Dropping the old per-collapse
tlb_remove_table_sync_one() IPI presumably pays for it, but that was not
measured and is not claimed here.

There may be a way out -- a PMD migration entry over the table during
the window, so the CPU never caches a walk to shoot down -- but that
means teaching every pmd-level walker a new kind of entry, and I have
not tried it.

PMD collapse deposits a freshly allocated page table instead of
redepositing the detached one.  Whoever withdraws a deposited table
frees it immediately, with nothing to hold a lockless walker off first,
and under a read lock the detached table may still be traversed by
GUP-fast or an RCU pte walk.  It goes to pte_free_defer() instead,
exactly as retract_page_tables() does.  One transient table page per PMD
collapse buys the IPI's absence.

That cost goes away if zap_deposited_table() -- the only site that frees
a deposited table outright, the others redeposit it or repopulate the
PMD with it -- used pte_free_defer().  The deposit would no longer have
to be quiescent and the detached table could go straight back.  It would
defer every THP zap's table free, and I have not tried it.

A shared source now costs an extra copy.  The freeze needs every page
exclusive to this mm, so the fault-in pass breaks CoW first -- an
allocation and a copy -- and the collapse then copies that page into the
destination; the old mechanism copied a shared page straight into the
new folio and broke the sharing that way.  It is bounded by
max_ptes_shared, which khugepaged holds at zero below the PMD order, so
in practice this is PMD-order collapse and MADV_COLLAPSE.

Size
====

mm/ grows by 1915 lines net: 4475 added against 2560 deleted.

That is not a claim that this is less code, but it is less than it
looks.  khugepaged.c goes from 3283 lines to 908.  The new engine is
4052 lines across mm/collapse.c and mm/collapse.h, of which 1344 --
about a third -- are comments, which is where the pipeline's invariants
are written down.  What replaces three install paths with their own
isolate/copy/rollback is one engine and one contract.

Testing
=======

Both matrices run the mm selftests plus a race harness, on the
validation config: KASAN, lockdep, PROVE_LOCKING, DEBUG_VM and
PAGE_TABLE_CHECK, 16G of guest memory, swap active so the swap-in
prepass is exercised rather than skipped.

  x86-64        433 pass, 0 fail, 12 skip
  arm64/64K     581 pass, 0 fail, 18 skip

dmesg clean on both.  The arm64 skips are a pre-existing shmem
MADV_COLLAPSE -EINVAL on 64K pages, confirmed against the base by A/B.

Every one of the 57 patches builds with no new warnings; !NUMA and !MMU
(arm nommu) build clean.  SMP=n does not build, for the reason in the
dependencies section above.

The race harness also gets longer soaks -- 1800s per driver mode, with
memory pressure and swap -- and the engine is fuzzed with syzkaller on a
KCOV+KASAN build.  That found two bugs the selftests could not reach: a
teardown that dropped rmap while the source was still frozen, where
removing an mlocked mapping munlocks and munlock_folio() takes a
reference a frozen folio forbids; and a whole-table MADV_DONTNEED racing
the copy window under CONFIG_PT_RECLAIM, which freed the table and left
the sources frozen and locked.  Both are fixed, and both gained coverage
-- the mlocked case is patch 53.

Kiryl Shutsemau (Meta) (57):
  mm: add pte_folio()
  mm: add pte_none_or_zero()
  mm/collapse: add collapse.h for the shared collapse state
  mm/collapse: rename mthp_present_ptes to eligible_ptes
  mm/collapse: state what a collapse may do in the policy
  mm/collapse: move the smallest collapse order to collapse.h
  mm/collapse: sketch the new anonymous collapse engine
  mm/collapse: scan a table for what a collapse could use
  mm/collapse: collect candidate windows into a round
  mm/collapse: run a round and feed the outcomes back
  mm/collapse: sketch the passes of a round
  mm/collapse: allocate a destination per candidate
  mm/collapse: revalidate a round against the VMA
  mm/collapse: fault the sources in before the freeze
  mm/collapse: check what a candidate would freeze
  mm/collapse: freeze the sources behind migration entries
  mm/collapse: copy the sources into the destinations
  mm/collapse: install the destinations at PTE level
  mm/collapse: install a PMD leaf as the terminal layer
  mm/collapse: put the sources back
  mm/collapse: settle whatever the round reached
  mm/collapse: walk a table with a selection cursor
  mm/collapse: give a refused region a second chance
  mm/collapse: report each candidate's outcome to tracing
  mm/collapse: collapse anonymous memory with the new engine
  mm/collapse: give collapse_single_pmd() the range to work on
  mm/collapse: scan the windows a VMA can hold
  mm/collapse: remove the mechanism the engine replaces
  mm/collapse: move what a collapse is judged on into collapse.c
  mm/collapse: name the max_ptes ceiling after collapse
  mm/khugepaged: count collapses where khugepaged makes them
  mm/collapse: move the file collapse into collapse.c
  mm/collapse: split collapse into a scan and a run
  mm/collapse: implement MADV_COLLAPSE in madvise.c
  mm/madvise: drop MADV_COLLAPSE's redundant mm reference
  mm/collapse: report what the scan found
  mm/collapse: report what the fault-in pass paid
  mm/collapse: report the round, and what it made faulters wait
  mm/collapse: name the file collapse's tracepoints after collapse
  mm/collapse: remove the tracepoints of the mechanism that is gone
  mm/collapse: give collapse its own trace header
  mm/collapse: allow error injection into the freeze
  mm/khugepaged: check the scan budget before the work, not after
  mm/khugepaged: hold the address space open across a scan
  mm/collapse: take a per-VMA read lock for the round
  mm/khugepaged: scan under a per-VMA read lock
  mm/madvise: collapse under a per-VMA read lock
  mm/collapse: assert the mm reference the engine relies on
  mm/khugepaged: drop the mmap_lock barrier from __khugepaged_exit()
  selftests/mm: attribute collapses by candidate event alone
  selftests/mm: cover collapse inside a sub-PMD VMA
  selftests/mm: cover a hole-y window in a sub-PMD VMA
  selftests/mm: cover collapse of mlocked ranges
  selftests/mm: cover collapse beside a MADV_FREE'd page
  selftests/mm: cover collapse beside a pinned page
  selftests/mm: cover the scaled max_ptes_shared limit
  MAINTAINERS: add an entry for collapse

 MAINTAINERS                                   |   19 +-
 fs/proc/task_mmu.c                            |    4 +-
 include/linux/huge_mm.h                       |    9 -
 include/linux/mm.h                            |   14 +
 include/linux/pgtable.h                       |   17 +
 .../events/{huge_memory.h => collapse.h}      |  176 +-
 kernel/bpf/btf.c                              |    8 +-
 mm/Makefile                                   |    2 +-
 mm/collapse.c                                 | 3808 +++++++++++++++++
 mm/collapse.h                                 |  244 ++
 mm/hugetlb.c                                  |    8 +-
 mm/khugepaged.c                               | 2711 +-----------
 mm/madvise.c                                  |  251 +-
 mm/migrate_device.c                           |    9 +-
 mm/mremap.c                                   |    2 +-
 tools/testing/selftests/mm/khugepaged.c       |  346 ++
 tools/testing/selftests/mm/khugepaged_race.c  |   29 +-
 .../selftests/mm/khugepaged_sync_check.c      |   65 +-
 tools/testing/selftests/mm/vm_util.c          |    2 +-
 19 files changed, 5022 insertions(+), 2702 deletions(-)
 rename include/trace/events/{huge_memory.h => collapse.h} (60%)
 create mode 100644 mm/collapse.c
 create mode 100644 mm/collapse.h


base-commit: 8b76faf42c5d342b3bf0b1fd97bdaf603ee57354
-- 
2.54.0

[RFC PATCH 01/57] mm: add pte_folio()

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Callers that want the folio behind a present PTE spell it out as
page_folio(pte_page(pte)).

Add pte_folio() as the folio companion to pte_page(), and convert the
callers in fs/proc/task_mmu.c and mm/hugetlb.c.

Preparation for the anonymous collapse engine, which reads the folio
behind a PTE in several places.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 fs/proc/task_mmu.c |  4 ++--
 include/linux/mm.h | 14 ++++++++++++++
 mm/hugetlb.c       |  8 ++++----
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/fs/proc/task_mmu.c b/fs/proc/task_mmu.c
index 5c54aebe2118..459c779b8141 100644
--- a/fs/proc/task_mmu.c
+++ b/fs/proc/task_mmu.c
@@ -1277,7 +1277,7 @@ static int smaps_hugetlb_range(pte_t *pte, unsigned long hmask,
 	ptl = huge_pte_lock(hstate_vma(vma), walk->mm, pte);
 	ptent = huge_ptep_get(walk->mm, addr, pte);
 	if (pte_present(ptent)) {
-		folio = page_folio(pte_page(ptent));
+		folio = pte_folio(ptent);
 		present = true;
 	} else {
 		const softleaf_t entry = softleaf_from_pte(ptent);
@@ -2227,7 +2227,7 @@ static int pagemap_hugetlb_range(pte_t *ptep, unsigned long hmask,
 	ptl = huge_pte_lock(hstate_vma(vma), walk->mm, ptep);
 	pte = huge_ptep_get(walk->mm, addr, ptep);
 	if (pte_present(pte)) {
-		struct folio *folio = page_folio(pte_page(pte));
+		struct folio *folio = pte_folio(pte);
 
 		if (!folio_test_anon(folio))
 			flags |= PM_FILE;
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 0829e0d3b2d1..eb44e3dfee09 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -2681,6 +2681,20 @@ static inline pte_t mk_pte(const struct page *page, pgprot_t pgprot)
 	return pfn_pte(page_to_pfn(page), pgprot);
 }
 
+/**
+ * pte_folio - Return the folio mapped by a present PTE.
+ * @pte: A present page table entry.
+ *
+ * The folio companion to pte_page(); only meaningful for a present PTE
+ * that maps a struct-page-backed folio.
+ *
+ * Return: The folio containing the page @pte maps.
+ */
+static inline struct folio *pte_folio(pte_t pte)
+{
+	return page_folio(pte_page(pte));
+}
+
 /**
  * folio_mk_pte - Create a PTE for this folio
  * @folio: The folio to create a PTE for
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index dded1768193a..bceab8e14118 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -5280,7 +5280,7 @@ void __unmap_hugepage_range(struct mmu_gather *tlb, struct vm_area_struct *vma,
 		 * are about to unmap is the actual folio of interest.
 		 */
 		if (folio_provided) {
-			if (folio != page_folio(pte_page(pte))) {
+			if (folio != pte_folio(pte)) {
 				spin_unlock(ptl);
 				continue;
 			}
@@ -5291,7 +5291,7 @@ void __unmap_hugepage_range(struct mmu_gather *tlb, struct vm_area_struct *vma,
 			 */
 			set_vma_resv_flags(vma, HPAGE_RESV_UNMAPPED);
 		} else {
-			folio = page_folio(pte_page(pte));
+			folio = pte_folio(pte);
 		}
 
 		pte = huge_ptep_get_and_clear(mm, address, ptep, sz);
@@ -5514,7 +5514,7 @@ static vm_fault_t hugetlb_wp(struct vm_fault *vmf)
 		return 0;
 	}
 
-	old_folio = page_folio(pte_page(pte));
+	old_folio = pte_folio(pte);
 
 	delayacct_wpcopy_start();
 
@@ -6189,7 +6189,7 @@ vm_fault_t hugetlb_fault(struct mm_struct *mm, struct vm_area_struct *vma,
 			 * checks whether we can re-use the folio exclusively
 			 * for us in case we are the only user of it.
 			 */
-			folio = page_folio(pte_page(vmf.orig_pte));
+			folio = pte_folio(vmf.orig_pte);
 			if (folio_test_anon(folio) && !folio_trylock(folio)) {
 				need_wait_lock = true;
 				goto out_ptl;
-- 
2.54.0

[RFC PATCH 02/57] mm: add pte_none_or_zero()

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A PTE that is none and one that maps the shared zeropage both stand for
a page of zeroes the mapping does not own.  Code that cares only about
the contents can treat the two alike.

Move khugepaged's local helper for that test to pgtable.h, below the
is_zero_pfn() it is built on.

migrate_vma_insert_page() open-codes the same test on the slot it is
about to fill.  Convert it.  It still tells none from the zeropage, but
only to decide whether there is an old mapping to flush.

No functional change intended.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/linux/pgtable.h | 17 +++++++++++++++++
 mm/khugepaged.c         |  7 -------
 mm/migrate_device.c     |  9 ++-------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index 8c093c119e5a..bbee6d31f015 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -2064,6 +2064,23 @@ static inline struct page *_zero_page(unsigned long addr)
 
 #ifdef CONFIG_MMU
 
+/**
+ * pte_none_or_zero - Does this PTE map nothing, or the shared zeropage?
+ * @pte: The page table entry to test.
+ *
+ * A PTE that is none and one that maps the shared zeropage both stand for a
+ * page of zeroes the mapping does not own, so code that only cares about the
+ * contents can treat them alike.
+ *
+ * Return: %true if @pte is none or maps the shared zeropage.
+ */
+static inline bool pte_none_or_zero(pte_t pte)
+{
+	if (pte_none(pte))
+		return true;
+	return pte_present(pte) && is_zero_pfn(pte_pfn(pte));
+}
+
 #ifndef CONFIG_TRANSPARENT_HUGEPAGE
 static inline int pmd_trans_huge(pmd_t pmd)
 {
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 5a06e3942e88..5f7126cf42f5 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -348,13 +348,6 @@ struct attribute_group khugepaged_attr_group = {
 };
 #endif /* CONFIG_SYSFS */
 
-static bool pte_none_or_zero(pte_t pte)
-{
-	if (pte_none(pte))
-		return true;
-	return pte_present(pte) && is_zero_pfn(pte_pfn(pte));
-}
-
 /**
  * collapse_max_ptes_none - Calculate maximum allowed empty PTEs or PTEs mapping
  * the shared zeropage for the given collapse operation.
diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 9a346162c688..60afa556b994 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -1067,14 +1067,9 @@ static void migrate_vma_insert_page(struct migrate_vma *migrate,
 	if (check_stable_address_space(mm))
 		goto unlock_abort;
 
-	if (pte_present(orig_pte)) {
-		unsigned long pfn = pte_pfn(orig_pte);
-
-		if (!is_zero_pfn(pfn))
-			goto unlock_abort;
-		flush = true;
-	} else if (!pte_none(orig_pte))
+	if (!pte_none_or_zero(orig_pte))
 		goto unlock_abort;
+	flush = pte_present(orig_pte);
 
 	/*
 	 * Check for userfaultfd but do not deliver the fault. Instead,
-- 
2.54.0

[RFC PATCH 03/57] mm/collapse: add collapse.h for the shared collapse state

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Preparation for building the new collapse engine in its own file.  The
engine and khugepaged.c need to agree on what a collapse result is and
what state a scan carries.

Move enum scan_result and struct collapse_control into a new
mm/collapse.h.

No functional change intended.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   | 60 +++++++++++++++++++++++++++++++++++++++++++++++++
 mm/khugepaged.c | 52 +-----------------------------------------
 2 files changed, 61 insertions(+), 51 deletions(-)
 create mode 100644 mm/collapse.h

diff --git a/mm/collapse.h b/mm/collapse.h
new file mode 100644
index 000000000000..26dbac7beddd
--- /dev/null
+++ b/mm/collapse.h
@@ -0,0 +1,60 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __MM_COLLAPSE_H
+#define __MM_COLLAPSE_H
+
+#include <linux/nodemask.h>
+#include <linux/pgtable.h>
+#include <linux/types.h>
+
+enum scan_result {
+	SCAN_FAIL,
+	SCAN_SUCCEED,
+	SCAN_NO_PTE_TABLE,
+	SCAN_PMD_MAPPED,
+	SCAN_EXCEED_NONE_PTE,
+	SCAN_EXCEED_SWAP_PTE,
+	SCAN_EXCEED_SHARED_PTE,
+	SCAN_PTE_NON_PRESENT,
+	SCAN_PTE_UFFD,
+	SCAN_PTE_MAPPED_HUGEPAGE,
+	SCAN_LACK_REFERENCED_PAGE,
+	SCAN_PAGE_NULL,
+	SCAN_SCAN_ABORT,
+	SCAN_PAGE_COUNT,
+	SCAN_PAGE_LRU,
+	SCAN_PAGE_LOCK,
+	SCAN_PAGE_ANON,
+	SCAN_PAGE_LAZYFREE,
+	SCAN_PAGE_COMPOUND,
+	SCAN_ANY_PROCESS,
+	SCAN_VMA_NULL,
+	SCAN_VMA_CHECK,
+	SCAN_ADDRESS_RANGE,
+	SCAN_DEL_PAGE_LRU,
+	SCAN_ALLOC_HUGE_PAGE_FAIL,
+	SCAN_CGROUP_CHARGE_FAIL,
+	SCAN_TRUNCATED,
+	SCAN_PAGE_HAS_PRIVATE,
+	SCAN_STORE_FAILED,
+	SCAN_COPY_MC,
+	SCAN_PAGE_FILLED,
+	SCAN_PAGE_DIRTY_OR_WRITEBACK,
+};
+
+struct collapse_control {
+	bool is_khugepaged;
+
+	/* Num pages scanned per node */
+	u32 node_load[MAX_NUMNODES];
+
+	/* Num pages scanned (see khugepaged_pages_to_scan) */
+	unsigned int progress;
+
+	/* nodemask for allocation fallback */
+	nodemask_t alloc_nmask;
+
+	/* Each bit represents a single occupied (!none/zero) page. */
+	DECLARE_BITMAP(mthp_present_ptes, MAX_PTRS_PER_PTE);
+};
+
+#endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 5f7126cf42f5..804b1d35f52a 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -26,45 +26,11 @@
 #include <linux/cleanup.h>
 
 #include <asm/tlb.h>
+#include "collapse.h"
 #include "internal.h"
 #include "page_alloc.h"
 #include "mm_slot.h"
 
-enum scan_result {
-	SCAN_FAIL,
-	SCAN_SUCCEED,
-	SCAN_NO_PTE_TABLE,
-	SCAN_PMD_MAPPED,
-	SCAN_EXCEED_NONE_PTE,
-	SCAN_EXCEED_SWAP_PTE,
-	SCAN_EXCEED_SHARED_PTE,
-	SCAN_PTE_NON_PRESENT,
-	SCAN_PTE_UFFD,
-	SCAN_PTE_MAPPED_HUGEPAGE,
-	SCAN_LACK_REFERENCED_PAGE,
-	SCAN_PAGE_NULL,
-	SCAN_SCAN_ABORT,
-	SCAN_PAGE_COUNT,
-	SCAN_PAGE_LRU,
-	SCAN_PAGE_LOCK,
-	SCAN_PAGE_ANON,
-	SCAN_PAGE_LAZYFREE,
-	SCAN_PAGE_COMPOUND,
-	SCAN_ANY_PROCESS,
-	SCAN_VMA_NULL,
-	SCAN_VMA_CHECK,
-	SCAN_ADDRESS_RANGE,
-	SCAN_DEL_PAGE_LRU,
-	SCAN_ALLOC_HUGE_PAGE_FAIL,
-	SCAN_CGROUP_CHARGE_FAIL,
-	SCAN_TRUNCATED,
-	SCAN_PAGE_HAS_PRIVATE,
-	SCAN_STORE_FAILED,
-	SCAN_COPY_MC,
-	SCAN_PAGE_FILLED,
-	SCAN_PAGE_DIRTY_OR_WRITEBACK,
-};
-
 #define CREATE_TRACE_POINTS
 #include <trace/events/huge_memory.h>
 
@@ -103,22 +69,6 @@ static struct kmem_cache *mm_slot_cache __ro_after_init;
 
 #define KHUGEPAGED_MIN_MTHP_ORDER	2
 
-struct collapse_control {
-	bool is_khugepaged;
-
-	/* Num pages scanned per node */
-	u32 node_load[MAX_NUMNODES];
-
-	/* Num pages scanned (see khugepaged_pages_to_scan) */
-	unsigned int progress;
-
-	/* nodemask for allocation fallback */
-	nodemask_t alloc_nmask;
-
-	/* Each bit represents a single occupied (!none/zero) page. */
-	DECLARE_BITMAP(mthp_present_ptes, MAX_PTRS_PER_PTE);
-};
-
 /**
  * struct khugepaged_scan - cursor for scanning
  * @mm_head: the head of the mm list to scan
-- 
2.54.0

[RFC PATCH 04/57] mm/collapse: rename mthp_present_ptes to eligible_ptes

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Neither half of the name holds.  A bit is set only after the PTE has
passed every check the scan makes: uffd, lazyfree, anonymity and sharing
among them.  Presence is the first of those criteria, not the whole of
it.

mthp_collapse() then reads the bitmap starting at the PMD order, so the
bitmap is not specific to mTHP either.

Name the bitmap for what a set bit means: the scan accepted that PTE as
a collapse source.

No functional change.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   |  4 ++--
 mm/khugepaged.c | 13 ++++++-------
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 26dbac7beddd..9c82e71533df 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -53,8 +53,8 @@ struct collapse_control {
 	/* nodemask for allocation fallback */
 	nodemask_t alloc_nmask;
 
-	/* Each bit represents a single occupied (!none/zero) page. */
-	DECLARE_BITMAP(mthp_present_ptes, MAX_PTRS_PER_PTE);
+	/* Each bit marks a PTE the scan accepted as a collapse source */
+	DECLARE_BITMAP(eligible_ptes, MAX_PTRS_PER_PTE);
 };
 
 #endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 804b1d35f52a..a12aafae8d9c 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -576,7 +576,7 @@ static void collapse_control_init_scan(struct collapse_control *cc)
 {
 	memset(cc->node_load, 0, sizeof(cc->node_load));
 	nodes_clear(cc->alloc_nmask);
-	bitmap_zero(cc->mthp_present_ptes, MAX_PTRS_PER_PTE);
+	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
 }
 
 static void release_pte_folio(struct folio *folio)
@@ -1437,8 +1437,8 @@ static unsigned int max_order_from_offset(unsigned int offset)
  * mthp_collapse() consumes the bitmap that is generated during
  * collapse_scan_pmd() to determine what regions and mTHP orders fit best.
  *
- * Each bit in cc->mthp_present_ptes represents a single occupied (!none/zero)
- * page. We start at the PMD order and check if it is eligible for collapse;
+ * Each bit in cc->eligible_ptes marks a PTE the scan accepted as a collapse
+ * source. We start at the PMD order and check if it is eligible for collapse;
  * if not, we check the left and right halves of the PTE page table we are
  * examining at a lower order.
  *
@@ -1469,12 +1469,12 @@ static enum scan_result mthp_collapse(struct mm_struct *mm,
 			goto next_order;
 
 		max_ptes_none = collapse_max_ptes_none(cc, NULL, order);
-		nr_occupied_ptes = bitmap_weight_from(cc->mthp_present_ptes, offset,
+		nr_occupied_ptes = bitmap_weight_from(cc->eligible_ptes, offset,
 						      offset + nr_ptes);
 
 		/*
 		 * Swap PTEs accepted during the scan are counted in @unmapped,
-		 * not in the present-PTE bitmap. Account them for the PMD-order
+		 * not in the eligible bitmap. Account them for the PMD-order
 		 * candidate.
 		 */
 		if (is_pmd_order(order))
@@ -1682,8 +1682,7 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 			}
 		}
 
-		/* Set bit for occupied pages */
-		__set_bit(i, cc->mthp_present_ptes);
+		__set_bit(i, cc->eligible_ptes);
 		/*
 		 * Record which node the original page is from and save this
 		 * information to cc->node_load[].
-- 
2.54.0

[RFC PATCH 05/57] mm/collapse: state what a collapse may do in the policy

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Tests scattered through the collapse path decide what a collapse is
allowed to do by asking whether khugepaged started it.  Between them
they settle:

 - which VMAs are eligible, and how hard to try for a folio;
 - how many empty, swapped-out or shared PTEs a window may contain, and
   whether a sub-PMD window is held to a stricter rule than a PMD;
 - whether a range has to look used, and whether a MADV_FREE'd page is
   left alone;
 - whether the PMD is mapped as part of the request, and whether dirty
   pages are worth writing back and retrying.

None of those is a fact about khugepaged.  Each is something the caller
decided before asking, and the collapse code should not have to look up
who called to find out.

Add struct collapse_policy for the caller to fill: khugepaged from its
own settings, MADV_COLLAPSE from the fact that a user asked explicitly.
Every test becomes a read of a field.

khugepaged fills the policy once per scan pass, MADV_COLLAPSE once per
call.  That is the one change in behaviour: the tunables are sampled
once per pass rather than on every call, so a table scanned early in a
pass and one scanned late are judged alike.

cc->is_khugepaged stays, with a single reader left: the daemon's
collapse counter, which is bookkeeping and not policy.

collapse_file() also drops a NULL check on the collapse_control.  It has
one call site, reached only from collapse_single_pmd(), which
dereferences cc unconditionally, so the check was already dead.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   |  47 ++++++++++++++++++++++
 mm/khugepaged.c | 105 +++++++++++++++++++++++++++---------------------
 2 files changed, 107 insertions(+), 45 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 9c82e71533df..44f52ea5bbb8 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -2,6 +2,7 @@
 #ifndef __MM_COLLAPSE_H
 #define __MM_COLLAPSE_H
 
+#include <linux/mm.h>
 #include <linux/nodemask.h>
 #include <linux/pgtable.h>
 #include <linux/types.h>
@@ -41,7 +42,53 @@ enum scan_result {
 	SCAN_PAGE_DIRTY_OR_WRITEBACK,
 };
 
+/*
+ * What a collapse is allowed to do, decided by whoever asked for it, so the
+ * code doing it need not ask who its caller is: khugepaged fills this in from
+ * its own settings, MADV_COLLAPSE from the fact that a user asked explicitly.
+ */
+struct collapse_policy {
+	/* Limits, stated per PMD; HPAGE_PMD_NR means "no limit" */
+	unsigned int max_ptes_none;
+	unsigned int max_ptes_swap;
+	unsigned int max_ptes_shared;
+
+	/*
+	 * Hold a sub-PMD window to a stricter rule than a PMD: no swapped-out
+	 * and no shared PTEs at all, and max_ptes_none as
+	 * collapse_max_ptes_none() scales it.  khugepaged holds mTHP collapse
+	 * to that; an explicit request does not.
+	 */
+	bool strict_sub_pmd;
+
+	/*
+	 * Collapse only where it looks worth doing: require some sign the range
+	 * is in use, and leave clean lazyfree folios for reclaim rather than
+	 * collapsing them into a folio that is not lazyfree.  A user who asked
+	 * for a collapse gets one either way.
+	 */
+	bool skip_lazyfree;
+	bool require_referenced;
+
+	/*
+	 * Finish the job rather than leaving it half done for a fault to pick
+	 * up: map the PMD over a file collapse before returning, and write
+	 * dirty pages back and retry once instead of refusing them.  Both cost
+	 * latency the caller has asked to pay.
+	 */
+	bool install_pmd;
+	bool writeback_dirty;
+
+	/* How hard to try for a destination folio */
+	gfp_t gfp;
+
+	/* Which VMAs are eligible, as thp_vma_allowable_orders() spells it */
+	enum tva_type tva_type;
+};
+
 struct collapse_control {
+	struct collapse_policy policy;
+
 	bool is_khugepaged;
 
 	/* Num pages scanned per node */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index a12aafae8d9c..eebc044a930e 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -310,15 +310,12 @@ struct attribute_group khugepaged_attr_group = {
 static unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order)
 {
-	const unsigned int max_ptes_none = khugepaged_max_ptes_none;
+	const unsigned int max_ptes_none = cc->policy.max_ptes_none;
 
 	if (vma && userfaultfd_armed(vma))
 		return 0;
-	/* for MADV_COLLAPSE, allow any empty/shared zeropage PTEs */
-	if (!cc->is_khugepaged)
-		return HPAGE_PMD_NR;
-	/* for PMD collapse, respect the user defined maximum */
-	if (is_pmd_order(order))
+	/* The limit as given, at the PMD order and wherever it is not capped */
+	if (is_pmd_order(order) || !cc->policy.strict_sub_pmd)
 		return max_ptes_none;
 	/*
 	 * for mTHP collapse with the sysctl value set to KHUGEPAGED_MAX_PTES_LIMIT,
@@ -350,19 +347,12 @@ static unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
 		unsigned int order)
 {
 	/*
-	 * For MADV_COLLAPSE, do not restrict the number of PTEs that map shared
-	 * anonymous pages.
+	 * A sub-PMD window held to the strict rule takes no shared page at all:
+	 * an mTHP is not worth the CoW-breaking.
 	 */
-	if (!cc->is_khugepaged)
-		return HPAGE_PMD_NR;
-	/*
-	 * for mTHP collapse do not allow collapsing anonymous memory pages that
-	 * are shared between processes.
-	 */
-	if (!is_pmd_order(order))
+	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
 		return 0;
-	/* for PMD collapse, respect the user defined maximum */
-	return khugepaged_max_ptes_shared;
+	return cc->policy.max_ptes_shared;
 }
 
 /**
@@ -378,16 +368,12 @@ static unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
 		unsigned int order)
 {
 	/*
-	 * For MADV_COLLAPSE, do not restrict the number PTEs entries or
-	 * pagecache entries that are non-present.
+	 * A sub-PMD window held to the strict rule takes nothing non-present:
+	 * reading pages back to build an mTHP is not worth the latency.
 	 */
-	if (!cc->is_khugepaged)
-		return HPAGE_PMD_NR;
-	/* for mTHP collapse do not allow any non-present PTEs or pagecache entries */
-	if (!is_pmd_order(order))
+	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
 		return 0;
-	/* for PMD collapse, respect the user defined maximum */
-	return khugepaged_max_ptes_swap;
+	return cc->policy.max_ptes_swap;
 }
 
 int hugepage_madvise(struct vm_area_struct *vma,
@@ -686,7 +672,7 @@ static enum scan_result __collapse_huge_page_isolate(struct vm_area_struct *vma,
 		 * If the vma has the VM_DROPPABLE flag, the collapse will
 		 * preserve the lazyfree property without needing to skip.
 		 */
-		if (cc->is_khugepaged && !(vma->vm_flags & VM_DROPPABLE) &&
+		if (cc->policy.skip_lazyfree && !(vma->vm_flags & VM_DROPPABLE) &&
 		    folio_test_lazyfree(folio) && !pte_dirty(pteval)) {
 			result = SCAN_PAGE_LAZYFREE;
 			goto out;
@@ -775,12 +761,12 @@ static enum scan_result __collapse_huge_page_isolate(struct vm_area_struct *vma,
 		if (folio_test_large(folio))
 			list_add_tail(&folio->lru, compound_pagelist);
 next:
-		if (cc->is_khugepaged &&
+		if (cc->policy.require_referenced &&
 		    folio_pte_referenced(folio, vma, addr, pteval))
 			referenced++;
 	}
 
-	if (unlikely(cc->is_khugepaged && !referenced)) {
+	if (unlikely(cc->policy.require_referenced && !referenced)) {
 		result = SCAN_LACK_REFERENCED_PAGE;
 	} else {
 		result = SCAN_SUCCEED;
@@ -984,6 +970,36 @@ static inline gfp_t alloc_hugepage_khugepaged_gfpmask(void)
 	return khugepaged_defrag() ? GFP_TRANSHUGE : GFP_TRANSHUGE_LIGHT;
 }
 
+/* khugepaged collapses on its own initiative, so it obeys its own settings. */
+static void collapse_policy_khugepaged(struct collapse_policy *p)
+{
+	p->max_ptes_none = READ_ONCE(khugepaged_max_ptes_none);
+	p->max_ptes_swap = READ_ONCE(khugepaged_max_ptes_swap);
+	p->max_ptes_shared = READ_ONCE(khugepaged_max_ptes_shared);
+	p->strict_sub_pmd = true;
+	p->skip_lazyfree = true;
+	p->require_referenced = true;
+	p->install_pmd = false;
+	p->writeback_dirty = false;
+	p->gfp = alloc_hugepage_khugepaged_gfpmask();
+	p->tva_type = TVA_KHUGEPAGED;
+}
+
+/* MADV_COLLAPSE was asked for explicitly, so it is not held to those. */
+static void collapse_policy_forced(struct collapse_policy *p)
+{
+	p->max_ptes_none = HPAGE_PMD_NR;
+	p->max_ptes_swap = HPAGE_PMD_NR;
+	p->max_ptes_shared = HPAGE_PMD_NR;
+	p->strict_sub_pmd = false;
+	p->skip_lazyfree = false;
+	p->require_referenced = false;
+	p->install_pmd = true;
+	p->writeback_dirty = true;
+	p->gfp = GFP_TRANSHUGE;
+	p->tva_type = TVA_FORCED_COLLAPSE;
+}
+
 #ifdef CONFIG_NUMA
 static int collapse_find_target_node(struct collapse_control *cc)
 {
@@ -1021,8 +1037,7 @@ static enum scan_result hugepage_vma_revalidate(struct mm_struct *mm, unsigned l
 		struct collapse_control *cc, unsigned int order)
 {
 	struct vm_area_struct *vma;
-	enum tva_type type = cc->is_khugepaged ? TVA_KHUGEPAGED :
-				 TVA_FORCED_COLLAPSE;
+	enum tva_type type = cc->policy.tva_type;
 
 	if (unlikely(collapse_test_exit_or_disable(mm)))
 		return SCAN_ANY_PROCESS;
@@ -1205,8 +1220,7 @@ static enum scan_result __collapse_huge_page_swapin(struct mm_struct *mm,
 static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_struct *mm,
 		struct collapse_control *cc, unsigned int order)
 {
-	gfp_t gfp = (cc->is_khugepaged ? alloc_hugepage_khugepaged_gfpmask() :
-		     GFP_TRANSHUGE);
+	gfp_t gfp = cc->policy.gfp;
 	int node = collapse_find_target_node(cc);
 	struct folio *folio;
 
@@ -1559,7 +1573,7 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 	const unsigned int max_ptes_shared = collapse_max_ptes_shared(cc, HPAGE_PMD_ORDER);
 	const unsigned int max_ptes_swap = collapse_max_ptes_swap(cc, HPAGE_PMD_ORDER);
 	unsigned int max_ptes_none = collapse_max_ptes_none(cc, vma, HPAGE_PMD_ORDER);
-	enum tva_type tva_flags = cc->is_khugepaged ? TVA_KHUGEPAGED : TVA_FORCED_COLLAPSE;
+	enum tva_type tva_flags = cc->policy.tva_type;
 	pmd_t *pmd;
 	pte_t *pte, *_pte, pteval;
 	int i;
@@ -1658,7 +1672,7 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 		 * If the vma has the VM_DROPPABLE flag, the collapse will
 		 * preserve the lazyfree property without needing to skip.
 		 */
-		if (cc->is_khugepaged && !(vma->vm_flags & VM_DROPPABLE) &&
+		if (cc->policy.skip_lazyfree && !(vma->vm_flags & VM_DROPPABLE) &&
 		    folio_test_lazyfree(folio) && !pte_dirty(pteval)) {
 			result = SCAN_PAGE_LAZYFREE;
 			goto out_unmap;
@@ -1716,11 +1730,11 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 			goto out_unmap;
 		}
 
-		if (cc->is_khugepaged &&
+		if (cc->policy.require_referenced &&
 		    folio_pte_referenced(folio, vma, addr, pteval))
 			referenced++;
 	}
-	if (cc->is_khugepaged &&
+	if (cc->policy.require_referenced &&
 		   (!referenced ||
 		    (unmapped && referenced < HPAGE_PMD_NR / 2))) {
 		result = SCAN_LACK_REFERENCED_PAGE;
@@ -2572,11 +2586,11 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 	xas_unlock_irq(&xas);
 
 	/*
-	 * Remove pte page tables, so we can re-fault the page as huge.
-	 * If MADV_COLLAPSE, adjust result to call try_collapse_pte_mapped_thp().
+	 * Remove pte page tables, so we can re-fault the page as huge.  A caller
+	 * that wants the PMD mapped now is told to go and do that.
 	 */
 	retract_page_tables(mapping, start);
-	if (cc && !cc->is_khugepaged)
+	if (cc->policy.install_pmd)
 		result = SCAN_PTE_MAPPED_HUGEPAGE;
 	folio_unlock(new_folio);
 
@@ -2760,11 +2774,8 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 retry:
 	result = collapse_scan_file(mm, addr, file, pgoff, cc);
 
-	/*
-	 * For MADV_COLLAPSE, when encountering dirty pages, try to writeback,
-	 * then retry the collapse one time.
-	 */
-	if (!cc->is_khugepaged && result == SCAN_PAGE_DIRTY_OR_WRITEBACK &&
+	/* Dirty pages are worth a writeback and one more try, if asked for */
+	if (cc->policy.writeback_dirty && result == SCAN_PAGE_DIRTY_OR_WRITEBACK &&
 	    !triggered_wb && mapping_can_writeback(file->f_mapping)) {
 		const loff_t lstart = (loff_t)pgoff << PAGE_SHIFT;
 		const loff_t lend = lstart + HPAGE_PMD_SIZE - 1;
@@ -2781,7 +2792,7 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 			result = SCAN_ANY_PROCESS;
 		else
 			result = try_collapse_pte_mapped_thp(mm, addr,
-							     !cc->is_khugepaged);
+							cc->policy.install_pmd);
 		if (result == SCAN_PMD_MAPPED)
 			result = SCAN_SUCCEED;
 		mmap_read_unlock(mm);
@@ -2931,6 +2942,9 @@ static void khugepaged_do_scan(struct collapse_control *cc)
 
 	lru_add_drain_all();
 
+	/* One policy for the whole pass, so every table is judged the same */
+	collapse_policy_khugepaged(&cc->policy);
+
 	cc->progress = 0;
 	while (true) {
 		cond_resched();
@@ -3159,6 +3173,7 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	if (!cc)
 		return -ENOMEM;
 	cc->is_khugepaged = false;
+	collapse_policy_forced(&cc->policy);
 	cc->progress = 0;
 
 	mmgrab(mm);
-- 
2.54.0

[RFC PATCH 06/57] mm/collapse: move the smallest collapse order to collapse.h

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The floor on the order a collapse will build is a property of collapse,
not of khugepaged.  MADV_COLLAPSE reaches the same code and is held to
the same floor.

Move KHUGEPAGED_MIN_MTHP_ORDER to the shared header as
COLLAPSE_MIN_MTHP_ORDER.  Preparation for the collapse engine, which
sizes its per-window arrays from the same floor.

No functional change intended.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   | 3 +++
 mm/khugepaged.c | 6 ++----
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 44f52ea5bbb8..1e969292edcb 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -7,6 +7,9 @@
 #include <linux/pgtable.h>
 #include <linux/types.h>
 
+/* The smallest order a collapse will build, and so the finest window it cuts */
+#define COLLAPSE_MIN_MTHP_ORDER		2
+
 enum scan_result {
 	SCAN_FAIL,
 	SCAN_SUCCEED,
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index eebc044a930e..f31689bf75a6 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -67,8 +67,6 @@ static DEFINE_READ_MOSTLY_HASHTABLE(mm_slots_hash, MM_SLOTS_HASH_BITS);
 
 static struct kmem_cache *mm_slot_cache __ro_after_init;
 
-#define KHUGEPAGED_MIN_MTHP_ORDER	2
-
 /**
  * struct khugepaged_scan - cursor for scanning
  * @mm_head: the head of the mm list to scan
@@ -1540,8 +1538,8 @@ static enum scan_result mthp_collapse(struct mm_struct *mm,
 		 * any smaller order enabled. When at the smallest order
 		 * we must always move to the next offset.
 		 */
-		if (order > KHUGEPAGED_MIN_MTHP_ORDER &&
-			(enabled_orders & GENMASK(order - 1, 0))) {
+		if (order > COLLAPSE_MIN_MTHP_ORDER &&
+		    (enabled_orders & GENMASK(order - 1, 0))) {
 			order--;
 			continue;
 		}
-- 
2.54.0

[RFC PATCH 07/57] mm/collapse: sketch the new anonymous collapse engine

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Anonymous collapse is held back in two ways that its current shape cannot
be patched out of.

Functionally, an mTHP is collapsed only when the whole PMD-sized window
qualifies: collapse_scan_pmd() reaches mthp_collapse() only on
SCAN_SUCCEED.  One PTE that disqualifies itself -- uffd-armed, not
anonymous, clean lazyfree, off the LRU, pinned -- takes the whole table
with it.  A 2M range with a single such page yields nothing, even where
the half beside it would collapse perfectly well.

For scalability, collapse_huge_page() holds mmap_write across a collapse,
and anon_vma_lock_write() with it, then does the same again for the next
huge page.  Every collapse stops every fault in the address space, one
huge page at a time.

The new engine addresses both.  It quiesces its sources the way migration
does: migration entries in their PTEs, then a frozen refcount.  That is
enough to make the copy safe without the exclusive lock, so the engine
runs under mmap_read throughout.

It also carries a batch of windows through each step together, rather
than one window through all of them.  Its verdict is per window rather
than per table, so a table that cannot become one huge page still yields
the largest windows inside it.

Lay the design out first and fill it in afterwards.  What arrives here is
the shape of the engine: the two entry points, and a comment at the top
of mm/collapse.c mapping the whole call tree.  Every step below them is a
stub, filled in before the anonymous path is pointed at the engine.

collapse_scan_anon_pmd() finds the table, asks the VMA which orders it
allows, scans it, and leaves in the collapse_control what a collapse
could use.  collapse_anon_pmd() takes that range and cuts it into
windows.

The line between the two is where the lock goes.  A scan only reads a VMA
and a page table, so it keeps the mmap_lock it was called under and hands
it back.  A collapse allocates, copies and flushes, so it is called
without the lock and takes it again for each round of its own.  Nothing
then has to tell a caller whether the lock survived a call, and a caller
that finds nothing to collapse never gives the lock up at all.

Nothing calls either entry point yet; the anonymous path keeps using the
mechanism the engine replaces.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/Makefile     |   2 +-
 mm/collapse.c   | 165 ++++++++++++++++++++++++++++++++++++++++++++++++
 mm/collapse.h   |  24 +++++++
 mm/khugepaged.c |   4 +-
 4 files changed, 192 insertions(+), 3 deletions(-)
 create mode 100644 mm/collapse.c

diff --git a/mm/Makefile b/mm/Makefile
index e7245cb88c66..2bef749a5c21 100644
--- a/mm/Makefile
+++ b/mm/Makefile
@@ -98,7 +98,7 @@ obj-$(CONFIG_MEMTEST)		+= memtest.o
 obj-$(CONFIG_MIGRATION) += migrate.o
 obj-$(CONFIG_NUMA) += memory-tiers.o
 obj-$(CONFIG_DEVICE_MIGRATION) += migrate_device.o
-obj-$(CONFIG_TRANSPARENT_HUGEPAGE) += huge_memory.o khugepaged.o
+obj-$(CONFIG_TRANSPARENT_HUGEPAGE) += collapse.o huge_memory.o khugepaged.o
 obj-$(CONFIG_PAGE_COUNTER) += page_counter.o
 obj-$(CONFIG_LIVEUPDATE_MEMFD) += memfd_luo.o
 obj-$(CONFIG_MEMCG_V1) += memcontrol-v1.o
diff --git a/mm/collapse.c b/mm/collapse.c
new file mode 100644
index 000000000000..0e6c3c68b44c
--- /dev/null
+++ b/mm/collapse.c
@@ -0,0 +1,165 @@
+// SPDX-License-Identifier: GPL-2.0
+#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
+
+#include <linux/bitops.h>
+#include <linux/highmem.h>
+#include <linux/huge_mm.h>
+#include <linux/hugetlb.h>	/* x86 flush_tlb_range() uses hstate_vma() */
+#include <linux/leafops.h>
+#include <linux/memcontrol.h>
+#include <linux/mm.h>
+#include <linux/mmu_notifier.h>
+#include <linux/pagemap.h>
+#include <linux/pgalloc.h>
+#include <linux/rmap.h>
+#include <linux/sched.h>
+#include <linux/sizes.h>
+#include <linux/slab.h>
+#include <linux/swap.h>
+#include <linux/userfaultfd_k.h>
+
+#include <asm/tlb.h>
+#include "collapse.h"
+#include "internal.h"
+
+/*
+ * Anonymous collapse, in rounds.
+ *
+ * The folios mapped across a window of PTEs become one folio of that window's
+ * order, with the sources quiesced by the two barriers migration uses --
+ * migration entries in their PTEs, then a frozen refcount -- so the copy itself
+ * needs no lock.  The engine runs under mmap_read throughout.
+ *
+ * A round carries a batch of candidate windows through the passes together,
+ * rather than carrying one window through the whole collapse.  [ptl] and
+ * [pmd lock] mark a pass that takes that lock and drops it again, so no
+ * page-table lock is ever held across passes; the source folio locks are the
+ * exception, held from freeze to putback.  [rcu] marks a pass that takes no
+ * page-table lock at all and reads the table racily, which only the scan does.
+ *
+ * Allocation happens on both sides of the freeze, and which side comes first
+ * matters.  collapse_provision() tries first, inside the window and after the
+ * freeze, with reclaim masked out of the gfp: the sources are frozen by then,
+ * so a faulter on one of them waits for this allocation.  A candidate the
+ * allocator has nothing ready for is not failed -- it goes back to selection,
+ * and collapse_reserve() allocates for it before the next round freezes
+ * anything, outside the window, where reclaim costs khugepaged its own
+ * progress and nobody else's wait.
+ *
+ * That second chance needs the policy's gfp to allow reclaim at all.  When it
+ * does not, a retry would miss the same way, so the first miss is the answer.
+ *
+ * collapse_scan_anon_pmd()            judge one PTE table's worth of a VMA
+ * `- collapse_scan_table()            [rcu] a bit per PTE a collapse can use
+ *
+ * collapse_anon_pmd()                 cut windows from those bits, run them
+ * |- collapse_next_candidate()        the next window worth attempting
+ * `- collapse_run_batch()             run the batch, then classify it
+ *    |- collapse_round()              below
+ *    `- collapse_classify_result()    carry on / lower / abandon
+ *       `- collapse_push_retry()      queue it for a lower order
+ *
+ * collapse_round()                    one batch of candidates
+ * |- collapse_reserve()               second try for what the last round
+ * |                                   missed, with reclaim; sleeps
+ * |- collapse_deposit()               a page table per PMD-order candidate
+ * |- collapse_revalidate()            check the VMA and the table survived
+ * |- collapse_faultin()               make the sources present and exclusive;
+ * |                                   sleeps, and may leave the lock dropped
+ * |- collapse_freeze()                raise the barriers [ptl], flush the TLB
+ * |- collapse_provision()             first try for every other destination,
+ * |                                   without reclaim: the sources are frozen
+ * |- collapse_copy()                  copy into the destinations; sleeps
+ * |- collapse_install()               publish them [ptl], or [pmd lock] and a
+ * |                                   second TLB flush at the PMD order
+ * |- collapse_putback()               lower the barriers, in order
+ * `- collapse_finish()                settle whatever the round reached
+ *
+ * Every slot of a candidate is a real source, a hole (pte_none, zero-filled
+ * and re-verified still-none at install), or the zeropage (cleared at freeze,
+ * zero-filled).  Sources come in "spans" -- consecutive PTEs mapping
+ * consecutive pages of one folio -- so partially mapped and compound sources
+ * collapse too: any order below the window's is a source, and a PMD candidate
+ * takes even a PTE-mapped THP of its own order.
+ *
+ * Nothing calls any of this yet: the anon path still uses the mechanism it
+ * replaces, and is switched over once both halves are complete.
+ */
+
+/*
+ * Scan the PTEs between @start and @end and record what a collapse could use: a
+ * bit in cc->eligible_ptes for every PTE that may be a source.  Returns
+ * SCAN_SUCCEED when every PTE in the range qualified, otherwise the reason one
+ * did not, and narrows cc->select_orders to what is still worth trying here.
+ */
+static enum scan_result collapse_scan_table(struct vm_area_struct *vma,
+					    pmd_t *pmd, unsigned long start,
+					    unsigned long end,
+					    struct collapse_control *cc)
+{
+	return SCAN_SUCCEED;
+}
+
+/* Everything a table is judged on starts empty for each table */
+static void collapse_anon_scan_init(struct collapse_control *cc)
+{
+	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
+	memset(cc->node_load, 0, sizeof(cc->node_load));
+	nodes_clear(cc->alloc_nmask);
+
+	cc->select_orders = 0;
+	cc->nr_collapsed = 0;
+}
+
+/*
+ * Judge one table's worth of @vma, leaving in @cc what a collapse could use:
+ * which orders are still worth attempting, and why the table was turned down if
+ * some order was.  Holds mmap_lock throughout -- it only reads -- and a caller
+ * that acts on what it found hands the range to collapse_anon_pmd() afterwards,
+ * without the lock.
+ */
+static enum scan_result __maybe_unused
+collapse_scan_anon_pmd(struct vm_area_struct *vma, unsigned long start,
+		       unsigned long end, struct collapse_control *cc)
+{
+	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
+	struct mm_struct *mm = vma->vm_mm;
+	pmd_t *pmd;
+
+	/* One table's worth at most, not empty, and inside the VMA */
+	VM_WARN_ON_ONCE(end > pmd_addr + HPAGE_PMD_SIZE || start >= end);
+	VM_WARN_ON_ONCE(start < vma->vm_start || end > vma->vm_end);
+
+	cc->scan_refusal = find_pmd_or_thp_or_none(mm, pmd_addr, &pmd);
+	if (cc->scan_refusal != SCAN_SUCCEED) {
+		cc->progress++;
+		return cc->scan_refusal;
+	}
+
+	/* Cleared only once a table has turned out to be there */
+	collapse_anon_scan_init(cc);
+
+	cc->select_orders = collapse_possible_orders(vma, vma->vm_flags,
+						     cc->policy.tva_type);
+	if (!cc->select_orders) {
+		cc->scan_refusal = SCAN_VMA_CHECK;
+		return cc->scan_refusal;
+	}
+
+	/* The scan narrows select_orders to whatever is left worth trying */
+	cc->scan_refusal = collapse_scan_table(vma, pmd, start, end, cc);
+
+	return cc->scan_refusal;
+}
+
+/*
+ * Cut the table into candidate windows and collapse what fits, from the
+ * largest order downwards.  Returns what the table yielded: a collapse, or
+ * the reason it did not.
+ */
+static enum scan_result __maybe_unused
+collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
+		  struct collapse_control *cc)
+{
+	return SCAN_FAIL;
+}
diff --git a/mm/collapse.h b/mm/collapse.h
index 1e969292edcb..e2af4c47cb60 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -105,6 +105,30 @@ struct collapse_control {
 
 	/* Each bit marks a PTE the scan accepted as a collapse source */
 	DECLARE_BITMAP(eligible_ptes, MAX_PTRS_PER_PTE);
+
+	/* Orders still worth attempting in the table being scanned */
+	unsigned long select_orders;
+
+	/* PTEs collapsed in it so far */
+	unsigned int nr_collapsed;
+
+	/*
+	 * Why the scan would not take all of the table, or SCAN_SUCCEED if it
+	 * took every order it was offered.  Not the opposite of what the scan
+	 * selected: a table can be worth collapsing at one order and refused at
+	 * another, so a scan that found work still has a reason to report, and
+	 * the collapse reports it when it salvages nothing.
+	 */
+	enum scan_result scan_refusal;
 };
 
+/*
+ * Defined in khugepaged.c, which still uses them itself.
+ * TODO: move each into collapse.c once its last khugepaged.c user is gone.
+ */
+unsigned long collapse_possible_orders(struct vm_area_struct *vma,
+		vm_flags_t vm_flags, enum tva_type tva_flags);
+enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
+		unsigned long address, pmd_t **pmd);
+
 #endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index f31689bf75a6..26d25093260b 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -498,7 +498,7 @@ void __khugepaged_enter(struct mm_struct *mm)
  * Check what orders are possible based on the vma and collapse type.
  * This is used to determine if mTHP collapse is a viable option.
  */
-static unsigned long collapse_possible_orders(struct vm_area_struct *vma,
+unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags)
 {
 	unsigned long orders;
@@ -1090,7 +1090,7 @@ static inline enum scan_result check_pmd_state(pmd_t *pmd)
 	return SCAN_SUCCEED;
 }
 
-static enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
+enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd)
 {
 	*pmd = mm_find_pmd(mm, address);
-- 
2.54.0

[RFC PATCH 08/57] mm/collapse: scan a table for what a collapse could use

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the scan.  Walk the range and set a bit in cc->eligible_ptes for
every PTE a collapse may take as a source: present, anonymous, not
uffd-armed, on the LRU and unlocked.  The bit is set last, so a PTE that
failed anything leaves it clear.

The walk takes no page table lock.  What it produces is advice: the
freeze settles every question the scan asks, by re-reading the table
under the lock and freezing each source to the count it expects.  A racy
read can only cost a candidate the freeze then refuses, or miss one the
next pass finds.  What it buys is that a fault in the range does not wait
for a walk of the whole table.

pte_offset_map() holds rcu_read_lock() until pte_unmap(), which keeps the
table from being freed underneath the walk.  mmap_lock keeps the VMA
attached, without which free_pgtables() could free it without waiting for
RCU at all.

The verdict is two-sided, which is the point:

 - A PTE that disqualifies only itself leaves the bitmap clear there and
   drops the PMD order, since a PMD candidate needs the whole table.
   Selection still gets the smaller windows that avoid it.
 - What refuses the table as a unit -- a limit the whole range exceeds,
   or sources spread across nodes too distant for one folio to serve --
   leaves no order eligible at all.

Limits on swapped-out and shared PTEs are stated per PMD and scaled to
what was actually scanned, so a partial table is held to the same density
as a whole one.

A folio whose reference count its mappings do not account for -- a GUP
pin, say -- is left to the freeze rather than refused here.
folio_expected_ref_count() wants a folio that cannot change order while
it is read.  This walk holds no page table lock and no folio lock, so a
folio splitting underneath it would have its count read for the wrong
size.  A reference of its own would not help: that stops a folio being
freed, not split.

Whether a range has to look used at all is the caller's policy, so only a
caller that asks gathers the young/referenced evidence.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 239 +++++++++++++++++++++++++++++++++++++++++++++++-
 mm/collapse.h   |   7 ++
 mm/khugepaged.c |   8 +-
 3 files changed, 249 insertions(+), 5 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 0e6c3c68b44c..66931ef6a6d0 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -86,6 +86,20 @@
  * replaces, and is switched over once both halves are complete.
  */
 
+/*
+ * Is @count past a limit stated per PMD, when only part of a table was scanned?
+ * Scale the comparison to the table so a partial scan is held to the same
+ * density as a whole one.
+ */
+static bool collapse_exceeds_limit(unsigned int count, unsigned int max_per_pmd,
+				   unsigned long start, unsigned long end)
+{
+	const unsigned long nr_scanned = (end - start) >> PAGE_SHIFT;
+
+	return (unsigned long)count * HPAGE_PMD_NR >
+	       (unsigned long)max_per_pmd * nr_scanned;
+}
+
 /*
  * Scan the PTEs between @start and @end and record what a collapse could use: a
  * bit in cc->eligible_ptes for every PTE that may be a source.  Returns
@@ -97,7 +111,230 @@ static enum scan_result collapse_scan_table(struct vm_area_struct *vma,
 					    unsigned long end,
 					    struct collapse_control *cc)
 {
-	return SCAN_SUCCEED;
+	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
+	unsigned int max_ptes_none, max_ptes_swap, max_ptes_shared;
+	int none_or_zero = 0, shared = 0, referenced = 0, unmapped = 0;
+	enum scan_result result, pmd_result = SCAN_SUCCEED;
+	unsigned int first_offset;
+	unsigned long addr;
+	pte_t *pte;
+	int i;
+
+	max_ptes_none = collapse_max_ptes_none(cc, vma, HPAGE_PMD_ORDER);
+	max_ptes_swap = collapse_max_ptes_swap(cc, HPAGE_PMD_ORDER);
+	max_ptes_shared = collapse_max_ptes_shared(cc, HPAGE_PMD_ORDER);
+
+	/*
+	 * No page table lock: what this builds is advice, and the freeze settles
+	 * every question it asks by re-reading the table under the lock and
+	 * freezing each source to the count it expects.  A racy read can only
+	 * cost a candidate that the freeze then refuses, or miss one that the
+	 * next pass finds.  What it buys is that a fault in this range does not
+	 * wait for a scan of the whole table.
+	 *
+	 * pte_offset_map() holds rcu_read_lock() until pte_unmap(), which is
+	 * what keeps the table itself from being freed underneath the walk;
+	 * mmap_lock keeps the VMA attached, without which free_pgtables() could
+	 * free it without waiting for RCU at all.  Nothing below here sleeps.
+	 */
+	pte = pte_offset_map(pmd, start);
+	if (!pte) {
+		cc->progress++;
+		result = SCAN_NO_PTE_TABLE;
+		goto out_no_table;
+	}
+
+	/*
+	 * The bitmap and the selection offsets stay relative to the table:
+	 * natural-alignment math needs the table-absolute position, not the
+	 * position within an arbitrarily placed VMA.
+	 */
+	first_offset = (start - pmd_addr) >> PAGE_SHIFT;
+	for (i = first_offset, addr = start; addr < end;
+	     i++, addr += PAGE_SIZE) {
+		pte_t pteval = ptep_get(pte + (i - first_offset));
+		struct folio *folio;
+		struct page *page;
+		int node;
+
+		cc->progress++;
+
+		if (pte_none_or_zero(pteval)) {
+			if (++none_or_zero > max_ptes_none &&
+			    pmd_result == SCAN_SUCCEED) {
+				pmd_result = SCAN_EXCEED_NONE_PTE;
+				count_vm_event(THP_SCAN_EXCEED_NONE_PTE);
+				count_mthp_stat(HPAGE_PMD_ORDER,
+						MTHP_STAT_COLLAPSE_EXCEED_NONE);
+			}
+			continue;
+		}
+		if (!pte_present(pteval)) {
+			unmapped++;
+			if (collapse_exceeds_limit(unmapped, max_ptes_swap,
+						   start, end)) {
+				result = SCAN_EXCEED_SWAP_PTE;
+				count_vm_event(THP_SCAN_EXCEED_SWAP_PTE);
+				count_mthp_stat(HPAGE_PMD_ORDER,
+						MTHP_STAT_COLLAPSE_EXCEED_SWAP);
+				goto out_table_refused;
+			}
+			/* Swap entries armed with uffd-wp are refused too */
+			if (pte_swp_uffd_any(pteval) &&
+			    pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PTE_UFFD;
+			continue;
+		}
+		if (pte_uffd(pteval)) {
+			/*
+			 * The huge PMD could be marked write protected when any
+			 * of the small ones is, but that could deliver
+			 * userfaults outside the registered range.  Keep it
+			 * simple and refuse the PTE.
+			 */
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PTE_UFFD;
+			continue;
+		}
+
+		page = vm_normal_page(vma, addr, pteval);
+		if (unlikely(!page) || unlikely(is_zone_device_page(page))) {
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PAGE_NULL;
+			continue;
+		}
+		folio = page_folio(page);
+
+		/*
+		 * A VM_DROPPABLE VMA keeps the lazyfree property across the
+		 * collapse, so there is nothing to preserve by skipping.
+		 */
+		if (cc->policy.skip_lazyfree &&
+		    !(vma->vm_flags & VM_DROPPABLE) &&
+		    folio_test_lazyfree(folio) && !pte_dirty(pteval)) {
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PAGE_LAZYFREE;
+			continue;
+		}
+
+		if (!folio_test_anon(folio)) {
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PAGE_ANON;
+			continue;
+		}
+
+		/*
+		 * A page counts as shared if any part of its folio is, which
+		 * bounds the cost of CoW-breaking rather than the count of it:
+		 * collapse_faultin() unshares on !PageAnonExclusive(), a broader
+		 * test -- a page whose fork co-mapper has exited is
+		 * single-mapped, so not counted here, yet stays non-exclusive
+		 * until a write reuses it.  Those are the cheap ones, reused in
+		 * place.  A page that has to be copied is one this test catches,
+		 * so the limit does bound the copying it is there to bound.
+		 */
+		if (folio_maybe_mapped_shared(folio)) {
+			shared++;
+			if (collapse_exceeds_limit(shared, max_ptes_shared,
+						   start, end)) {
+				result = SCAN_EXCEED_SHARED_PTE;
+				count_vm_event(THP_SCAN_EXCEED_SHARED_PTE);
+				count_mthp_stat(HPAGE_PMD_ORDER,
+						MTHP_STAT_COLLAPSE_EXCEED_SHARED);
+				goto out_table_refused;
+			}
+		}
+
+		/*
+		 * Which node the sources are on decides where the destination is
+		 * allocated: the one with the most of them wins.
+		 */
+		node = folio_nid(folio);
+		if (collapse_scan_abort(node, cc)) {
+			result = SCAN_SCAN_ABORT;
+			goto out_table_refused;
+		}
+		cc->node_load[node]++;
+
+		/*
+		 * Usually a folio somebody else is already isolating, whose
+		 * reference the freeze would refuse anyway.  Not exact: one
+		 * still on a per-CPU add batch reads the same, and the freeze
+		 * drains those before it starts.
+		 */
+		if (!folio_test_lru(folio)) {
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PAGE_LRU;
+			continue;
+		}
+		if (folio_test_locked(folio)) {
+			if (pmd_result == SCAN_SUCCEED)
+				pmd_result = SCAN_PAGE_LOCK;
+			continue;
+		}
+
+		/*
+		 * A folio whose reference count its mappings do not account for
+		 * -- a GUP pin, say -- is refused by the freeze, not here.
+		 * folio_expected_ref_count() wants a folio that cannot change
+		 * order while it is read, and this walk holds no page table lock
+		 * and no folio lock, so a folio splitting underneath it would
+		 * have the count read for the wrong size.  A reference of our
+		 * own would not help: it stops the folio being freed, not split.
+		 *
+		 * So leave it to the freeze, which reads the table under the
+		 * lock and settles the question by freezing each source to the
+		 * count it expects.  What it costs is a window selected here and
+		 * refused there.
+		 */
+
+		/*
+		 * Every check passed: this PTE can be a collapse source.  The
+		 * bit is set last, so a disqualified PTE leaves it clear.
+		 */
+		__set_bit(i, cc->eligible_ptes);
+
+		/*
+		 * Whether a range has to look used at all is the caller's
+		 * policy, so only a caller that asks gathers the evidence.
+		 */
+		if (cc->policy.require_referenced &&
+		    (pte_young(pteval) || folio_test_young(folio) ||
+		     folio_test_referenced(folio) ||
+		     mmu_notifier_test_young(vma->vm_mm, addr)))
+			referenced++;
+	}
+
+	if (cc->policy.require_referenced &&
+	    (!referenced || (unmapped && referenced < HPAGE_PMD_NR / 2)))
+		result = SCAN_LACK_REFERENCED_PAGE;
+	else
+		result = pmd_result;
+	pte_unmap(pte);
+	goto out;
+
+out_table_refused:
+	/*
+	 * The table is refused as a unit -- a limit the whole range exceeds, or
+	 * pages on nodes too distant for one folio to serve them all -- so no
+	 * window inside it is eligible either.
+	 */
+	pte_unmap(pte);
+out_no_table:
+	cc->select_orders = 0;
+out:
+	/*
+	 * A PMD candidate needs the whole table, so anything that disqualified a
+	 * single PTE rules it out.  Smaller windows that avoid the offending
+	 * PTEs are still collapsible, so drop just that order and leave the rest
+	 * to selection -- dropping it also lowers the order selection roots its
+	 * windows at.  MADV_COLLAPSE has no other order enabled, so it is left
+	 * with none.
+	 */
+	if (result != SCAN_SUCCEED)
+		cc->select_orders &= ~BIT(HPAGE_PMD_ORDER);
+
+	return result;
 }
 
 /* Everything a table is judged on starts empty for each table */
diff --git a/mm/collapse.h b/mm/collapse.h
index e2af4c47cb60..ad88b91d9a72 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -130,5 +130,12 @@ unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags);
 enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd);
+bool collapse_scan_abort(int nid, struct collapse_control *cc);
+unsigned int collapse_max_ptes_none(struct collapse_control *cc,
+		struct vm_area_struct *vma, unsigned int order);
+unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
+		unsigned int order);
+unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
+		unsigned int order);
 
 #endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 26d25093260b..9823884a83c9 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -305,7 +305,7 @@ struct attribute_group khugepaged_attr_group = {
  *
  * Return: Maximum number of empty/shared zeropage PTEs for the collapse operation
  */
-static unsigned int collapse_max_ptes_none(struct collapse_control *cc,
+unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order)
 {
 	const unsigned int max_ptes_none = cc->policy.max_ptes_none;
@@ -341,7 +341,7 @@ static unsigned int collapse_max_ptes_none(struct collapse_control *cc,
  * Return: Maximum number of PTEs that map shared anonymous pages for the
  * collapse operation
  */
-static unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
+unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
 		unsigned int order)
 {
 	/*
@@ -362,7 +362,7 @@ static unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
  * Return: Maximum number of non-present PTEs or the maximum allowed non-present
  * pagecache entries for the collapse operation.
  */
-static unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
+unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
 		unsigned int order)
 {
 	/*
@@ -934,7 +934,7 @@ static struct collapse_control khugepaged_collapse_control = {
 	.is_khugepaged = true,
 };
 
-static bool collapse_scan_abort(int nid, struct collapse_control *cc)
+bool collapse_scan_abort(int nid, struct collapse_control *cc)
 {
 	int i;
 
-- 
2.54.0

[RFC PATCH 09/57] mm/collapse: collect candidate windows into a round

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

collapse_anon_pmd() is the half of a table's collapse that follows the
scan: cut windows out of the PTEs the scan accepted, and run them.  It
runs them a round at a time, so a round needs somewhere to be collected.

Add that array to collapse_control.  Its size is the number of windows
one table holds at the smallest order a collapse builds, or as many as
the byte cap allows, whichever is fewer.

A round is capped because it holds destination folios that are allocated
but not yet installed, and because a faulter on any source inside it
waits for the round to finish.  A dense table is collapsed as several
rounds rather than one.

The array is too large for the stack.  khugepaged takes it when the
daemon starts, so an allocation failure is reported to the sysfs write
that enabled khugepaged rather than surfacing inside the daemon;
MADV_COLLAPSE takes one per call.

collapse_anon_pmd() then gets its shape.  Take candidates from selection
until the round is full or selection is done, run the round, and stop
once selection has nothing left and the round is empty.  A candidate the
full round could not take stays pending for the next one, so nothing is
dropped at the boundary.

Selection and the batch run are stubs here, so the loop collects nothing
and the range still yields nothing.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 156 +++++++++++++++++++++++++++++++++++++++++++++++-
 mm/collapse.h   |   9 +++
 mm/khugepaged.c |  24 +++++++-
 3 files changed, 185 insertions(+), 4 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 66931ef6a6d0..6dae5e35e61d 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -86,6 +86,52 @@
  * replaces, and is switched over once both halves are complete.
  */
 
+/*
+ * Cap on the memory a round may hold in flight: destination folios allocated
+ * but not yet installed, the fault latency of anything inside a candidate being
+ * collapsed, and memcg charge pressure all scale with it.  A dense table is
+ * collapsed as several rounds rather than one.
+ */
+#define COLLAPSE_BATCH_BYTES		SZ_32M
+
+/* Windows in one table at the finest order collapse cuts */
+#define COLLAPSE_TABLE_WINDOWS	(HPAGE_PMD_NR >> COLLAPSE_MIN_MTHP_ORDER)
+
+/*
+ * How many candidates a round can hold, fixed by the table geometry: the byte
+ * cap decides it at the smallest order collapse builds, but never more than the
+ * windows one table has at that order.
+ */
+#define COLLAPSE_MAX_CANDIDATES						\
+	min(COLLAPSE_BATCH_BYTES >> (PAGE_SHIFT + COLLAPSE_MIN_MTHP_ORDER), \
+	    COLLAPSE_TABLE_WINDOWS)
+
+/*
+ * A candidate is an (addr, order) window selected for collapse.  Selection
+ * counts in PTE offsets -- the bitmap it reads and the alignment it honours are
+ * indexed that way -- while the passes that run a candidate work in addresses,
+ * like the page tables and VMAs they touch.  This is where the two meet.
+ */
+struct collapse_candidate {
+	unsigned long addr;
+	unsigned int order;
+};
+
+void collapse_control_release(struct collapse_control *cc)
+{
+	kfree(cc->candidates);
+	cc->candidates = NULL;
+}
+
+int collapse_control_init(struct collapse_control *cc)
+{
+	cc->nr_candidates = 0;
+	cc->candidates = kmalloc_objs(*cc->candidates, COLLAPSE_MAX_CANDIDATES);
+	if (!cc->candidates)
+		return -ENOMEM;
+	return 0;
+}
+
 /*
  * Is @count past a limit stated per PMD, when only part of a table was scanned?
  * Scale the comparison to the table so a partial scan is held to the same
@@ -389,6 +435,75 @@ collapse_scan_anon_pmd(struct vm_area_struct *vma, unsigned long start,
 	return cc->scan_refusal;
 }
 
+/* Point the selection cursor at [start, end) of the table, in PTE offsets */
+static void collapse_selection_init(struct collapse_control *cc,
+				    unsigned int start, unsigned int end)
+{
+}
+
+/*
+ * The next window worth attempting, as an (offset, order) pair.  False when
+ * selection is exhausted, which is what ends the range.
+ *
+ * A candidate is only ever an (offset, order) pair: the scan that recorded the
+ * eligible PTEs has dropped the ptl, so anything else -- folio pointers in
+ * particular -- would be stale by construction.
+ */
+static bool collapse_next_candidate(struct collapse_control *cc,
+				    unsigned int *offset, unsigned int *order)
+{
+	return false;
+}
+
+/*
+ * Run and classify the collected batch.  Returns false when a candidate's
+ * outcome abandons the table.
+ */
+static bool collapse_run_batch(struct mm_struct *mm, unsigned long pmd_addr,
+			       struct collapse_control *cc)
+{
+	/* collapse_anon_pmd() only runs a round it has put something in */
+	VM_WARN_ON_ONCE(!cc->nr_candidates);
+
+	cc->nr_candidates = 0;
+	return true;
+}
+
+/*
+ * One more candidate of @order would either overflow the array or push what the
+ * round holds past the byte cap.  An empty round takes whatever it is offered:
+ * a single candidate is above the cap all by itself once a PMD is (512M with
+ * 64K pages), and refusing it would collapse nothing at all.
+ */
+static bool collapse_batch_full(struct collapse_control *cc,
+				unsigned long bytes, unsigned int order)
+{
+	if (!cc->nr_candidates)
+		return false;
+
+	return cc->nr_candidates == COLLAPSE_MAX_CANDIDATES ||
+	       bytes + (PAGE_SIZE << order) > COLLAPSE_BATCH_BYTES;
+}
+
+/*
+ * Take the next array slot for the window at @addr.  A slot may still hold a
+ * previous round's values, so every field is set here.
+ */
+static void collapse_add_candidate(struct collapse_control *cc,
+				   unsigned long addr, unsigned int order)
+{
+	struct collapse_candidate *cand;
+
+	/* collapse_batch_full() has already made room */
+	if (WARN_ON_ONCE(cc->nr_candidates >= COLLAPSE_MAX_CANDIDATES))
+		return;
+
+	cand = &cc->candidates[cc->nr_candidates];
+	cc->nr_candidates++;
+	cand->addr = addr;
+	cand->order = order;
+}
+
 /*
  * Cut the table into candidate windows and collapse what fits, from the
  * largest order downwards.  Returns what the table yielded: a collapse, or
@@ -398,5 +513,44 @@ static enum scan_result __maybe_unused
 collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 		  struct collapse_control *cc)
 {
-	return SCAN_FAIL;
+	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
+	unsigned int offset, order;
+	unsigned long bytes = 0;
+	bool pending = false;
+	bool cont = true;
+
+	collapse_selection_init(cc, (start - pmd_addr) >> PAGE_SHIFT,
+				(end - pmd_addr) >> PAGE_SHIFT);
+
+	while (cont) {
+		if (!pending)
+			pending = collapse_next_candidate(cc, &offset, &order);
+
+		if (!pending || collapse_batch_full(cc, bytes, order)) {
+			/*
+			 * Selection is exhausted and the round is empty: the
+			 * range is done.  Without this a flush of an empty
+			 * round would return, collect nothing, and come
+			 * straight back here.
+			 */
+			if (!cc->nr_candidates)
+				break;
+
+			cont = collapse_run_batch(mm, pmd_addr, cc);
+			bytes = 0;
+			continue;
+		}
+
+		/*
+		 * The round holds no resources until it is run, so
+		 * collecting costs nothing but the array slot.  A candidate the
+		 * full round could not take is kept pending for the next one.
+		 */
+		collapse_add_candidate(cc, pmd_addr + offset * PAGE_SIZE, order);
+
+		bytes += PAGE_SIZE << order;
+		pending = false;
+	}
+
+	return cc->nr_collapsed ? SCAN_SUCCEED : SCAN_FAIL;
 }
diff --git a/mm/collapse.h b/mm/collapse.h
index ad88b91d9a72..1159ed39b9eb 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -10,6 +10,8 @@
 /* The smallest order a collapse will build, and so the finest window it cuts */
 #define COLLAPSE_MIN_MTHP_ORDER		2
 
+struct collapse_candidate;
+
 enum scan_result {
 	SCAN_FAIL,
 	SCAN_SUCCEED,
@@ -120,8 +122,15 @@ struct collapse_control {
 	 * the collapse reports it when it salvages nothing.
 	 */
 	enum scan_result scan_refusal;
+
+	/* The candidate windows collected for the current round */
+	struct collapse_candidate *candidates;
+	unsigned int nr_candidates;
 };
 
+int collapse_control_init(struct collapse_control *cc);
+void collapse_control_release(struct collapse_control *cc);
+
 /*
  * Defined in khugepaged.c, which still uses them itself.
  * TODO: move each into collapse.c once its last khugepaged.c user is gone.
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 9823884a83c9..43f6107c953a 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -3079,12 +3079,22 @@ int start_stop_khugepaged(void)
 	guard(mutex)(&khugepaged_mutex);
 	if (hugepage_enabled()) {
 		if (!khugepaged_thread) {
-			struct task_struct *new_thread = kthread_run(khugepaged,
-								     NULL,
-								     "khugepaged");
+			struct task_struct *new_thread;
+			int err;
 
+			/*
+			 * The engine collapses out of its candidate array, so
+			 * take it before starting the thread that needs it: a
+			 * failure surfaces here rather than in the daemon.
+			 */
+			err = collapse_control_init(&khugepaged_collapse_control);
+			if (err)
+				return err;
+
+			new_thread = kthread_run(khugepaged, NULL, "khugepaged");
 			if (IS_ERR(new_thread)) {
 				pr_err("khugepaged: kthread_run(khugepaged) failed\n");
+				collapse_control_release(&khugepaged_collapse_control);
 				return PTR_ERR(new_thread);
 			}
 
@@ -3096,6 +3106,7 @@ int start_stop_khugepaged(void)
 	} else if (khugepaged_thread) {
 		kthread_stop(khugepaged_thread);
 		khugepaged_thread = NULL;
+		collapse_control_release(&khugepaged_collapse_control);
 	}
 	set_recommended_min_free_kbytes();
 	return 0;
@@ -3154,6 +3165,7 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	enum scan_result last_fail = SCAN_FAIL;
 	int thps = 0;
 	bool mmap_unlocked = false;
+	int err;
 
 	BUG_ON(vma->vm_start > start);
 	BUG_ON(vma->vm_end < end);
@@ -3173,6 +3185,11 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	cc->is_khugepaged = false;
 	collapse_policy_forced(&cc->policy);
 	cc->progress = 0;
+	err = collapse_control_init(cc);
+	if (err) {
+		kfree(cc);
+		return err;
+	}
 
 	mmgrab(mm);
 	lru_add_drain_all();
@@ -3231,6 +3248,7 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 out_nolock:
 	mmap_assert_locked(mm);
 	mmdrop(mm);
+	collapse_control_release(cc);
 	kfree(cc);
 
 	return thps == ((hend - hstart) >> HPAGE_PMD_SHIFT) ? 0
-- 
2.54.0

[RFC PATCH 10/57] mm/collapse: run a round and feed the outcomes back

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A candidate a round attempts either collapsed or did not, and if it did
not there is a reason.  Selection needs those outcomes to decide what
comes next: carry on past the window, try the same region at a lower
order, or give the table up.

Fill in collapse_run_batch(): run the round, then walk the batch handing
each candidate's result to classification.

The walk covers the whole batch.  A pass that refuses one candidate marks
it and carries on rather than truncating the round, so every candidate
has a result of its own to hand back.  Only an outcome that condemns the
table cuts the walk short, and then nothing of that table re-enters
selection.

The round and the classification it feeds are both stubs, so nothing is
attempted and nothing is decided.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/mm/collapse.c b/mm/collapse.c
index 6dae5e35e61d..ad9e5a447854 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -115,8 +115,16 @@
 struct collapse_candidate {
 	unsigned long addr;
 	unsigned int order;
+	enum scan_result result;
 };
 
+/* Where a candidate sits in the table, in the PTE offsets selection counts in */
+static unsigned int candidate_offset(const struct collapse_candidate *cand,
+				     unsigned long pmd_addr)
+{
+	return (cand->addr - pmd_addr) >> PAGE_SHIFT;
+}
+
 void collapse_control_release(struct collapse_control *cc)
 {
 	kfree(cc->candidates);
@@ -132,6 +140,17 @@ int collapse_control_init(struct collapse_control *cc)
 	return 0;
 }
 
+/*
+ * Carry one batch of candidates through the passes.  Every candidate comes back
+ * with a result of its own: the passes before the freeze mark what they refuse
+ * and carry on, each pass after it works on what the last left, so no failure
+ * truncates the round.
+ */
+static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
+			   struct collapse_control *cc)
+{
+}
+
 /*
  * Is @count past a limit stated per PMD, when only part of a table was scanned?
  * Scale the comparison to the table so a partial scan is held to the same
@@ -455,6 +474,18 @@ static bool collapse_next_candidate(struct collapse_control *cc,
 	return false;
 }
 
+/*
+ * Feed one candidate's outcome back into selection: its region is done, it
+ * re-enters the retry store at a lower order, or the table is abandoned.
+ * Returns false in that last case.
+ */
+static bool collapse_classify_result(struct collapse_control *cc,
+				     unsigned int offset, unsigned int order,
+				     enum scan_result result)
+{
+	return true;
+}
+
 /*
  * Run and classify the collected batch.  Returns false when a candidate's
  * outcome abandons the table.
@@ -462,9 +493,30 @@ static bool collapse_next_candidate(struct collapse_control *cc,
 static bool collapse_run_batch(struct mm_struct *mm, unsigned long pmd_addr,
 			       struct collapse_control *cc)
 {
+	unsigned int i;
+
 	/* collapse_anon_pmd() only runs a round it has put something in */
 	VM_WARN_ON_ONCE(!cc->nr_candidates);
 
+	collapse_round(mm, pmd_addr, cc);
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		unsigned int offset = candidate_offset(cand, pmd_addr);
+
+		if (!collapse_classify_result(cc, offset, cand->order,
+					      cand->result)) {
+			/*
+			 * The table is abandoned: the candidates behind this one
+			 * keep their results and are left unclassified, so
+			 * nothing more of this table enters selection, and the
+			 * abandoning result clears what earlier ones left there.
+			 */
+			cc->nr_candidates = 0;
+			return false;
+		}
+	}
+
 	cc->nr_candidates = 0;
 	return true;
 }
@@ -502,6 +554,7 @@ static void collapse_add_candidate(struct collapse_control *cc,
 	cc->nr_candidates++;
 	cand->addr = addr;
 	cand->order = order;
+	cand->result = SCAN_FAIL;
 }
 
 /*
-- 
2.54.0

[RFC PATCH 11/57] mm/collapse: sketch the passes of a round

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A round is a sequence of passes over the same batch, and their order is
most of the design.  It falls into three parts:

 - before any lock, the allocations that may sleep: destination folios,
   and the page table a PMD-order candidate deposits;
 - under mmap_read, revalidation and fault-in, which may have to give the
   lock up;
 - from the freeze onwards, a stretch that has to run to completion.

Each pass in that last part works on what the one before it left, and
every barrier raised has to be lowered again.  The destinations still
missing are asked for there too, without reclaim: a faulter on a frozen
source would wait for the allocation.

Lay that sequence out, with every pass a stub but one.  collapse_round()
takes mmap_read for the middle of it: a collapse is called without the
lock, and takes its own for each round.  It brackets the frozen window in
one mmu-notifier invalidate over the whole batch.

That invalidate needs a span before any pass has a body, so
collapse_revalidate() settles it from the start, in cc->batch_start and
cc->batch_end.  The span is taken over the candidates rather than off the
ends of the array: a region refused at one order can re-enter selection
at a lower one, so a round is not address-ordered and candidates[0] need
not be the lowest.

A fault-in that had to sleep comes back with the lock dropped, reported
as SCAN_LOCK_DROPPED.  The round takes the lock again and runs the pass
afresh, a bounded number of times, rather than sending the batch back to
selection.

The stubs do nothing, so the round does nothing.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h |   1 +
 mm/collapse.c                      | 207 +++++++++++++++++++++++++++++
 mm/collapse.h                      |   9 ++
 3 files changed, 217 insertions(+)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 5a48c5406cce..778f5a56956c 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -24,6 +24,7 @@
 	EM( SCAN_PAGE_COUNT,		"not_suitable_page_count")	\
 	EM( SCAN_PAGE_LRU,		"page_not_in_lru")		\
 	EM( SCAN_PAGE_LOCK,		"page_locked")			\
+	EM( SCAN_LOCK_DROPPED,		"lock_dropped")			\
 	EM( SCAN_PAGE_ANON,		"page_not_anon")		\
 	EM( SCAN_PAGE_LAZYFREE,		"page_lazyfree")		\
 	EM( SCAN_PAGE_COMPOUND,		"page_compound")		\
diff --git a/mm/collapse.c b/mm/collapse.c
index ad9e5a447854..25c0f72a9a68 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -94,6 +94,14 @@
  */
 #define COLLAPSE_BATCH_BYTES		SZ_32M
 
+/*
+ * How many times a round runs the fault-in pass.  A fault that has to wait drops
+ * the lock, and running the pass again costs a walk of the batch but buys at
+ * least one completed swap-in; readahead brings a cluster in at a time, so this
+ * covers a PMD's default max_ptes_swap.
+ */
+#define COLLAPSE_FAULTIN_PASSES	8
+
 /* Windows in one table at the finest order collapse cuts */
 #define COLLAPSE_TABLE_WINDOWS	(HPAGE_PMD_NR >> COLLAPSE_MIN_MTHP_ORDER)
 
@@ -118,6 +126,21 @@ struct collapse_candidate {
 	enum scan_result result;
 };
 
+static unsigned long candidate_start(const struct collapse_candidate *cand)
+{
+	return cand->addr;
+}
+
+static unsigned long candidate_size(const struct collapse_candidate *cand)
+{
+	return PAGE_SIZE << cand->order;
+}
+
+static unsigned long candidate_end(const struct collapse_candidate *cand)
+{
+	return candidate_start(cand) + candidate_size(cand);
+}
+
 /* Where a candidate sits in the table, in the PTE offsets selection counts in */
 static unsigned int candidate_offset(const struct collapse_candidate *cand,
 				     unsigned long pmd_addr)
@@ -140,6 +163,134 @@ int collapse_control_init(struct collapse_control *cc)
 	return 0;
 }
 
+/*
+ * The scan and the allocation both dropped mmap_lock, so nothing seen before it
+ * can be trusted: find the VMA and the PTE table again, and check they still
+ * allow every provisioned candidate.
+ *
+ * This is also where the batch's span is settled, for the invalidate the round
+ * issues over it.
+ */
+static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
+					    unsigned long pmd_addr,
+					    struct collapse_control *cc,
+					    pmd_t **pmdp)
+{
+	unsigned int i;
+
+	cc->batch_start = ULONG_MAX;
+	cc->batch_end = 0;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+
+		cc->batch_start = min(cc->batch_start, candidate_start(cand));
+		cc->batch_end = max(cc->batch_end, candidate_end(cand));
+	}
+
+	return SCAN_SUCCEED;
+}
+
+/*
+ * Make every source the round needs present and exclusively owned by this mm,
+ * by faulting it in as an ordinary access would.  Sleeps, and drops mmap_lock on
+ * failure, since a fault may have to be retried with it released.
+ *
+ * Anything faulted in lands on a per-CPU LRU batch, holding a reference the
+ * freeze cannot account for, so the freeze drains those batches before it
+ * starts.
+ */
+static enum scan_result collapse_faultin(struct vm_area_struct *vma,
+					 struct collapse_control *cc,
+					 pmd_t *pmd)
+{
+	return SCAN_SUCCEED;
+}
+
+/*
+ * Raise the two barriers on the sources of every candidate: migration entries in
+ * their PTEs, then a frozen refcount.  Takes the table's ptl once for the whole
+ * batch, and flushes the TLB once before dropping it.  A candidate whose sources
+ * moved is dropped here.
+ */
+static void collapse_freeze(struct vm_area_struct *vma,
+			    struct collapse_control *cc, pmd_t *pmd)
+{
+}
+
+/*
+ * Allocate ahead of the freeze for the candidates whose light allocation missed
+ * last round.  This is where reclaim belongs: nothing is held or frozen, so a
+ * long compaction costs only khugepaged's own progress -- which is why the
+ * mechanism this replaces allocated here too.  Having asked the allocator to try
+ * hard, a miss now is a failure.
+ */
+static void collapse_reserve(struct mm_struct *mm, struct collapse_control *cc)
+{
+}
+
+/*
+ * Secure the page table the PMD terminal layer deposits.  This stays ahead of the
+ * freeze because pte_alloc_one() allocates with GFP_PGTABLE_USER and takes no gfp
+ * to strip: order-0 or not, it may reclaim and sleep, which is what the window
+ * exists to keep out.  The destination folio has a light gfp to fall back on and
+ * so can be deferred; this has none.
+ */
+static void collapse_deposit(struct mm_struct *mm, struct collapse_control *cc)
+{
+}
+
+/*
+ * Give the frozen candidates that still need one a destination folio, without
+ * reclaim: a faulter on their sources would wait for it.
+ *
+ * A miss here is not a failure, as long as a retry could do better: the
+ * candidate keeps its freeze and asks for the reclaiming gfp, which
+ * collapse_reserve() uses before the next round freezes anything.  When the
+ * policy forbids reclaim there is nothing better to retry with, so the miss is
+ * the answer, and a smaller order over the same region is the better next move.
+ */
+static void collapse_provision(struct mm_struct *mm,
+			       struct collapse_control *cc)
+{
+}
+
+/*
+ * Copy the frozen sources into their destinations.  Nothing can reach either
+ * side, so this needs no page-table lock, and it sleeps.
+ */
+static void collapse_copy(struct vm_area_struct *vma,
+			  struct collapse_control *cc)
+{
+}
+
+/* Publish each destination folio in place of the sources it replaces */
+static void collapse_install(struct vm_area_struct *vma,
+			     struct collapse_control *cc, pmd_t *pmd)
+{
+}
+
+/*
+ * Lower the barriers the freeze raised, on the sources of an installed candidate
+ * and on those of one that got no further.
+ */
+static void collapse_putback(struct vm_area_struct *vma,
+			     struct collapse_control *cc)
+{
+}
+
+/*
+ * Settle whatever the round reached: account what was installed, release what
+ * was not, and give every candidate the result selection will classify.  Returns
+ * how many candidates were installed.
+ */
+static unsigned int collapse_finish(struct mm_struct *mm,
+				    struct collapse_control *cc,
+				    enum scan_result result)
+{
+	return 0;
+}
+
 /*
  * Carry one batch of candidates through the passes.  Every candidate comes back
  * with a result of its own: the passes before the freeze mark what they refuse
@@ -149,6 +300,62 @@ int collapse_control_init(struct collapse_control *cc)
 static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 			   struct collapse_control *cc)
 {
+	unsigned int passes = COLLAPSE_FAULTIN_PASSES;
+	struct mmu_notifier_range range;
+	struct vm_area_struct *vma;
+	enum scan_result result;
+	pmd_t *pmd;
+
+	collapse_reserve(mm, cc);
+	collapse_deposit(mm, cc);
+
+retry:
+	mmap_read_lock(mm);
+
+	vma = find_vma(mm, pmd_addr);
+	if (!vma) {
+		result = SCAN_VMA_NULL;
+		goto out_unlock;
+	}
+
+	result = collapse_revalidate(vma, pmd_addr, cc, &pmd);
+	if (result != SCAN_SUCCEED)
+		goto out_unlock;
+
+	result = collapse_faultin(vma, cc, pmd);
+	/*
+	 * A fault dropped the lock to wait, as a swap-in does.  The swap-in it
+	 * started is still running and the walk skips whatever has arrived, so
+	 * take the lock again rather than send the batch back to selection.  The
+	 * VMA and the table are looked up afresh: both may have changed.
+	 */
+	if (result == SCAN_LOCK_DROPPED && --passes)
+		goto retry;
+	if (result != SCAN_SUCCEED)
+		goto out;	/* the callee released mmap_lock */
+
+	/* One invalidate window spans the batch, as collapse_revalidate() left it */
+	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
+				cc->batch_start, cc->batch_end);
+	mmu_notifier_invalidate_range_start(&range);
+
+	/*
+	 * None of these can fail as a whole: the freeze takes the sources it
+	 * can and drops the candidates it cannot, and each pass after it works
+	 * on what the one before left, so every barrier raised is lowered again.
+	 */
+	collapse_freeze(vma, cc, pmd);
+	collapse_provision(mm, cc);
+	collapse_copy(vma, cc);
+	collapse_install(vma, cc, pmd);
+	collapse_putback(vma, cc);
+
+	mmu_notifier_invalidate_range_end(&range);
+
+out_unlock:
+	mmap_read_unlock(mm);
+out:
+	collapse_finish(mm, cc, result);
 }
 
 /*
diff --git a/mm/collapse.h b/mm/collapse.h
index 1159ed39b9eb..c61db86dc6c2 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -29,6 +29,7 @@ enum scan_result {
 	SCAN_PAGE_COUNT,
 	SCAN_PAGE_LRU,
 	SCAN_PAGE_LOCK,
+	SCAN_LOCK_DROPPED,
 	SCAN_PAGE_ANON,
 	SCAN_PAGE_LAZYFREE,
 	SCAN_PAGE_COMPOUND,
@@ -126,6 +127,14 @@ struct collapse_control {
 	/* The candidate windows collected for the current round */
 	struct collapse_candidate *candidates;
 	unsigned int nr_candidates;
+
+	/*
+	 * What the candidates the round still means to freeze span, settled by
+	 * collapse_revalidate() as it walks them.  A round is not necessarily
+	 * address-ordered, so this cannot be read off the ends of the array.
+	 */
+	unsigned long batch_start;
+	unsigned long batch_end;
 };
 
 int collapse_control_init(struct collapse_control *cc);
-- 
2.54.0

[RFC PATCH 12/57] mm/collapse: allocate a destination per candidate

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the allocation, which happens on both sides of the freeze.

A destination is a folio of the candidate's order, charged to the memcg,
with the memcg's deferred-split list entry taken up front while sleeping
is still allowed: the PMD-order install would otherwise need one under
the pmd lock.

collapse_alloc() does all of that for one candidate with the gfp it is
handed, and counts nothing when it fails: what a miss means is up to the
caller.

collapse_provision() is the caller inside the window.  The sources are
frozen by then and a faulter on any of them is waiting, so it asks
without __GFP_DIRECT_RECLAIM: reclaim entered there would be paid for by
that faulter.

A candidate the allocator cannot spare one for is declined rather than
failed.  It keeps its freeze and records SCAN_ALLOC_LIGHT_MISS, which
asks for the reclaiming gfp so a later round can allocate for it before
freezing anything.  Where the policy forbids reclaim there is nothing
better to retry with, so the miss is the verdict: the real result is
recorded and the failure counters fire.

collapse_reserve() honours those requests, before the round takes any
lock.  This is where reclaim belongs: nothing is held or frozen, so a
long compaction costs only khugepaged's own progress, which is why the
mechanism being replaced allocated here too.  Having asked the allocator
to try hard, a miss there is a failure.

Nothing sets cand->reclaim yet, so collapse_reserve() has nothing to do.
The request comes from the selection side, which queues a region refused
at one order for another attempt.

The page table a PMD-order candidate deposits cannot be deferred the same
way.  pte_alloc_one() allocates with GFP_PGTABLE_USER and takes no gfp to
strip, so it may reclaim and sleep whatever the order asked for.
collapse_deposit() secures it ahead of the freeze and refuses the
candidate when it cannot.  A PMD-order window is a whole table, so there
is at most one such candidate and it is the first.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h |   3 +-
 mm/collapse.c                      | 126 +++++++++++++++++++++++++++++
 mm/collapse.h                      |   2 +
 mm/khugepaged.c                    |   4 +-
 4 files changed, 132 insertions(+), 3 deletions(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 778f5a56956c..68693eba82ef 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -40,7 +40,8 @@
 	EM( SCAN_STORE_FAILED,		"store_failed")			\
 	EM( SCAN_COPY_MC,		"copy_poisoned_page")		\
 	EM( SCAN_PAGE_FILLED,		"page_filled")			\
-	EMe(SCAN_PAGE_DIRTY_OR_WRITEBACK, "page_dirty_or_writeback")
+	EM( SCAN_PAGE_DIRTY_OR_WRITEBACK, "page_dirty_or_writeback")	\
+	EMe(SCAN_ALLOC_LIGHT_MISS,	"alloc_light_miss")
 
 #undef EM
 #undef EMe
diff --git a/mm/collapse.c b/mm/collapse.c
index 25c0f72a9a68..58c8d83f3468 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -17,6 +17,7 @@
 #include <linux/slab.h>
 #include <linux/swap.h>
 #include <linux/userfaultfd_k.h>
+#include <linux/vmstat.h>
 
 #include <asm/tlb.h>
 #include "collapse.h"
@@ -114,6 +115,12 @@
 	min(COLLAPSE_BATCH_BYTES >> (PAGE_SHIFT + COLLAPSE_MIN_MTHP_ORDER), \
 	    COLLAPSE_TABLE_WINDOWS)
 
+/* How far a candidate got, and so what a failure has to undo for it */
+enum collapse_candidate_state {
+	CAND_SELECTED,		/* collected; nothing held on its behalf yet */
+	CAND_SKIPPED,		/* refused; nothing of it left to undo */
+};
+
 /*
  * A candidate is an (addr, order) window selected for collapse.  Selection
  * counts in PTE offsets -- the bitmap it reads and the alignment it honours are
@@ -123,7 +130,12 @@
 struct collapse_candidate {
 	unsigned long addr;
 	unsigned int order;
+	/* The light allocation missed last round: this one may reclaim for it */
+	bool reclaim;
+	enum collapse_candidate_state state;
 	enum scan_result result;
+	struct folio *new_folio;
+	pgtable_t deposit;		/* PMD order: fresh table to deposit */
 };
 
 static unsigned long candidate_start(const struct collapse_candidate *cand)
@@ -218,6 +230,43 @@ static void collapse_freeze(struct vm_area_struct *vma,
 {
 }
 
+/*
+ * Allocate one candidate's destination with @gfp: a folio of its order, charged,
+ * with the memcg's deferred-split list heads in place so the install cannot need
+ * to allocate under the pmd lock.  Those heads cost only the first collapse in a
+ * memcg.
+ *
+ * A failure counts nothing and changes nothing: what a miss means is the caller's
+ * policy.
+ */
+static enum scan_result collapse_alloc(struct mm_struct *mm,
+				       struct collapse_control *cc,
+				       struct collapse_candidate *cand,
+				       gfp_t gfp)
+{
+	struct folio *folio;
+
+	folio = __folio_alloc(gfp, cand->order, collapse_find_target_node(cc),
+			      &cc->alloc_nmask);
+	if (!folio)
+		return SCAN_ALLOC_HUGE_PAGE_FAIL;
+
+	if (unlikely(mem_cgroup_charge(folio, mm, gfp)) ||
+	    folio_memcg_alloc_deferred(folio)) {
+		folio_put(folio);
+		return SCAN_CGROUP_CHARGE_FAIL;
+	}
+
+	if (is_pmd_order(cand->order)) {
+		count_vm_event(THP_COLLAPSE_ALLOC);
+		count_memcg_folio_events(folio, THP_COLLAPSE_ALLOC, 1);
+	}
+	count_mthp_stat(cand->order, MTHP_STAT_COLLAPSE_ALLOC);
+	cand->new_folio = folio;
+
+	return SCAN_SUCCEED;
+}
+
 /*
  * Allocate ahead of the freeze for the candidates whose light allocation missed
  * last round.  This is where reclaim belongs: nothing is held or frozen, so a
@@ -227,6 +276,31 @@ static void collapse_freeze(struct vm_area_struct *vma,
  */
 static void collapse_reserve(struct mm_struct *mm, struct collapse_control *cc)
 {
+	unsigned int i;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		enum scan_result result;
+
+		if (!cand->reclaim)
+			continue;
+		cand->reclaim = false;
+
+		result = collapse_alloc(mm, cc, cand, cc->policy.gfp);
+		if (result == SCAN_SUCCEED)
+			continue;
+
+		if (result == SCAN_ALLOC_HUGE_PAGE_FAIL) {
+			/* Asked the allocator to try hard and it still missed */
+			if (is_pmd_order(cand->order))
+				count_vm_event(THP_COLLAPSE_ALLOC_FAILED);
+			count_mthp_stat(cand->order,
+					MTHP_STAT_COLLAPSE_ALLOC_FAILED);
+		}
+
+		cand->state = CAND_SKIPPED;
+		cand->result = result;
+	}
 }
 
 /*
@@ -235,9 +309,28 @@ static void collapse_reserve(struct mm_struct *mm, struct collapse_control *cc)
  * to strip: order-0 or not, it may reclaim and sleep, which is what the window
  * exists to keep out.  The destination folio has a light gfp to fall back on and
  * so can be deferred; this has none.
+ *
+ * A round is one table and a PMD-order window is the whole of it, so such a
+ * candidate cannot share a round: if there is one it is the only one, and it is
+ * candidates[0].  This secures one page table, never a batch of them.
  */
 static void collapse_deposit(struct mm_struct *mm, struct collapse_control *cc)
 {
+	struct collapse_candidate *cand = &cc->candidates[0];
+
+	if (!is_pmd_order(cand->order))
+		return;
+
+	VM_WARN_ON_ONCE(cc->nr_candidates != 1);
+
+	if (cand->state != CAND_SELECTED)
+		return;
+
+	cand->deposit = pte_alloc_one(mm);
+	if (!cand->deposit) {
+		cand->state = CAND_SKIPPED;
+		cand->result = SCAN_ALLOC_HUGE_PAGE_FAIL;
+	}
 }
 
 /*
@@ -253,6 +346,35 @@ static void collapse_deposit(struct mm_struct *mm, struct collapse_control *cc)
 static void collapse_provision(struct mm_struct *mm,
 			       struct collapse_control *cc)
 {
+	const gfp_t gfp = cc->policy.gfp & ~__GFP_DIRECT_RECLAIM;
+	const bool may_retry = gfp != cc->policy.gfp;
+	unsigned int i;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		enum scan_result result;
+
+		if (cand->state != CAND_SELECTED || cand->new_folio)
+			continue;
+
+		result = collapse_alloc(mm, cc, cand, gfp);
+		if (result == SCAN_SUCCEED)
+			continue;
+
+		if (may_retry) {
+			/* A charge miss too: charging may reclaim when allowed */
+			cand->result = SCAN_ALLOC_LIGHT_MISS;
+		} else {
+			/* The gfp a retry would use, so this is the answer */
+			if (result == SCAN_ALLOC_HUGE_PAGE_FAIL) {
+				if (is_pmd_order(cand->order))
+					count_vm_event(THP_COLLAPSE_ALLOC_FAILED);
+				count_mthp_stat(cand->order,
+						MTHP_STAT_COLLAPSE_ALLOC_FAILED);
+			}
+			cand->result = result;
+		}
+	}
 }
 
 /*
@@ -761,7 +883,11 @@ static void collapse_add_candidate(struct collapse_control *cc,
 	cc->nr_candidates++;
 	cand->addr = addr;
 	cand->order = order;
+	cand->reclaim = false;
+	cand->state = CAND_SELECTED;
 	cand->result = SCAN_FAIL;
+	cand->new_folio = NULL;
+	cand->deposit = NULL;
 }
 
 /*
diff --git a/mm/collapse.h b/mm/collapse.h
index c61db86dc6c2..feb2e0d57339 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -46,6 +46,7 @@ enum scan_result {
 	SCAN_COPY_MC,
 	SCAN_PAGE_FILLED,
 	SCAN_PAGE_DIRTY_OR_WRITEBACK,
+	SCAN_ALLOC_LIGHT_MISS,
 };
 
 /*
@@ -148,6 +149,7 @@ unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags);
 enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd);
+int collapse_find_target_node(struct collapse_control *cc);
 bool collapse_scan_abort(int nid, struct collapse_control *cc);
 unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order);
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 43f6107c953a..50b520961b9b 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -999,7 +999,7 @@ static void collapse_policy_forced(struct collapse_policy *p)
 }
 
 #ifdef CONFIG_NUMA
-static int collapse_find_target_node(struct collapse_control *cc)
+int collapse_find_target_node(struct collapse_control *cc)
 {
 	int nid, target_node = 0, max_value = 0;
 
@@ -1018,7 +1018,7 @@ static int collapse_find_target_node(struct collapse_control *cc)
 	return target_node;
 }
 #else
-static int collapse_find_target_node(struct collapse_control *cc)
+int collapse_find_target_node(struct collapse_control *cc)
 {
 	return 0;
 }
-- 
2.54.0

[RFC PATCH 13/57] mm/collapse: revalidate a round against the VMA

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the pass that re-establishes what the round is working on.
Selection ran under mmap_lock and the allocation ran without it, so by
the time the round takes the lock back the address space may have changed
underneath it.

Check that the mm is not exiting and has not had THP disabled, that the
VMA the round looked up is still anonymous with an anon_vma, and find the
PTE table again in case it became a huge PMD or went away.

Then re-check each candidate on its own.  A VMA that shrank, or was
replaced by a smaller one, may no longer hold a window that fitted when
it was selected, and per-size enablement may have been turned off for its
order since.  Such a candidate is dropped and the rest of the round goes
on without it.

The check is per candidate rather than over the batch because a window is
aligned to its own order: thp_vma_suitable_order() on each one is the
containment check in full.

What the walk leaves is what the round goes on to freeze, so it also
settles the batch's span, in cc->batch_start and cc->batch_end, for the
one invalidate the round issues.  A candidate the walk dropped is not in
the span, and a round left with no candidates has no span and nothing to
run.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 49 ++++++++++++++++++++++++++++++++++++++++++++-----
 mm/collapse.h   | 11 +++++++++++
 mm/khugepaged.c | 11 -----------
 3 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 58c8d83f3468..1367ade721f7 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -177,18 +177,36 @@ int collapse_control_init(struct collapse_control *cc)
 
 /*
  * The scan and the allocation both dropped mmap_lock, so nothing seen before it
- * can be trusted: find the VMA and the PTE table again, and check they still
- * allow every provisioned candidate.
+ * can be trusted: check the VMA the round just looked up and the PTE table
+ * again, and that they still allow every provisioned candidate.
  *
- * This is also where the batch's span is settled, for the invalidate the round
- * issues over it.
+ * The VMA was found by address, so it need not be the one the scan saw, nor
+ * still cover everything the round collected -- thp_vma_suitable_order() asks
+ * that of each candidate, since a window is aligned to its own order.  A VMA
+ * that shrank under a candidate therefore refuses that candidate and no more,
+ * like every other pass.
+ *
+ * What survives is what the round goes on to freeze, so this is also where the
+ * batch's span is settled, for the invalidate the round issues over it.
  */
 static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 					    unsigned long pmd_addr,
 					    struct collapse_control *cc,
 					    pmd_t **pmdp)
 {
-	unsigned int i;
+	struct mm_struct *mm = vma->vm_mm;
+	enum scan_result result;
+	unsigned int i, nr_live = 0;
+
+	if (unlikely(collapse_test_exit_or_disable(mm)))
+		return SCAN_ANY_PROCESS;
+
+	if (!vma->anon_vma || !vma_is_anonymous(vma))
+		return SCAN_PAGE_ANON;
+
+	result = find_pmd_or_thp_or_none(mm, pmd_addr, pmdp);
+	if (result != SCAN_SUCCEED)
+		return result;
 
 	cc->batch_start = ULONG_MAX;
 	cc->batch_end = 0;
@@ -196,10 +214,31 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 	for (i = 0; i < cc->nr_candidates; i++) {
 		struct collapse_candidate *cand = &cc->candidates[i];
 
+		if (cand->state != CAND_SELECTED)
+			continue;
+
+		/*
+		 * The window has to still fit the VMA, which may have shrunk or
+		 * been replaced, and its order to still be one the VMA allows.
+		 */
+		if (!thp_vma_suitable_order(vma, cand->addr, cand->order) ||
+		    !thp_vma_allowable_orders(vma, vma->vm_flags,
+					      cc->policy.tva_type,
+					      BIT(cand->order))) {
+			cand->state = CAND_SKIPPED;
+			cand->result = SCAN_VMA_CHECK;
+			continue;
+		}
+
 		cc->batch_start = min(cc->batch_start, candidate_start(cand));
 		cc->batch_end = max(cc->batch_end, candidate_end(cand));
+		nr_live++;
 	}
 
+	/* Nothing the VMA still allows: no span to invalidate, nothing to run */
+	if (!nr_live)
+		return SCAN_VMA_CHECK;
+
 	return SCAN_SUCCEED;
 }
 
diff --git a/mm/collapse.h b/mm/collapse.h
index feb2e0d57339..0d6f77a7233b 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -138,6 +138,17 @@ struct collapse_control {
 	unsigned long batch_end;
 };
 
+static inline int collapse_test_exit(struct mm_struct *mm)
+{
+	return atomic_read(&mm->mm_users) == 0;
+}
+
+static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
+{
+	return collapse_test_exit(mm) ||
+		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
+}
+
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
 
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 50b520961b9b..1244e161beae 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -421,17 +421,6 @@ void __init khugepaged_destroy(void)
 	kmem_cache_destroy(mm_slot_cache);
 }
 
-static inline int collapse_test_exit(struct mm_struct *mm)
-{
-	return atomic_read(&mm->mm_users) == 0;
-}
-
-static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
-{
-	return collapse_test_exit(mm) ||
-		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
-}
-
 static inline bool anon_hpage_enabled(void)
 {
 	if (READ_ONCE(huge_anon_orders_always))
-- 
2.54.0

[RFC PATCH 14/57] mm/collapse: fault the sources in before the freeze

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the pass that makes the sources fit to freeze.  The freeze takes
the PTEs as it finds them and cannot fault, so before it starts every slot
a candidate covers has to be a hole, the zeropage, or a present page this
mm owns exclusively.

Walk each candidate slot by slot and let the fault path do the work: a
swap entry is read back in, a page shared with a fork child is unshared in
place.

Exclusivity is tested with PageAnonExclusive() rather than by asking
whether the folio looks shared.  They are not the same test.  A page whose
fork co-mapper has exited is mapped once and looks unshared, but stays
non-exclusive until some write reuses it, so a sharing test would skip the
unshare on exactly the pages that need one.

Each address gets a few tries, since an unshare can lose a race with a
co-mapper and a swap read can be interrupted; the PTE is re-read after
every fault.  What is still unfit after that is left to the freeze, which
refuses it.

Swap-in is refused outright below the PMD order, where reading pages back
to build an mTHP is not worth the latency.  That verdict is against the
one candidate, which is skipped so the rest of the batch can go on.

The pass sleeps, and on failure it returns with mmap_lock already dropped,
since the fault path may drop it and the caller cannot tell which case
happened.

A fault that has to wait for a swap read is one of those: it drops the
lock and returns VM_FAULT_RETRY, which says nothing about the window it
was working on.  Report it as SCAN_LOCK_DROPPED, so the round can take the
lock and run the pass again rather than treat it as a refusal.  Reporting
it as SCAN_PAGE_LOCK made it indistinguishable in a trace from a folio
someone else had locked.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 130 +++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 129 insertions(+), 1 deletion(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 1367ade721f7..4ec02071f588 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -153,6 +153,11 @@ static unsigned long candidate_end(const struct collapse_candidate *cand)
 	return candidate_start(cand) + candidate_size(cand);
 }
 
+static unsigned int candidate_nr_pages(const struct collapse_candidate *cand)
+{
+	return 1U << cand->order;
+}
+
 /* Where a candidate sits in the table, in the PTE offsets selection counts in */
 static unsigned int candidate_offset(const struct collapse_candidate *cand,
 				     unsigned long pmd_addr)
@@ -242,6 +247,93 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 	return SCAN_SUCCEED;
 }
 
+/*
+ * Faults one address may take before the freeze is left to judge it.  More than
+ * one because the unshare can race a co-mapper re-sharing the page, and a swap
+ * read can be interrupted; each try re-reads the PTE to see where it stands.
+ */
+#define COLLAPSE_FAULTIN_TRIES	3
+
+/*
+ * Bring one address to a state the freeze will accept: present, and exclusive if
+ * it is anonymous.  Returns with mmap_lock dropped on every failure, because the
+ * fault path may drop it and the caller cannot tell which case it is in.
+ *
+ * SCAN_EXCEED_SWAP_PTE is the exception: it is a verdict on this candidate
+ * rather than on the round, nothing was faulted to reach it, and it keeps the
+ * lock so the caller can refuse this candidate and carry on with the rest.
+ */
+static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
+					      struct collapse_candidate *cand,
+					      pmd_t *pmd, unsigned long addr)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	const unsigned int flags = FAULT_FLAG_ALLOW_RETRY | FAULT_FLAG_UNSHARE |
+		(mm != current->mm ? FAULT_FLAG_REMOTE : 0);
+	unsigned int tries;
+
+	for (tries = 0; tries <= COLLAPSE_FAULTIN_TRIES; tries++) {
+		struct page *page;
+		pte_t ptent, *pte;
+		vm_fault_t ret;
+
+		pte = pte_offset_map(pmd, addr);
+		if (!pte) {
+			mmap_read_unlock(mm);
+			return SCAN_NO_PTE_TABLE;
+		}
+		ptent = ptep_get_lockless(pte);
+		pte_unmap(pte);
+
+		/* A hole or the zeropage is population's business */
+		if (pte_none_or_zero(ptent))
+			break;
+
+		if (pte_present(ptent)) {
+			page = vm_normal_page(vma, addr, ptent);
+
+			/*
+			 * PageAnonExclusive is the invariant the freeze relies
+			 * on, and the only exact test for it.  Testing sharing
+			 * with folio_maybe_mapped_shared() is not the same: a
+			 * page whose fork co-mapper has gone away is
+			 * single-mapped, yet stays non-exclusive until a write
+			 * reuses it, so sharing would skip the unshare on
+			 * exactly the pages that need it.  Unsharing one of
+			 * those is cheap -- it reuses the page in place and
+			 * just sets the bit.
+			 */
+			if (!page || !folio_test_anon(page_folio(page)) ||
+			    PageAnonExclusive(page))
+				break;		/* already exclusive */
+		} else if (!is_pmd_order(cand->order)) {
+			/* Sub-PMD collapse does not fault swap in */
+			count_mthp_stat(cand->order,
+					MTHP_STAT_COLLAPSE_EXCEED_SWAP);
+			return SCAN_EXCEED_SWAP_PTE;
+		}
+
+		if (tries == COLLAPSE_FAULTIN_TRIES)
+			break;		/* the freeze refuses it if still unfit */
+
+		/* Only swap or shared PTEs reach here; the rest broke out */
+		ret = handle_mm_fault(vma, addr, flags, NULL);
+		/*
+		 * Not a verdict on this window: the fault dropped the lock to
+		 * wait, which is what a swap-in normally does.  Distinct from
+		 * SCAN_PAGE_LOCK, a folio someone else holds locked.
+		 */
+		if (ret & VM_FAULT_RETRY)
+			return SCAN_LOCK_DROPPED;
+		if (ret & VM_FAULT_ERROR) {
+			mmap_read_unlock(mm);
+			return SCAN_FAIL;
+		}
+	}
+
+	return SCAN_SUCCEED;
+}
+
 /*
  * Make every source the round needs present and exclusively owned by this mm,
  * by faulting it in as an ordinary access would.  Sleeps, and drops mmap_lock on
@@ -255,7 +347,43 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 					 struct collapse_control *cc,
 					 pmd_t *pmd)
 {
-	return SCAN_SUCCEED;
+	enum scan_result result = SCAN_SUCCEED;
+	unsigned int i;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		unsigned long addr;
+		unsigned int j;
+
+		if (cand->state != CAND_SELECTED)
+			continue;
+
+		for (j = 0, addr = cand->addr;
+		     j < candidate_nr_pages(cand);
+		     j++, addr += PAGE_SIZE) {
+			enum scan_result r;
+
+			r = collapse_faultin_addr(vma, cand, pmd, addr);
+			/*
+			 * The one failure that judges this candidate rather
+			 * than the round, and so the one that leaves the lock
+			 * in our hands: refuse it and go on to the next.
+			 * Failing the round here would lower the order of every
+			 * candidate it carries, a verdict nobody reached.
+			 */
+			if (r == SCAN_EXCEED_SWAP_PTE) {
+				cand->state = CAND_SKIPPED;
+				cand->result = r;
+				break;
+			}
+			if (r != SCAN_SUCCEED) {
+				result = r;
+				goto out;
+			}
+		}
+	}
+out:
+	return result;
 }
 
 /*
-- 
2.54.0

[RFC PATCH 15/57] mm/collapse: check what a candidate would freeze

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The freeze takes folio locks, rewrites PTEs and flushes the TLB, and any
of that has to be undone slot by slot if the candidate turns out unfit --
while faulters on those sources wait.  So it decides first and acts
second.

This is the deciding half: walk every slot a candidate covers, under the
table's ptl, and answer whether all of it can be frozen.  It touches
nothing, so a refusal costs the round only the walk.

The walk goes in source spans, a span being consecutive PTEs mapping
consecutive pages of one folio.  No layout is refused for its shape:
where a span ends, the next slot starts one of its own, which is what
lets partially mapped and compound sources collapse.  A slot may also be
a hole or the zeropage, both of which the destination just zero-fills.

What a span has to satisfy, beyond being present, anonymous and not
uffd-armed:

 - Every live mapping of its folio is this span.  The freeze is
   whole-folio, so a live PTE anywhere else would race a zap whose
   folio_put() underflows the frozen count.  Under the ptl this is exact,
   since fork -- the only way an exclusive anon folio gains mappings --
   takes mmap_write.

 - Every page of it is PageAnonExclusive().  A shared folio has no
   refcount the freeze can pin down without the other mappers' ptls.

 - It is not MADV_FREE'd, unless the caller asked for the collapse.
   Copying a lazyfree page into a folio that is not lazyfree would quietly
   make memory the user offered up undroppable again, which is why the
   policy carries that choice.

Sub-PMD candidates also refuse folios already at or above their own order,
there being nothing to gain; a PMD candidate takes them, that being the
PTE-mapped-THP re-collapse case.

SCAN_PAGE_NOT_EXCLUSIVE joins enum scan_result and the trace symbol list.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h |   1 +
 mm/collapse.c                      | 177 +++++++++++++++++++++++++++++
 mm/collapse.h                      |   1 +
 3 files changed, 179 insertions(+)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 68693eba82ef..ff938ac9c43c 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -41,6 +41,7 @@
 	EM( SCAN_COPY_MC,		"copy_poisoned_page")		\
 	EM( SCAN_PAGE_FILLED,		"page_filled")			\
 	EM( SCAN_PAGE_DIRTY_OR_WRITEBACK, "page_dirty_or_writeback")	\
+	EM( SCAN_PAGE_NOT_EXCLUSIVE,	"page_not_exclusive")		\
 	EMe(SCAN_ALLOC_LIGHT_MISS,	"alloc_light_miss")
 
 #undef EM
diff --git a/mm/collapse.c b/mm/collapse.c
index 4ec02071f588..c75d91cb9d48 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -386,6 +386,145 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 	return result;
 }
 
+/*
+ * How many slots a source span starting at @first may cover: the pages left in
+ * its folio, capped at @max.  Every freeze-side walker bounds spans with this,
+ * so per-span batching of clears, locks and freezes cannot reach a slot the span
+ * does not cover.
+ */
+static unsigned int collapse_span_max(pte_t first, unsigned int max)
+{
+	struct page *page = pte_page(first);
+	struct folio *folio = page_folio(page);
+	unsigned int left = folio_nr_pages(folio) - folio_page_idx(folio, page);
+
+	return min(max, left);
+}
+
+/*
+ * Can this candidate's sources be frozen?  Every slot is checked and nothing is
+ * touched, so a refusal costs the round nothing but the walk.
+ *
+ * The walk is in source spans: a span is consecutive PTEs mapping consecutive
+ * pages of one folio, and it ends wherever the next PTE stops being the folio's
+ * next page.  No layout is refused for its shape -- the next slot simply starts
+ * its own span -- so partially mapped and compound sources collapse too.
+ *
+ * Caller holds mmap_read and the table's ptl.
+ */
+static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
+						 struct collapse_control *cc,
+						 struct collapse_candidate *cand,
+						 pte_t *pte)
+{
+	const unsigned int nr_pages = candidate_nr_pages(cand);
+	unsigned long addr;
+	unsigned int i;
+
+	for (i = 0, addr = cand->addr; i < nr_pages;) {
+		pte_t ptent = ptep_get(pte + i);
+		unsigned int nr, nr_max, k;
+		struct folio *folio;
+		struct page *page;
+
+		if (!pte_present(ptent)) {
+			/* Holes are population; swap and markers are not */
+			if (pte_none(ptent)) {
+				i++;
+				addr += PAGE_SIZE;
+				continue;
+			}
+			return SCAN_PTE_NON_PRESENT;
+		}
+		if (pte_uffd(ptent))
+			return SCAN_PTE_UFFD;
+
+		/* The zeropage zero-fills like a hole, and has no normal page */
+		if (is_zero_pfn(pte_pfn(ptent))) {
+			i++;
+			addr += PAGE_SIZE;
+			continue;
+		}
+		page = vm_normal_page(vma, addr, ptent);
+		if (!page || unlikely(is_zone_device_page(page)))
+			return SCAN_PAGE_NULL;
+
+		folio = page_folio(page);
+		if (!folio_test_anon(folio))
+			return SCAN_PAGE_ANON;
+
+		/*
+		 * Collapsing a MADV_FREE'd page would copy it into a folio that
+		 * is not lazyfree, quietly making memory the user offered up
+		 * undroppable again.
+		 */
+		if (cc->policy.skip_lazyfree &&
+		    !(vma->vm_flags & VM_DROPPABLE) &&
+		    folio_test_lazyfree(folio) && !pte_dirty(ptent))
+			return SCAN_PAGE_LAZYFREE;
+
+		/*
+		 * A sub-PMD candidate refuses folios of its own order and above:
+		 * collapsing those would gain nothing.  A PMD candidate accepts
+		 * every order up to its own -- the PTE-mapped-THP re-collapse
+		 * class.
+		 */
+		if (folio_order(folio) >= cand->order &&
+		    !is_pmd_order(cand->order))
+			return SCAN_PTE_MAPPED_HUGEPAGE;
+
+		/*
+		 * Exclusive anon only: the expected refcount of a shared folio
+		 * cannot be pinned down without its other mappers' ptls.
+		 * Swapcache membership is fine -- folio_expected_ref_count()
+		 * accounts those references.
+		 */
+		if (folio_maybe_mapped_shared(folio))
+			return SCAN_PAGE_NOT_EXCLUSIVE;
+
+		nr_max = collapse_span_max(ptent, nr_pages - i);
+		for (nr = 1; nr < nr_max; nr++) {
+			pte_t tail = ptep_get(pte + i + nr);
+
+			if (!pte_present(tail) ||
+			    pte_pfn(tail) != pte_pfn(ptent) + nr)
+				break;
+			if (pte_uffd(tail))
+				return SCAN_PTE_UFFD;
+		}
+
+		/*
+		 * Every live mapping of the folio must be this span: the freeze
+		 * is whole-folio, and a live PTE left anywhere else loses to a
+		 * racing zap -- its rmap drop is paired with a folio_put() that
+		 * would underflow the frozen count.  The check is race-free
+		 * under our ptl: in-window PTEs are ours, fork (the only way
+		 * exclusive anon gains mappings) takes mmap_write, and a folio
+		 * whose mappings all sit under this ptl cannot lose one either.
+		 * This also refuses a folio scattered across several spans of
+		 * the window, whose mapcount exceeds any single span.
+		 */
+		if (folio_mapcount(folio) != nr)
+			return SCAN_PAGE_COUNT;
+
+		/*
+		 * Every page of the span must be exclusive: the freeze accounts
+		 * only references it can see, and a non-exclusive page may be
+		 * unshared under us.  collapse_faultin() should have arranged
+		 * this; enforce it here, where it is depended on.
+		 */
+		for (k = 0; k < nr; k++) {
+			if (!PageAnonExclusive(pte_page(ptep_get(pte + i + k))))
+				return SCAN_PAGE_NOT_EXCLUSIVE;
+		}
+
+		i += nr;
+		addr += nr * PAGE_SIZE;
+	}
+
+	return SCAN_SUCCEED;
+}
+
 /*
  * Raise the two barriers on the sources of every candidate: migration entries in
  * their PTEs, then a frozen refcount.  Takes the table's ptl once for the whole
@@ -395,6 +534,44 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 static void collapse_freeze(struct vm_area_struct *vma,
 			    struct collapse_control *cc, pmd_t *pmd)
 {
+	struct mm_struct *mm = vma->vm_mm;
+	pte_t *pte, *table;
+	spinlock_t *ptl;
+	unsigned int i;
+
+	pte = pte_offset_map_lock(mm, pmd, cc->candidates[0].addr, &ptl);
+	if (!pte) {
+		for (i = 0; i < cc->nr_candidates; i++) {
+			struct collapse_candidate *cand = &cc->candidates[i];
+
+			if (cand->state != CAND_SELECTED)
+				continue;
+			cand->state = CAND_SKIPPED;
+			cand->result = SCAN_NO_PTE_TABLE;
+		}
+		return;
+	}
+
+	/*
+	 * Index each candidate from the table base, not relative to
+	 * candidates[0]: a round is not necessarily address-ordered, so
+	 * candidates[0] need not be the lowest.  They all share one table.
+	 */
+	table = pte - pte_index(cc->candidates[0].addr);
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		pte_t *cand_pte = table + pte_index(cand->addr);
+
+		if (cand->state != CAND_SELECTED)
+			continue;
+
+		cand->result = collapse_check_candidate(vma, cc, cand, cand_pte);
+		if (cand->result != SCAN_SUCCEED)
+			cand->state = CAND_SKIPPED;
+	}
+
+	pte_unmap_unlock(pte, ptl);
 }
 
 /*
diff --git a/mm/collapse.h b/mm/collapse.h
index 0d6f77a7233b..747168104a72 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -46,6 +46,7 @@ enum scan_result {
 	SCAN_COPY_MC,
 	SCAN_PAGE_FILLED,
 	SCAN_PAGE_DIRTY_OR_WRITEBACK,
+	SCAN_PAGE_NOT_EXCLUSIVE,
 	SCAN_ALLOC_LIGHT_MISS,
 };
 
-- 
2.54.0

[RFC PATCH 16/57] mm/collapse: freeze the sources behind migration entries

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The acting half.  For every span the checking half accepted:

 1. take a reference on the source folio and lock it;
 2. replace its PTEs with migration entries.  This closes the userspace
    side: faults and GUP-slow now wait on that folio lock, taken before
    the first entry becomes visible;
 3. freeze the folio to folio_expected_ref_count() + 1.  This closes the
    kernel side: folio_try_get() fails.

The two barriers rise in that order because it is reachability order, and
together they are what lets the copy run with no lock at all.

Writeback is the one case the freeze cannot catch, so such a folio is
refused up front.  PG_writeback holds no reference of its own, so the
frozen count is exactly right and the freeze succeeds -- then
folio_end_writeback() takes a reference outright and BUGs on it, or frees
it under the copy.

From the freeze to the putback the round holds every source folio's lock
at once, and folio locks have no global order.  The engine only ever
folio_trylock()s, and unfreezes rather than blocks on refusal, so it is
never the waiting edge of a cycle.

One ranged TLB flush covers everything that froze, before the ptl is
dropped.  Until it completes a CPU with a stale entry could still write a
source through the old mapping, and that path never consults a refcount.

A candidate that cannot finish restores what it displaced, unfreezes,
unlocks and drops out; its neighbours carry on.  The restore needs no
flush of its own: what it puts back is identical to whatever a stale
entry holds.

It restores slot by slot.  The PTEs of one span agree on the PFN and
nothing else, so a partial CoW, or a clear_refs write-protect undone one
page at a time, leaves permissions the first PTE cannot stand for.  Dirty
accumulated over the span goes to the folio instead, the way unmap does.

The displaced values live in a pool sized to a whole table, since one
candidate can displace that much and the byte cap bounds a round, not a
candidate.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 335 ++++++++++++++++++++++++++++++++++++++++++++++++--
 mm/collapse.h |   3 +
 2 files changed, 330 insertions(+), 8 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index c75d91cb9d48..cf2b9b3640ae 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -115,10 +115,20 @@
 	min(COLLAPSE_BATCH_BYTES >> (PAGE_SHIFT + COLLAPSE_MIN_MTHP_ORDER), \
 	    COLLAPSE_TABLE_WINDOWS)
 
+/*
+ * The saved-PTE pool spans a whole table.  The byte cap bounds what a round
+ * holds, but not what one candidate does: a sub-PMD order goes up to
+ * HPAGE_PMD_NR/2 pages -- 256M at order 12 with 64K pages -- and displaces all
+ * of its PTEs in one shot regardless.  So the pool has to fit the largest span
+ * of displaced PTEs a table can hold, which is the table itself.
+ */
+#define COLLAPSE_SAVED_PTES	HPAGE_PMD_NR
+
 /* How far a candidate got, and so what a failure has to undo for it */
 enum collapse_candidate_state {
 	CAND_SELECTED,		/* collected; nothing held on its behalf yet */
 	CAND_SKIPPED,		/* refused; nothing of it left to undo */
+	CAND_FROZEN,		/* sources displaced and frozen */
 };
 
 /*
@@ -136,6 +146,7 @@ struct collapse_candidate {
 	enum scan_result result;
 	struct folio *new_folio;
 	pgtable_t deposit;		/* PMD order: fresh table to deposit */
+	pte_t *saved_ptes;		/* its slice of collapse_control::saved_ptes */
 };
 
 static unsigned long candidate_start(const struct collapse_candidate *cand)
@@ -168,15 +179,20 @@ static unsigned int candidate_offset(const struct collapse_candidate *cand,
 void collapse_control_release(struct collapse_control *cc)
 {
 	kfree(cc->candidates);
+	kfree(cc->saved_ptes);
 	cc->candidates = NULL;
+	cc->saved_ptes = NULL;
 }
 
 int collapse_control_init(struct collapse_control *cc)
 {
 	cc->nr_candidates = 0;
 	cc->candidates = kmalloc_objs(*cc->candidates, COLLAPSE_MAX_CANDIDATES);
-	if (!cc->candidates)
+	cc->saved_ptes = kmalloc_objs(*cc->saved_ptes, COLLAPSE_SAVED_PTES);
+	if (!cc->candidates || !cc->saved_ptes) {
+		collapse_control_release(cc);
 		return -ENOMEM;
+	}
 	return 0;
 }
 
@@ -401,6 +417,99 @@ static unsigned int collapse_span_max(pte_t first, unsigned int max)
 	return min(max, left);
 }
 
+/*
+ * Length of the source span at slot @i, read from the saved PTEs rather than the
+ * table: once frozen the slots hold migration entries, so a rollback re-derives
+ * the freeze's spans from what it displaced.
+ */
+static unsigned int collapse_saved_span_len(struct collapse_candidate *cand,
+					    unsigned int i, unsigned int bound)
+{
+	pte_t first = cand->saved_ptes[i];
+	unsigned int nr, nr_max;
+
+	nr_max = collapse_span_max(first, bound - i);
+	for (nr = 1; nr < nr_max; nr++) {
+		pte_t saved = cand->saved_ptes[i + nr];
+
+		if (pte_none_or_zero(saved) ||
+		    pte_pfn(saved) != pte_pfn(first) + nr)
+			break;
+	}
+	return nr;
+}
+
+/*
+ * Undo a freeze that could not complete: restore the displaced PTE values over
+ * the candidate's migration entries, then unfreeze, unlock and release the
+ * source folios.
+ *
+ * How far the freeze got:
+ *
+ *  - @nr_saved slots were displaced, in PTEs;
+ *  - @nr_frozen of those belong to folios that were also frozen.
+ *
+ * Each slot restores by class: a hole was never modified, a cleared zeropage is
+ * stored back plainly, and a source's saved value goes back as it was.  All are
+ * plain stores -- writing over a non-present entry has no hardware A/D race.
+ *
+ * Slot by slot, not one set_ptes() over the span: the PTEs of one folio need
+ * not agree on more than the PFN, so the first one's permissions are not the
+ * span's.
+ *
+ * Deliberately no TLB flush: the restored translation is identical to anything
+ * a stale TLB entry may hold, so every stale entry is benign.  This reads like
+ * a missing flush; it is not.
+ *
+ * Caller holds the table's ptl -- the same uninterrupted hold the freeze ran
+ * under.
+ */
+static void collapse_unfreeze_candidate(struct mm_struct *mm,
+					struct collapse_candidate *cand,
+					pte_t *pte, unsigned int nr_saved,
+					unsigned int nr_frozen)
+{
+	unsigned long addr = cand->addr;
+	unsigned int i = 0;
+
+	while (i < nr_saved) {
+		pte_t saved = cand->saved_ptes[i];
+		struct folio *folio;
+		unsigned int nr, k;
+
+		if (pte_none(saved)) {
+			/* Hole: nothing was touched */
+			i++;
+			addr += PAGE_SIZE;
+			continue;
+		}
+		if (is_zero_pfn(pte_pfn(saved))) {
+			/* Cleared zeropage: plain non-present -> present store */
+			set_pte_at(mm, addr, pte + i, saved);
+			i++;
+			addr += PAGE_SIZE;
+			continue;
+		}
+
+		folio = pte_folio(saved);
+		nr = collapse_saved_span_len(cand, i, nr_saved);
+
+		for (k = 0; k < nr; k++) {
+			set_pte_at(mm, addr + k * PAGE_SIZE, pte + i + k,
+				   cand->saved_ptes[i + k]);
+		}
+		if (i < nr_frozen) {
+			folio_ref_unfreeze(folio,
+					   folio_expected_ref_count(folio) + 1);
+		}
+		folio_unlock(folio);
+		folio_put(folio);
+
+		i += nr;
+		addr += nr * PAGE_SIZE;
+	}
+}
+
 /*
  * Can this candidate's sources be frozen?  Every slot is checked and nothing is
  * touched, so a refusal costs the round nothing but the walk.
@@ -525,6 +634,193 @@ static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
 	return SCAN_SUCCEED;
 }
 
+/*
+ * Freeze one candidate's sources, span by span, raising both quiescence
+ * barriers in reachability order:
+ *
+ *  1. the span's PTEs become migration entries.  Faults and GUP-slow now wait
+ *     on the source folio's lock, taken before the first entry is visible.
+ *  2. the folio is frozen to its expected reference count, so folio_try_get()
+ *     fails for anyone taking a speculative reference.
+ *
+ * All or nothing: a failure part way through unwinds what it displaced and
+ * leaves the table as it was found.
+ *
+ * Neither barrier deflects a path that takes its reference outright rather
+ * than speculatively.  Such a source has to be refused before the freeze, not
+ * survive it -- see the writeback test below.
+ *
+ * A round holds every source folio's lock at once, from freeze to putback, and
+ * folio locks have no global order.  That cannot deadlock: folio_trylock() is
+ * the engine's only acquisition and a refusal unfreezes instead of blocking, so
+ * the engine is never the waiting edge of a cycle.  Nothing between freeze and
+ * putback waits on anything that could wait on us -- allocation and charging
+ * happen earlier, and the copy only copies.  The install does take the ptl
+ * while holding these folio locks, which is the safe order: a faulter on one of
+ * our migration entries cannot sleep on the folio lock under a spinlock, so it
+ * drops the ptl first.  Do not add a blocking lock or a sleeping allocation
+ * between freeze and putback.
+ *
+ * On entry:
+ *
+ *  - mmap_read is held, and the table's ptl for the whole freeze;
+ *  - collapse_check_candidate() has accepted the candidate under that same ptl
+ *    hold;
+ *  - the round is covered by an mmu_notifier_invalidate_range_start() issued
+ *    outside the ptl.
+ *
+ * collapse_freeze() issues the ranged TLB flush over everything that froze
+ * before dropping the ptl.  No copy may run before it completes.
+ */
+static enum scan_result collapse_freeze_candidate(struct mm_struct *mm,
+		struct collapse_candidate *cand, pte_t *pte)
+{
+	const unsigned int nr_pages = candidate_nr_pages(cand);
+	unsigned int nr_saved = 0, nr_frozen = 0;
+	enum scan_result result;
+	struct folio *folio;
+	unsigned long addr;
+	unsigned int i;
+
+	for (i = 0, addr = cand->addr; i < nr_pages;) {
+		pte_t ptent = ptep_get(pte + i);
+		unsigned int nr, nr_max, k;
+		pte_t rep;
+
+		if (pte_none(ptent)) {
+			/* Hole: nothing to freeze; install verifies it stayed one */
+			cand->saved_ptes[i] = ptent;
+			nr_saved = ++i;
+			addr += PAGE_SIZE;
+			continue;
+		}
+		if (is_zero_pfn(pte_pfn(ptent))) {
+			/*
+			 * Clear the zeropage mapping now, covered by the round's
+			 * ranged flush: overwriting a live PTE at install would
+			 * be a valid->valid transition, breaking arm64's
+			 * break-before-make.  The zeropage has neither rmap nor
+			 * per-map references -- the saved value alone undoes it.
+			 */
+			cand->saved_ptes[i] =
+				ptep_get_and_clear(mm, addr, pte + i);
+			nr_saved = ++i;
+			addr += PAGE_SIZE;
+			continue;
+		}
+
+		folio = pte_folio(ptent);
+
+		/*
+		 * A folio revisited by a second span of this round is already
+		 * ours and frozen at its first span: folio_get() on a zero count
+		 * is a bug, and try-get fails cleanly.  Scrambled layouts
+		 * (mremap) construct this; nothing else can hold a folio frozen
+		 * while its PTE is live under our ptl, so it is not transient.
+		 */
+		if (!folio_try_get(folio)) {
+			result = SCAN_PAGE_COUNT;
+			goto unfreeze;
+		}
+		if (!folio_trylock(folio)) {
+			folio_put(folio);
+			result = SCAN_PAGE_LOCK;
+			goto unfreeze;
+		}
+
+		/*
+		 * Never freeze a folio under writeback.  PG_writeback holds no
+		 * reference of its own -- the swapcache reference keeps the folio
+		 * alive, and everything that would drop it waits for the flag --
+		 * so folio_end_writeback() plain folio_get()s a folio it may
+		 * assume is alive: a BUG on a frozen one, or with
+		 * CONFIG_DEBUG_VM off, a free under our copy.
+		 *
+		 * Unlike every other hazard here, the freeze does not catch it.
+		 * folio_expected_ref_count() counts the swapcache reference, so
+		 * the count is exactly right and the freeze succeeds.  Nor can
+		 * "is it in the swapcache" stand in for this test: that would
+		 * refuse the pages the fault-in pass just swapped in.
+		 *
+		 * Reachable even though writeback starts on an unmapped folio: a
+		 * re-fault from the swapcache maps it back before the bio
+		 * completes, and folio_free_swap() will not drop the cache entry
+		 * under writeback.  Testing once is enough -- writeback starts
+		 * only under the folio lock, which we hold from here through
+		 * putback.
+		 */
+		if (folio_test_writeback(folio)) {
+			folio_unlock(folio);
+			folio_put(folio);
+			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+			goto unfreeze;
+		}
+
+		/* Each slot's own value: a span agrees on the PFN, not the rest */
+		cand->saved_ptes[i] = ptent;
+		nr_max = collapse_span_max(ptent, nr_pages - i);
+		for (nr = 1; nr < nr_max; nr++) {
+			pte_t tail = ptep_get(pte + i + nr);
+
+			if (!pte_present(tail) ||
+			    pte_pfn(tail) != pte_pfn(ptent) + nr)
+				break;
+			cand->saved_ptes[i + nr] = tail;
+		}
+
+		/*
+		 * The clear is the GUP-fast linearization point: a grab landing
+		 * before it elevates the refcount and the freeze below fails
+		 * (the candidate unfreezes); one landing after fails its PTE
+		 * re-read and retries.  Clear and store sit adjacent under one
+		 * uninterrupted ptl hold, batched per span
+		 * (get_and_clear_full_ptes() unfolds contpte), so the transient
+		 * none window is invisible to installers, which all take the ptl.
+		 */
+		rep = get_and_clear_full_ptes(mm, addr, pte + i, nr, 0);
+
+		/*
+		 * Dirty from the clear -- including any the hardware set since
+		 * the reads above -- goes to the folio, the way unmap does,
+		 * rather than onto PTEs that never had it.  Young needs no such
+		 * care: a migration entry drops it either way.
+		 */
+		if (pte_dirty(rep))
+			folio_mark_dirty(folio);
+
+		for (k = 0; k < nr; k++) {
+			pte_t saved = cand->saved_ptes[i + k];
+			swp_entry_t entry;
+			pte_t swp_pte;
+
+			entry = make_readable_migration_entry(pte_pfn(saved));
+			swp_pte = swp_entry_to_pte(entry);
+			if (pte_soft_dirty(saved))
+				swp_pte = pte_swp_mksoft_dirty(swp_pte);
+			set_pte_at(mm, addr + k * PAGE_SIZE, pte + i + k,
+				   swp_pte);
+		}
+		nr_saved = i + nr;
+
+		if (!folio_ref_freeze(folio,
+				      folio_expected_ref_count(folio) + 1)) {
+			result = SCAN_PAGE_COUNT;
+			goto unfreeze;
+		}
+		nr_frozen = nr_saved;
+
+		i += nr;
+		addr += nr * PAGE_SIZE;
+	}
+
+	cand->state = CAND_FROZEN;
+	return SCAN_SUCCEED;
+
+unfreeze:
+	collapse_unfreeze_candidate(mm, cand, pte, nr_saved, nr_frozen);
+	return result;
+}
+
 /*
  * Raise the two barriers on the sources of every candidate: migration entries in
  * their PTEs, then a frozen refcount.  Takes the table's ptl once for the whole
@@ -534,11 +830,15 @@ static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
 static void collapse_freeze(struct vm_area_struct *vma,
 			    struct collapse_control *cc, pmd_t *pmd)
 {
+	unsigned long flush_start = ULONG_MAX, flush_end = 0;
 	struct mm_struct *mm = vma->vm_mm;
 	pte_t *pte, *table;
 	spinlock_t *ptl;
 	unsigned int i;
 
+	/* Pending per-CPU folio batches hold references that fail the freeze */
+	lru_add_drain();
+
 	pte = pte_offset_map_lock(mm, pmd, cc->candidates[0].addr, &ptl);
 	if (!pte) {
 		for (i = 0; i < cc->nr_candidates; i++) {
@@ -562,15 +862,27 @@ static void collapse_freeze(struct vm_area_struct *vma,
 	for (i = 0; i < cc->nr_candidates; i++) {
 		struct collapse_candidate *cand = &cc->candidates[i];
 		pte_t *cand_pte = table + pte_index(cand->addr);
+		enum scan_result result;
 
 		if (cand->state != CAND_SELECTED)
 			continue;
 
-		cand->result = collapse_check_candidate(vma, cc, cand, cand_pte);
-		if (cand->result != SCAN_SUCCEED)
+		result = collapse_check_candidate(vma, cc, cand, cand_pte);
+		if (result == SCAN_SUCCEED)
+			result = collapse_freeze_candidate(mm, cand, cand_pte);
+
+		cand->result = result;
+		if (result != SCAN_SUCCEED) {
 			cand->state = CAND_SKIPPED;
+			continue;
+		}
+
+		flush_start = min(flush_start, candidate_start(cand));
+		flush_end = max(flush_end, candidate_end(cand));
 	}
 
+	if (flush_end)
+		flush_tlb_range(vma, flush_start, flush_end);
 	pte_unmap_unlock(pte, ptl);
 }
 
@@ -698,7 +1010,7 @@ static void collapse_provision(struct mm_struct *mm,
 		struct collapse_candidate *cand = &cc->candidates[i];
 		enum scan_result result;
 
-		if (cand->state != CAND_SELECTED || cand->new_folio)
+		if (cand->state != CAND_FROZEN || cand->new_folio)
 			continue;
 
 		result = collapse_alloc(mm, cc, cand, gfp);
@@ -1200,13 +1512,14 @@ static bool collapse_run_batch(struct mm_struct *mm, unsigned long pmd_addr,
  * a single candidate is above the cap all by itself once a PMD is (512M with
  * 64K pages), and refusing it would collapse nothing at all.
  */
-static bool collapse_batch_full(struct collapse_control *cc,
+static bool collapse_batch_full(struct collapse_control *cc, unsigned int slots,
 				unsigned long bytes, unsigned int order)
 {
 	if (!cc->nr_candidates)
 		return false;
 
 	return cc->nr_candidates == COLLAPSE_MAX_CANDIDATES ||
+	       slots + (1U << order) > COLLAPSE_SAVED_PTES ||
 	       bytes + (PAGE_SIZE << order) > COLLAPSE_BATCH_BYTES;
 }
 
@@ -1215,7 +1528,8 @@ static bool collapse_batch_full(struct collapse_control *cc,
  * previous round's values, so every field is set here.
  */
 static void collapse_add_candidate(struct collapse_control *cc,
-				   unsigned long addr, unsigned int order)
+				   unsigned long addr, unsigned int order,
+				   pte_t *saved_ptes)
 {
 	struct collapse_candidate *cand;
 
@@ -1232,6 +1546,7 @@ static void collapse_add_candidate(struct collapse_control *cc,
 	cand->result = SCAN_FAIL;
 	cand->new_folio = NULL;
 	cand->deposit = NULL;
+	cand->saved_ptes = saved_ptes;
 }
 
 /*
@@ -1246,6 +1561,7 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
 	unsigned int offset, order;
 	unsigned long bytes = 0;
+	unsigned int slots = 0;
 	bool pending = false;
 	bool cont = true;
 
@@ -1256,7 +1572,7 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 		if (!pending)
 			pending = collapse_next_candidate(cc, &offset, &order);
 
-		if (!pending || collapse_batch_full(cc, bytes, order)) {
+		if (!pending || collapse_batch_full(cc, slots, bytes, order)) {
 			/*
 			 * Selection is exhausted and the round is empty: the
 			 * range is done.  Without this a flush of an empty
@@ -1267,6 +1583,7 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 				break;
 
 			cont = collapse_run_batch(mm, pmd_addr, cc);
+			slots = 0;
 			bytes = 0;
 			continue;
 		}
@@ -1276,8 +1593,10 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 		 * collecting costs nothing but the array slot.  A candidate the
 		 * full round could not take is kept pending for the next one.
 		 */
-		collapse_add_candidate(cc, pmd_addr + offset * PAGE_SIZE, order);
+		collapse_add_candidate(cc, pmd_addr + offset * PAGE_SIZE, order,
+				       cc->saved_ptes + slots);
 
+		slots += 1U << order;
 		bytes += PAGE_SIZE << order;
 		pending = false;
 	}
diff --git a/mm/collapse.h b/mm/collapse.h
index 747168104a72..3256c45ee228 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -137,6 +137,9 @@ struct collapse_control {
 	 */
 	unsigned long batch_start;
 	unsigned long batch_end;
+
+	/* PTE values the round displaced, carved up between its candidates */
+	pte_t *saved_ptes;
 };
 
 static inline int collapse_test_exit(struct mm_struct *mm)
-- 
2.54.0

[RFC PATCH 17/57] mm/collapse: copy the sources into the destinations

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the copy.  For each candidate that both froze and has a
destination folio, the page mapped at slot k of its window becomes page k
of that folio.  A slot with no source -- a hole, or a zeropage the freeze
cleared -- is zero-filled instead.

Neither condition implies the other.  A reserve before the lock can leave
a folio on a candidate that then fails to freeze, and the provision
inside the window can decline one for a candidate that froze.

Nothing else is carried across: the destination's PTE bits are not
derived from the sources, and the install rebuilds them the way a fault
would.

A machine check reading a source is the only failure the copy can report,
and only where the architecture provides an MC-safe copy.  The candidate
records SCAN_COPY_MC and the install undoes it, that being where the ptl
the undoing needs is held.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/mm/collapse.c b/mm/collapse.c
index cf2b9b3640ae..a3882d897d11 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1040,6 +1040,46 @@ static void collapse_provision(struct mm_struct *mm,
 static void collapse_copy(struct vm_area_struct *vma,
 			  struct collapse_control *cc)
 {
+	unsigned int i;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		const unsigned int nr_pages = candidate_nr_pages(cand);
+		unsigned long addr = cand->addr;
+		unsigned int k;
+
+		/* A folio does not imply a freeze: reserve runs before the lock */
+		if (cand->state != CAND_FROZEN)
+			continue;
+
+		/* A freeze does not imply a folio: provision may have declined */
+		if (!cand->new_folio)
+			continue;
+
+		/* Each source lands where its address puts it: slot k, page k */
+		for (k = 0; k < nr_pages; k++, addr += PAGE_SIZE) {
+			struct page *dst = folio_page(cand->new_folio, k);
+			struct page *src;
+
+			/* No source: a hole, or a zeropage the freeze cleared */
+			if (pte_none_or_zero(cand->saved_ptes[k])) {
+				clear_user_highpage(dst, addr);
+				continue;
+			}
+
+			src = pte_page(cand->saved_ptes[k]);
+
+			/*
+			 * A machine check on a source is the only way this
+			 * fails, and the install is what undoes the candidate:
+			 * that is where the ptl the undoing needs is held.
+			 */
+			if (copy_mc_user_highpage(dst, src, addr, vma)) {
+				cand->result = SCAN_COPY_MC;
+				break;
+			}
+		}
+	}
 }
 
 /* Publish each destination folio in place of the sources it replaces */
-- 
2.54.0

[RFC PATCH 18/57] mm/collapse: install the destinations at PTE level

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the install for sub-PMD candidates, and with it the two things
every install needs: a verify, and an abort.

The verify decides whether the window is still the round's to publish.
Under the ptl, every slot that had a source must still hold the round's
migration entry, and every slot that had none must still be none.  A hole
some fault refilled is not the round's to overwrite.  Verify and install
share one ptl hold, so a verified candidate cannot lose a slot before it
is published.

The abort undoes a frozen candidate that cannot be published: the copy
took a machine check, the provision pass could not spare it a
destination, or the verify refused it.  Slot by slot:

 - a slot still holding the round's migration entry is restored from the
   saved value.  No TLB flush: the translation is identical.
 - a slot that does not is left exactly as found, since restoring it
   would resurrect memory the user zapped.  Its rmap is dropped here,
   because the zapper fixed up rss for what it cleared but could not drop
   the rmap a frozen source keeps.

The table itself going away is that same rule at whole-table scale.  A
racing MADV_DONTNEED over the whole table, and the empty-table reclaim
behind it, can free the table between freeze and install.  Every slot
then reads as foreign, and what is left to undo is exactly the half of
the teardown a zapper cannot do for a frozen source.

Publishing is the fault path's own helper, so the destination gets rmap
and LRU insertion before its PTEs, fresh bits from the VMA, contpte
painting, and no TLB flush -- every transition is non-present to present.
Slots that had no source become anon memory no zap ever accounted for, so
rss is corrected by hand.

PMD-order candidates need a terminal layer of their own, a stub here.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 226 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 226 insertions(+)

diff --git a/mm/collapse.c b/mm/collapse.c
index a3882d897d11..842adc30aeb0 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -129,6 +129,7 @@ enum collapse_candidate_state {
 	CAND_SELECTED,		/* collected; nothing held on its behalf yet */
 	CAND_SKIPPED,		/* refused; nothing of it left to undo */
 	CAND_FROZEN,		/* sources displaced and frozen */
+	CAND_INSTALLED,		/* the destination is mapped */
 };
 
 /*
@@ -1082,10 +1083,235 @@ static void collapse_copy(struct vm_area_struct *vma,
 	}
 }
 
+/*
+ * Undo one frozen slot: restore the saved PTE if our migration entry is still
+ * there, or drop the rmap the freeze took if a racing zap already replaced it.
+ * Returns true when the slot was zapped -- its mapping reference is then ours to
+ * release.
+ */
+static bool collapse_abort_slot(struct vm_area_struct *vma, struct folio *folio,
+				pte_t *slot, unsigned long addr, pte_t saved)
+{
+	/* No table left: the slot cannot still be holding our entry */
+	if (slot) {
+		softleaf_t entry = softleaf_from_pte(ptep_get(slot));
+
+		if (softleaf_is_migration(entry) &&
+		    softleaf_to_pfn(entry) == pte_pfn(saved)) {
+			set_pte_at(vma->vm_mm, addr, slot, saved);
+			return false;
+		}
+	}
+	folio_remove_rmap_pte(folio, pte_page(saved), vma);
+	return true;
+}
+
+/*
+ * Abort one frozen candidate at install time: it took a machine check during the
+ * copy, or some of its slots no longer hold our migration entries.  mmap_read
+ * (held freeze..putback) blocks fork, mremap and munmap, and faults wait on the
+ * migration entries -- but madvise-class operations run under mmap_read too, so a
+ * concurrent MADV_DONTNEED may have zapped frozen slots, and a fault may have
+ * refilled a zapped one.
+ *
+ * Slots still holding our entries are restored from the saved values (no TLB
+ * flush: identical translation).  Foreign slots are left exactly as found --
+ * restoring them would resurrect memory the user zapped -- but their rmap is
+ * dropped here: the zapper fixed up rss for the slots it cleared, yet could not
+ * drop the rmap a frozen source keeps, unlike a migrating one, which unmaps at
+ * freeze time.  Slots with no source follow the same rule with no rmap to drop: a
+ * cleared zeropage is restored only while its slot is still none, and a hole was
+ * never touched at all.
+ *
+ * @pte is NULL when the table itself is gone: a racing whole-table MADV_DONTNEED
+ * zapped every entry, frozen slots included, and the empty-table reclaim
+ * (CONFIG_PT_RECLAIM) freed it, clearing the pmd under the pmd lock and the pte
+ * ptl, neither of which excludes it between our freeze and install.  Every slot
+ * then reads as foreign, which is exactly right: nothing of ours survives to
+ * restore or verify, and what is left is the half of the teardown the zapper
+ * cannot perform for a frozen source -- the kept rmap, the freeze, the folio
+ * locks and the references.  The caller holds no page-table lock in that case,
+ * there being no table to lock.
+ */
+static void collapse_abort_candidate(struct vm_area_struct *vma,
+				     struct collapse_candidate *cand,
+				     pte_t *pte)
+{
+	const unsigned int nr_pages = candidate_nr_pages(cand);
+	struct mm_struct *mm = vma->vm_mm;
+	unsigned long addr = cand->addr;
+	unsigned int i, nr;
+
+	for (i = 0; i < nr_pages; i += nr, addr += nr * PAGE_SIZE) {
+		pte_t saved = cand->saved_ptes[i];
+		unsigned int k, nr_dropped;
+		struct folio *folio;
+
+		nr = 1;		/* skip stride; a span overrides it */
+		if (pte_none(saved))
+			continue;
+		if (is_zero_pfn(pte_pfn(saved))) {
+			if (pte && pte_none(ptep_get(pte + i)))
+				set_pte_at(mm, addr, pte + i, saved);
+			continue;
+		}
+
+		folio = pte_folio(saved);
+		nr = collapse_saved_span_len(cand, i, nr_pages);
+
+		/*
+		 * Unfreeze before any rmap drop: rmap removal munlocks under
+		 * VM_LOCKED, and munlock_folio() takes a reference a frozen folio
+		 * forbids.  The expected count still holds every slot's mapping
+		 * reference; restored slots keep theirs, and the zapped slots'
+		 * references become ours to drop with the rmap.
+		 */
+		folio_ref_unfreeze(folio, folio_expected_ref_count(folio) + 1);
+
+		nr_dropped = 0;
+		for (k = 0; k < nr; k++) {
+			pte_t *slot = pte ? pte + i + k : NULL;
+
+			nr_dropped += collapse_abort_slot(vma, folio, slot,
+							  addr + k * PAGE_SIZE,
+							  cand->saved_ptes[i + k]);
+		}
+
+		folio_unlock(folio);
+		folio_put_refs(folio, nr_dropped + 1);
+	}
+
+	/*
+	 * Not installed; collapse_finish() releases the destination, which has to
+	 * wait for the ptl to be dropped.
+	 */
+	cand->state = CAND_SKIPPED;
+}
+
+/*
+ * Nothing may have shifted under the round: every source slot must still hold our
+ * migration entry, and every slot with no source must still be none -- the freeze
+ * cleared the zeropage ones, so both read as none by then, and a slot some fault
+ * has refilled, with a page or with a zeropage, is not ours to overwrite.
+ * @nr_populated returns how many source-less slots the install is about to make
+ * present, which is rss no zap ever accounted for.
+ */
+static bool collapse_verify_candidate(struct collapse_candidate *cand,
+				      pte_t *pte, unsigned int *nr_populated)
+{
+	const unsigned int nr_pages = candidate_nr_pages(cand);
+	unsigned int k, populated = 0;
+
+	for (k = 0; k < nr_pages; k++) {
+		pte_t live = ptep_get(pte + k);
+		softleaf_t entry;
+
+		if (pte_none_or_zero(cand->saved_ptes[k])) {
+			if (!pte_none(live))
+				return false;
+			populated++;
+			continue;
+		}
+
+		entry = softleaf_from_pte(live);
+		if (!softleaf_is_migration(entry) ||
+		    softleaf_to_pfn(entry) != pte_pfn(cand->saved_ptes[k]))
+			return false;
+	}
+	*nr_populated = populated;
+	return true;
+}
+
+/*
+ * The PMD terminal layer: verify, detach the table, deposit a fresh one and
+ * install the leaf, as one atomic section under the pmd lock.  A pmd_none window
+ * never exists -- faults stay held at pte level by the migration entries
+ * throughout -- which is what lets PMD collapse run under mmap_read like the rest
+ * of the engine.
+ */
+static void collapse_install_pmd(struct vm_area_struct *vma,
+				 struct collapse_control *cc, pmd_t *pmd)
+{
+}
+
 /* Publish each destination folio in place of the sources it replaces */
 static void collapse_install(struct vm_area_struct *vma,
 			     struct collapse_control *cc, pmd_t *pmd)
 {
+	struct mm_struct *mm = vma->vm_mm;
+	pte_t *pte, *table;
+	spinlock_t *ptl;
+	unsigned int i;
+
+	if (is_pmd_order(cc->candidates[0].order)) {
+		/* A PMD candidate fills the slot pool: always alone */
+		VM_WARN_ON_ONCE(cc->nr_candidates != 1);
+		collapse_install_pmd(vma, cc, pmd);
+		return;
+	}
+
+	pte = pte_offset_map_lock(mm, pmd, cc->candidates[0].addr, &ptl);
+	if (!pte) {
+		/*
+		 * Table gone under us (see collapse_abort_candidate() on @pte).
+		 * Tear down every frozen candidate -- stranding them would leak
+		 * frozen, locked sources.
+		 */
+		for (i = 0; i < cc->nr_candidates; i++) {
+			struct collapse_candidate *cand = &cc->candidates[i];
+
+			if (cand->state != CAND_FROZEN)
+				continue;
+
+			cand->result = SCAN_NO_PTE_TABLE;
+			collapse_abort_candidate(vma, cand, NULL);
+		}
+		return;
+	}
+	table = pte - pte_index(cc->candidates[0].addr);
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		pte_t *cand_pte = table + pte_index(cand->addr);
+		unsigned int nr_populated;
+
+		if (cand->state != CAND_FROZEN)
+			continue;
+
+		if (cand->result != SCAN_SUCCEED) {
+			/* Machine check during the copy */
+			collapse_abort_candidate(vma, cand, cand_pte);
+			continue;
+		}
+
+		/* No destination: the provision pass could not spare one */
+		if (!cand->new_folio) {
+			collapse_abort_candidate(vma, cand, cand_pte);
+			continue;
+		}
+
+		if (!collapse_verify_candidate(cand, cand_pte, &nr_populated)) {
+			cand->result = SCAN_PTE_NON_PRESENT;
+			collapse_abort_candidate(vma, cand, cand_pte);
+			continue;
+		}
+
+		/*
+		 * The smp_wmb() in __folio_mark_uptodate() orders the copied
+		 * data before the set_ptes() that publishes it.
+		 */
+		__folio_mark_uptodate(cand->new_folio);
+		map_anon_folio_pte_nopf(cand->new_folio, cand_pte, vma,
+					cand->addr, /*uffd_wp=*/ false);
+
+		/* Slots with no source gain anon memory that no zap accounted */
+		if (nr_populated)
+			add_mm_counter(mm, MM_ANONPAGES, nr_populated);
+		cand->new_folio = NULL;	/* ownership: the mappings */
+		cand->state = CAND_INSTALLED;
+	}
+
+	pte_unmap_unlock(pte, ptl);
 }
 
 /*
-- 
2.54.0

[RFC PATCH 19/57] mm/collapse: install a PMD leaf as the terminal layer

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the PMD install.  Under the pmd lock, with the pte ptl nested
inside it: verify, detach the table with pmdp_collapse_flush(), deposit a
fresh one and map the leaf.

That is one atomic section, so no pmd_none() window ever exists: faults
stay held down at pte level by the migration entries throughout.  It is
what lets PMD collapse run under mmap_read like everything else here.

Two things force that nesting, which is the one the tree already uses to
reinstall a table.  A racing zap of a frozen entry takes the pte ptl, so
the verify has to hold it.  And the table must not come apart between
verify and detach, which is the pmd lock's job.

Nothing leaves the section early, aborts included.  An abort only
restores PTEs and would need no pmd-level exclusion of its own, except
that its pte pointer came from pte_offset_map_rw_nolock(), whose caller
must establish that the pmd is stable.

The deposited table is the freshly allocated one, never the table just
detached.  A deposited table has to be quiescent, because whoever
withdraws it frees it immediately with nothing to hold a lockless walker
off first, and a table that has never been reachable is quiescent by
construction.

The detached one is not: GUP-fast and RCU pte walks that read the old PMD
may still be inside it, and on broadcast-TLBI architectures the flush
expels nobody.  Quiescing it would need an IPI, which has nowhere to go
here -- outside the pmd lock it opens the pmd_none() window this design
does not have, inside it is a broadcast under a spinlock.  So the
detached table goes to pte_free_defer(), which holds the free until those
walkers finish.  One transient table page per PMD collapse is the cost.

No anon_vma_lock_write() is taken, unlike the mechanism being replaced:

 - rmap walks on the sources are unreachable, their refcounts frozen and
   their folio locks held from freeze to putback;
 - non-rmap pte walkers see migration entries;
 - pmd-level observers see either the old table or the leaf, never an
   intermediate;
 - fork, mremap and munmap take mmap_write, which the mmap_read held here
   excludes.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 118 insertions(+)

diff --git a/mm/collapse.c b/mm/collapse.c
index 842adc30aeb0..ab7476471b8d 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1232,6 +1232,124 @@ static bool collapse_verify_candidate(struct collapse_candidate *cand,
 static void collapse_install_pmd(struct vm_area_struct *vma,
 				 struct collapse_control *cc, pmd_t *pmd)
 {
+	struct collapse_candidate *cand = &cc->candidates[0];
+	struct mm_struct *mm = vma->vm_mm;
+	spinlock_t *pmd_ptl, *pte_ptl;
+	pgtable_t old_table = NULL;
+	unsigned int nr_populated;
+	pmd_t old_pmd, pmdval;
+	pte_t *pte;
+
+	if (cand->state != CAND_FROZEN)
+		return;
+
+	/* No destination: the provision pass could not spare one */
+	if (!cand->new_folio) {
+		pte = pte_offset_map_lock(mm, pmd, cand->addr, &pte_ptl);
+		collapse_abort_candidate(vma, cand, pte);
+		if (pte)
+			pte_unmap_unlock(pte, pte_ptl);
+		return;
+	}
+
+	/*
+	 * The pte ptl nests inside the pmd lock, the nesting the tree already
+	 * uses for reinstalling a table: a racing zap of a frozen entry takes
+	 * the pte ptl, so the verify must hold it, and the table must not come
+	 * apart between verify and detach.  pmd_same() rechecks are unnecessary,
+	 * the pmd lock being held across the whole section.
+	 */
+	pmd_ptl = pmd_lock(mm, pmd);
+	pte = pte_offset_map_rw_nolock(mm, pmd, cand->addr, &pmdval, &pte_ptl);
+	if (!pte) {
+		/* Table gone under us; see collapse_abort_candidate() on @pte */
+		spin_unlock(pmd_ptl);
+		cand->result = SCAN_NO_PTE_TABLE;
+		collapse_abort_candidate(vma, cand, NULL);
+		return;
+	}
+	if (pte_ptl != pmd_ptl)
+		spin_lock_nested(pte_ptl, SINGLE_DEPTH_NESTING);
+
+	/*
+	 * Every exit is inside that section, the aborts as much as the install.
+	 * An abort needs no pmd-level exclusion of its own; it only restores
+	 * PTEs.  But the table it works on came from pte_offset_map_rw_nolock(),
+	 * which leaves its caller to establish that the pmd is stable, and the
+	 * held pmd lock is what does that here.
+	 */
+	if (cand->result != SCAN_SUCCEED) {
+		/* Machine check during the copy */
+		collapse_abort_candidate(vma, cand, pte);
+		goto out_unlock;
+	}
+
+	if (!collapse_verify_candidate(cand, pte, &nr_populated)) {
+		cand->result = SCAN_PTE_NON_PRESENT;
+		collapse_abort_candidate(vma, cand, pte);
+		goto out_unlock;
+	}
+
+	/*
+	 * Nothing fallible sits past here.  No anon_vma_lock_write either: rmap
+	 * walks on the sources are unreachable -- refcounts frozen, folio locks
+	 * held from freeze to putback -- non-rmap pte walkers see migration
+	 * entries, pmd-level observers see the old table or the leaf and never an
+	 * intermediate, and fork, mremap and munmap take mmap_write, which our
+	 * mmap_read excludes.
+	 *
+	 * The flush inside pmdp_collapse_flush() is the round's second over this
+	 * range: the freeze displaced every leaf here and flushed before dropping
+	 * the ptl, and the verify above proved nothing has been mapped since.
+	 * What it covers is the paging-structure caches -- a CPU may still hold
+	 * the pmd-to-table link, for a table that is about to be freed -- which
+	 * is why the helper shoots down a pte range rather than a pmd.
+	 */
+	old_pmd = pmdp_collapse_flush(vma, cand->addr, pmd);
+	old_table = pmd_pgtable(old_pmd);
+
+	/*
+	 * The smp_wmb() in __folio_mark_uptodate() orders the copied data before
+	 * the install below publishes it.
+	 */
+	__folio_mark_uptodate(cand->new_folio);
+
+	/*
+	 * Deposit a freshly allocated table, not the one just detached: a
+	 * deposited table has to be quiescent, because whoever withdraws it frees
+	 * it immediately (zap_huge_pmd()) with nothing to hold a lockless walker
+	 * off first.  A table that has never been reachable is quiescent by
+	 * construction, which is why collapse_alloc() secured one.
+	 *
+	 * The detached table is not.  GUP-fast and RCU pte walks that read the
+	 * old PMD before pmdp_collapse_flush() may still be inside it, and on
+	 * broadcast-TLBI arches that flush expels nobody.  Quiescing it would
+	 * take an IPI (tlb_remove_table_sync_one()), which has nowhere to go
+	 * here: outside the pmd lock it opens a pmd_none window a fault can fill,
+	 * inside it is a broadcast under a spinlock.  So it goes to
+	 * pte_free_defer(), which holds the free until those walkers finish, as
+	 * retract_page_tables() does.  One transient table page per PMD collapse
+	 * is what that costs.
+	 */
+	pgtable_trans_huge_deposit(mm, pmd, cand->deposit);
+	map_anon_folio_pmd_nopf(cand->new_folio, pmd, vma, cand->addr);
+
+	/* Slots with no source gain anon memory that no zap accounted */
+	if (nr_populated)
+		add_mm_counter(mm, MM_ANONPAGES, nr_populated);
+	cand->deposit = NULL;
+	cand->new_folio = NULL;	/* ownership: the mapping */
+	cand->state = CAND_INSTALLED;
+
+out_unlock:
+	if (pte_ptl != pmd_ptl)
+		spin_unlock(pte_ptl);
+	pte_unmap(pte);
+	spin_unlock(pmd_ptl);
+
+	/* The deposit balanced the detached table, so the count is already right */
+	if (old_table)
+		pte_free_defer(mm, old_table);
 }
 
 /* Publish each destination folio in place of the sources it replaces */
-- 
2.54.0

[RFC PATCH 20/57] mm/collapse: put the sources back

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the putback: for every installed candidate, lower the barriers
the freeze raised, span by span.

This is also what wakes the faulters the collapse held up.  They sleep on
a source folio's lock; once it is dropped they refault and find present
PTEs pointing at the new folio.

The order within a span is important:

 - Unfreeze first.  Rmap removal munlocks under VM_LOCKED, and
   munlock_folio() takes a reference a frozen folio forbids.

 - Then drop the rmap.  Until it is gone the expected count still holds
   the span's mapping references; afterwards they belong to the round, so
   every folio_remove_rmap_ptes() is paired with a folio_put_refs() for
   the same slots.

 - Then unlock, which is the wake.  Holding the lock until here keeps
   lock-taking rmap walkers out, and the window it leaves -- a live folio
   with no PTEs -- is one any teardown of a mapped folio passes through.

 - Drop the references strictly last, the round's included.  Waiters wait
   without a reference of their own, so the round's has to outlive the
   unlock.

The stale swapcache entry goes too: the copy has replaced what it
described.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/mm/collapse.c b/mm/collapse.c
index ab7476471b8d..f65f413339bf 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1439,6 +1439,56 @@ static void collapse_install(struct vm_area_struct *vma,
 static void collapse_putback(struct vm_area_struct *vma,
 			     struct collapse_control *cc)
 {
+	unsigned int i;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+		const unsigned int nr_pages = candidate_nr_pages(cand);
+		unsigned int k = 0;
+
+		if (cand->state != CAND_INSTALLED)
+			continue;
+
+		while (k < nr_pages) {
+			struct folio *folio;
+			unsigned int nr;
+
+			/* A slot with no source has nothing to put back */
+			if (pte_none_or_zero(cand->saved_ptes[k])) {
+				k++;
+				continue;
+			}
+
+			folio = pte_folio(cand->saved_ptes[k]);
+			nr = collapse_saved_span_len(cand, k, nr_pages);
+
+			/*
+			 * Unfreeze before the rmap drop: rmap removal munlocks
+			 * under VM_LOCKED, and munlock_folio() takes a reference
+			 * a frozen folio forbids.  The expected count still
+			 * holds the span's mapping references; once the rmap is
+			 * gone they are ours to drop, so every
+			 * folio_remove_rmap_ptes() is paired with a
+			 * folio_put_refs() for the same slots.  The folio lock
+			 * is held until the wake below, so lock-taking rmap
+			 * walkers stay excluded, and the stale-rmap window this
+			 * leaves -- live folio, no PTEs -- is one any teardown of
+			 * a mapped folio passes through.
+			 */
+			folio_ref_unfreeze(folio,
+					   folio_expected_ref_count(folio) + 1);
+			folio_remove_rmap_ptes(folio,
+					       pte_page(cand->saved_ptes[k]),
+					       nr, vma);
+			folio_unlock(folio);
+
+			/* The copy replaced it; drop the stale swap entry */
+			free_swap_cache(folio);
+			folio_put_refs(folio, nr + 1);
+
+			k += nr;
+		}
+	}
 }
 
 /*
-- 
2.54.0

[RFC PATCH 21/57] mm/collapse: settle whatever the round reached

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the last pass.  A candidate that never froze is recorded as given
up on, with whatever result ended the round.  Anything no pass took
ownership of -- a destination folio, a table meant for deposit -- goes
back.  The count of installed candidates is what the round reports.

Holding no lock here is the point.  Every refusal before this happens
under a page-table lock: the freeze unwinds under the ptl it took, and
the install aborts under that ptl or the pmd lock.

Dropping the last reference to a folio, and the memcg uncharge behind it,
is not spinlock work.  So a refused candidate keeps its folio and its
table until the round is over, and this is where they are released.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index f65f413339bf..2da1f8ddcca8 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1500,7 +1500,30 @@ static unsigned int collapse_finish(struct mm_struct *mm,
 				    struct collapse_control *cc,
 				    enum scan_result result)
 {
-	return 0;
+	unsigned int i, nr_installed = 0;
+
+	for (i = 0; i < cc->nr_candidates; i++) {
+		struct collapse_candidate *cand = &cc->candidates[i];
+
+		/* Never froze: the round gave up before it got that far */
+		if (cand->state == CAND_SELECTED) {
+			cand->state = CAND_SKIPPED;
+			cand->result = result;
+		}
+
+		if (cand->new_folio) {
+			folio_put(cand->new_folio);
+			cand->new_folio = NULL;
+		}
+		if (cand->deposit) {
+			pte_free(mm, cand->deposit);
+			cand->deposit = NULL;
+		}
+		if (cand->state == CAND_INSTALLED)
+			nr_installed++;
+	}
+
+	return nr_installed;
 }
 
 /*
-- 
2.54.0

[RFC PATCH 22/57] mm/collapse: walk a table with a selection cursor

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the half of selection that emits candidates: a cursor over the
table, handing out the largest window that fits where it stands.

Two things bound the order at any point.  A huge page has to be naturally
aligned, so the cursor's own offset caps it -- at offset 4 nothing above
order 2 can start -- and the largest enabled order caps it too.

A window qualifies when enough of it is eligible: the scan's bits counted
over the window, against the max_ptes_none limit for that order.  One
that does not qualify drops to the next enabled order below, which need
not be half of it, since a sparse set of enabled sizes may skip several.
When no smaller order is left, the cursor steps over the region.

Only the scan's bitmap is read, so a clear bit is either a hole or a PTE
the scan disqualified.  Occupancy here means what a collapse could use,
not what is present.

Non-present PTEs the scan accepted are the exception.  They are counted
apart, in cc->scan_unmapped, and added back only for a PMD candidate,
which faults them in; a smaller window leaves them as holes, sub-PMD
collapse not reading swap.

The cursor advances at emission and never rewinds.  A round is collected
before it is run, so within a round every attempt is assumed to succeed.
Nothing here gives a refused region a second chance.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 126 ++++++++++++++++++++++++++++++++++++++++++++++++
 mm/collapse.h   |  10 ++++
 mm/khugepaged.c |   2 +-
 3 files changed, 137 insertions(+), 1 deletion(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 2da1f8ddcca8..258bb9cc32c5 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1841,6 +1841,7 @@ static enum scan_result collapse_scan_table(struct vm_area_struct *vma,
 	if (result != SCAN_SUCCEED)
 		cc->select_orders &= ~BIT(HPAGE_PMD_ORDER);
 
+	cc->scan_unmapped = unmapped;
 	return result;
 }
 
@@ -1852,6 +1853,7 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
 	nodes_clear(cc->alloc_nmask);
 
 	cc->select_orders = 0;
+	cc->scan_unmapped = 0;
 	cc->nr_collapsed = 0;
 }
 
@@ -1896,10 +1898,116 @@ collapse_scan_anon_pmd(struct vm_area_struct *vma, unsigned long start,
 	return cc->scan_refusal;
 }
 
+/*
+ * Selection cuts the table into candidate windows and feeds them to rounds.  A
+ * window is cut at the largest enabled order that fits and qualifies -- the PMD
+ * order, when the whole table qualified -- and a region that does not qualify is
+ * probed at the next enabled order below, which need not be half of it: a sparse
+ * set of enabled sizes may skip several.
+ *
+ * Only cc->eligible_ptes is read, so a clear bit is either a hole or a PTE the
+ * scan disqualified: a window's occupancy is what a collapse could use, not what
+ * is present.
+ */
+
+/*
+ * Largest order a window may be rooted at: the largest enabled one.
+ * select_orders is fixed for the table, and the caller checked it is not empty,
+ * so this is well-defined for the whole walk.
+ */
+static unsigned int collapse_root_order(struct collapse_control *cc)
+{
+	return __fls(cc->select_orders);
+}
+
+/*
+ * The next enabled order below @order, or 0 when there is none.  select_orders
+ * never carries an order below COLLAPSE_MIN_MTHP_ORDER -- THP_ORDERS_ALL_ANON
+ * masks orders 0 and 1 -- so __fls() honours that floor by itself.  Order 0 has
+ * no bits below it to mask and has to answer 0 outright: a walk that ascended
+ * instead would emit a window at an offset it is not aligned for.
+ */
+static unsigned int collapse_lower_order(struct collapse_control *cc,
+					 unsigned int order)
+{
+	unsigned long lower;
+
+	if (!order)
+		return 0;
+
+	lower = cc->select_orders & GENMASK(order - 1, 0);
+	return lower ? __fls(lower) : 0;
+}
+
 /* Point the selection cursor at [start, end) of the table, in PTE offsets */
 static void collapse_selection_init(struct collapse_control *cc,
 				    unsigned int start, unsigned int end)
 {
+	cc->select_start = start;
+	cc->select_end = end;
+	cc->select_offset = start;
+	cc->select_order = min(max_order_from_offset(start),
+			       collapse_root_order(cc));
+}
+
+/*
+ * Advance past the region [select_offset, select_offset + nr_ptes) and determine
+ * the highest order that can be attempted next.  Since huge pages must be
+ * naturally aligned, it is limited by the alignment of the new offset: after an
+ * order-2 mTHP at offset 0 the offset becomes 4, and __ffs(4) == 2, so the next
+ * attempt starts at order 2.
+ */
+static void collapse_selection_advance(struct collapse_control *cc,
+				       unsigned int nr_ptes)
+{
+	cc->select_offset += nr_ptes;
+	cc->select_order = min(max_order_from_offset(cc->select_offset),
+			       collapse_root_order(cc));
+}
+
+/*
+ * The window at the cursor did not qualify.  Drop to the next smaller enabled
+ * order over the same region, or -- when no smaller order remains -- give the
+ * region up and advance the cursor past it.
+ */
+static void collapse_selection_reject(struct collapse_control *cc)
+{
+	unsigned int lower = collapse_lower_order(cc, cc->select_order);
+
+	if (lower)
+		cc->select_order = lower;
+	else
+		collapse_selection_advance(cc, 1U << cc->select_order);
+}
+
+/* Is the window at @offset one a collapse of @order should be attempted on? */
+static bool collapse_window_eligible(struct collapse_control *cc,
+				     unsigned int offset, unsigned int order)
+{
+	unsigned int nr_ptes = 1U << order;
+	unsigned int max_ptes_none, nr_eligible_ptes;
+
+	if (!test_bit(order, &cc->select_orders))
+		return false;
+
+	/* The window must lie inside the scanned range */
+	if (offset < cc->select_start || offset + nr_ptes > cc->select_end)
+		return false;
+
+	max_ptes_none = collapse_max_ptes_none(cc, NULL, order);
+	nr_eligible_ptes = bitmap_weight_from(cc->eligible_ptes, offset,
+					      offset + nr_ptes);
+
+	/*
+	 * Swap PTEs the scan accepted are counted in cc->scan_unmapped, not in
+	 * the bitmap.  collapse_faultin() reads them in for a PMD candidate, so
+	 * there they do become sources; a smaller window leaves them as holes,
+	 * sub-PMD collapse not faulting swap in.
+	 */
+	if (is_pmd_order(order))
+		nr_eligible_ptes += cc->scan_unmapped;
+
+	return nr_eligible_ptes >= nr_ptes - max_ptes_none;
 }
 
 /*
@@ -1913,6 +2021,24 @@ static void collapse_selection_init(struct collapse_control *cc,
 static bool collapse_next_candidate(struct collapse_control *cc,
 				    unsigned int *offset, unsigned int *order)
 {
+	while (cc->select_offset < cc->select_end) {
+		if (!collapse_window_eligible(cc, cc->select_offset,
+					      cc->select_order)) {
+			collapse_selection_reject(cc);
+			continue;
+		}
+
+		/*
+		 * The cursor advances past the window at emission: a round is
+		 * collected before it is run, so within a round every attempt is
+		 * assumed to succeed.
+		 */
+		*offset = cc->select_offset;
+		*order = cc->select_order;
+		collapse_selection_advance(cc, 1U << cc->select_order);
+		return true;
+	}
+
 	return false;
 }
 
diff --git a/mm/collapse.h b/mm/collapse.h
index 3256c45ee228..94b796271843 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -114,6 +114,15 @@ struct collapse_control {
 	/* Orders still worth attempting in the table being scanned */
 	unsigned long select_orders;
 
+	/* Non-present PTEs the scan accepted, which no bitmap bit marks */
+	unsigned int scan_unmapped;
+
+	/* Where selection has got to in the table, and at what order */
+	unsigned int select_start;
+	unsigned int select_end;
+	unsigned int select_offset;
+	unsigned int select_order;
+
 	/* PTEs collapsed in it so far */
 	unsigned int nr_collapsed;
 
@@ -166,6 +175,7 @@ enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd);
 int collapse_find_target_node(struct collapse_control *cc);
 bool collapse_scan_abort(int nid, struct collapse_control *cc);
+unsigned int max_order_from_offset(unsigned int offset);
 unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order);
 unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 1244e161beae..c7c933e819e2 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -1426,7 +1426,7 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 }
 
 /* Return the highest naturally aligned order that fits at @offset within a PMD. */
-static unsigned int max_order_from_offset(unsigned int offset)
+unsigned int max_order_from_offset(unsigned int offset)
 {
 	if (offset == 0)
 		return HPAGE_PMD_ORDER;
-- 
2.54.0

[RFC PATCH 23/57] mm/collapse: give a refused region a second chance

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the rest of selection: a store of regions to re-enter, and the
classification that decides what goes in it.

Each attempted candidate's outcome is one of four:

 - it collapsed, or was already a huge page.  The region is done; the
   cursor stepped past it at emission.
 - it was refused for something a smaller window might avoid.  The whole
   region goes back at the next enabled order down: selection cannot tell
   which slot refused, so it re-probes the region rather than guessing.
 - only its in-window allocation missed.  The region goes back at the
   same order, marked so the round that picks it up allocates with
   reclaim before freezing anything.
 - the outcome condemns the table.  Selection stops there and drops
   everything queued.

A queued region is walked exactly as the table is: the largest order its
offset's alignment allows, capped by the region's own, descending through
the enabled orders until one qualifies, then stepping past what it
emitted.  Walking rather than shrinking one window is what keeps the tail
of a region in play, which is often where the collapsible part is.

The store is a stack, and the classify loop feeding it walks the batch in
emission order, so entries pop in the order they were refused rather than
by address.  A round drawn from two of them is not address-ordered.

Selection terminates because the pushes that tile a region strictly
descend, and the one that keeps the order cannot repeat for a region: it
comes back asking for reclaim, and a miss the allocator was asked to work
for is a failure, which descends.

An allocation failure is only reported as one when nothing smaller is
left to try.  The caller answers such a failure by backing off for a
while, and a failure at a large order is no reason to: one PMD is 512M
with 64K pages, so that attempt fails as a matter of course, while the
order the region settles for allocates fine.

collapse_anon_pmd() can now say what the table yielded, in order of
precedence:

 - a collapse;
 - an allocation failure, which the caller answers by backing off.  It
   outranks both refusals below, being the only result acted on;
 - the scan's own refusal, when selection never got as far as refusing a
   window;
 - the last reason a window was refused.

All are per table, so the reset that opens a table clears them, the retry
store included.

A dropped lock is classified with the outcomes that abandon the table,
not with those that try a smaller order.  The fault-in pass has already
taken the lock again as many times as it may, so what is left says
nothing about any window, and demoting every candidate the round was
carrying would be a verdict no pass reached.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 201 +++++++++++++++++++++++++++++++++++++++++++++++---
 mm/collapse.h |  11 +++
 2 files changed, 203 insertions(+), 9 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 258bb9cc32c5..9b73ebff1103 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -124,6 +124,16 @@
  */
 #define COLLAPSE_SAVED_PTES	HPAGE_PMD_NR
 
+/*
+ * Capacity of the retry store: the most regions a table can hold at once.  Live
+ * entries cover disjoint regions -- a region is one candidate's extent, and an
+ * extent is consumed from the cursor or from one entry, never from two -- and
+ * the smallest a producer pushes is one window at the smallest order.  Not
+ * bounded by what a round pushes: the stack is drained from the top, so an entry
+ * below a live one outlives the round that pushed it.
+ */
+#define COLLAPSE_RETRY_STORE_SIZE	COLLAPSE_TABLE_WINDOWS
+
 /* How far a candidate got, and so what a failure has to undo for it */
 enum collapse_candidate_state {
 	CAND_SELECTED,		/* collected; nothing held on its behalf yet */
@@ -132,6 +142,21 @@ enum collapse_candidate_state {
 	CAND_INSTALLED,		/* the destination is mapped */
 };
 
+/*
+ * A region queued to re-enter selection, walked like the table itself: @offset is
+ * the next window to probe, @end one past the region, and @order the largest to
+ * try -- below the order that just failed, so the same window cannot be emitted
+ * twice.  Walking the region rather than shrinking one window keeps its tail,
+ * which is often where the collapsible window is.
+ */
+struct collapse_retry {
+	unsigned int offset;
+	unsigned int end;
+	unsigned int order;
+	/* The light allocation missed here: the next attempt may reclaim */
+	bool reclaim;
+};
+
 /*
  * A candidate is an (addr, order) window selected for collapse.  Selection
  * counts in PTE offsets -- the bitmap it reads and the alignment it honours are
@@ -181,16 +206,20 @@ void collapse_control_release(struct collapse_control *cc)
 {
 	kfree(cc->candidates);
 	kfree(cc->saved_ptes);
+	kfree(cc->retries);
 	cc->candidates = NULL;
 	cc->saved_ptes = NULL;
+	cc->retries = NULL;
 }
 
 int collapse_control_init(struct collapse_control *cc)
 {
 	cc->nr_candidates = 0;
+	cc->nr_retries = 0;
 	cc->candidates = kmalloc_objs(*cc->candidates, COLLAPSE_MAX_CANDIDATES);
 	cc->saved_ptes = kmalloc_objs(*cc->saved_ptes, COLLAPSE_SAVED_PTES);
-	if (!cc->candidates || !cc->saved_ptes) {
+	cc->retries = kmalloc_objs(*cc->retries, COLLAPSE_RETRY_STORE_SIZE);
+	if (!cc->candidates || !cc->saved_ptes || !cc->retries) {
 		collapse_control_release(cc);
 		return -ENOMEM;
 	}
@@ -1855,6 +1884,9 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
 	cc->select_orders = 0;
 	cc->scan_unmapped = 0;
 	cc->nr_collapsed = 0;
+	cc->select_result = SCAN_FAIL;
+	cc->smallest_alloc_failed = false;
+	cc->nr_retries = 0;
 }
 
 /*
@@ -2010,6 +2042,48 @@ static bool collapse_window_eligible(struct collapse_control *cc,
 	return nr_eligible_ptes >= nr_ptes - max_ptes_none;
 }
 
+/*
+ * Queue the region [@offset, @end) to re-enter selection at @order.  Two
+ * producers push, both in collapse_classify_result(): a refused region, tiled at
+ * the next enabled order down because selection cannot tell which slot refused;
+ * and a region whose in-window allocation missed, at an unchanged order, asking
+ * for reclaim next time.
+ *
+ * The store is a stack, and the classify loop that feeds it walks the batch by
+ * ascending address, so entries pop in the order they were refused rather than
+ * by address: a round drawn from two of them descends.  Nothing may take
+ * candidates[0] for the lowest -- what a round spans is cc->batch_start and
+ * cc->batch_end, taken over its candidates by collapse_revalidate().
+ *
+ * Selection terminates because the tiling producer strictly descends, and the
+ * unchanged-order one cannot fire twice for a region: its retry arrives with
+ * reclaim set, so the next miss is a failure that descends.
+ *
+ * The store is sized for the most regions a table can hold, so this cannot
+ * overflow; losing an entry would cost a region its lower-order attempt, so it
+ * asserts rather than fails.
+ */
+static void collapse_push_retry(struct collapse_control *cc, unsigned int offset,
+				unsigned int end, unsigned int order,
+				bool reclaim)
+{
+	struct collapse_retry *retry;
+
+	if (cc->nr_retries >= COLLAPSE_RETRY_STORE_SIZE) {
+		VM_WARN_ON_ONCE(1);
+		return;
+	}
+
+	retry = &cc->retries[cc->nr_retries];
+
+	retry->offset = offset;
+	retry->end = end;
+	retry->order = order;
+	retry->reclaim = reclaim;
+
+	cc->nr_retries++;
+}
+
 /*
  * The next window worth attempting, as an (offset, order) pair.  False when
  * selection is exhausted, which is what ends the range.
@@ -2019,8 +2093,42 @@ static bool collapse_window_eligible(struct collapse_control *cc,
  * particular -- would be stale by construction.
  */
 static bool collapse_next_candidate(struct collapse_control *cc,
-				    unsigned int *offset, unsigned int *order)
+				    unsigned int *offset, unsigned int *order,
+				    bool *reclaim)
 {
+	while (cc->nr_retries) {
+		struct collapse_retry *r = &cc->retries[cc->nr_retries - 1];
+		unsigned int try, smallest;
+
+		if (r->offset >= r->end) {
+			cc->nr_retries--;
+			continue;
+		}
+
+		/*
+		 * The same walk as the table's own: the largest order the
+		 * offset's alignment allows, capped by the region's, descending
+		 * through the enabled orders until one fits.  If nothing fits
+		 * here, step over the smallest window tried and carry on --
+		 * which is what keeps the region's tail in play.
+		 */
+		try = min(max_order_from_offset(r->offset), r->order);
+		smallest = try;
+		while (try && !collapse_window_eligible(cc, r->offset, try)) {
+			smallest = try;
+			try = collapse_lower_order(cc, try);
+		}
+
+		if (try) {
+			*offset = r->offset;
+			*order = try;
+			*reclaim = r->reclaim;
+			r->offset += 1U << try;
+			return true;
+		}
+		r->offset += 1U << smallest;
+	}
+
 	while (cc->select_offset < cc->select_end) {
 		if (!collapse_window_eligible(cc, cc->select_offset,
 					      cc->select_order)) {
@@ -2035,6 +2143,7 @@ static bool collapse_next_candidate(struct collapse_control *cc,
 		 */
 		*offset = cc->select_offset;
 		*order = cc->select_order;
+		*reclaim = false;
 		collapse_selection_advance(cc, 1U << cc->select_order);
 		return true;
 	}
@@ -2051,7 +2160,68 @@ static bool collapse_classify_result(struct collapse_control *cc,
 				     unsigned int offset, unsigned int order,
 				     enum scan_result result)
 {
-	return true;
+	unsigned int lower;
+
+	switch (result) {
+	/* Done with the region: the cursor moved past it at emission */
+	case SCAN_SUCCEED:
+		cc->nr_collapsed += 1U << order;
+		fallthrough;
+	case SCAN_PTE_MAPPED_HUGEPAGE:
+		return true;
+	/* Only the light allocation missed: the same order, allowed to reclaim */
+	case SCAN_ALLOC_LIGHT_MISS:
+		collapse_push_retry(cc, offset, offset + (1U << order), order,
+				    /*reclaim=*/ true);
+		return true;
+	/* A smaller order over the same region might still fit */
+	case SCAN_ALLOC_HUGE_PAGE_FAIL:
+		/*
+		 * Only a failure with nothing left below it says the allocator
+		 * cannot serve this collapse.  A failure at a large order says
+		 * nothing about what the region will settle for -- one PMD is
+		 * 512M with 64K pages, so that attempt fails as a matter of
+		 * course -- and the caller answers an allocation failure by
+		 * backing off for a while.
+		 */
+		if (!collapse_lower_order(cc, order))
+			cc->smallest_alloc_failed = true;
+		fallthrough;
+	case SCAN_LACK_REFERENCED_PAGE:
+	case SCAN_EXCEED_NONE_PTE:
+	case SCAN_EXCEED_SWAP_PTE:
+	case SCAN_EXCEED_SHARED_PTE:
+	case SCAN_PAGE_LOCK:
+	case SCAN_PAGE_COUNT:
+	case SCAN_PAGE_NOT_EXCLUSIVE:
+	case SCAN_PAGE_NULL:
+	case SCAN_DEL_PAGE_LRU:
+	case SCAN_PTE_NON_PRESENT:
+	case SCAN_PTE_UFFD:
+	case SCAN_PAGE_LAZYFREE:
+	case SCAN_PAGE_DIRTY_OR_WRITEBACK:
+		cc->select_result = result;
+		lower = collapse_lower_order(cc, order);
+		if (lower) {
+			/* The whole failed region re-enters, as one entry */
+			collapse_push_retry(cc, offset, offset + (1U << order),
+					    lower, /*reclaim=*/ false);
+		}
+		return true;
+	/*
+	 * Nothing further is worth attempting in this table.  A dropped lock
+	 * belongs here rather than above: it says nothing about any window, so
+	 * lowering the order of every candidate the round was carrying would be
+	 * a verdict nobody reached.  The next scan finds the table again.
+	 */
+	case SCAN_LOCK_DROPPED:
+	case SCAN_PMD_MAPPED:
+	default:
+		cc->select_result = result;
+		cc->select_offset = cc->select_end;
+		cc->nr_retries = 0;
+		return false;
+	}
 }
 
 /*
@@ -2112,7 +2282,7 @@ static bool collapse_batch_full(struct collapse_control *cc, unsigned int slots,
  */
 static void collapse_add_candidate(struct collapse_control *cc,
 				   unsigned long addr, unsigned int order,
-				   pte_t *saved_ptes)
+				   bool reclaim, pte_t *saved_ptes)
 {
 	struct collapse_candidate *cand;
 
@@ -2124,7 +2294,7 @@ static void collapse_add_candidate(struct collapse_control *cc,
 	cc->nr_candidates++;
 	cand->addr = addr;
 	cand->order = order;
-	cand->reclaim = false;
+	cand->reclaim = reclaim;
 	cand->state = CAND_SELECTED;
 	cand->result = SCAN_FAIL;
 	cand->new_folio = NULL;
@@ -2145,7 +2315,7 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 	unsigned int offset, order;
 	unsigned long bytes = 0;
 	unsigned int slots = 0;
-	bool pending = false;
+	bool pending = false, reclaim = false;
 	bool cont = true;
 
 	collapse_selection_init(cc, (start - pmd_addr) >> PAGE_SHIFT,
@@ -2153,7 +2323,8 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 
 	while (cont) {
 		if (!pending)
-			pending = collapse_next_candidate(cc, &offset, &order);
+			pending = collapse_next_candidate(cc, &offset, &order,
+							  &reclaim);
 
 		if (!pending || collapse_batch_full(cc, slots, bytes, order)) {
 			/*
@@ -2177,12 +2348,24 @@ collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
 		 * full round could not take is kept pending for the next one.
 		 */
 		collapse_add_candidate(cc, pmd_addr + offset * PAGE_SIZE, order,
-				       cc->saved_ptes + slots);
+				       reclaim, cc->saved_ptes + slots);
 
 		slots += 1U << order;
 		bytes += PAGE_SIZE << order;
 		pending = false;
 	}
 
-	return cc->nr_collapsed ? SCAN_SUCCEED : SCAN_FAIL;
+	if (cc->nr_collapsed)
+		return SCAN_SUCCEED;
+	/*
+	 * Report an allocation failure over any refusal, the scan's included: it
+	 * is the one outcome the caller acts on, by backing off rather than
+	 * scanning on.
+	 */
+	if (cc->smallest_alloc_failed)
+		return SCAN_ALLOC_HUGE_PAGE_FAIL;
+	/* Nothing salvaged and nothing to wait for: say what was refused */
+	if (cc->scan_refusal != SCAN_SUCCEED)
+		return cc->scan_refusal;
+	return cc->select_result;
 }
diff --git a/mm/collapse.h b/mm/collapse.h
index 94b796271843..3803f5a89087 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -11,6 +11,7 @@
 #define COLLAPSE_MIN_MTHP_ORDER		2
 
 struct collapse_candidate;
+struct collapse_retry;
 
 enum scan_result {
 	SCAN_FAIL,
@@ -135,6 +136,16 @@ struct collapse_control {
 	 */
 	enum scan_result scan_refusal;
 
+	/* Why the last window was refused */
+	enum scan_result select_result;
+
+	/* A region ran out of orders to try because none could be allocated */
+	bool smallest_alloc_failed;
+
+	/* Regions waiting to re-enter selection at a lower order */
+	struct collapse_retry *retries;
+	unsigned int nr_retries;
+
 	/* The candidate windows collected for the current round */
 	struct collapse_candidate *candidates;
 	unsigned int nr_candidates;
-- 
2.54.0

[RFC PATCH 24/57] mm/collapse: report each candidate's outcome to tracing

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The engine decides per candidate, and every one of those decisions is
currently invisible: the mechanism it is about to replace reports through
mm_collapse_huge_page_isolate, which the engine never calls.  Switching
the anonymous path over without something in its place would take
existing tracing with it.

Add one tracepoint, mm_collapse_candidate: a window's address and order,
the pass that reached a verdict on it, and what that verdict was.

Every candidate a pass judged produces exactly one -- the pass that
refused it, or the install for one that made it.  A candidate the round
gave up on before any pass judged it produces none.

That is enough to follow a round: which windows were attempted, and which
ones the batch dropped and where.  It is also what a scan of the trace
buffer can attribute to an address.

It goes in the huge_memory trace system, next to the events it stands in
for, so a consumer enabling that system keeps seeing collapses.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 40 ++++++++++++++++++++++++++++++
 mm/collapse.c                      | 32 +++++++++++++++++++++++-
 mm/collapse.h                      | 13 ++++++++++
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index ff938ac9c43c..86131845b761 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -44,12 +44,21 @@
 	EM( SCAN_PAGE_NOT_EXCLUSIVE,	"page_not_exclusive")		\
 	EMe(SCAN_ALLOC_LIGHT_MISS,	"alloc_light_miss")
 
+#define COLLAPSE_PASS_STATUS						\
+	EM( COLLAPSE_PASS_ALLOC,	"alloc")			\
+	EM( COLLAPSE_PASS_REVALIDATE,	"revalidate")			\
+	EM( COLLAPSE_PASS_FAULTIN,	"faultin")			\
+	EM( COLLAPSE_PASS_FREEZE,	"freeze")			\
+	EM( COLLAPSE_PASS_COPY,		"copy")				\
+	EMe(COLLAPSE_PASS_INSTALL,	"install")
+
 #undef EM
 #undef EMe
 #define EM(a, b)	TRACE_DEFINE_ENUM(a);
 #define EMe(a, b)	TRACE_DEFINE_ENUM(a);
 
 SCAN_STATUS
+COLLAPSE_PASS_STATUS
 
 #undef EM
 #undef EMe
@@ -117,6 +126,37 @@ TRACE_EVENT(mm_collapse_huge_page,
 		__entry->order)
 );
 
+TRACE_EVENT(mm_collapse_candidate,
+
+	TP_PROTO(struct mm_struct *mm, unsigned long addr, unsigned int order,
+		 int pass, int result),
+
+	TP_ARGS(mm, addr, order, pass, result),
+
+	TP_STRUCT__entry(
+		__field(struct mm_struct *, mm)
+		__field(unsigned long, addr)
+		__field(unsigned int, order)
+		__field(int, pass)
+		__field(int, result)
+	),
+
+	TP_fast_assign(
+		__entry->mm = mm;
+		__entry->addr = addr;
+		__entry->order = order;
+		__entry->pass = pass;
+		__entry->result = result;
+	),
+
+	TP_printk("mm=%p, addr=0x%lx, order=%u, pass=%s, result=%s",
+		__entry->mm,
+		__entry->addr,
+		__entry->order,
+		__print_symbolic(__entry->pass, COLLAPSE_PASS_STATUS),
+		__print_symbolic(__entry->result, SCAN_STATUS))
+);
+
 TRACE_EVENT(mm_collapse_huge_page_isolate,
 
 	TP_PROTO(struct folio *folio, int none_or_zero,
diff --git a/mm/collapse.c b/mm/collapse.c
index 9b73ebff1103..91ff20138a8e 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -20,6 +20,7 @@
 #include <linux/vmstat.h>
 
 #include <asm/tlb.h>
+#include <trace/events/huge_memory.h>
 #include "collapse.h"
 #include "internal.h"
 
@@ -195,6 +196,14 @@ static unsigned int candidate_nr_pages(const struct collapse_candidate *cand)
 	return 1U << cand->order;
 }
 
+static void collapse_trace_candidate(struct mm_struct *mm,
+				     const struct collapse_candidate *cand,
+				     enum collapse_pass pass)
+{
+	trace_mm_collapse_candidate(mm, cand->addr, cand->order, pass,
+				    cand->result);
+}
+
 /* Where a candidate sits in the table, in the PTE offsets selection counts in */
 static unsigned int candidate_offset(const struct collapse_candidate *cand,
 				     unsigned long pmd_addr)
@@ -278,6 +287,8 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 					      BIT(cand->order))) {
 			cand->state = CAND_SKIPPED;
 			cand->result = SCAN_VMA_CHECK;
+			collapse_trace_candidate(mm, cand,
+						 COLLAPSE_PASS_REVALIDATE);
 			continue;
 		}
 
@@ -420,6 +431,8 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 			if (r == SCAN_EXCEED_SWAP_PTE) {
 				cand->state = CAND_SKIPPED;
 				cand->result = r;
+				collapse_trace_candidate(vma->vm_mm, cand,
+							 COLLAPSE_PASS_FAULTIN);
 				break;
 			}
 			if (r != SCAN_SUCCEED) {
@@ -878,6 +891,7 @@ static void collapse_freeze(struct vm_area_struct *vma,
 				continue;
 			cand->state = CAND_SKIPPED;
 			cand->result = SCAN_NO_PTE_TABLE;
+			collapse_trace_candidate(mm, cand, COLLAPSE_PASS_FREEZE);
 		}
 		return;
 	}
@@ -904,6 +918,7 @@ static void collapse_freeze(struct vm_area_struct *vma,
 		cand->result = result;
 		if (result != SCAN_SUCCEED) {
 			cand->state = CAND_SKIPPED;
+			collapse_trace_candidate(mm, cand, COLLAPSE_PASS_FREEZE);
 			continue;
 		}
 
@@ -986,6 +1001,7 @@ static void collapse_reserve(struct mm_struct *mm, struct collapse_control *cc)
 
 		cand->state = CAND_SKIPPED;
 		cand->result = result;
+		collapse_trace_candidate(mm, cand, COLLAPSE_PASS_ALLOC);
 	}
 }
 
@@ -1016,6 +1032,7 @@ static void collapse_deposit(struct mm_struct *mm, struct collapse_control *cc)
 	if (!cand->deposit) {
 		cand->state = CAND_SKIPPED;
 		cand->result = SCAN_ALLOC_HUGE_PAGE_FAIL;
+		collapse_trace_candidate(mm, cand, COLLAPSE_PASS_ALLOC);
 	}
 }
 
@@ -1060,6 +1077,8 @@ static void collapse_provision(struct mm_struct *mm,
 			}
 			cand->result = result;
 		}
+
+		collapse_trace_candidate(mm, cand, COLLAPSE_PASS_ALLOC);
 	}
 }
 
@@ -1106,6 +1125,8 @@ static void collapse_copy(struct vm_area_struct *vma,
 			 */
 			if (copy_mc_user_highpage(dst, src, addr, vma)) {
 				cand->result = SCAN_COPY_MC;
+				collapse_trace_candidate(vma->vm_mm, cand,
+							 COLLAPSE_PASS_COPY);
 				break;
 			}
 		}
@@ -1294,6 +1315,7 @@ static void collapse_install_pmd(struct vm_area_struct *vma,
 		/* Table gone under us; see collapse_abort_candidate() on @pte */
 		spin_unlock(pmd_ptl);
 		cand->result = SCAN_NO_PTE_TABLE;
+		collapse_trace_candidate(mm, cand, COLLAPSE_PASS_INSTALL);
 		collapse_abort_candidate(vma, cand, NULL);
 		return;
 	}
@@ -1315,6 +1337,7 @@ static void collapse_install_pmd(struct vm_area_struct *vma,
 
 	if (!collapse_verify_candidate(cand, pte, &nr_populated)) {
 		cand->result = SCAN_PTE_NON_PRESENT;
+		collapse_trace_candidate(mm, cand, COLLAPSE_PASS_INSTALL);
 		collapse_abort_candidate(vma, cand, pte);
 		goto out_unlock;
 	}
@@ -1411,6 +1434,8 @@ static void collapse_install(struct vm_area_struct *vma,
 				continue;
 
 			cand->result = SCAN_NO_PTE_TABLE;
+			collapse_trace_candidate(mm, cand,
+						 COLLAPSE_PASS_INSTALL);
 			collapse_abort_candidate(vma, cand, NULL);
 		}
 		return;
@@ -1439,6 +1464,8 @@ static void collapse_install(struct vm_area_struct *vma,
 
 		if (!collapse_verify_candidate(cand, cand_pte, &nr_populated)) {
 			cand->result = SCAN_PTE_NON_PRESENT;
+			collapse_trace_candidate(mm, cand,
+						 COLLAPSE_PASS_INSTALL);
 			collapse_abort_candidate(vma, cand, cand_pte);
 			continue;
 		}
@@ -1548,8 +1575,11 @@ static unsigned int collapse_finish(struct mm_struct *mm,
 			pte_free(mm, cand->deposit);
 			cand->deposit = NULL;
 		}
-		if (cand->state == CAND_INSTALLED)
+		if (cand->state == CAND_INSTALLED) {
 			nr_installed++;
+			collapse_trace_candidate(mm, cand,
+						 COLLAPSE_PASS_INSTALL);
+		}
 	}
 
 	return nr_installed;
diff --git a/mm/collapse.h b/mm/collapse.h
index 3803f5a89087..34de3ebb05e3 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -13,6 +13,19 @@
 struct collapse_candidate;
 struct collapse_retry;
 
+/*
+ * Which pass of a round reached a verdict on a candidate.  Only collapse.c
+ * produces these; the trace header khugepaged.c builds names them.
+ */
+enum collapse_pass {
+	COLLAPSE_PASS_ALLOC,
+	COLLAPSE_PASS_REVALIDATE,
+	COLLAPSE_PASS_FAULTIN,
+	COLLAPSE_PASS_FREEZE,
+	COLLAPSE_PASS_COPY,
+	COLLAPSE_PASS_INSTALL,
+};
+
 enum scan_result {
 	SCAN_FAIL,
 	SCAN_SUCCEED,
-- 
2.54.0

[RFC PATCH 25/57] mm/collapse: collapse anonymous memory with the new engine

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Point the anonymous path at the engine.  Everything it needs is in place,
so this is the whole switch: collapse_single_pmd() calls
collapse_scan_anon_pmd() and then collapse_anon_pmd(), where it used to
call collapse_scan_pmd().  The scan runs under the mmap_read the caller
already holds; the collapse is called after dropping it, and takes the
lock itself for each round.

Both callers hand the engine a PMD-aligned address with the whole table
inside the VMA: khugepaged walks [ALIGN(vm_start), ALIGN_DOWN(vm_end)) a
table at a time, and MADV_COLLAPSE aligns its range inwards the same way.
So the range passed is always the table.  The engine accepts a narrower
one, which nothing asks for yet.

Two things userspace sees change:

 - A collapse runs under mmap_read rather than holding mmap_write for its
   duration, so faults elsewhere in the address space are no longer
   stopped while it works.
 - A table that cannot become one huge page still yields the largest
   windows inside it, where before a single disqualified PTE gave up the
   whole table.

The result the caller gets is the engine's, and it still acts on an
allocation failure by backing off.

The mechanism this replaces is left in place, now unreferenced.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 15 ++++++---------
 mm/collapse.h   |  5 +++++
 mm/khugepaged.c | 14 ++++++++++++--
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 91ff20138a8e..df3760e3918b 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -83,9 +83,6 @@
  * consecutive pages of one folio -- so partially mapped and compound sources
  * collapse too: any order below the window's is a source, and a PMD candidate
  * takes even a PTE-mapped THP of its own order.
- *
- * Nothing calls any of this yet: the anon path still uses the mechanism it
- * replaces, and is switched over once both halves are complete.
  */
 
 /*
@@ -1926,9 +1923,9 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
  * that acts on what it found hands the range to collapse_anon_pmd() afterwards,
  * without the lock.
  */
-static enum scan_result __maybe_unused
-collapse_scan_anon_pmd(struct vm_area_struct *vma, unsigned long start,
-		       unsigned long end, struct collapse_control *cc)
+enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
+					unsigned long start, unsigned long end,
+					struct collapse_control *cc)
 {
 	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
 	struct mm_struct *mm = vma->vm_mm;
@@ -2337,9 +2334,9 @@ static void collapse_add_candidate(struct collapse_control *cc,
  * largest order downwards.  Returns what the table yielded: a collapse, or
  * the reason it did not.
  */
-static enum scan_result __maybe_unused
-collapse_anon_pmd(struct mm_struct *mm, unsigned long start, unsigned long end,
-		  struct collapse_control *cc)
+enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
+				   unsigned long end,
+				   struct collapse_control *cc)
 {
 	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
 	unsigned int offset, order;
diff --git a/mm/collapse.h b/mm/collapse.h
index 34de3ebb05e3..50a9d59bbf03 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -186,6 +186,11 @@ static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
 		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
 }
 
+enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
+		unsigned long start, unsigned long end,
+		struct collapse_control *cc);
+enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
+		unsigned long end, struct collapse_control *cc);
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
 
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index c7c933e819e2..0662d08f7c60 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -1553,7 +1553,8 @@ static enum scan_result mthp_collapse(struct mm_struct *mm,
 	return last_result;
 }
 
-static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
+static enum scan_result __maybe_unused
+collapse_scan_pmd(struct mm_struct *mm,
 		struct vm_area_struct *vma, unsigned long start_addr,
 		bool *lock_dropped, struct collapse_control *cc)
 {
@@ -2749,7 +2750,16 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 	mmap_assert_locked(mm);
 
 	if (vma_is_anonymous(vma)) {
-		result = collapse_scan_pmd(mm, vma, addr, lock_dropped, cc);
+		result = collapse_scan_anon_pmd(vma, addr, addr + HPAGE_PMD_SIZE,
+						cc);
+		if (!cc->select_orders)
+			goto end;
+
+		/* collapse_anon_pmd() takes mmap_lock itself, where it needs it */
+		mmap_read_unlock(mm);
+		*lock_dropped = true;
+
+		result = collapse_anon_pmd(mm, addr, addr + HPAGE_PMD_SIZE, cc);
 		goto end;
 	}
 
-- 
2.54.0

[RFC PATCH 26/57] mm/collapse: give collapse_single_pmd() the range to work on

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

collapse_single_pmd() derives the end of its range from its start: one
PMD, always.  Both of its callers already know the range they mean.

Take the end as an argument and pass it to the scan and the collapse,
both of which already accept a partial table.  Both callers pass what the
function computed for itself.

Preparation for scanning a VMA that holds less than a whole table.

No functional change intended.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/khugepaged.c | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 0662d08f7c60..d1e031ed3e6f 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2738,8 +2738,8 @@ static enum scan_result collapse_scan_file(struct mm_struct *mm,
  * the results.
  */
 static enum scan_result collapse_single_pmd(unsigned long addr,
-		struct vm_area_struct *vma, bool *lock_dropped,
-		struct collapse_control *cc)
+		unsigned long end, struct vm_area_struct *vma,
+		bool *lock_dropped, struct collapse_control *cc)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	bool triggered_wb = false;
@@ -2750,8 +2750,7 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 	mmap_assert_locked(mm);
 
 	if (vma_is_anonymous(vma)) {
-		result = collapse_scan_anon_pmd(vma, addr, addr + HPAGE_PMD_SIZE,
-						cc);
+		result = collapse_scan_anon_pmd(vma, addr, end, cc);
 		if (!cc->select_orders)
 			goto end;
 
@@ -2759,7 +2758,7 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 		mmap_read_unlock(mm);
 		*lock_dropped = true;
 
-		result = collapse_anon_pmd(mm, addr, addr + HPAGE_PMD_SIZE, cc);
+		result = collapse_anon_pmd(mm, addr, end, cc);
 		goto end;
 	}
 
@@ -2872,6 +2871,8 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 				  hend);
 
 			*result = collapse_single_pmd(khugepaged_scan.address,
+						      khugepaged_scan.address +
+						      HPAGE_PMD_SIZE,
 						      vma, &lock_dropped, cc);
 			/* move to next address */
 			khugepaged_scan.address += HPAGE_PMD_SIZE;
@@ -3211,7 +3212,8 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 			hend = min(hend, vma->vm_end & HPAGE_PMD_MASK);
 		}
 
-		result = collapse_single_pmd(addr, vma, &mmap_unlocked, cc);
+		result = collapse_single_pmd(addr, addr + HPAGE_PMD_SIZE, vma,
+					     &mmap_unlocked, cc);
 
 		switch (result) {
 		case SCAN_SUCCEED:
-- 
2.54.0

[RFC PATCH 27/57] mm/collapse: scan the windows a VMA can hold

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

khugepaged covers each VMA in whole PTE tables, from its first PMD-aligned
address to its last.  A VMA smaller than a table is never scanned at all,
and in a larger one everything outside its PMD-aligned span is skipped.

That was the right shape when a collapse was always a PMD.  It keeps mTHP
collapse away from every range that is not PMD-shaped, which is most of
what an mTHP is for.

The gap is widest where a PMD is largest.  On arm64 with 64K base pages a
PMD is 512M, so the old walk reached only VMAs big enough and aligned well
enough to hold one.  A 2M mTHP -- order 5 there -- was unreachable in
anything smaller, however many such windows the VMA had room for.

Root the coverage at windows of the largest order the VMA allows, and hand
the range on one table at a time, clamped to the VMA.  The engine already
accepts a partial table; this is the first caller that gives it one.

The cursor is no longer PMD-aligned, so the assert that said it was goes.
The bound beside it goes too: the range is clamped to the VMA where it is
computed, leaving nothing for it to catch.

For a VMA whose largest allowed order is the PMD order -- every file VMA,
and any anonymous VMA with only PMD-order THP enabled -- the walk is
exactly what it was.  Where smaller orders are enabled, khugepaged now
reaches VMAs a table would not fit in, and the edges of VMAs it used to
leave out.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/khugepaged.c | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index d1e031ed3e6f..895183d92fb8 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2838,44 +2838,56 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 
 	vma_iter_init(&vmi, mm, khugepaged_scan.address);
 	for_each_vma(vmi, vma) {
-		unsigned long hstart, hend;
+		unsigned long hstart, hend, window;
+		unsigned long orders;
 
 		cond_resched();
 		if (unlikely(collapse_test_exit_or_disable(mm))) {
 			cc->progress++;
 			break;
 		}
-		if (!collapse_possible(vma, vma->vm_flags, TVA_KHUGEPAGED)) {
+		orders = collapse_possible_orders(vma, vma->vm_flags,
+						  TVA_KHUGEPAGED);
+		if (!orders) {
 			cc->progress++;
 			continue;
 		}
-		hstart = ALIGN(vma->vm_start, HPAGE_PMD_SIZE);
-		hend = ALIGN_DOWN(vma->vm_end, HPAGE_PMD_SIZE);
+
+		/*
+		 * Coverage is rooted at windows of the largest order the VMA
+		 * allows: below the PMD order that reaches VMAs a whole table
+		 * would not fit in, and parts of a VMA that a whole table would
+		 * leave out.
+		 */
+		window = PAGE_SIZE << __fls(orders);
+		hstart = ALIGN(vma->vm_start, window);
+		hend = ALIGN_DOWN(vma->vm_end, window);
 		if (khugepaged_scan.address > hend) {
 			cc->progress++;
 			continue;
 		}
 		if (khugepaged_scan.address < hstart)
 			khugepaged_scan.address = hstart;
-		VM_BUG_ON(khugepaged_scan.address & ~HPAGE_PMD_MASK);
 
 		while (khugepaged_scan.address < hend) {
+			unsigned long pmd_addr, range_end;
 			bool lock_dropped = false;
 
+			/* One table's worth at most, and never past the VMA */
+			pmd_addr = khugepaged_scan.address & HPAGE_PMD_MASK;
+			range_end = min(hend, pmd_addr + HPAGE_PMD_SIZE);
+
 			cond_resched();
 			if (unlikely(collapse_test_exit_or_disable(mm)))
 				goto breakouterloop;
 
-			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart ||
-				  khugepaged_scan.address + HPAGE_PMD_SIZE >
-				  hend);
+			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart);
 
 			*result = collapse_single_pmd(khugepaged_scan.address,
-						      khugepaged_scan.address +
-						      HPAGE_PMD_SIZE,
-						      vma, &lock_dropped, cc);
+						      range_end, vma,
+						      &lock_dropped, cc);
 			/* move to next address */
-			khugepaged_scan.address += HPAGE_PMD_SIZE;
+			khugepaged_scan.address = range_end;
 			if (lock_dropped)
 				/*
 				 * We released mmap_lock so break loop.  Note
-- 
2.54.0

[RFC PATCH 28/57] mm/collapse: remove the mechanism the engine replaces

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Nothing reaches the old anonymous collapse any more: the entry point was
rewired to the engine, and every function below it lost its last caller.

Delete the chain: the scan, the mTHP order walk, the collapse itself,
isolation, swap-in, the copy with its success and failure paths, the PTE
release helpers, folio_pte_referenced() and the pmd-still-valid check.

What stays is what the file paths and MADV_COLLAPSE still call:
alloc_charge_folio() for a file collapse's destination,
hugepage_vma_revalidate() for the VMA check after MADV_COLLAPSE drops
mmap_lock, and count_collapse_event() and collapse_control_init_scan()
for the file scan.

Four tracepoints lose their only emitter here: mm_khugepaged_scan_pmd,
mm_collapse_huge_page, mm_collapse_huge_page_isolate and
mm_collapse_huge_page_swapin.  Their definitions stay, now without an
emitter, and the engine reports through mm_collapse_candidate.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/khugepaged.c | 984 +-----------------------------------------------
 1 file changed, 9 insertions(+), 975 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 895183d92fb8..6203473f4953 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -545,370 +545,6 @@ void __khugepaged_exit(struct mm_struct *mm)
 	}
 }
 
-static void collapse_control_init_scan(struct collapse_control *cc)
-{
-	memset(cc->node_load, 0, sizeof(cc->node_load));
-	nodes_clear(cc->alloc_nmask);
-	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
-}
-
-static void release_pte_folio(struct folio *folio)
-{
-	node_stat_mod_folio(folio,
-			NR_ISOLATED_ANON + folio_is_file_lru(folio),
-			-folio_nr_pages(folio));
-	folio_unlock(folio);
-	folio_putback_lru(folio);
-}
-
-static void release_pte_pages(pte_t *pte, pte_t *_pte,
-		struct list_head *compound_pagelist)
-{
-	struct folio *folio, *tmp;
-
-	while (--_pte >= pte) {
-		pte_t pteval = ptep_get(_pte);
-		unsigned long pfn;
-
-		if (pte_none(pteval))
-			continue;
-		VM_WARN_ON_ONCE(!pte_present(pteval));
-		pfn = pte_pfn(pteval);
-		if (is_zero_pfn(pfn))
-			continue;
-		folio = pfn_folio(pfn);
-		if (folio_test_large(folio))
-			continue;
-		release_pte_folio(folio);
-	}
-
-	list_for_each_entry_safe(folio, tmp, compound_pagelist, lru) {
-		list_del(&folio->lru);
-		release_pte_folio(folio);
-	}
-}
-
-/*
- * folio_pte_referenced() - Check if a folio or its PTE mapping was recently used
- *
- * Return: true if recent access was observed through either the folio state
- * or the current PTE mapping.
- */
-static inline bool folio_pte_referenced(struct folio *folio,
-		struct vm_area_struct *vma, unsigned long addr, pte_t pteval)
-{
-	/* The folio was referenced previously ... */
-	if (folio_test_young(folio) || folio_test_referenced(folio))
-		return true;
-	/* ... or the PTE mapping was recently used */
-	return pte_young(pteval) || mmu_notifier_test_young(vma->vm_mm, addr);
-}
-
-static void count_collapse_event(unsigned int order, enum vm_event_item vm_event,
-		enum mthp_stat_item mthp_event)
-{
-	if (is_pmd_order(order))
-		count_vm_event(vm_event);
-	count_mthp_stat(order, mthp_event);
-}
-
-static enum scan_result __collapse_huge_page_isolate(struct vm_area_struct *vma,
-		unsigned long start_addr, pte_t *pte, struct collapse_control *cc,
-		unsigned int order, struct list_head *compound_pagelist)
-{
-	const unsigned int max_ptes_none = collapse_max_ptes_none(cc, vma, order);
-	const unsigned int max_ptes_shared = collapse_max_ptes_shared(cc, order);
-	const unsigned long nr_pages = 1UL << order;
-	struct page *page = NULL;
-	struct folio *folio = NULL;
-	unsigned long addr = start_addr;
-	pte_t *_pte;
-	int none_or_zero = 0, shared = 0, referenced = 0;
-	enum scan_result result = SCAN_FAIL;
-
-	for (_pte = pte; _pte < pte + nr_pages;
-	     _pte++, addr += PAGE_SIZE) {
-		pte_t pteval = ptep_get(_pte);
-		if (pte_none_or_zero(pteval)) {
-			if (++none_or_zero > max_ptes_none) {
-				result = SCAN_EXCEED_NONE_PTE;
-				count_collapse_event(order, THP_SCAN_EXCEED_NONE_PTE,
-						     MTHP_STAT_COLLAPSE_EXCEED_NONE);
-				goto out;
-			}
-			continue;
-		}
-		if (!pte_present(pteval)) {
-			result = SCAN_PTE_NON_PRESENT;
-			goto out;
-		}
-		if (pte_uffd(pteval)) {
-			result = SCAN_PTE_UFFD;
-			goto out;
-		}
-		page = vm_normal_page(vma, addr, pteval);
-		if (unlikely(!page) || unlikely(is_zone_device_page(page))) {
-			result = SCAN_PAGE_NULL;
-			goto out;
-		}
-
-		folio = page_folio(page);
-		VM_BUG_ON_FOLIO(!folio_test_anon(folio), folio);
-
-		/*
-		 * If the vma has the VM_DROPPABLE flag, the collapse will
-		 * preserve the lazyfree property without needing to skip.
-		 */
-		if (cc->policy.skip_lazyfree && !(vma->vm_flags & VM_DROPPABLE) &&
-		    folio_test_lazyfree(folio) && !pte_dirty(pteval)) {
-			result = SCAN_PAGE_LAZYFREE;
-			goto out;
-		}
-
-		/* See collapse_scan_pmd(). */
-		if (folio_maybe_mapped_shared(folio)) {
-			/*
-			 * TODO: Support shared pages without leading to further
-			 * mTHP collapses. Currently bringing in new pages via
-			 * shared may cause a future higher order collapse on a
-			 * rescan of the same range.
-			 */
-			if (++shared > max_ptes_shared) {
-				result = SCAN_EXCEED_SHARED_PTE;
-				count_collapse_event(order, THP_SCAN_EXCEED_SHARED_PTE,
-						     MTHP_STAT_COLLAPSE_EXCEED_SHARED);
-				goto out;
-			}
-		}
-		/*
-		 * TODO: In some cases of partially-mapped folios, we'd actually
-		 * want to collapse.
-		 */
-		if (!is_pmd_order(order) && folio_order(folio) >= order) {
-			result = SCAN_PTE_MAPPED_HUGEPAGE;
-			goto out;
-		}
-
-		if (folio_test_large(folio)) {
-			struct folio *f;
-
-			/*
-			 * Check if we have dealt with the compound page
-			 * already
-			 */
-			list_for_each_entry(f, compound_pagelist, lru) {
-				if (folio == f)
-					goto next;
-			}
-		}
-
-		/*
-		 * We can do it before folio_isolate_lru because the
-		 * folio can't be freed from under us. NOTE: folio lock
-		 * is needed to serialize against split_huge_page()
-		 * when invoked from the VM.
-		 */
-		if (!folio_trylock(folio)) {
-			result = SCAN_PAGE_LOCK;
-			goto out;
-		}
-
-		/*
-		 * Check if the page has any GUP (or other external) pins.
-		 *
-		 * The page table that maps the page has been already unlinked
-		 * from the page table tree and this process cannot get
-		 * an additional pin on the page.
-		 *
-		 * New pins can come later if the page is shared across fork,
-		 * but not from this process. The other process cannot write to
-		 * the page, only trigger CoW.
-		 */
-		if (folio_expected_ref_count(folio) != folio_ref_count(folio)) {
-			folio_unlock(folio);
-			result = SCAN_PAGE_COUNT;
-			goto out;
-		}
-
-		/*
-		 * Isolate the folio to avoid collapsing a hugepage
-		 * currently in use by the VM.
-		 */
-		if (!folio_isolate_lru(folio)) {
-			folio_unlock(folio);
-			result = SCAN_DEL_PAGE_LRU;
-			goto out;
-		}
-		node_stat_mod_folio(folio,
-				NR_ISOLATED_ANON + folio_is_file_lru(folio),
-				folio_nr_pages(folio));
-		VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
-		VM_BUG_ON_FOLIO(folio_test_lru(folio), folio);
-
-		if (folio_test_large(folio))
-			list_add_tail(&folio->lru, compound_pagelist);
-next:
-		if (cc->policy.require_referenced &&
-		    folio_pte_referenced(folio, vma, addr, pteval))
-			referenced++;
-	}
-
-	if (unlikely(cc->policy.require_referenced && !referenced)) {
-		result = SCAN_LACK_REFERENCED_PAGE;
-	} else {
-		result = SCAN_SUCCEED;
-		trace_mm_collapse_huge_page_isolate(folio, none_or_zero,
-						    referenced, result, order);
-		return result;
-	}
-out:
-	release_pte_pages(pte, _pte, compound_pagelist);
-	trace_mm_collapse_huge_page_isolate(folio, none_or_zero,
-					    referenced, result, order);
-	return result;
-}
-
-static void __collapse_huge_page_copy_succeeded(pte_t *pte,
-		struct vm_area_struct *vma, unsigned long address,
-		spinlock_t *ptl, unsigned int order,
-		struct list_head *compound_pagelist)
-{
-	const unsigned long nr_pages = 1UL << order;
-	unsigned long end = address + (PAGE_SIZE * nr_pages);
-	struct folio *src, *tmp;
-	pte_t pteval;
-	pte_t *_pte;
-	unsigned int nr_ptes;
-
-	for (_pte = pte; _pte < pte + nr_pages; _pte += nr_ptes,
-	     address += nr_ptes * PAGE_SIZE) {
-		nr_ptes = 1;
-		pteval = ptep_get(_pte);
-		if (pte_none_or_zero(pteval)) {
-			add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
-			if (pte_none(pteval))
-				continue;
-			/*
-			 * ptl mostly unnecessary.
-			 */
-			spin_lock(ptl);
-			ptep_clear(vma->vm_mm, address, _pte);
-			spin_unlock(ptl);
-			ksm_might_unmap_zero_page(vma->vm_mm, pteval);
-		} else {
-			struct page *src_page = pte_page(pteval);
-
-			src = page_folio(src_page);
-
-			if (folio_test_large(src)) {
-				unsigned int max_nr_ptes = (end - address) >> PAGE_SHIFT;
-
-				nr_ptes = folio_pte_batch(src, _pte, pteval, max_nr_ptes);
-			} else {
-				release_pte_folio(src);
-			}
-
-			/*
-			 * ptl mostly unnecessary, but preempt has to
-			 * be disabled to update the per-cpu stats
-			 * inside folio_remove_rmap_pte().
-			 */
-			spin_lock(ptl);
-			clear_ptes(vma->vm_mm, address, _pte, nr_ptes);
-			folio_remove_rmap_ptes(src, src_page, nr_ptes, vma);
-			spin_unlock(ptl);
-			free_swap_cache(src);
-			folio_put_refs(src, nr_ptes);
-		}
-	}
-
-	list_for_each_entry_safe(src, tmp, compound_pagelist, lru) {
-		list_del(&src->lru);
-		node_stat_sub_folio(src, NR_ISOLATED_ANON +
-				folio_is_file_lru(src));
-		folio_unlock(src);
-		free_swap_cache(src);
-		folio_putback_lru(src);
-	}
-}
-
-static void __collapse_huge_page_copy_failed(pte_t *pte,
-		pmd_t *pmd, pmd_t orig_pmd, struct vm_area_struct *vma,
-		unsigned int order, struct list_head *compound_pagelist)
-{
-	const unsigned long nr_pages = 1UL << order;
-	spinlock_t *pmd_ptl;
-
-	/*
-	 * Re-establish the PMD to point to the original page table
-	 * entry. Restoring PMD needs to be done prior to releasing
-	 * pages. Since pages are still isolated and locked here,
-	 * acquiring anon_vma_lock_write() is unnecessary.
-	 */
-	pmd_ptl = pmd_lock(vma->vm_mm, pmd);
-	pmd_populate(vma->vm_mm, pmd, pmd_pgtable(orig_pmd));
-	spin_unlock(pmd_ptl);
-	/*
-	 * Release both raw and compound pages isolated
-	 * in __collapse_huge_page_isolate.
-	 */
-	release_pte_pages(pte, pte + nr_pages, compound_pagelist);
-}
-
-/*
- * __collapse_huge_page_copy - attempts to copy memory contents from raw
- * pages to a hugepage. Cleans up the raw pages if copying succeeds;
- * otherwise restores the original page table and releases isolated raw pages.
- * Returns SCAN_SUCCEED if copying succeeds, otherwise returns SCAN_COPY_MC.
- *
- * @pte: starting of the PTEs to copy from
- * @folio: the new hugepage to copy contents to
- * @pmd: pointer to the new hugepage's PMD
- * @orig_pmd: the original raw pages' PMD
- * @vma: the original raw pages' virtual memory area
- * @address: starting address to copy
- * @ptl: lock on raw pages' PTEs
- * @compound_pagelist: list that stores compound pages
- */
-static enum scan_result __collapse_huge_page_copy(pte_t *pte, struct folio *folio,
-		pmd_t *pmd, pmd_t orig_pmd, struct vm_area_struct *vma,
-		unsigned long address, spinlock_t *ptl, unsigned int order,
-		struct list_head *compound_pagelist)
-{
-	const unsigned long nr_pages = 1UL << order;
-	unsigned int i;
-	enum scan_result result = SCAN_SUCCEED;
-
-	/*
-	 * Copying pages' contents is subject to memory poison at any iteration.
-	 */
-	for (i = 0; i < nr_pages; i++) {
-		pte_t pteval = ptep_get(pte + i);
-		struct page *page = folio_page(folio, i);
-		unsigned long src_addr = address + i * PAGE_SIZE;
-		struct page *src_page;
-
-		if (pte_none_or_zero(pteval)) {
-			clear_user_highpage(page, src_addr);
-			continue;
-		}
-		src_page = pte_page(pteval);
-		if (copy_mc_user_highpage(page, src_page, src_addr, vma) > 0) {
-			result = SCAN_COPY_MC;
-			break;
-		}
-	}
-
-	if (likely(result == SCAN_SUCCEED))
-		__collapse_huge_page_copy_succeeded(pte, vma, address, ptl,
-						    order, compound_pagelist);
-	else
-		__collapse_huge_page_copy_failed(pte, pmd, orig_pmd, vma,
-						 order, compound_pagelist);
-
-	return result;
-}
-
 static void khugepaged_alloc_sleep(void)
 {
 	DEFINE_WAIT(wait);
@@ -1089,119 +725,19 @@ enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 	return check_pmd_state(*pmd);
 }
 
-static enum scan_result check_pmd_still_valid(struct mm_struct *mm,
-		unsigned long address, pmd_t *pmd)
+static void count_collapse_event(unsigned int order, enum vm_event_item vm_event,
+		enum mthp_stat_item mthp_event)
 {
-	pmd_t *new_pmd;
-	enum scan_result result = find_pmd_or_thp_or_none(mm, address, &new_pmd);
-
-	if (result != SCAN_SUCCEED)
-		return result;
-	if (new_pmd != pmd)
-		return SCAN_FAIL;
-	return SCAN_SUCCEED;
+	if (is_pmd_order(order))
+		count_vm_event(vm_event);
+	count_mthp_stat(order, mthp_event);
 }
 
-/*
- * Bring missing pages in from swap, to complete THP collapse.
- * Only done if collapse_scan_pmd() believes it is worthwhile.
- *
- * For mTHP orders the function bails on the first swap entry, because
- * faulting pages back in during collapse could re-populate PTEs that
- * push a later scan over the threshold for a higher-order collapse.
- *
- * Called and returns without pte mapped or spinlocks held.
- * Returns result: if not SCAN_SUCCEED, mmap_lock has been released.
- */
-static enum scan_result __collapse_huge_page_swapin(struct mm_struct *mm,
-		struct vm_area_struct *vma, unsigned long start_addr,
-		pmd_t *pmd, int referenced, unsigned int order)
+static void collapse_control_init_scan(struct collapse_control *cc)
 {
-	int swapped_in = 0;
-	vm_fault_t ret = 0;
-	unsigned long addr, end = start_addr + (PAGE_SIZE << order);
-	enum scan_result result;
-	pte_t *pte = NULL;
-	spinlock_t *ptl;
-
-	for (addr = start_addr; addr < end; addr += PAGE_SIZE) {
-		struct vm_fault vmf = {
-			.vma = vma,
-			.address = addr,
-			.pgoff = linear_page_index(vma, addr),
-			.flags = FAULT_FLAG_ALLOW_RETRY,
-			.pmd = pmd,
-		};
-
-		if (!pte++) {
-			/*
-			 * Here the ptl is only used to check pte_same() in
-			 * do_swap_page(), so readonly version is enough.
-			 */
-			pte = pte_offset_map_ro_nolock(mm, pmd, addr, &ptl);
-			if (!pte) {
-				mmap_read_unlock(mm);
-				result = SCAN_NO_PTE_TABLE;
-				goto out;
-			}
-		}
-
-		vmf.orig_pte = ptep_get_lockless(pte);
-		if (pte_none(vmf.orig_pte) ||
-		    pte_present(vmf.orig_pte))
-			continue;
-
-		/*
-		 * TODO: Support swapin without leading to further mTHP
-		 * collapses. Currently bringing in new pages via swapin may
-		 * cause a future higher order collapse on a rescan of the same
-		 * range.
-		 */
-		if (!is_pmd_order(order)) {
-			count_mthp_stat(order, MTHP_STAT_COLLAPSE_EXCEED_SWAP);
-			pte_unmap(pte);
-			mmap_read_unlock(mm);
-			result = SCAN_EXCEED_SWAP_PTE;
-			goto out;
-		}
-
-		vmf.pte = pte;
-		vmf.ptl = ptl;
-		ret = do_swap_page(&vmf);
-		/* Which unmaps pte (after perhaps re-checking the entry) */
-		pte = NULL;
-
-		/*
-		 * do_swap_page() returns VM_FAULT_RETRY with released mmap_lock.
-		 * Note we treat VM_FAULT_RETRY as VM_FAULT_ERROR here because
-		 * we do not retry here and swap entry will remain in pagetable
-		 * resulting in later failure.
-		 */
-		if (ret & VM_FAULT_RETRY) {
-			/* Likely, but not guaranteed, that page lock failed */
-			result = SCAN_PAGE_LOCK;
-			goto out;
-		}
-		if (ret & VM_FAULT_ERROR) {
-			mmap_read_unlock(mm);
-			result = SCAN_FAIL;
-			goto out;
-		}
-		swapped_in++;
-	}
-
-	if (pte)
-		pte_unmap(pte);
-
-	/* Drain LRU cache to remove extra pin on the swapped in pages */
-	if (swapped_in)
-		lru_add_drain();
-
-	result = SCAN_SUCCEED;
-out:
-	trace_mm_collapse_huge_page_swapin(mm, swapped_in, referenced, result,
-					   order);
-	return result;
+	memset(cc->node_load, 0, sizeof(cc->node_load));
+	nodes_clear(cc->alloc_nmask);
+	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
 }
 
 static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_struct *mm,
@@ -1234,197 +770,6 @@ static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_stru
 	return SCAN_SUCCEED;
 }
 
-/*
- * collapse_huge_page() expects the mmap_lock to be unlocked before entering and
- * will always return with the lock unlocked, to avoid holding the mmap_lock
- * while allocating a THP, as that could trigger direct reclaim/compaction.
- * Note that the VMA must be rechecked after grabbing the mmap_lock again.
- */
-static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long start_addr,
-		int referenced, int unmapped, struct collapse_control *cc,
-		unsigned int order)
-{
-	const unsigned long pmd_addr = start_addr & HPAGE_PMD_MASK;
-	const unsigned long end_addr = start_addr + (PAGE_SIZE << order);
-	LIST_HEAD(compound_pagelist);
-	pmd_t *pmd, _pmd;
-	pte_t *pte = NULL;
-	pgtable_t pgtable;
-	struct folio *folio;
-	spinlock_t *pmd_ptl, *pte_ptl;
-	enum scan_result result = SCAN_FAIL;
-	struct vm_area_struct *vma;
-	struct mmu_notifier_range range;
-	bool anon_vma_locked = false;
-
-	result = alloc_charge_folio(&folio, mm, cc, order);
-	if (result != SCAN_SUCCEED)
-		goto out_nolock;
-
-	if (folio_memcg_alloc_deferred(folio)) {
-		result = SCAN_ALLOC_HUGE_PAGE_FAIL;
-		goto out_nolock;
-	}
-
-	mmap_read_lock(mm);
-	result = hugepage_vma_revalidate(mm, pmd_addr, /*expect_anon=*/ true,
-					 &vma, cc, order);
-	if (result != SCAN_SUCCEED) {
-		mmap_read_unlock(mm);
-		goto out_nolock;
-	}
-
-	result = find_pmd_or_thp_or_none(mm, pmd_addr, &pmd);
-	if (result != SCAN_SUCCEED) {
-		mmap_read_unlock(mm);
-		goto out_nolock;
-	}
-
-	if (unmapped) {
-		/*
-		 * __collapse_huge_page_swapin() will return with mmap_lock
-		 * released when it fails. So we jump out_nolock directly in
-		 * that case.  Continuing to collapse causes inconsistency.
-		 */
-		result = __collapse_huge_page_swapin(mm, vma, start_addr, pmd,
-						     referenced, order);
-		if (result != SCAN_SUCCEED)
-			goto out_nolock;
-	}
-
-	mmap_read_unlock(mm);
-	/*
-	 * Prevent all access to pagetables with the exception of
-	 * gup_fast later handled by the pmdp_collapse_flush() and the VM
-	 * handled by the anon_vma lock + folio lock.
-	 *
-	 * UFFDIO_MOVE is prevented to race as well thanks to the
-	 * mmap_lock.
-	 */
-	mmap_write_lock(mm);
-	result = hugepage_vma_revalidate(mm, pmd_addr, /*expect_anon=*/ true,
-					 &vma, cc, order);
-	if (result != SCAN_SUCCEED)
-		goto out_up_write;
-	/* check if the pmd is still valid */
-	vma_start_write(vma);
-	result = check_pmd_still_valid(mm, pmd_addr, pmd);
-	if (result != SCAN_SUCCEED)
-		goto out_up_write;
-
-	anon_vma_lock_write(vma->anon_vma);
-	anon_vma_locked = true;
-
-	/*
-	 * Only notify about the PTE range we will actually modify. While we
-	 * temporary unmap the whole PTE table for mTHP collapse, we'll remap
-	 * it later, leaving other PTEs effectively unmodified. The locks we
-	 * hold prevent anybody from stumbling over such temporarily unmapped
-	 * PTE tables.
-	 */
-	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm, start_addr,
-				end_addr);
-	mmu_notifier_invalidate_range_start(&range);
-
-	pmd_ptl = pmd_lock(mm, pmd); /* probably unnecessary */
-	/*
-	 * This removes any huge TLB entry from the CPU so we won't allow
-	 * huge and small TLB entries for the same virtual address to
-	 * avoid the risk of CPU bugs in that area.
-	 *
-	 * Parallel GUP-fast is fine since GUP-fast will back off when
-	 * it detects PMD is changed.
-	 */
-	_pmd = pmdp_collapse_flush(vma, pmd_addr, pmd);
-	spin_unlock(pmd_ptl);
-	mmu_notifier_invalidate_range_end(&range);
-	tlb_remove_table_sync_one();
-
-	pte = pte_offset_map_lock(mm, &_pmd, start_addr, &pte_ptl);
-	if (pte) {
-		result = __collapse_huge_page_isolate(vma, start_addr, pte, cc,
-						      order, &compound_pagelist);
-		spin_unlock(pte_ptl);
-	} else {
-		result = SCAN_NO_PTE_TABLE;
-	}
-
-	if (unlikely(result != SCAN_SUCCEED)) {
-		spin_lock(pmd_ptl);
-		VM_WARN_ON_ONCE(!pmd_none(*pmd));
-		/*
-		 * We can only use set_pmd_at() when establishing
-		 * hugepmds and never for establishing regular pmds that
-		 * points to regular pagetables. Use pmd_populate() for that
-		 */
-		pmd_populate(mm, pmd, pmd_pgtable(_pmd));
-		spin_unlock(pmd_ptl);
-		goto out_up_write;
-	}
-
-	/*
-	 * For PMD collapse all pages are isolated and locked so anon_vma
-	 * rmap can't run anymore. For mTHP collapse the PMD entry has been
-	 * removed and not all pages are isolated and locked, so we must hold
-	 * the lock to prevent neighboring folios from attempting to access
-	 * this PMD until its reinstalled.
-	 */
-	if (is_pmd_order(order)) {
-		anon_vma_unlock_write(vma->anon_vma);
-		anon_vma_locked = false;
-	}
-
-	result = __collapse_huge_page_copy(pte, folio, pmd, _pmd,
-					   vma, start_addr, pte_ptl,
-					   order, &compound_pagelist);
-	if (unlikely(result != SCAN_SUCCEED))
-		goto out_up_write;
-
-	/*
-	 * The smp_wmb() inside __folio_mark_uptodate() ensures the
-	 * copy_huge_page writes become visible before the set_pmd_at()
-	 * write.
-	 */
-	__folio_mark_uptodate(folio);
-	spin_lock(pmd_ptl);
-	VM_WARN_ON_ONCE(!pmd_none(*pmd));
-	if (is_pmd_order(order)) {
-		pgtable = pmd_pgtable(_pmd);
-		pgtable_trans_huge_deposit(mm, pmd, pgtable);
-		map_anon_folio_pmd_nopf(folio, pmd, vma, pmd_addr);
-	} else {
-		/*
-		 * Some architectures (e.g. MIPS) walk the live page table in
-		 * their implementation. update_mmu_cache_range() must be called
-		 * with a valid page table hierarchy and the PTE lock held.
-		 * Acquire it nested inside pmd_ptl when they are distinct locks.
-		 */
-		if (pte_ptl != pmd_ptl)
-			spin_lock_nested(pte_ptl, SINGLE_DEPTH_NESTING);
-		pmd_populate(mm, pmd, pmd_pgtable(_pmd));
-		map_anon_folio_pte_nopf(folio, pte, vma, start_addr,
-					  /*uffd_wp=*/ false);
-		if (pte_ptl != pmd_ptl)
-			spin_unlock(pte_ptl);
-	}
-	spin_unlock(pmd_ptl);
-
-	folio = NULL;
-
-	result = SCAN_SUCCEED;
-out_up_write:
-	if (pte)
-		pte_unmap(pte);
-	if (anon_vma_locked)
-		anon_vma_unlock_write(vma->anon_vma);
-	mmap_write_unlock(mm);
-out_nolock:
-	if (folio)
-		folio_put(folio);
-	trace_mm_collapse_huge_page(mm, result == SCAN_SUCCEED, result, order);
-	return result;
-}
-
 /* Return the highest naturally aligned order that fits at @offset within a PMD. */
 unsigned int max_order_from_offset(unsigned int offset)
 {
@@ -1434,317 +779,6 @@ unsigned int max_order_from_offset(unsigned int offset)
 	return min_t(unsigned int, __ffs(offset), HPAGE_PMD_ORDER);
 }
 
-/*
- * mthp_collapse() consumes the bitmap that is generated during
- * collapse_scan_pmd() to determine what regions and mTHP orders fit best.
- *
- * Each bit in cc->eligible_ptes marks a PTE the scan accepted as a collapse
- * source. We start at the PMD order and check if it is eligible for collapse;
- * if not, we check the left and right halves of the PTE page table we are
- * examining at a lower order.
- *
- * For each of these, we determine how many PTE entries are occupied in the
- * range of PTE entries we propose to collapse, then we compare this to a
- * threshold number of PTE entries which would need to be occupied for a
- * collapse to be permitted at that order (accounting for max_ptes_none).
- *
- * If a collapse is permitted, we attempt to collapse the PTE range into a
- * mTHP.
- */
-static enum scan_result mthp_collapse(struct mm_struct *mm,
-		unsigned long address, int referenced, int unmapped,
-		struct collapse_control *cc, unsigned long enabled_orders)
-{
-	unsigned int nr_occupied_ptes, nr_ptes, max_ptes_none;
-	enum scan_result last_result = SCAN_FAIL;
-	int collapsed = 0;
-	bool alloc_failed = false;
-	unsigned long collapse_address;
-	unsigned int offset = 0;
-	unsigned int order = HPAGE_PMD_ORDER;
-
-	while (offset < HPAGE_PMD_NR) {
-		nr_ptes = 1UL << order;
-
-		if (!test_bit(order, &enabled_orders))
-			goto next_order;
-
-		max_ptes_none = collapse_max_ptes_none(cc, NULL, order);
-		nr_occupied_ptes = bitmap_weight_from(cc->eligible_ptes, offset,
-						      offset + nr_ptes);
-
-		/*
-		 * Swap PTEs accepted during the scan are counted in @unmapped,
-		 * not in the eligible bitmap. Account them for the PMD-order
-		 * candidate.
-		 */
-		if (is_pmd_order(order))
-			nr_occupied_ptes += unmapped;
-
-		if (nr_occupied_ptes >= nr_ptes - max_ptes_none) {
-			enum scan_result ret;
-
-			collapse_address = address + offset * PAGE_SIZE;
-			ret = collapse_huge_page(mm, collapse_address, referenced,
-						 unmapped, cc, order);
-
-			switch (ret) {
-			/* Cases where we continue to next collapse candidate */
-			case SCAN_SUCCEED:
-				collapsed += nr_ptes;
-				fallthrough;
-			case SCAN_PTE_MAPPED_HUGEPAGE:
-				goto next_offset;
-			/* Cases where lower orders might still succeed */
-			case SCAN_ALLOC_HUGE_PAGE_FAIL:
-				alloc_failed = true;
-				fallthrough;
-			case SCAN_LACK_REFERENCED_PAGE:
-			case SCAN_EXCEED_NONE_PTE:
-			case SCAN_EXCEED_SWAP_PTE:
-			case SCAN_EXCEED_SHARED_PTE:
-			case SCAN_PAGE_LOCK:
-			case SCAN_PAGE_COUNT:
-			case SCAN_PAGE_NULL:
-			case SCAN_DEL_PAGE_LRU:
-			case SCAN_PTE_NON_PRESENT:
-			case SCAN_PTE_UFFD:
-			case SCAN_PAGE_LAZYFREE:
-				last_result = ret;
-				goto next_order;
-			/* Cases where no further collapse is possible */
-			case SCAN_PMD_MAPPED:
-				fallthrough;
-			default:
-				last_result = ret;
-				goto done;
-			}
-		}
-
-next_order:
-		/*
-		 * Continue with the next smaller order if there is still
-		 * any smaller order enabled. When at the smallest order
-		 * we must always move to the next offset.
-		 */
-		if (order > COLLAPSE_MIN_MTHP_ORDER &&
-		    (enabled_orders & GENMASK(order - 1, 0))) {
-			order--;
-			continue;
-		}
-next_offset:
-		/*
-		 * Advance past the region we just processed and determine the
-		 * highest order we can attempt next. Since huge pages must be
-		 * naturally aligned, the max order we can attempt next is
-		 * limited by the alignment of the new offset.
-		 * E.g. if we collapsed a order-2 mTHP at offset 0, offset
-		 * becomes 4 and __ffs(4) == 2, so the next attempt starts at
-		 * order 2.
-		 */
-		offset += nr_ptes;
-		order = max_order_from_offset(offset);
-	}
-done:
-	if (collapsed)
-		return SCAN_SUCCEED;
-	if (alloc_failed)
-		return SCAN_ALLOC_HUGE_PAGE_FAIL;
-	return last_result;
-}
-
-static enum scan_result __maybe_unused
-collapse_scan_pmd(struct mm_struct *mm,
-		struct vm_area_struct *vma, unsigned long start_addr,
-		bool *lock_dropped, struct collapse_control *cc)
-{
-	const unsigned int max_ptes_shared = collapse_max_ptes_shared(cc, HPAGE_PMD_ORDER);
-	const unsigned int max_ptes_swap = collapse_max_ptes_swap(cc, HPAGE_PMD_ORDER);
-	unsigned int max_ptes_none = collapse_max_ptes_none(cc, vma, HPAGE_PMD_ORDER);
-	enum tva_type tva_flags = cc->policy.tva_type;
-	pmd_t *pmd;
-	pte_t *pte, *_pte, pteval;
-	int i;
-	int none_or_zero = 0, shared = 0, referenced = 0;
-	enum scan_result result = SCAN_FAIL;
-	struct page *page = NULL;
-	struct folio *folio = NULL;
-	unsigned long addr;
-	unsigned long enabled_orders;
-	spinlock_t *ptl;
-	int node = NUMA_NO_NODE, unmapped = 0;
-
-	VM_BUG_ON(start_addr & ~HPAGE_PMD_MASK);
-
-	result = find_pmd_or_thp_or_none(mm, start_addr, &pmd);
-	if (result != SCAN_SUCCEED) {
-		cc->progress++;
-		goto out;
-	}
-
-	collapse_control_init_scan(cc);
-
-	enabled_orders = collapse_possible_orders(vma, vma->vm_flags, tva_flags);
-
-	/*
-	 * If PMD is the only enabled order, enforce max_ptes_none, otherwise
-	 * scan all pages to populate the bitmap for mTHP collapse. The bitmap
-	 * is then checked again in mthp_collapse() for each attempted order.
-	 */
-	if (enabled_orders != BIT(HPAGE_PMD_ORDER))
-		max_ptes_none = KHUGEPAGED_MAX_PTES_LIMIT;
-
-	pte = pte_offset_map_lock(mm, pmd, start_addr, &ptl);
-	if (!pte) {
-		cc->progress++;
-		result = SCAN_NO_PTE_TABLE;
-		goto out;
-	}
-
-	for (i = 0; i < HPAGE_PMD_NR; i++) {
-		_pte = pte + i;
-		addr = start_addr + i * PAGE_SIZE;
-		pteval = ptep_get(_pte);
-
-		cc->progress++;
-
-		if (pte_none_or_zero(pteval)) {
-			if (++none_or_zero > max_ptes_none) {
-				result = SCAN_EXCEED_NONE_PTE;
-				count_collapse_event(HPAGE_PMD_ORDER, THP_SCAN_EXCEED_NONE_PTE,
-						     MTHP_STAT_COLLAPSE_EXCEED_NONE);
-				goto out_unmap;
-			}
-			continue;
-		}
-		if (!pte_present(pteval)) {
-			if (++unmapped > max_ptes_swap) {
-				result = SCAN_EXCEED_SWAP_PTE;
-				count_collapse_event(HPAGE_PMD_ORDER, THP_SCAN_EXCEED_SWAP_PTE,
-						     MTHP_STAT_COLLAPSE_EXCEED_SWAP);
-				goto out_unmap;
-			}
-			/*
-			 * Always be strict with uffd-wp
-			 * enabled swap entries.  Please see
-			 * comment below for pte_uffd().
-			 */
-			if (pte_swp_uffd_any(pteval)) {
-				result = SCAN_PTE_UFFD;
-				goto out_unmap;
-			}
-			continue;
-		}
-		if (pte_uffd(pteval)) {
-			/*
-			 * Don't collapse the page if any of the small
-			 * PTEs are armed with uffd write protection.
-			 * Here we can also mark the new huge pmd as
-			 * write protected if any of the small ones is
-			 * marked but that could bring unknown
-			 * userfault messages that falls outside of
-			 * the registered range.  So, just be simple.
-			 */
-			result = SCAN_PTE_UFFD;
-			goto out_unmap;
-		}
-
-		page = vm_normal_page(vma, addr, pteval);
-		if (unlikely(!page) || unlikely(is_zone_device_page(page))) {
-			result = SCAN_PAGE_NULL;
-			goto out_unmap;
-		}
-		folio = page_folio(page);
-
-		/*
-		 * If the vma has the VM_DROPPABLE flag, the collapse will
-		 * preserve the lazyfree property without needing to skip.
-		 */
-		if (cc->policy.skip_lazyfree && !(vma->vm_flags & VM_DROPPABLE) &&
-		    folio_test_lazyfree(folio) && !pte_dirty(pteval)) {
-			result = SCAN_PAGE_LAZYFREE;
-			goto out_unmap;
-		}
-
-		if (!folio_test_anon(folio)) {
-			result = SCAN_PAGE_ANON;
-			goto out_unmap;
-		}
-
-		/*
-		 * We treat a single page as shared if any part of the THP
-		 * is shared.
-		 */
-		if (folio_maybe_mapped_shared(folio)) {
-			if (++shared > max_ptes_shared) {
-				result = SCAN_EXCEED_SHARED_PTE;
-				count_collapse_event(HPAGE_PMD_ORDER, THP_SCAN_EXCEED_SHARED_PTE,
-						     MTHP_STAT_COLLAPSE_EXCEED_SHARED);
-				goto out_unmap;
-			}
-		}
-
-		__set_bit(i, cc->eligible_ptes);
-		/*
-		 * Record which node the original page is from and save this
-		 * information to cc->node_load[].
-		 * Khugepaged will allocate hugepage from the node has the max
-		 * hit record.
-		 */
-		node = folio_nid(folio);
-		if (collapse_scan_abort(node, cc)) {
-			result = SCAN_SCAN_ABORT;
-			goto out_unmap;
-		}
-		cc->node_load[node]++;
-		if (!folio_test_lru(folio)) {
-			result = SCAN_PAGE_LRU;
-			goto out_unmap;
-		}
-		if (folio_test_locked(folio)) {
-			result = SCAN_PAGE_LOCK;
-			goto out_unmap;
-		}
-
-		/*
-		 * Check if the page has any GUP (or other external) pins.
-		 *
-		 * Here the check is racy, but such case is ephemeral and
-		 * we could always retry collapse later. Anyway the same
-		 * check will be done again later the risk seems low.
-		 */
-		if (folio_expected_ref_count(folio) != folio_ref_count(folio)) {
-			result = SCAN_PAGE_COUNT;
-			goto out_unmap;
-		}
-
-		if (cc->policy.require_referenced &&
-		    folio_pte_referenced(folio, vma, addr, pteval))
-			referenced++;
-	}
-	if (cc->policy.require_referenced &&
-		   (!referenced ||
-		    (unmapped && referenced < HPAGE_PMD_NR / 2))) {
-		result = SCAN_LACK_REFERENCED_PAGE;
-	} else {
-		result = SCAN_SUCCEED;
-	}
-out_unmap:
-	pte_unmap_unlock(pte, ptl);
-	if (result == SCAN_SUCCEED) {
-		/* collapse_huge_page() expects the lock to be dropped before calling */
-		mmap_read_unlock(mm);
-		result = mthp_collapse(mm, start_addr, referenced,
-				       unmapped, cc, enabled_orders);
-		/* mmap_lock was released above, set lock_dropped */
-		*lock_dropped = true;
-	}
-out:
-	trace_mm_khugepaged_scan_pmd(mm, folio, referenced,
-				     none_or_zero, result, unmapped);
-	return result;
-}
-
 static void collect_mm_slot(struct mm_slot *slot)
 {
 	struct mm_struct *mm = slot->mm;
-- 
2.54.0

[RFC PATCH 29/57] mm/collapse: move what a collapse is judged on into collapse.c

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

collapse.c had to call back into khugepaged.c for nine things: the PMD
state checks, the orders a VMA allows, the per-order limits, the NUMA
heuristics and the window alignment.  The engine therefore depended on
the daemon, which is backwards -- khugepaged is one caller of a collapse,
not the place a collapse lives.

Move them across.  max_order_from_offset() and collapse_max_ptes_shared()
have no caller left outside collapse.c and land static; the rest stay
exported while khugepaged.c's file paths still ask for them.

KHUGEPAGED_MAX_PTES_LIMIT goes to collapse.h with them, being both the
ceiling the tunables accept and the value they read as "no limit".

Every function body is byte for byte what it was, and the dependency now
points one way.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 187 +++++++++++++++++++++++++++++++++++++++++++++++
 mm/collapse.h   |  11 +--
 mm/khugepaged.c | 189 ------------------------------------------------
 3 files changed, 191 insertions(+), 196 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index df3760e3918b..43a4b771bbe1 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -113,6 +113,193 @@
 	min(COLLAPSE_BATCH_BYTES >> (PAGE_SHIFT + COLLAPSE_MIN_MTHP_ORDER), \
 	    COLLAPSE_TABLE_WINDOWS)
 
+static inline enum scan_result check_pmd_state(pmd_t *pmd)
+{
+	pmd_t pmde = pmdp_get_lockless(pmd);
+
+	if (pmd_none(pmde))
+		return SCAN_NO_PTE_TABLE;
+
+	/*
+	 * The folio may be under migration when khugepaged is trying to
+	 * collapse it. Migration success or failure will eventually end
+	 * up with a present PMD mapping a folio again.
+	 */
+	if (pmd_is_migration_entry(pmde))
+		return SCAN_PMD_MAPPED;
+	if (!pmd_present(pmde))
+		return SCAN_NO_PTE_TABLE;
+	if (pmd_trans_huge(pmde))
+		return SCAN_PMD_MAPPED;
+	if (pmd_bad(pmde))
+		return SCAN_NO_PTE_TABLE;
+	return SCAN_SUCCEED;
+}
+
+enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
+		unsigned long address, pmd_t **pmd)
+{
+	*pmd = mm_find_pmd(mm, address);
+	if (!*pmd)
+		return SCAN_NO_PTE_TABLE;
+
+	return check_pmd_state(*pmd);
+}
+
+/*
+ * Check what orders are possible based on the vma and collapse type.
+ * This is used to determine if mTHP collapse is a viable option.
+ */
+unsigned long collapse_possible_orders(struct vm_area_struct *vma,
+		vm_flags_t vm_flags, enum tva_type tva_flags)
+{
+	unsigned long orders;
+
+	/* If khugepaged is scanning an anonymous vma, allow mTHP collapse */
+	if ((tva_flags == TVA_KHUGEPAGED) && vma_is_anonymous(vma))
+		orders = THP_ORDERS_ALL_ANON;
+	else
+		orders = BIT(HPAGE_PMD_ORDER);
+
+	return thp_vma_allowable_orders(vma, vm_flags, tva_flags, orders);
+}
+
+/* Return the highest naturally aligned order that fits at @offset within a PMD. */
+static unsigned int max_order_from_offset(unsigned int offset)
+{
+	if (offset == 0)
+		return HPAGE_PMD_ORDER;
+
+	return min_t(unsigned int, __ffs(offset), HPAGE_PMD_ORDER);
+}
+
+/**
+ * collapse_max_ptes_none - Calculate maximum allowed empty PTEs or PTEs mapping
+ * the shared zeropage for the given collapse operation.
+ * @cc: The collapse control struct
+ * @vma: The vma to check for userfaultfd
+ * @order: The folio order being collapsed to
+ *
+ * Return: Maximum number of empty/shared zeropage PTEs for the collapse operation
+ */
+unsigned int collapse_max_ptes_none(struct collapse_control *cc,
+		struct vm_area_struct *vma, unsigned int order)
+{
+	const unsigned int max_ptes_none = cc->policy.max_ptes_none;
+
+	if (vma && userfaultfd_armed(vma))
+		return 0;
+	/* The limit as given, at the PMD order and wherever it is not capped */
+	if (is_pmd_order(order) || !cc->policy.strict_sub_pmd)
+		return max_ptes_none;
+	/*
+	 * for mTHP collapse with the sysctl value set to KHUGEPAGED_MAX_PTES_LIMIT,
+	 * scale the maximum number of PTEs to the order of the collapse.
+	 */
+	if (max_ptes_none == KHUGEPAGED_MAX_PTES_LIMIT)
+		return (1 << order) - 1;
+	/*
+	 * For mTHP collapse of values other than 0 or KHUGEPAGED_MAX_PTES_LIMIT,
+	 * emit a warning and return 0.
+	 */
+	if (max_ptes_none)
+		pr_warn_once("mTHP collapse does not support max_ptes_none"
+		     " values other than 0 or %u, defaulting to 0.\n",
+		     KHUGEPAGED_MAX_PTES_LIMIT);
+	return 0;
+}
+
+/**
+ * collapse_max_ptes_shared - Calculate maximum allowed PTEs that map shared
+ * anonymous pages for the given collapse operation.
+ * @cc: The collapse control struct
+ * @order: The folio order being collapsed to
+ *
+ * Return: Maximum number of PTEs that map shared anonymous pages for the
+ * collapse operation
+ */
+static unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
+		unsigned int order)
+{
+	/*
+	 * A sub-PMD window held to the strict rule takes no shared page at all:
+	 * an mTHP is not worth the CoW-breaking.
+	 */
+	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
+		return 0;
+	return cc->policy.max_ptes_shared;
+}
+
+/**
+ * collapse_max_ptes_swap - Calculate the maximum allowed non-present PTEs or the
+ * maximum allowed non-present pagecache entries for the given collapse operation.
+ * @cc: The collapse control struct
+ * @order: The folio order being collapsed to
+ *
+ * Return: Maximum number of non-present PTEs or the maximum allowed non-present
+ * pagecache entries for the collapse operation.
+ */
+unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
+		unsigned int order)
+{
+	/*
+	 * A sub-PMD window held to the strict rule takes nothing non-present:
+	 * reading pages back to build an mTHP is not worth the latency.
+	 */
+	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
+		return 0;
+	return cc->policy.max_ptes_swap;
+}
+
+bool collapse_scan_abort(int nid, struct collapse_control *cc)
+{
+	int i;
+
+	/*
+	 * If node_reclaim_mode is disabled, then no extra effort is made to
+	 * allocate memory locally.
+	 */
+	if (!node_reclaim_enabled())
+		return false;
+
+	/* If there is a count for this node already, it must be acceptable */
+	if (cc->node_load[nid])
+		return false;
+
+	for (i = 0; i < MAX_NUMNODES; i++) {
+		if (!cc->node_load[i])
+			continue;
+		if (node_distance(nid, i) > node_reclaim_distance)
+			return true;
+	}
+	return false;
+}
+
+#ifdef CONFIG_NUMA
+int collapse_find_target_node(struct collapse_control *cc)
+{
+	int nid, target_node = 0, max_value = 0;
+
+	/* find first node with max normal pages hit */
+	for (nid = 0; nid < MAX_NUMNODES; nid++)
+		if (cc->node_load[nid] > max_value) {
+			max_value = cc->node_load[nid];
+			target_node = nid;
+		}
+
+	for_each_online_node(nid) {
+		if (max_value == cc->node_load[nid])
+			node_set(nid, cc->alloc_nmask);
+	}
+
+	return target_node;
+}
+#else
+int collapse_find_target_node(struct collapse_control *cc)
+{
+	return 0;
+}
+#endif
 /*
  * The saved-PTE pool spans a whole table.  The byte cap bounds what a round
  * holds, but not what one candidate does: a sub-PMD order goes up to
diff --git a/mm/collapse.h b/mm/collapse.h
index 50a9d59bbf03..2fdb6418653a 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -7,6 +7,9 @@
 #include <linux/pgtable.h>
 #include <linux/types.h>
 
+/* Ceiling the max_ptes_* tunables accept, and the value meaning "no limit" */
+#define KHUGEPAGED_MAX_PTES_LIMIT	(HPAGE_PMD_NR - 1)
+
 /* The smallest order a collapse will build, and so the finest window it cuts */
 #define COLLAPSE_MIN_MTHP_ORDER		2
 
@@ -194,22 +197,16 @@ enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
 
-/*
- * Defined in khugepaged.c, which still uses them itself.
- * TODO: move each into collapse.c once its last khugepaged.c user is gone.
- */
 unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags);
+enum scan_result check_pmd_state(pmd_t *pmd);
 enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd);
 int collapse_find_target_node(struct collapse_control *cc);
 bool collapse_scan_abort(int nid, struct collapse_control *cc);
-unsigned int max_order_from_offset(unsigned int offset);
 unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order);
 unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
 		unsigned int order);
-unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
-		unsigned int order);
 
 #endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 6203473f4953..26c0e961ac9f 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -57,7 +57,6 @@ static DECLARE_WAIT_QUEUE_HEAD(khugepaged_wait);
  *
  * Note that these are only respected if collapse was initiated by khugepaged.
  */
-#define KHUGEPAGED_MAX_PTES_LIMIT (HPAGE_PMD_NR - 1)
 unsigned int khugepaged_max_ptes_none __read_mostly;
 static unsigned int khugepaged_max_ptes_swap __read_mostly;
 static unsigned int khugepaged_max_ptes_shared __read_mostly;
@@ -296,84 +295,6 @@ struct attribute_group khugepaged_attr_group = {
 };
 #endif /* CONFIG_SYSFS */
 
-/**
- * collapse_max_ptes_none - Calculate maximum allowed empty PTEs or PTEs mapping
- * the shared zeropage for the given collapse operation.
- * @cc: The collapse control struct
- * @vma: The vma to check for userfaultfd
- * @order: The folio order being collapsed to
- *
- * Return: Maximum number of empty/shared zeropage PTEs for the collapse operation
- */
-unsigned int collapse_max_ptes_none(struct collapse_control *cc,
-		struct vm_area_struct *vma, unsigned int order)
-{
-	const unsigned int max_ptes_none = cc->policy.max_ptes_none;
-
-	if (vma && userfaultfd_armed(vma))
-		return 0;
-	/* The limit as given, at the PMD order and wherever it is not capped */
-	if (is_pmd_order(order) || !cc->policy.strict_sub_pmd)
-		return max_ptes_none;
-	/*
-	 * for mTHP collapse with the sysctl value set to KHUGEPAGED_MAX_PTES_LIMIT,
-	 * scale the maximum number of PTEs to the order of the collapse.
-	 */
-	if (max_ptes_none == KHUGEPAGED_MAX_PTES_LIMIT)
-		return (1 << order) - 1;
-	/*
-	 * For mTHP collapse of values other than 0 or KHUGEPAGED_MAX_PTES_LIMIT,
-	 * emit a warning and return 0.
-	 */
-	if (max_ptes_none)
-		pr_warn_once("mTHP collapse does not support max_ptes_none"
-		     " values other than 0 or %u, defaulting to 0.\n",
-		     KHUGEPAGED_MAX_PTES_LIMIT);
-	return 0;
-}
-
-/**
- * collapse_max_ptes_shared - Calculate maximum allowed PTEs that map shared
- * anonymous pages for the given collapse operation.
- * @cc: The collapse control struct
- * @order: The folio order being collapsed to
- *
- * Return: Maximum number of PTEs that map shared anonymous pages for the
- * collapse operation
- */
-unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
-		unsigned int order)
-{
-	/*
-	 * A sub-PMD window held to the strict rule takes no shared page at all:
-	 * an mTHP is not worth the CoW-breaking.
-	 */
-	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
-		return 0;
-	return cc->policy.max_ptes_shared;
-}
-
-/**
- * collapse_max_ptes_swap - Calculate the maximum allowed non-present PTEs or the
- * maximum allowed non-present pagecache entries for the given collapse operation.
- * @cc: The collapse control struct
- * @order: The folio order being collapsed to
- *
- * Return: Maximum number of non-present PTEs or the maximum allowed non-present
- * pagecache entries for the collapse operation.
- */
-unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
-		unsigned int order)
-{
-	/*
-	 * A sub-PMD window held to the strict rule takes nothing non-present:
-	 * reading pages back to build an mTHP is not worth the latency.
-	 */
-	if (!is_pmd_order(order) && cc->policy.strict_sub_pmd)
-		return 0;
-	return cc->policy.max_ptes_swap;
-}
-
 int hugepage_madvise(struct vm_area_struct *vma,
 		     vm_flags_t *vm_flags, int advice)
 {
@@ -483,24 +404,6 @@ void __khugepaged_enter(struct mm_struct *mm)
 		wake_up_interruptible(&khugepaged_wait);
 }
 
-/*
- * Check what orders are possible based on the vma and collapse type.
- * This is used to determine if mTHP collapse is a viable option.
- */
-unsigned long collapse_possible_orders(struct vm_area_struct *vma,
-		vm_flags_t vm_flags, enum tva_type tva_flags)
-{
-	unsigned long orders;
-
-	/* If khugepaged is scanning an anonymous vma, allow mTHP collapse */
-	if ((tva_flags == TVA_KHUGEPAGED) && vma_is_anonymous(vma))
-		orders = THP_ORDERS_ALL_ANON;
-	else
-		orders = BIT(HPAGE_PMD_ORDER);
-
-	return thp_vma_allowable_orders(vma, vm_flags, tva_flags, orders);
-}
-
 static bool collapse_possible(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags)
 {
@@ -559,30 +462,6 @@ static struct collapse_control khugepaged_collapse_control = {
 	.is_khugepaged = true,
 };
 
-bool collapse_scan_abort(int nid, struct collapse_control *cc)
-{
-	int i;
-
-	/*
-	 * If node_reclaim_mode is disabled, then no extra effort is made to
-	 * allocate memory locally.
-	 */
-	if (!node_reclaim_enabled())
-		return false;
-
-	/* If there is a count for this node already, it must be acceptable */
-	if (cc->node_load[nid])
-		return false;
-
-	for (i = 0; i < MAX_NUMNODES; i++) {
-		if (!cc->node_load[i])
-			continue;
-		if (node_distance(nid, i) > node_reclaim_distance)
-			return true;
-	}
-	return false;
-}
-
 #define khugepaged_defrag()					\
 	(transparent_hugepage_flags &				\
 	 (1<<TRANSPARENT_HUGEPAGE_DEFRAG_KHUGEPAGED_FLAG))
@@ -623,32 +502,6 @@ static void collapse_policy_forced(struct collapse_policy *p)
 	p->tva_type = TVA_FORCED_COLLAPSE;
 }
 
-#ifdef CONFIG_NUMA
-int collapse_find_target_node(struct collapse_control *cc)
-{
-	int nid, target_node = 0, max_value = 0;
-
-	/* find first node with max normal pages hit */
-	for (nid = 0; nid < MAX_NUMNODES; nid++)
-		if (cc->node_load[nid] > max_value) {
-			max_value = cc->node_load[nid];
-			target_node = nid;
-		}
-
-	for_each_online_node(nid) {
-		if (max_value == cc->node_load[nid])
-			node_set(nid, cc->alloc_nmask);
-	}
-
-	return target_node;
-}
-#else
-int collapse_find_target_node(struct collapse_control *cc)
-{
-	return 0;
-}
-#endif
-
 /*
  * If mmap_lock temporarily dropped, revalidate vma
  * after taking the mmap_lock again.
@@ -692,39 +545,6 @@ static enum scan_result hugepage_vma_revalidate(struct mm_struct *mm, unsigned l
 	return SCAN_SUCCEED;
 }
 
-static inline enum scan_result check_pmd_state(pmd_t *pmd)
-{
-	pmd_t pmde = pmdp_get_lockless(pmd);
-
-	if (pmd_none(pmde))
-		return SCAN_NO_PTE_TABLE;
-
-	/*
-	 * The folio may be under migration when khugepaged is trying to
-	 * collapse it. Migration success or failure will eventually end
-	 * up with a present PMD mapping a folio again.
-	 */
-	if (pmd_is_migration_entry(pmde))
-		return SCAN_PMD_MAPPED;
-	if (!pmd_present(pmde))
-		return SCAN_NO_PTE_TABLE;
-	if (pmd_trans_huge(pmde))
-		return SCAN_PMD_MAPPED;
-	if (pmd_bad(pmde))
-		return SCAN_NO_PTE_TABLE;
-	return SCAN_SUCCEED;
-}
-
-enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
-		unsigned long address, pmd_t **pmd)
-{
-	*pmd = mm_find_pmd(mm, address);
-	if (!*pmd)
-		return SCAN_NO_PTE_TABLE;
-
-	return check_pmd_state(*pmd);
-}
-
 static void count_collapse_event(unsigned int order, enum vm_event_item vm_event,
 		enum mthp_stat_item mthp_event)
 {
@@ -770,15 +590,6 @@ static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_stru
 	return SCAN_SUCCEED;
 }
 
-/* Return the highest naturally aligned order that fits at @offset within a PMD. */
-unsigned int max_order_from_offset(unsigned int offset)
-{
-	if (offset == 0)
-		return HPAGE_PMD_ORDER;
-
-	return min_t(unsigned int, __ffs(offset), HPAGE_PMD_ORDER);
-}
-
 static void collect_mm_slot(struct mm_slot *slot)
 {
 	struct mm_struct *mm = slot->mm;
-- 
2.54.0

[RFC PATCH 30/57] mm/collapse: name the max_ptes ceiling after collapse

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The constant lives in collapse.h and is read by the limit helpers there,
so KHUGEPAGED_MAX_PTES_LIMIT names it after one of its callers rather
than after what it is.

Rename it to COLLAPSE_MAX_PTES_LIMIT, matching the other constant the
header defines.  khugepaged keeps using it to bound what its tunables
accept.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 8 ++++----
 mm/collapse.h   | 2 +-
 mm/khugepaged.c | 8 ++++----
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 43a4b771bbe1..ae7c2777b279 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -193,19 +193,19 @@ unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 	if (is_pmd_order(order) || !cc->policy.strict_sub_pmd)
 		return max_ptes_none;
 	/*
-	 * for mTHP collapse with the sysctl value set to KHUGEPAGED_MAX_PTES_LIMIT,
+	 * for mTHP collapse with the sysctl value set to COLLAPSE_MAX_PTES_LIMIT,
 	 * scale the maximum number of PTEs to the order of the collapse.
 	 */
-	if (max_ptes_none == KHUGEPAGED_MAX_PTES_LIMIT)
+	if (max_ptes_none == COLLAPSE_MAX_PTES_LIMIT)
 		return (1 << order) - 1;
 	/*
-	 * For mTHP collapse of values other than 0 or KHUGEPAGED_MAX_PTES_LIMIT,
+	 * For mTHP collapse of values other than 0 or COLLAPSE_MAX_PTES_LIMIT,
 	 * emit a warning and return 0.
 	 */
 	if (max_ptes_none)
 		pr_warn_once("mTHP collapse does not support max_ptes_none"
 		     " values other than 0 or %u, defaulting to 0.\n",
-		     KHUGEPAGED_MAX_PTES_LIMIT);
+		     COLLAPSE_MAX_PTES_LIMIT);
 	return 0;
 }
 
diff --git a/mm/collapse.h b/mm/collapse.h
index 2fdb6418653a..94c11051f06a 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -8,7 +8,7 @@
 #include <linux/types.h>
 
 /* Ceiling the max_ptes_* tunables accept, and the value meaning "no limit" */
-#define KHUGEPAGED_MAX_PTES_LIMIT	(HPAGE_PMD_NR - 1)
+#define COLLAPSE_MAX_PTES_LIMIT		(HPAGE_PMD_NR - 1)
 
 /* The smallest order a collapse will build, and so the finest window it cuts */
 #define COLLAPSE_MIN_MTHP_ORDER		2
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 26c0e961ac9f..8c770e251c22 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -214,7 +214,7 @@ static ssize_t max_ptes_none_store(struct kobject *kobj,
 	unsigned long max_ptes_none;
 
 	err = kstrtoul(buf, 10, &max_ptes_none);
-	if (err || max_ptes_none > KHUGEPAGED_MAX_PTES_LIMIT)
+	if (err || max_ptes_none > COLLAPSE_MAX_PTES_LIMIT)
 		return -EINVAL;
 
 	khugepaged_max_ptes_none = max_ptes_none;
@@ -239,7 +239,7 @@ static ssize_t max_ptes_swap_store(struct kobject *kobj,
 	unsigned long max_ptes_swap;
 
 	err  = kstrtoul(buf, 10, &max_ptes_swap);
-	if (err || max_ptes_swap > KHUGEPAGED_MAX_PTES_LIMIT)
+	if (err || max_ptes_swap > COLLAPSE_MAX_PTES_LIMIT)
 		return -EINVAL;
 
 	khugepaged_max_ptes_swap = max_ptes_swap;
@@ -265,7 +265,7 @@ static ssize_t max_ptes_shared_store(struct kobject *kobj,
 	unsigned long max_ptes_shared;
 
 	err  = kstrtoul(buf, 10, &max_ptes_shared);
-	if (err || max_ptes_shared > KHUGEPAGED_MAX_PTES_LIMIT)
+	if (err || max_ptes_shared > COLLAPSE_MAX_PTES_LIMIT)
 		return -EINVAL;
 
 	khugepaged_max_ptes_shared = max_ptes_shared;
@@ -330,7 +330,7 @@ int __init khugepaged_init(void)
 		return -ENOMEM;
 
 	khugepaged_pages_to_scan = HPAGE_PMD_NR * 8;
-	khugepaged_max_ptes_none = KHUGEPAGED_MAX_PTES_LIMIT;
+	khugepaged_max_ptes_none = COLLAPSE_MAX_PTES_LIMIT;
 	khugepaged_max_ptes_swap = HPAGE_PMD_NR / 8;
 	khugepaged_max_ptes_shared = HPAGE_PMD_NR / 2;
 
-- 
2.54.0

[RFC PATCH 31/57] mm/khugepaged: count collapses where khugepaged makes them

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

collapse_single_pmd() bumps khugepaged_pages_collapsed for its caller, and
tests cc->is_khugepaged to know whether it should: the counter belongs to
the daemon, and MADV_COLLAPSE must not touch it.  So the one thing the
shared path still asks about its caller is bookkeeping, not policy.

Count in khugepaged's own walk instead, at the call it already makes.  The
question goes away, and cc->is_khugepaged with it -- nothing else read it.
current_is_khugepaged() is a different test, on the task rather than on
the request.

Preparation for moving the dispatcher into collapse.c, from where a static
in khugepaged.c is out of reach.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   | 2 --
 mm/khugepaged.c | 9 +++------
 2 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 94c11051f06a..11f51c6ea444 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -114,8 +114,6 @@ struct collapse_policy {
 struct collapse_control {
 	struct collapse_policy policy;
 
-	bool is_khugepaged;
-
 	/* Num pages scanned per node */
 	u32 node_load[MAX_NUMNODES];
 
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 8c770e251c22..907ed1131460 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -458,9 +458,7 @@ static void khugepaged_alloc_sleep(void)
 	remove_wait_queue(&khugepaged_wait, &wait);
 }
 
-static struct collapse_control khugepaged_collapse_control = {
-	.is_khugepaged = true,
-};
+static struct collapse_control khugepaged_collapse_control;
 
 #define khugepaged_defrag()					\
 	(transparent_hugepage_flags &				\
@@ -1639,8 +1637,6 @@ static enum scan_result collapse_single_pmd(unsigned long addr,
 		mmap_read_unlock(mm);
 	}
 end:
-	if (cc->is_khugepaged && result == SCAN_SUCCEED)
-		++khugepaged_pages_collapsed;
 	return result;
 }
 
@@ -1731,6 +1727,8 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			*result = collapse_single_pmd(khugepaged_scan.address,
 						      range_end, vma,
 						      &lock_dropped, cc);
+			if (*result == SCAN_SUCCEED)
+				++khugepaged_pages_collapsed;
 			/* move to next address */
 			khugepaged_scan.address = range_end;
 			if (lock_dropped)
@@ -2039,7 +2037,6 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	cc = kmalloc_obj(*cc);
 	if (!cc)
 		return -ENOMEM;
-	cc->is_khugepaged = false;
 	collapse_policy_forced(&cc->policy);
 	cc->progress = 0;
 	err = collapse_control_init(cc);
-- 
2.54.0

[RFC PATCH 32/57] mm/collapse: move the file collapse into collapse.c

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Collapse is split across two files with no boundary to speak of: the
engine in collapse.c, the file and shmem half in khugepaged.c, helpers
reaching in both directions, and the function that chooses between them
sitting with the daemon.

Move the rest of the mechanism over -- the file collapse and its scan, the
PTE-mapped-THP recollapse and the table retraction it needs, and the
allocator they share -- and with them collapse_single_pmd(), which
dispatches on the VMA.

khugepaged.c keeps what is actually khugepaged: the tunables, the daemon
and its scan budget, the mm_slot bookkeeping, and the walk that decides
which ranges to offer.

Five helpers lose their last caller outside collapse.c and go static, as
do the engine's two entries now that the dispatcher reaches them from the
same file.  check_pmd_state() was already static inline, so its
declaration was only ever dead.  collapse_single_pmd() takes their place
in collapse.h, keeping the interface it has today: a VMA to dispatch on
and an out-param saying whether the lock survived.

Pure motion: every function arrives in collapse.c exactly as it left
khugepaged.c.  Beyond that the diff has only the four includes the moved
code needs, and the header declarations that changed hands.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 1095 ++++++++++++++++++++++++++++++++++++++++++++++-
 mm/collapse.h   |   17 +-
 mm/khugepaged.c | 1075 ----------------------------------------------
 3 files changed, 1090 insertions(+), 1097 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index ae7c2777b279..21bfbc038044 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -1,7 +1,9 @@
 // SPDX-License-Identifier: GPL-2.0
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
 
+#include <linux/backing-dev.h>
 #include <linux/bitops.h>
+#include <linux/dax.h>
 #include <linux/highmem.h>
 #include <linux/huge_mm.h>
 #include <linux/hugetlb.h>	/* x86 flush_tlb_range() uses hstate_vma() */
@@ -11,8 +13,10 @@
 #include <linux/mmu_notifier.h>
 #include <linux/pagemap.h>
 #include <linux/pgalloc.h>
+#include <linux/rcupdate_wait.h>
 #include <linux/rmap.h>
 #include <linux/sched.h>
+#include <linux/shmem_fs.h>
 #include <linux/sizes.h>
 #include <linux/slab.h>
 #include <linux/swap.h>
@@ -136,7 +140,7 @@ static inline enum scan_result check_pmd_state(pmd_t *pmd)
 	return SCAN_SUCCEED;
 }
 
-enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
+static enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
 		unsigned long address, pmd_t **pmd)
 {
 	*pmd = mm_find_pmd(mm, address);
@@ -182,7 +186,7 @@ static unsigned int max_order_from_offset(unsigned int offset)
  *
  * Return: Maximum number of empty/shared zeropage PTEs for the collapse operation
  */
-unsigned int collapse_max_ptes_none(struct collapse_control *cc,
+static unsigned int collapse_max_ptes_none(struct collapse_control *cc,
 		struct vm_area_struct *vma, unsigned int order)
 {
 	const unsigned int max_ptes_none = cc->policy.max_ptes_none;
@@ -239,7 +243,7 @@ static unsigned int collapse_max_ptes_shared(struct collapse_control *cc,
  * Return: Maximum number of non-present PTEs or the maximum allowed non-present
  * pagecache entries for the collapse operation.
  */
-unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
+static unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
 		unsigned int order)
 {
 	/*
@@ -251,7 +255,7 @@ unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
 	return cc->policy.max_ptes_swap;
 }
 
-bool collapse_scan_abort(int nid, struct collapse_control *cc)
+static bool collapse_scan_abort(int nid, struct collapse_control *cc)
 {
 	int i;
 
@@ -276,7 +280,7 @@ bool collapse_scan_abort(int nid, struct collapse_control *cc)
 }
 
 #ifdef CONFIG_NUMA
-int collapse_find_target_node(struct collapse_control *cc)
+static int collapse_find_target_node(struct collapse_control *cc)
 {
 	int nid, target_node = 0, max_value = 0;
 
@@ -295,7 +299,7 @@ int collapse_find_target_node(struct collapse_control *cc)
 	return target_node;
 }
 #else
-int collapse_find_target_node(struct collapse_control *cc)
+static int collapse_find_target_node(struct collapse_control *cc)
 {
 	return 0;
 }
@@ -2110,7 +2114,7 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
  * that acts on what it found hands the range to collapse_anon_pmd() afterwards,
  * without the lock.
  */
-enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
+static enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
 					unsigned long start, unsigned long end,
 					struct collapse_control *cc)
 {
@@ -2521,7 +2525,7 @@ static void collapse_add_candidate(struct collapse_control *cc,
  * largest order downwards.  Returns what the table yielded: a collapse, or
  * the reason it did not.
  */
-enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
+static enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
 				   unsigned long end,
 				   struct collapse_control *cc)
 {
@@ -2583,3 +2587,1078 @@ enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
 		return cc->scan_refusal;
 	return cc->select_result;
 }
+
+static void count_collapse_event(unsigned int order, enum vm_event_item vm_event,
+		enum mthp_stat_item mthp_event)
+{
+	if (is_pmd_order(order))
+		count_vm_event(vm_event);
+	count_mthp_stat(order, mthp_event);
+}
+
+static void collapse_control_init_scan(struct collapse_control *cc)
+{
+	memset(cc->node_load, 0, sizeof(cc->node_load));
+	nodes_clear(cc->alloc_nmask);
+	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
+}
+
+static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_struct *mm,
+		struct collapse_control *cc, unsigned int order)
+{
+	gfp_t gfp = cc->policy.gfp;
+	int node = collapse_find_target_node(cc);
+	struct folio *folio;
+
+	folio = __folio_alloc(gfp, order, node, &cc->alloc_nmask);
+	if (!folio) {
+		*foliop = NULL;
+		count_collapse_event(order, THP_COLLAPSE_ALLOC_FAILED,
+				     MTHP_STAT_COLLAPSE_ALLOC_FAILED);
+		return SCAN_ALLOC_HUGE_PAGE_FAIL;
+	}
+
+	count_collapse_event(order, THP_COLLAPSE_ALLOC, MTHP_STAT_COLLAPSE_ALLOC);
+
+	if (unlikely(mem_cgroup_charge(folio, mm, gfp))) {
+		folio_put(folio);
+		*foliop = NULL;
+		return SCAN_CGROUP_CHARGE_FAIL;
+	}
+
+	if (is_pmd_order(order))
+		count_memcg_folio_events(folio, THP_COLLAPSE_ALLOC, 1);
+
+	*foliop = folio;
+	return SCAN_SUCCEED;
+}
+
+/* folio must be locked, and mmap_lock must be held */
+static enum scan_result set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
+		pmd_t *pmdp, struct folio *folio, struct page *page)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	struct vm_fault vmf = {
+		.vma = vma,
+		.address = addr,
+		.flags = 0,
+	};
+	pgd_t *pgdp;
+	p4d_t *p4dp;
+	pud_t *pudp;
+
+	mmap_assert_locked(vma->vm_mm);
+
+	if (!pmdp) {
+		pgdp = pgd_offset(mm, addr);
+		p4dp = p4d_alloc(mm, pgdp, addr);
+		if (!p4dp)
+			return SCAN_FAIL;
+		pudp = pud_alloc(mm, p4dp, addr);
+		if (!pudp)
+			return SCAN_FAIL;
+		pmdp = pmd_alloc(mm, pudp, addr);
+		if (!pmdp)
+			return SCAN_FAIL;
+	}
+
+	vmf.pmd = pmdp;
+	if (do_set_pmd(&vmf, folio, page))
+		return SCAN_FAIL;
+
+	folio_get(folio);
+	return SCAN_SUCCEED;
+}
+
+static enum scan_result try_collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
+		bool install_pmd)
+{
+	enum scan_result result = SCAN_FAIL;
+	int nr_mapped_ptes = 0;
+	unsigned int nr_batch_ptes;
+	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 i;
+
+	mmap_assert_locked(mm);
+
+	/* First check VMA found, in case page tables are being torn down */
+	if (!vma || !vma->vm_file ||
+	    !range_in_vma(vma, haddr, haddr + HPAGE_PMD_SIZE))
+		return SCAN_VMA_CHECK;
+
+	/* Fast check before locking page if already PMD-mapped */
+	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
+	if (result == SCAN_PMD_MAPPED)
+		return result;
+
+	/*
+	 * If we are here, we've succeeded in replacing all the native pages
+	 * in the page cache with a single hugepage. If a mm were to fault-in
+	 * this memory (mapped by a suitably aligned VMA), we'd get the hugepage
+	 * and map it by a PMD, regardless of sysfs THP settings. As such, let's
+	 * analogously elide sysfs THP settings here and force collapse.
+	 */
+	if (!thp_vma_allowable_order(vma, vma->vm_flags, TVA_FORCED_COLLAPSE, PMD_ORDER))
+		return SCAN_VMA_CHECK;
+
+	/*
+	 * Keep pmd pgtable while the uffd bit is in use; see comment in
+	 * retract_page_tables().
+	 */
+	if (userfaultfd_protected(vma))
+		return SCAN_PTE_UFFD;
+
+	folio = filemap_lock_folio(vma->vm_file->f_mapping,
+			       linear_page_index(vma, haddr));
+	if (IS_ERR(folio))
+		return SCAN_PAGE_NULL;
+
+	if (!is_pmd_order(folio_order(folio))) {
+		result = SCAN_PAGE_COMPOUND;
+		goto drop_folio;
+	}
+
+	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
+	switch (result) {
+	case SCAN_SUCCEED:
+		break;
+	case SCAN_NO_PTE_TABLE:
+		/*
+		 * All pte entries have been removed and pmd cleared.
+		 * Skip all the pte checks and just update the pmd mapping.
+		 */
+		goto maybe_install_pmd;
+	default:
+		goto drop_folio;
+	}
+
+	result = SCAN_FAIL;
+	start_pte = pte_offset_map_lock(mm, pmd, haddr, &ptl);
+	if (!start_pte)		/* mmap_lock + page lock should prevent this */
+		goto drop_folio;
+
+	/* step 1: check all mapped PTEs are to the right huge page */
+	for (i = 0, addr = haddr, pte = start_pte;
+	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
+		struct page *page;
+		pte_t ptent = ptep_get(pte);
+
+		/* empty pte, skip */
+		if (pte_none(ptent))
+			continue;
+
+		/* page swapped out, abort */
+		if (!pte_present(ptent)) {
+			result = SCAN_PTE_NON_PRESENT;
+			goto abort;
+		}
+
+		page = vm_normal_page(vma, addr, ptent);
+		if (WARN_ON_ONCE(page && is_zone_device_page(page)))
+			page = NULL;
+		/*
+		 * Note that uprobe, debugger, or MAP_PRIVATE may change the
+		 * page table, but the new page will not be a subpage of hpage.
+		 */
+		if (folio_page(folio, i) != page)
+			goto abort;
+	}
+
+	pte_unmap_unlock(start_pte, ptl);
+	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
+				haddr, haddr + HPAGE_PMD_SIZE);
+	mmu_notifier_invalidate_range_start(&range);
+	notified = true;
+
+	/*
+	 * pmd_lock covers a wider range than ptl, and (if split from mm's
+	 * page_table_lock) ptl nests inside pml. The less time we hold pml,
+	 * the better; but userfaultfd's mfill_atomic_pte() on a private VMA
+	 * inserts a valid as-if-COWed PTE without even looking up page cache.
+	 * So page lock of folio does not protect from it, so we must not drop
+	 * ptl before pgt_pmd is removed, so uffd private needs pml taken now.
+	 */
+	if (userfaultfd_armed(vma) && !(vma->vm_flags & VM_SHARED))
+		pml = pmd_lock(mm, pmd);
+
+	start_pte = pte_offset_map_rw_nolock(mm, pmd, haddr, &pgt_pmd, &ptl);
+	if (!start_pte)		/* mmap_lock + page lock should prevent this */
+		goto abort;
+	if (!pml)
+		spin_lock(ptl);
+	else if (ptl != pml)
+		spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
+
+	if (unlikely(!pmd_same(pgt_pmd, pmdp_get_lockless(pmd))))
+		goto abort;
+
+	/* step 2: clear page table and adjust rmap */
+	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
+	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
+	     pte += nr_batch_ptes) {
+		unsigned int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
+		struct page *page;
+		pte_t ptent = ptep_get(pte);
+
+		nr_batch_ptes = 1;
+
+		if (pte_none(ptent))
+			continue;
+		/*
+		 * We dropped ptl after the first scan, to do the mmu_notifier:
+		 * page lock stops more PTEs of the folio being faulted in, but
+		 * does not stop write faults COWing anon copies from existing
+		 * PTEs; and does not stop those being swapped out or migrated.
+		 */
+		if (!pte_present(ptent)) {
+			result = SCAN_PTE_NON_PRESENT;
+			goto abort;
+		}
+		page = vm_normal_page(vma, addr, ptent);
+
+		if (folio_page(folio, i) != page)
+			goto abort;
+
+		nr_batch_ptes = folio_pte_batch(folio, pte, ptent, max_nr_batch_ptes);
+
+		/*
+		 * Must clear entry, or a racing truncate may re-remove it.
+		 * TLB flush can be left until pmdp_collapse_flush() does it.
+		 * PTE dirty? Shmem page is already dirty; file is read-only.
+		 */
+		clear_ptes(mm, addr, pte, nr_batch_ptes);
+		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
+		nr_mapped_ptes += nr_batch_ptes;
+	}
+
+	if (!pml)
+		spin_unlock(ptl);
+
+	/* step 3: set proper refcount and mm_counters. */
+	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 */
+	if (!pml) {
+		pml = pmd_lock(mm, pmd);
+		if (ptl != pml) {
+			spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
+			if (unlikely(!pmd_same(pgt_pmd, pmdp_get_lockless(pmd)))) {
+				flush_tlb_mm(mm);
+				goto unlock;
+			}
+		}
+	}
+	pgt_pmd = pmdp_collapse_flush(vma, haddr, pmd);
+	pmdp_get_lockless_sync();
+	pte_unmap_unlock(start_pte, ptl);
+	if (ptl != pml)
+		spin_unlock(pml);
+
+	mmu_notifier_invalidate_range_end(&range);
+
+	mm_dec_nr_ptes(mm);
+	page_table_check_pte_clear_range(mm, haddr, pgt_pmd);
+	pte_free_defer(mm, pmd_pgtable(pgt_pmd));
+
+maybe_install_pmd:
+	/* step 5: install pmd entry */
+	result = install_pmd
+			? set_huge_pmd(vma, haddr, pmd, folio, &folio->page)
+			: SCAN_SUCCEED;
+	goto drop_folio;
+abort:
+	if (nr_mapped_ptes) {
+		flush_tlb_mm(mm);
+		folio_ref_sub(folio, nr_mapped_ptes);
+		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
+	}
+unlock:
+	if (start_pte)
+		pte_unmap_unlock(start_pte, ptl);
+	if (pml && pml != ptl)
+		spin_unlock(pml);
+	if (notified)
+		mmu_notifier_invalidate_range_end(&range);
+drop_folio:
+	folio_unlock(folio);
+	folio_put(folio);
+	return result;
+}
+
+/**
+ * collapse_pte_mapped_thp - Try to collapse a pte-mapped THP for mm at
+ * address haddr.
+ *
+ * @mm: process address space where collapse happens
+ * @addr: THP collapse address
+ * @install_pmd: If a huge PMD should be installed
+ *
+ * This function checks whether all the PTEs in the PMD are pointing to the
+ * right THP. If so, retract the page table so the THP can refault in with
+ * as pmd-mapped. Possibly install a huge PMD mapping the THP.
+ */
+void collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
+		bool install_pmd)
+{
+	try_collapse_pte_mapped_thp(mm, addr, install_pmd);
+}
+
+/* Can we retract page tables for this file-backed VMA? */
+static bool file_backed_vma_is_retractable(struct vm_area_struct *vma)
+{
+	/*
+	 * Check vma->anon_vma to exclude MAP_PRIVATE mappings that
+	 * got written to. These VMAs are likely not worth removing
+	 * page tables from, as PMD-mapping is likely to be split later.
+	 */
+	if (READ_ONCE(vma->anon_vma))
+		return false;
+
+	/*
+	 * When a vma is registered with uffd-wp or RWP, we cannot recycle
+	 * the page table because there may be pte markers installed.
+	 * VM_UFFD_RWP ranges similarly rely on per-PTE uffd state
+	 * and cannot be recycled to a shared PMD. Other vmas can still
+	 * have the same file mapped hugely, but skip this one: it will
+	 * always be mapped in small page size for these registrations.
+	 */
+	if (userfaultfd_protected(vma))
+		return false;
+
+	/*
+	 * If the VMA contains guard regions then we can't collapse it.
+	 *
+	 * This is set atomically on guard marker installation under mmap/VMA
+	 * read lock, and here we may not hold any VMA or mmap lock at all.
+	 *
+	 * This is therefore serialised on the PTE page table lock, which is
+	 * obtained on guard region installation after the flag is set, so this
+	 * check being performed under this lock excludes races.
+	 */
+	if (vma_test_atomic_flag(vma, VMA_MAYBE_GUARD_BIT))
+		return false;
+
+	return true;
+}
+
+static void retract_page_tables(struct address_space *mapping, pgoff_t pgoff)
+{
+	struct vm_area_struct *vma;
+
+	i_mmap_lock_read(mapping);
+	mapping_rmap_tree_foreach(vma, mapping, pgoff, pgoff) {
+		struct mmu_notifier_range range;
+		struct mm_struct *mm;
+		unsigned long addr;
+		pmd_t *pmd, pgt_pmd;
+		spinlock_t *pml;
+		spinlock_t *ptl;
+		bool success = false;
+
+		addr = vma->vm_start +
+			((pgoff - vma_start_pgoff(vma)) << PAGE_SHIFT);
+		if (addr & ~HPAGE_PMD_MASK ||
+		    vma->vm_end < addr + HPAGE_PMD_SIZE)
+			continue;
+
+		mm = vma->vm_mm;
+		if (find_pmd_or_thp_or_none(mm, addr, &pmd) != SCAN_SUCCEED)
+			continue;
+
+		if (collapse_test_exit(mm))
+			continue;
+
+		if (!file_backed_vma_is_retractable(vma))
+			continue;
+
+		/* PTEs were notified when unmapped; but now for the PMD? */
+		mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
+					addr, addr + HPAGE_PMD_SIZE);
+		mmu_notifier_invalidate_range_start(&range);
+
+		pml = pmd_lock(mm, pmd);
+		/*
+		 * The lock of new_folio is still held, we will be blocked in
+		 * the page fault path, which prevents the pte entries from
+		 * being set again. So even though the old empty PTE page may be
+		 * concurrently freed and a new PTE page is filled into the pmd
+		 * entry, it is still empty and can be removed.
+		 *
+		 * So here we only need to recheck if the state of pmd entry
+		 * still meets our requirements, rather than checking pmd_same()
+		 * like elsewhere.
+		 */
+		if (check_pmd_state(pmd) != SCAN_SUCCEED)
+			goto drop_pml;
+		ptl = pte_lockptr(mm, pmd);
+		if (ptl != pml)
+			spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
+
+		/*
+		 * Huge page lock is still held, so normally the page table must
+		 * remain empty; and we have already skipped anon_vma and
+		 * userfaultfd_wp() vmas.  But since the mmap_lock is not held,
+		 * it is still possible for a racing userfaultfd_ioctl() or
+		 * madvise() to have inserted ptes or markers.  Now that we hold
+		 * ptlock, repeating the retractable checks protects us from
+		 * races against the prior checks.
+		 */
+		if (likely(file_backed_vma_is_retractable(vma))) {
+			pgt_pmd = pmdp_collapse_flush(vma, addr, pmd);
+			pmdp_get_lockless_sync();
+			success = true;
+		}
+
+		if (ptl != pml)
+			spin_unlock(ptl);
+drop_pml:
+		spin_unlock(pml);
+
+		mmu_notifier_invalidate_range_end(&range);
+
+		if (success) {
+			mm_dec_nr_ptes(mm);
+			page_table_check_pte_clear_range(mm, addr, pgt_pmd);
+			pte_free_defer(mm, pmd_pgtable(pgt_pmd));
+		}
+	}
+	i_mmap_unlock_read(mapping);
+}
+
+/**
+ * collapse_file - collapse filemap/tmpfs/shmem pages into huge one.
+ *
+ * @mm: process address space where collapse happens
+ * @addr: virtual collapse start address
+ * @file: file that collapse on
+ * @start: collapse start address
+ * @cc: collapse context and scratchpad
+ *
+ * Basic scheme is simple, details are more complex:
+ *  - allocate and lock a new huge page;
+ *  - scan page cache, locking old pages
+ *    + swap/gup in pages if necessary;
+ *  - copy data to new page
+ *  - handle shmem holes
+ *    + re-validate that holes weren't filled by someone else
+ *    + check for userfaultfd
+ *  - finalize updates to the page cache;
+ *  - if replacing succeeds:
+ *    + unlock huge page;
+ *    + free old pages;
+ *  - if replacing failed;
+ *    + unlock old pages
+ *    + unlock and free huge page;
+ */
+static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
+		struct file *file, pgoff_t start, struct collapse_control *cc)
+{
+	struct address_space *mapping = file->f_mapping;
+	struct page *dst;
+	struct folio *folio, *tmp, *new_folio;
+	pgoff_t index = 0, end = start + HPAGE_PMD_NR;
+	LIST_HEAD(pagelist);
+	XA_STATE_ORDER(xas, &mapping->i_pages, start, HPAGE_PMD_ORDER);
+	enum scan_result result = SCAN_SUCCEED;
+	int nr_none = 0;
+	bool is_shmem = shmem_file(file);
+
+	/*
+	 * MADV_COLLAPSE ignores shmem huge config, so do not check shmem
+	 *
+	 * TODO: once shmem always calls mapping_set_large_folios() on its
+	 * mapping, the shmem check can be removed.
+	 */
+	VM_WARN_ON_ONCE(!is_shmem && !mapping_pmd_folio_support(mapping));
+	VM_WARN_ON_ONCE(start & (HPAGE_PMD_NR - 1));
+
+	result = alloc_charge_folio(&new_folio, mm, cc, HPAGE_PMD_ORDER);
+	if (result != SCAN_SUCCEED)
+		goto out;
+
+	mapping_set_update(&xas, mapping);
+
+	__folio_set_locked(new_folio);
+	if (is_shmem)
+		__folio_set_swapbacked(new_folio);
+	new_folio->index = start;
+	new_folio->mapping = mapping;
+
+	/*
+	 * Ensure we have slots for all the pages in the range.  This is
+	 * almost certainly a no-op because most of the pages must be present
+	 */
+	do {
+		xas_lock_irq(&xas);
+		xas_create_range(&xas);
+		if (!xas_error(&xas))
+			break;
+		xas_unlock_irq(&xas);
+		if (!xas_nomem(&xas, GFP_KERNEL)) {
+			result = SCAN_FAIL;
+			goto rollback;
+		}
+	} while (1);
+
+	for (index = start; index < end;) {
+		xas_set(&xas, index);
+		folio = xas_load(&xas);
+
+		VM_BUG_ON(index != xas.xa_index);
+		if (is_shmem) {
+			if (!folio) {
+				/*
+				 * Stop if extent has been truncated or
+				 * hole-punched, and is now completely
+				 * empty.
+				 */
+				if (index == start) {
+					if (!xas_next_entry(&xas, end - 1)) {
+						result = SCAN_TRUNCATED;
+						goto xa_locked;
+					}
+				}
+				nr_none++;
+				index++;
+				continue;
+			}
+
+			if (xa_is_value(folio) || !folio_test_uptodate(folio)) {
+				xas_unlock_irq(&xas);
+				/* swap in or instantiate fallocated page */
+				if (shmem_get_folio(mapping->host, index, 0,
+						&folio, SGP_NOALLOC)) {
+					result = SCAN_FAIL;
+					goto xa_unlocked;
+				}
+				/* drain lru cache to help folio_isolate_lru() */
+				lru_add_drain();
+			} else if (folio_trylock(folio)) {
+				folio_get(folio);
+				xas_unlock_irq(&xas);
+			} else {
+				result = SCAN_PAGE_LOCK;
+				goto xa_locked;
+			}
+		} else {	/* !is_shmem */
+			if (!folio || xa_is_value(folio)) {
+				xas_unlock_irq(&xas);
+				page_cache_sync_readahead(mapping, &file->f_ra,
+							  file, index,
+							  end - index);
+				/* drain lru cache to help folio_isolate_lru() */
+				lru_add_drain();
+				folio = filemap_lock_folio(mapping, index);
+				if (IS_ERR(folio)) {
+					result = SCAN_FAIL;
+					goto xa_unlocked;
+				}
+			} else if (folio_test_dirty(folio)) {
+				/*
+				 * This page is dirty because it hasn't
+				 * been flushed since first write.
+				 *
+				 * Trigger async flush for read-only files and
+				 * hope the writeback is done when khugepaged
+				 * revisits this page. Writable files can have
+				 * their folios dirty at any time; blindly
+				 * flushing them would cause undesirable
+				 * system-wide writeback.
+				 *
+				 * This is a one-off situation. We are not
+				 * forcing writeback in loop.
+				 */
+				xas_unlock_irq(&xas);
+				if (!inode_is_open_for_write(mapping->host))
+					filemap_flush(mapping);
+				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+				goto xa_unlocked;
+			} else if (folio_test_writeback(folio)) {
+				xas_unlock_irq(&xas);
+				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+				goto xa_unlocked;
+			} else if (folio_trylock(folio)) {
+				folio_get(folio);
+				xas_unlock_irq(&xas);
+			} else {
+				result = SCAN_PAGE_LOCK;
+				goto xa_locked;
+			}
+		}
+
+		/*
+		 * The folio must be locked, so we can drop the i_pages lock
+		 * without racing with truncate.
+		 */
+		VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
+
+		/* make sure the folio is up to date */
+		if (unlikely(!folio_test_uptodate(folio))) {
+			result = SCAN_FAIL;
+			goto out_unlock;
+		}
+
+		/*
+		 * If file was truncated then extended, or hole-punched, before
+		 * we locked the first folio, then a THP might be there already.
+		 * This will be discovered on the first iteration.
+		 */
+		if (is_pmd_order(folio_order(folio))) {
+			result = SCAN_PTE_MAPPED_HUGEPAGE;
+			goto out_unlock;
+		}
+
+		if (folio_mapping(folio) != mapping) {
+			result = SCAN_TRUNCATED;
+			goto out_unlock;
+		}
+
+		if (!is_shmem && (folio_test_dirty(folio) ||
+				  folio_test_writeback(folio))) {
+			/*
+			 * khugepaged only works on clean file-backed folios,
+			 * so this folio is dirty because it hasn't been flushed
+			 * since first write.
+			 */
+			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+			goto out_unlock;
+		}
+
+		if (!folio_isolate_lru(folio)) {
+			result = SCAN_DEL_PAGE_LRU;
+			goto out_unlock;
+		}
+
+		if (!filemap_release_folio(folio, GFP_KERNEL)) {
+			result = SCAN_PAGE_HAS_PRIVATE;
+			folio_putback_lru(folio);
+			goto out_unlock;
+		}
+
+		if (folio_mapped(folio))
+			try_to_unmap(folio,
+					TTU_IGNORE_MLOCK | TTU_BATCH_FLUSH);
+
+		xas_lock_irq(&xas);
+
+		VM_BUG_ON_FOLIO(folio != xa_load(xas.xa, index), folio);
+
+		/*
+		 * We control 2 + nr_pages references to the folio:
+		 *  - we hold a pin on it;
+		 *  - nr_pages reference from page cache;
+		 *  - one from lru_isolate_folio;
+		 * If those are the only references, then any new usage
+		 * of the folio will have to fetch it from the page
+		 * cache. That requires locking the folio to handle
+		 * truncate, so any new usage will be blocked until we
+		 * unlock folio after collapse/during rollback.
+		 */
+		if (folio_ref_count(folio) != 2 + folio_nr_pages(folio)) {
+			result = SCAN_PAGE_COUNT;
+			xas_unlock_irq(&xas);
+			folio_putback_lru(folio);
+			goto out_unlock;
+		}
+
+		/*
+		 * At this point, the folio is locked and unmapped. If the PTE
+		 * was dirty, try_to_unmap() has transferred the dirty bit to
+		 * the folio and we must not collapse it into a clean
+		 * file-backed folio.
+		 *
+		 * If the folio is clean here, no one can write it until we
+		 * drop the folio lock. A write through a stale TLB entry came
+		 * from a clean PTE and must fault because the PTE has been
+		 * cleared; the fault path has to take the folio lock before
+		 * installing a writable mapping. Buffered write paths also
+		 * have to take the folio lock before modifying file contents
+		 * without a mapping, typically via write_begin_get_folio().
+		 */
+		if (!is_shmem && folio_test_dirty(folio)) {
+			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
+			xas_unlock_irq(&xas);
+			folio_putback_lru(folio);
+			goto out_unlock;
+		}
+
+		/*
+		 * Accumulate the folios that are being collapsed.
+		 */
+		list_add_tail(&folio->lru, &pagelist);
+		index += folio_nr_pages(folio);
+		continue;
+out_unlock:
+		folio_unlock(folio);
+		folio_put(folio);
+		goto xa_unlocked;
+	}
+
+xa_locked:
+	xas_unlock_irq(&xas);
+xa_unlocked:
+
+	/*
+	 * If collapse is successful, flush must be done now before copying.
+	 * If collapse is unsuccessful, does flush actually need to be done?
+	 * Do it anyway, to clear the state.
+	 */
+	try_to_unmap_flush();
+
+	if (result == SCAN_SUCCEED && nr_none &&
+	    !shmem_charge(mapping->host, nr_none))
+		result = SCAN_FAIL;
+	if (result != SCAN_SUCCEED) {
+		nr_none = 0;
+		goto rollback;
+	}
+
+	/*
+	 * The old folios are locked, so they won't change anymore.
+	 */
+	index = start;
+	dst = folio_page(new_folio, 0);
+	list_for_each_entry(folio, &pagelist, lru) {
+		int i, nr_pages = folio_nr_pages(folio);
+
+		while (index < folio->index) {
+			clear_highpage(dst);
+			index++;
+			dst++;
+		}
+
+		for (i = 0; i < nr_pages; i++) {
+			if (copy_mc_highpage(dst, folio_page(folio, i)) > 0) {
+				result = SCAN_COPY_MC;
+				goto rollback;
+			}
+			index++;
+			dst++;
+		}
+	}
+	while (index < end) {
+		clear_highpage(dst);
+		index++;
+		dst++;
+	}
+
+	if (nr_none) {
+		struct vm_area_struct *vma;
+		int nr_none_check = 0;
+
+		i_mmap_lock_read(mapping);
+		xas_lock_irq(&xas);
+
+		xas_set(&xas, start);
+		for (index = start; index < end; index++) {
+			if (!xas_next(&xas)) {
+				xas_store(&xas, XA_RETRY_ENTRY);
+				if (xas_error(&xas)) {
+					result = SCAN_STORE_FAILED;
+					goto immap_locked;
+				}
+				nr_none_check++;
+			}
+		}
+
+		if (nr_none != nr_none_check) {
+			result = SCAN_PAGE_FILLED;
+			goto immap_locked;
+		}
+
+		/*
+		 * If userspace observed a missing page in a VMA with
+		 * a MODE_MISSING userfaultfd, then it might expect a
+		 * UFFD_EVENT_PAGEFAULT for that page. If so, we need to
+		 * roll back to avoid suppressing such an event. Since
+		 * wp/minor userfaultfds don't give userspace any
+		 * guarantees that the kernel doesn't fill a missing
+		 * page with a zero page, so they don't matter here.
+		 *
+		 * Any userfaultfds registered after this point will
+		 * not be able to observe any missing pages due to the
+		 * previously inserted retry entries.
+		 */
+		mapping_rmap_tree_foreach(vma, mapping, start, end) {
+			if (userfaultfd_missing(vma)) {
+				result = SCAN_EXCEED_NONE_PTE;
+				goto immap_locked;
+			}
+		}
+
+immap_locked:
+		i_mmap_unlock_read(mapping);
+		if (result != SCAN_SUCCEED) {
+			xas_set(&xas, start);
+			for (index = start; index < end; index++) {
+				if (xas_next(&xas) == XA_RETRY_ENTRY)
+					xas_store(&xas, NULL);
+			}
+
+			xas_unlock_irq(&xas);
+			goto rollback;
+		}
+	} else {
+		xas_lock_irq(&xas);
+	}
+
+	if (is_shmem) {
+		lruvec_stat_mod_folio(new_folio, NR_SHMEM, HPAGE_PMD_NR);
+		lruvec_stat_mod_folio(new_folio, NR_SHMEM_THPS, HPAGE_PMD_NR);
+	} else {
+		lruvec_stat_mod_folio(new_folio, NR_FILE_THPS, HPAGE_PMD_NR);
+	}
+	lruvec_stat_mod_folio(new_folio, NR_FILE_PAGES, HPAGE_PMD_NR);
+
+	/*
+	 * Mark new_folio as uptodate before inserting it into the
+	 * page cache so that it isn't mistaken for an fallocated but
+	 * unwritten page.
+	 */
+	folio_mark_uptodate(new_folio);
+	folio_ref_add(new_folio, HPAGE_PMD_NR - 1);
+
+	if (is_shmem)
+		folio_mark_dirty(new_folio);
+	folio_add_lru(new_folio);
+
+	/* Join all the small entries into a single multi-index entry. */
+	xas_set_order(&xas, start, HPAGE_PMD_ORDER);
+	xas_store(&xas, new_folio);
+	WARN_ON_ONCE(xas_error(&xas));
+	xas_unlock_irq(&xas);
+
+	/*
+	 * Remove pte page tables, so we can re-fault the page as huge.  A caller
+	 * that wants the PMD mapped now is told to go and do that.
+	 */
+	retract_page_tables(mapping, start);
+	if (cc->policy.install_pmd)
+		result = SCAN_PTE_MAPPED_HUGEPAGE;
+	folio_unlock(new_folio);
+
+	/*
+	 * The collapse has succeeded, so free the old folios.
+	 */
+	list_for_each_entry_safe(folio, tmp, &pagelist, lru) {
+		list_del(&folio->lru);
+		lruvec_stat_mod_folio(folio, NR_FILE_PAGES,
+				      -folio_nr_pages(folio));
+		if (is_shmem)
+			lruvec_stat_mod_folio(folio, NR_SHMEM,
+					      -folio_nr_pages(folio));
+		folio->mapping = NULL;
+		folio_clear_active(folio);
+		folio_clear_unevictable(folio);
+		folio_unlock(folio);
+		folio_put_refs(folio, 2 + folio_nr_pages(folio));
+	}
+
+	goto out;
+
+rollback:
+	/* Something went wrong: roll back page cache changes */
+	if (nr_none) {
+		xas_lock_irq(&xas);
+		mapping->nrpages -= nr_none;
+		xas_unlock_irq(&xas);
+		shmem_uncharge(mapping->host, nr_none);
+	}
+
+	list_for_each_entry_safe(folio, tmp, &pagelist, lru) {
+		list_del(&folio->lru);
+		folio_unlock(folio);
+		folio_putback_lru(folio);
+		folio_put(folio);
+	}
+
+	new_folio->mapping = NULL;
+
+	folio_unlock(new_folio);
+	folio_put(new_folio);
+out:
+	VM_BUG_ON(!list_empty(&pagelist));
+	trace_mm_khugepaged_collapse_file(mm, new_folio, index, addr, is_shmem, file, HPAGE_PMD_NR, result);
+	return result;
+}
+
+static enum scan_result collapse_scan_file(struct mm_struct *mm,
+		unsigned long addr, struct file *file, pgoff_t start,
+		struct collapse_control *cc)
+{
+	const unsigned int max_ptes_none = collapse_max_ptes_none(cc, NULL, HPAGE_PMD_ORDER);
+	const unsigned int max_ptes_swap = collapse_max_ptes_swap(cc, HPAGE_PMD_ORDER);
+	struct folio *folio = NULL;
+	struct address_space *mapping = file->f_mapping;
+	XA_STATE(xas, &mapping->i_pages, start);
+	int present, swap;
+	int node = NUMA_NO_NODE;
+	enum scan_result result = SCAN_SUCCEED;
+
+	present = 0;
+	swap = 0;
+	collapse_control_init_scan(cc);
+	rcu_read_lock();
+	xas_for_each(&xas, folio, start + HPAGE_PMD_NR - 1) {
+		if (xas_retry(&xas, folio))
+			continue;
+
+		if (xa_is_value(folio)) {
+			swap += 1 << xas_get_order(&xas);
+			if (swap > max_ptes_swap) {
+				result = SCAN_EXCEED_SWAP_PTE;
+				count_vm_event(THP_SCAN_EXCEED_SWAP_PTE);
+				break;
+			}
+			continue;
+		}
+
+		if (!folio_try_get(folio)) {
+			xas_reset(&xas);
+			continue;
+		}
+
+		if (unlikely(folio != xas_reload(&xas))) {
+			folio_put(folio);
+			xas_reset(&xas);
+			continue;
+		}
+
+		if (is_pmd_order(folio_order(folio))) {
+			result = SCAN_PTE_MAPPED_HUGEPAGE;
+			/*
+			 * PMD-sized THP implies that we can only try
+			 * retracting the PTE table.
+			 */
+			folio_put(folio);
+			break;
+		}
+
+		node = folio_nid(folio);
+		if (collapse_scan_abort(node, cc)) {
+			result = SCAN_SCAN_ABORT;
+			folio_put(folio);
+			break;
+		}
+		cc->node_load[node]++;
+
+		if (!folio_test_lru(folio)) {
+			result = SCAN_PAGE_LRU;
+			folio_put(folio);
+			break;
+		}
+
+		if (folio_expected_ref_count(folio) + 1 != folio_ref_count(folio)) {
+			result = SCAN_PAGE_COUNT;
+			folio_put(folio);
+			break;
+		}
+
+		/*
+		 * We probably should check if the folio is referenced
+		 * here, but nobody would transfer pte_young() to
+		 * folio_test_referenced() for us.  And rmap walk here
+		 * is just too costly...
+		 */
+
+		present += folio_nr_pages(folio);
+		folio_put(folio);
+
+		if (need_resched()) {
+			xas_pause(&xas);
+			cond_resched_rcu();
+		}
+	}
+	rcu_read_unlock();
+	if (result == SCAN_PTE_MAPPED_HUGEPAGE)
+		cc->progress++;
+	else
+		cc->progress += HPAGE_PMD_NR;
+
+	if (result == SCAN_SUCCEED) {
+		if (present < HPAGE_PMD_NR - max_ptes_none) {
+			result = SCAN_EXCEED_NONE_PTE;
+			count_vm_event(THP_SCAN_EXCEED_NONE_PTE);
+		} else {
+			result = collapse_file(mm, addr, file, start, cc);
+		}
+	}
+
+	trace_mm_khugepaged_scan_file(mm, folio, file, present, swap, result);
+	return result;
+}
+
+/*
+ * Try to collapse a single PMD starting at a PMD aligned addr, and return
+ * the results.
+ */
+enum scan_result collapse_single_pmd(unsigned long addr,
+		unsigned long end, struct vm_area_struct *vma,
+		bool *lock_dropped, struct collapse_control *cc)
+{
+	struct mm_struct *mm = vma->vm_mm;
+	bool triggered_wb = false;
+	enum scan_result result;
+	struct file *file;
+	pgoff_t pgoff;
+
+	mmap_assert_locked(mm);
+
+	if (vma_is_anonymous(vma)) {
+		result = collapse_scan_anon_pmd(vma, addr, end, cc);
+		if (!cc->select_orders)
+			goto end;
+
+		/* collapse_anon_pmd() takes mmap_lock itself, where it needs it */
+		mmap_read_unlock(mm);
+		*lock_dropped = true;
+
+		result = collapse_anon_pmd(mm, addr, end, cc);
+		goto end;
+	}
+
+	file = get_file(vma->vm_file);
+	pgoff = linear_page_index(vma, addr);
+
+	mmap_read_unlock(mm);
+	*lock_dropped = true;
+retry:
+	result = collapse_scan_file(mm, addr, file, pgoff, cc);
+
+	/* Dirty pages are worth a writeback and one more try, if asked for */
+	if (cc->policy.writeback_dirty && result == SCAN_PAGE_DIRTY_OR_WRITEBACK &&
+	    !triggered_wb && mapping_can_writeback(file->f_mapping)) {
+		const loff_t lstart = (loff_t)pgoff << PAGE_SHIFT;
+		const loff_t lend = lstart + HPAGE_PMD_SIZE - 1;
+
+		filemap_write_and_wait_range(file->f_mapping, lstart, lend);
+		triggered_wb = true;
+		goto retry;
+	}
+	fput(file);
+
+	if (result == SCAN_PTE_MAPPED_HUGEPAGE) {
+		mmap_read_lock(mm);
+		if (collapse_test_exit_or_disable(mm))
+			result = SCAN_ANY_PROCESS;
+		else
+			result = try_collapse_pte_mapped_thp(mm, addr,
+							cc->policy.install_pmd);
+		if (result == SCAN_PMD_MAPPED)
+			result = SCAN_SUCCEED;
+		mmap_read_unlock(mm);
+	}
+end:
+	return result;
+}
diff --git a/mm/collapse.h b/mm/collapse.h
index 11f51c6ea444..dc60806fb81e 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -187,24 +187,13 @@ static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
 		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
 }
 
-enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
-		unsigned long start, unsigned long end,
-		struct collapse_control *cc);
-enum scan_result collapse_anon_pmd(struct mm_struct *mm, unsigned long start,
-		unsigned long end, struct collapse_control *cc);
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
+enum scan_result collapse_single_pmd(unsigned long addr, unsigned long end,
+		struct vm_area_struct *vma, bool *lock_dropped,
+		struct collapse_control *cc);
 
 unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags);
-enum scan_result check_pmd_state(pmd_t *pmd);
-enum scan_result find_pmd_or_thp_or_none(struct mm_struct *mm,
-		unsigned long address, pmd_t **pmd);
-int collapse_find_target_node(struct collapse_control *cc);
-bool collapse_scan_abort(int nid, struct collapse_control *cc);
-unsigned int collapse_max_ptes_none(struct collapse_control *cc,
-		struct vm_area_struct *vma, unsigned int order);
-unsigned int collapse_max_ptes_swap(struct collapse_control *cc,
-		unsigned int order);
 
 #endif	/* __MM_COLLAPSE_H */
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 907ed1131460..b7fc93e11d6b 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -543,51 +543,6 @@ static enum scan_result hugepage_vma_revalidate(struct mm_struct *mm, unsigned l
 	return SCAN_SUCCEED;
 }
 
-static void count_collapse_event(unsigned int order, enum vm_event_item vm_event,
-		enum mthp_stat_item mthp_event)
-{
-	if (is_pmd_order(order))
-		count_vm_event(vm_event);
-	count_mthp_stat(order, mthp_event);
-}
-
-static void collapse_control_init_scan(struct collapse_control *cc)
-{
-	memset(cc->node_load, 0, sizeof(cc->node_load));
-	nodes_clear(cc->alloc_nmask);
-	bitmap_zero(cc->eligible_ptes, MAX_PTRS_PER_PTE);
-}
-
-static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_struct *mm,
-		struct collapse_control *cc, unsigned int order)
-{
-	gfp_t gfp = cc->policy.gfp;
-	int node = collapse_find_target_node(cc);
-	struct folio *folio;
-
-	folio = __folio_alloc(gfp, order, node, &cc->alloc_nmask);
-	if (!folio) {
-		*foliop = NULL;
-		count_collapse_event(order, THP_COLLAPSE_ALLOC_FAILED,
-				     MTHP_STAT_COLLAPSE_ALLOC_FAILED);
-		return SCAN_ALLOC_HUGE_PAGE_FAIL;
-	}
-
-	count_collapse_event(order, THP_COLLAPSE_ALLOC, MTHP_STAT_COLLAPSE_ALLOC);
-
-	if (unlikely(mem_cgroup_charge(folio, mm, gfp))) {
-		folio_put(folio);
-		*foliop = NULL;
-		return SCAN_CGROUP_CHARGE_FAIL;
-	}
-
-	if (is_pmd_order(order))
-		count_memcg_folio_events(folio, THP_COLLAPSE_ALLOC, 1);
-
-	*foliop = folio;
-	return SCAN_SUCCEED;
-}
-
 static void collect_mm_slot(struct mm_slot *slot)
 {
 	struct mm_struct *mm = slot->mm;
@@ -610,1036 +565,6 @@ static void collect_mm_slot(struct mm_slot *slot)
 	}
 }
 
-/* folio must be locked, and mmap_lock must be held */
-static enum scan_result set_huge_pmd(struct vm_area_struct *vma, unsigned long addr,
-		pmd_t *pmdp, struct folio *folio, struct page *page)
-{
-	struct mm_struct *mm = vma->vm_mm;
-	struct vm_fault vmf = {
-		.vma = vma,
-		.address = addr,
-		.flags = 0,
-	};
-	pgd_t *pgdp;
-	p4d_t *p4dp;
-	pud_t *pudp;
-
-	mmap_assert_locked(vma->vm_mm);
-
-	if (!pmdp) {
-		pgdp = pgd_offset(mm, addr);
-		p4dp = p4d_alloc(mm, pgdp, addr);
-		if (!p4dp)
-			return SCAN_FAIL;
-		pudp = pud_alloc(mm, p4dp, addr);
-		if (!pudp)
-			return SCAN_FAIL;
-		pmdp = pmd_alloc(mm, pudp, addr);
-		if (!pmdp)
-			return SCAN_FAIL;
-	}
-
-	vmf.pmd = pmdp;
-	if (do_set_pmd(&vmf, folio, page))
-		return SCAN_FAIL;
-
-	folio_get(folio);
-	return SCAN_SUCCEED;
-}
-
-static enum scan_result try_collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
-		bool install_pmd)
-{
-	enum scan_result result = SCAN_FAIL;
-	int nr_mapped_ptes = 0;
-	unsigned int nr_batch_ptes;
-	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 i;
-
-	mmap_assert_locked(mm);
-
-	/* First check VMA found, in case page tables are being torn down */
-	if (!vma || !vma->vm_file ||
-	    !range_in_vma(vma, haddr, haddr + HPAGE_PMD_SIZE))
-		return SCAN_VMA_CHECK;
-
-	/* Fast check before locking page if already PMD-mapped */
-	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
-	if (result == SCAN_PMD_MAPPED)
-		return result;
-
-	/*
-	 * If we are here, we've succeeded in replacing all the native pages
-	 * in the page cache with a single hugepage. If a mm were to fault-in
-	 * this memory (mapped by a suitably aligned VMA), we'd get the hugepage
-	 * and map it by a PMD, regardless of sysfs THP settings. As such, let's
-	 * analogously elide sysfs THP settings here and force collapse.
-	 */
-	if (!thp_vma_allowable_order(vma, vma->vm_flags, TVA_FORCED_COLLAPSE, PMD_ORDER))
-		return SCAN_VMA_CHECK;
-
-	/*
-	 * Keep pmd pgtable while the uffd bit is in use; see comment in
-	 * retract_page_tables().
-	 */
-	if (userfaultfd_protected(vma))
-		return SCAN_PTE_UFFD;
-
-	folio = filemap_lock_folio(vma->vm_file->f_mapping,
-			       linear_page_index(vma, haddr));
-	if (IS_ERR(folio))
-		return SCAN_PAGE_NULL;
-
-	if (!is_pmd_order(folio_order(folio))) {
-		result = SCAN_PAGE_COMPOUND;
-		goto drop_folio;
-	}
-
-	result = find_pmd_or_thp_or_none(mm, haddr, &pmd);
-	switch (result) {
-	case SCAN_SUCCEED:
-		break;
-	case SCAN_NO_PTE_TABLE:
-		/*
-		 * All pte entries have been removed and pmd cleared.
-		 * Skip all the pte checks and just update the pmd mapping.
-		 */
-		goto maybe_install_pmd;
-	default:
-		goto drop_folio;
-	}
-
-	result = SCAN_FAIL;
-	start_pte = pte_offset_map_lock(mm, pmd, haddr, &ptl);
-	if (!start_pte)		/* mmap_lock + page lock should prevent this */
-		goto drop_folio;
-
-	/* step 1: check all mapped PTEs are to the right huge page */
-	for (i = 0, addr = haddr, pte = start_pte;
-	     i < HPAGE_PMD_NR; i++, addr += PAGE_SIZE, pte++) {
-		struct page *page;
-		pte_t ptent = ptep_get(pte);
-
-		/* empty pte, skip */
-		if (pte_none(ptent))
-			continue;
-
-		/* page swapped out, abort */
-		if (!pte_present(ptent)) {
-			result = SCAN_PTE_NON_PRESENT;
-			goto abort;
-		}
-
-		page = vm_normal_page(vma, addr, ptent);
-		if (WARN_ON_ONCE(page && is_zone_device_page(page)))
-			page = NULL;
-		/*
-		 * Note that uprobe, debugger, or MAP_PRIVATE may change the
-		 * page table, but the new page will not be a subpage of hpage.
-		 */
-		if (folio_page(folio, i) != page)
-			goto abort;
-	}
-
-	pte_unmap_unlock(start_pte, ptl);
-	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
-				haddr, haddr + HPAGE_PMD_SIZE);
-	mmu_notifier_invalidate_range_start(&range);
-	notified = true;
-
-	/*
-	 * pmd_lock covers a wider range than ptl, and (if split from mm's
-	 * page_table_lock) ptl nests inside pml. The less time we hold pml,
-	 * the better; but userfaultfd's mfill_atomic_pte() on a private VMA
-	 * inserts a valid as-if-COWed PTE without even looking up page cache.
-	 * So page lock of folio does not protect from it, so we must not drop
-	 * ptl before pgt_pmd is removed, so uffd private needs pml taken now.
-	 */
-	if (userfaultfd_armed(vma) && !(vma->vm_flags & VM_SHARED))
-		pml = pmd_lock(mm, pmd);
-
-	start_pte = pte_offset_map_rw_nolock(mm, pmd, haddr, &pgt_pmd, &ptl);
-	if (!start_pte)		/* mmap_lock + page lock should prevent this */
-		goto abort;
-	if (!pml)
-		spin_lock(ptl);
-	else if (ptl != pml)
-		spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
-
-	if (unlikely(!pmd_same(pgt_pmd, pmdp_get_lockless(pmd))))
-		goto abort;
-
-	/* step 2: clear page table and adjust rmap */
-	for (i = 0, addr = haddr, pte = start_pte; i < HPAGE_PMD_NR;
-	     i += nr_batch_ptes, addr += nr_batch_ptes * PAGE_SIZE,
-	     pte += nr_batch_ptes) {
-		unsigned int max_nr_batch_ptes = (end - addr) >> PAGE_SHIFT;
-		struct page *page;
-		pte_t ptent = ptep_get(pte);
-
-		nr_batch_ptes = 1;
-
-		if (pte_none(ptent))
-			continue;
-		/*
-		 * We dropped ptl after the first scan, to do the mmu_notifier:
-		 * page lock stops more PTEs of the folio being faulted in, but
-		 * does not stop write faults COWing anon copies from existing
-		 * PTEs; and does not stop those being swapped out or migrated.
-		 */
-		if (!pte_present(ptent)) {
-			result = SCAN_PTE_NON_PRESENT;
-			goto abort;
-		}
-		page = vm_normal_page(vma, addr, ptent);
-
-		if (folio_page(folio, i) != page)
-			goto abort;
-
-		nr_batch_ptes = folio_pte_batch(folio, pte, ptent, max_nr_batch_ptes);
-
-		/*
-		 * Must clear entry, or a racing truncate may re-remove it.
-		 * TLB flush can be left until pmdp_collapse_flush() does it.
-		 * PTE dirty? Shmem page is already dirty; file is read-only.
-		 */
-		clear_ptes(mm, addr, pte, nr_batch_ptes);
-		folio_remove_rmap_ptes(folio, page, nr_batch_ptes, vma);
-		nr_mapped_ptes += nr_batch_ptes;
-	}
-
-	if (!pml)
-		spin_unlock(ptl);
-
-	/* step 3: set proper refcount and mm_counters. */
-	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 */
-	if (!pml) {
-		pml = pmd_lock(mm, pmd);
-		if (ptl != pml) {
-			spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
-			if (unlikely(!pmd_same(pgt_pmd, pmdp_get_lockless(pmd)))) {
-				flush_tlb_mm(mm);
-				goto unlock;
-			}
-		}
-	}
-	pgt_pmd = pmdp_collapse_flush(vma, haddr, pmd);
-	pmdp_get_lockless_sync();
-	pte_unmap_unlock(start_pte, ptl);
-	if (ptl != pml)
-		spin_unlock(pml);
-
-	mmu_notifier_invalidate_range_end(&range);
-
-	mm_dec_nr_ptes(mm);
-	page_table_check_pte_clear_range(mm, haddr, pgt_pmd);
-	pte_free_defer(mm, pmd_pgtable(pgt_pmd));
-
-maybe_install_pmd:
-	/* step 5: install pmd entry */
-	result = install_pmd
-			? set_huge_pmd(vma, haddr, pmd, folio, &folio->page)
-			: SCAN_SUCCEED;
-	goto drop_folio;
-abort:
-	if (nr_mapped_ptes) {
-		flush_tlb_mm(mm);
-		folio_ref_sub(folio, nr_mapped_ptes);
-		add_mm_counter(mm, mm_counter_file(folio), -nr_mapped_ptes);
-	}
-unlock:
-	if (start_pte)
-		pte_unmap_unlock(start_pte, ptl);
-	if (pml && pml != ptl)
-		spin_unlock(pml);
-	if (notified)
-		mmu_notifier_invalidate_range_end(&range);
-drop_folio:
-	folio_unlock(folio);
-	folio_put(folio);
-	return result;
-}
-
-/**
- * collapse_pte_mapped_thp - Try to collapse a pte-mapped THP for mm at
- * address haddr.
- *
- * @mm: process address space where collapse happens
- * @addr: THP collapse address
- * @install_pmd: If a huge PMD should be installed
- *
- * This function checks whether all the PTEs in the PMD are pointing to the
- * right THP. If so, retract the page table so the THP can refault in with
- * as pmd-mapped. Possibly install a huge PMD mapping the THP.
- */
-void collapse_pte_mapped_thp(struct mm_struct *mm, unsigned long addr,
-		bool install_pmd)
-{
-	try_collapse_pte_mapped_thp(mm, addr, install_pmd);
-}
-
-/* Can we retract page tables for this file-backed VMA? */
-static bool file_backed_vma_is_retractable(struct vm_area_struct *vma)
-{
-	/*
-	 * Check vma->anon_vma to exclude MAP_PRIVATE mappings that
-	 * got written to. These VMAs are likely not worth removing
-	 * page tables from, as PMD-mapping is likely to be split later.
-	 */
-	if (READ_ONCE(vma->anon_vma))
-		return false;
-
-	/*
-	 * When a vma is registered with uffd-wp or RWP, we cannot recycle
-	 * the page table because there may be pte markers installed.
-	 * VM_UFFD_RWP ranges similarly rely on per-PTE uffd state
-	 * and cannot be recycled to a shared PMD. Other vmas can still
-	 * have the same file mapped hugely, but skip this one: it will
-	 * always be mapped in small page size for these registrations.
-	 */
-	if (userfaultfd_protected(vma))
-		return false;
-
-	/*
-	 * If the VMA contains guard regions then we can't collapse it.
-	 *
-	 * This is set atomically on guard marker installation under mmap/VMA
-	 * read lock, and here we may not hold any VMA or mmap lock at all.
-	 *
-	 * This is therefore serialised on the PTE page table lock, which is
-	 * obtained on guard region installation after the flag is set, so this
-	 * check being performed under this lock excludes races.
-	 */
-	if (vma_test_atomic_flag(vma, VMA_MAYBE_GUARD_BIT))
-		return false;
-
-	return true;
-}
-
-static void retract_page_tables(struct address_space *mapping, pgoff_t pgoff)
-{
-	struct vm_area_struct *vma;
-
-	i_mmap_lock_read(mapping);
-	mapping_rmap_tree_foreach(vma, mapping, pgoff, pgoff) {
-		struct mmu_notifier_range range;
-		struct mm_struct *mm;
-		unsigned long addr;
-		pmd_t *pmd, pgt_pmd;
-		spinlock_t *pml;
-		spinlock_t *ptl;
-		bool success = false;
-
-		addr = vma->vm_start +
-			((pgoff - vma_start_pgoff(vma)) << PAGE_SHIFT);
-		if (addr & ~HPAGE_PMD_MASK ||
-		    vma->vm_end < addr + HPAGE_PMD_SIZE)
-			continue;
-
-		mm = vma->vm_mm;
-		if (find_pmd_or_thp_or_none(mm, addr, &pmd) != SCAN_SUCCEED)
-			continue;
-
-		if (collapse_test_exit(mm))
-			continue;
-
-		if (!file_backed_vma_is_retractable(vma))
-			continue;
-
-		/* PTEs were notified when unmapped; but now for the PMD? */
-		mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
-					addr, addr + HPAGE_PMD_SIZE);
-		mmu_notifier_invalidate_range_start(&range);
-
-		pml = pmd_lock(mm, pmd);
-		/*
-		 * The lock of new_folio is still held, we will be blocked in
-		 * the page fault path, which prevents the pte entries from
-		 * being set again. So even though the old empty PTE page may be
-		 * concurrently freed and a new PTE page is filled into the pmd
-		 * entry, it is still empty and can be removed.
-		 *
-		 * So here we only need to recheck if the state of pmd entry
-		 * still meets our requirements, rather than checking pmd_same()
-		 * like elsewhere.
-		 */
-		if (check_pmd_state(pmd) != SCAN_SUCCEED)
-			goto drop_pml;
-		ptl = pte_lockptr(mm, pmd);
-		if (ptl != pml)
-			spin_lock_nested(ptl, SINGLE_DEPTH_NESTING);
-
-		/*
-		 * Huge page lock is still held, so normally the page table must
-		 * remain empty; and we have already skipped anon_vma and
-		 * userfaultfd_wp() vmas.  But since the mmap_lock is not held,
-		 * it is still possible for a racing userfaultfd_ioctl() or
-		 * madvise() to have inserted ptes or markers.  Now that we hold
-		 * ptlock, repeating the retractable checks protects us from
-		 * races against the prior checks.
-		 */
-		if (likely(file_backed_vma_is_retractable(vma))) {
-			pgt_pmd = pmdp_collapse_flush(vma, addr, pmd);
-			pmdp_get_lockless_sync();
-			success = true;
-		}
-
-		if (ptl != pml)
-			spin_unlock(ptl);
-drop_pml:
-		spin_unlock(pml);
-
-		mmu_notifier_invalidate_range_end(&range);
-
-		if (success) {
-			mm_dec_nr_ptes(mm);
-			page_table_check_pte_clear_range(mm, addr, pgt_pmd);
-			pte_free_defer(mm, pmd_pgtable(pgt_pmd));
-		}
-	}
-	i_mmap_unlock_read(mapping);
-}
-
-/**
- * collapse_file - collapse filemap/tmpfs/shmem pages into huge one.
- *
- * @mm: process address space where collapse happens
- * @addr: virtual collapse start address
- * @file: file that collapse on
- * @start: collapse start address
- * @cc: collapse context and scratchpad
- *
- * Basic scheme is simple, details are more complex:
- *  - allocate and lock a new huge page;
- *  - scan page cache, locking old pages
- *    + swap/gup in pages if necessary;
- *  - copy data to new page
- *  - handle shmem holes
- *    + re-validate that holes weren't filled by someone else
- *    + check for userfaultfd
- *  - finalize updates to the page cache;
- *  - if replacing succeeds:
- *    + unlock huge page;
- *    + free old pages;
- *  - if replacing failed;
- *    + unlock old pages
- *    + unlock and free huge page;
- */
-static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
-		struct file *file, pgoff_t start, struct collapse_control *cc)
-{
-	struct address_space *mapping = file->f_mapping;
-	struct page *dst;
-	struct folio *folio, *tmp, *new_folio;
-	pgoff_t index = 0, end = start + HPAGE_PMD_NR;
-	LIST_HEAD(pagelist);
-	XA_STATE_ORDER(xas, &mapping->i_pages, start, HPAGE_PMD_ORDER);
-	enum scan_result result = SCAN_SUCCEED;
-	int nr_none = 0;
-	bool is_shmem = shmem_file(file);
-
-	/*
-	 * MADV_COLLAPSE ignores shmem huge config, so do not check shmem
-	 *
-	 * TODO: once shmem always calls mapping_set_large_folios() on its
-	 * mapping, the shmem check can be removed.
-	 */
-	VM_WARN_ON_ONCE(!is_shmem && !mapping_pmd_folio_support(mapping));
-	VM_WARN_ON_ONCE(start & (HPAGE_PMD_NR - 1));
-
-	result = alloc_charge_folio(&new_folio, mm, cc, HPAGE_PMD_ORDER);
-	if (result != SCAN_SUCCEED)
-		goto out;
-
-	mapping_set_update(&xas, mapping);
-
-	__folio_set_locked(new_folio);
-	if (is_shmem)
-		__folio_set_swapbacked(new_folio);
-	new_folio->index = start;
-	new_folio->mapping = mapping;
-
-	/*
-	 * Ensure we have slots for all the pages in the range.  This is
-	 * almost certainly a no-op because most of the pages must be present
-	 */
-	do {
-		xas_lock_irq(&xas);
-		xas_create_range(&xas);
-		if (!xas_error(&xas))
-			break;
-		xas_unlock_irq(&xas);
-		if (!xas_nomem(&xas, GFP_KERNEL)) {
-			result = SCAN_FAIL;
-			goto rollback;
-		}
-	} while (1);
-
-	for (index = start; index < end;) {
-		xas_set(&xas, index);
-		folio = xas_load(&xas);
-
-		VM_BUG_ON(index != xas.xa_index);
-		if (is_shmem) {
-			if (!folio) {
-				/*
-				 * Stop if extent has been truncated or
-				 * hole-punched, and is now completely
-				 * empty.
-				 */
-				if (index == start) {
-					if (!xas_next_entry(&xas, end - 1)) {
-						result = SCAN_TRUNCATED;
-						goto xa_locked;
-					}
-				}
-				nr_none++;
-				index++;
-				continue;
-			}
-
-			if (xa_is_value(folio) || !folio_test_uptodate(folio)) {
-				xas_unlock_irq(&xas);
-				/* swap in or instantiate fallocated page */
-				if (shmem_get_folio(mapping->host, index, 0,
-						&folio, SGP_NOALLOC)) {
-					result = SCAN_FAIL;
-					goto xa_unlocked;
-				}
-				/* drain lru cache to help folio_isolate_lru() */
-				lru_add_drain();
-			} else if (folio_trylock(folio)) {
-				folio_get(folio);
-				xas_unlock_irq(&xas);
-			} else {
-				result = SCAN_PAGE_LOCK;
-				goto xa_locked;
-			}
-		} else {	/* !is_shmem */
-			if (!folio || xa_is_value(folio)) {
-				xas_unlock_irq(&xas);
-				page_cache_sync_readahead(mapping, &file->f_ra,
-							  file, index,
-							  end - index);
-				/* drain lru cache to help folio_isolate_lru() */
-				lru_add_drain();
-				folio = filemap_lock_folio(mapping, index);
-				if (IS_ERR(folio)) {
-					result = SCAN_FAIL;
-					goto xa_unlocked;
-				}
-			} else if (folio_test_dirty(folio)) {
-				/*
-				 * This page is dirty because it hasn't
-				 * been flushed since first write.
-				 *
-				 * Trigger async flush for read-only files and
-				 * hope the writeback is done when khugepaged
-				 * revisits this page. Writable files can have
-				 * their folios dirty at any time; blindly
-				 * flushing them would cause undesirable
-				 * system-wide writeback.
-				 *
-				 * This is a one-off situation. We are not
-				 * forcing writeback in loop.
-				 */
-				xas_unlock_irq(&xas);
-				if (!inode_is_open_for_write(mapping->host))
-					filemap_flush(mapping);
-				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-				goto xa_unlocked;
-			} else if (folio_test_writeback(folio)) {
-				xas_unlock_irq(&xas);
-				result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-				goto xa_unlocked;
-			} else if (folio_trylock(folio)) {
-				folio_get(folio);
-				xas_unlock_irq(&xas);
-			} else {
-				result = SCAN_PAGE_LOCK;
-				goto xa_locked;
-			}
-		}
-
-		/*
-		 * The folio must be locked, so we can drop the i_pages lock
-		 * without racing with truncate.
-		 */
-		VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
-
-		/* make sure the folio is up to date */
-		if (unlikely(!folio_test_uptodate(folio))) {
-			result = SCAN_FAIL;
-			goto out_unlock;
-		}
-
-		/*
-		 * If file was truncated then extended, or hole-punched, before
-		 * we locked the first folio, then a THP might be there already.
-		 * This will be discovered on the first iteration.
-		 */
-		if (is_pmd_order(folio_order(folio))) {
-			result = SCAN_PTE_MAPPED_HUGEPAGE;
-			goto out_unlock;
-		}
-
-		if (folio_mapping(folio) != mapping) {
-			result = SCAN_TRUNCATED;
-			goto out_unlock;
-		}
-
-		if (!is_shmem && (folio_test_dirty(folio) ||
-				  folio_test_writeback(folio))) {
-			/*
-			 * khugepaged only works on clean file-backed folios,
-			 * so this folio is dirty because it hasn't been flushed
-			 * since first write.
-			 */
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			goto out_unlock;
-		}
-
-		if (!folio_isolate_lru(folio)) {
-			result = SCAN_DEL_PAGE_LRU;
-			goto out_unlock;
-		}
-
-		if (!filemap_release_folio(folio, GFP_KERNEL)) {
-			result = SCAN_PAGE_HAS_PRIVATE;
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
-		if (folio_mapped(folio))
-			try_to_unmap(folio,
-					TTU_IGNORE_MLOCK | TTU_BATCH_FLUSH);
-
-		xas_lock_irq(&xas);
-
-		VM_BUG_ON_FOLIO(folio != xa_load(xas.xa, index), folio);
-
-		/*
-		 * We control 2 + nr_pages references to the folio:
-		 *  - we hold a pin on it;
-		 *  - nr_pages reference from page cache;
-		 *  - one from lru_isolate_folio;
-		 * If those are the only references, then any new usage
-		 * of the folio will have to fetch it from the page
-		 * cache. That requires locking the folio to handle
-		 * truncate, so any new usage will be blocked until we
-		 * unlock folio after collapse/during rollback.
-		 */
-		if (folio_ref_count(folio) != 2 + folio_nr_pages(folio)) {
-			result = SCAN_PAGE_COUNT;
-			xas_unlock_irq(&xas);
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
-		/*
-		 * At this point, the folio is locked and unmapped. If the PTE
-		 * was dirty, try_to_unmap() has transferred the dirty bit to
-		 * the folio and we must not collapse it into a clean
-		 * file-backed folio.
-		 *
-		 * If the folio is clean here, no one can write it until we
-		 * drop the folio lock. A write through a stale TLB entry came
-		 * from a clean PTE and must fault because the PTE has been
-		 * cleared; the fault path has to take the folio lock before
-		 * installing a writable mapping. Buffered write paths also
-		 * have to take the folio lock before modifying file contents
-		 * without a mapping, typically via write_begin_get_folio().
-		 */
-		if (!is_shmem && folio_test_dirty(folio)) {
-			result = SCAN_PAGE_DIRTY_OR_WRITEBACK;
-			xas_unlock_irq(&xas);
-			folio_putback_lru(folio);
-			goto out_unlock;
-		}
-
-		/*
-		 * Accumulate the folios that are being collapsed.
-		 */
-		list_add_tail(&folio->lru, &pagelist);
-		index += folio_nr_pages(folio);
-		continue;
-out_unlock:
-		folio_unlock(folio);
-		folio_put(folio);
-		goto xa_unlocked;
-	}
-
-xa_locked:
-	xas_unlock_irq(&xas);
-xa_unlocked:
-
-	/*
-	 * If collapse is successful, flush must be done now before copying.
-	 * If collapse is unsuccessful, does flush actually need to be done?
-	 * Do it anyway, to clear the state.
-	 */
-	try_to_unmap_flush();
-
-	if (result == SCAN_SUCCEED && nr_none &&
-	    !shmem_charge(mapping->host, nr_none))
-		result = SCAN_FAIL;
-	if (result != SCAN_SUCCEED) {
-		nr_none = 0;
-		goto rollback;
-	}
-
-	/*
-	 * The old folios are locked, so they won't change anymore.
-	 */
-	index = start;
-	dst = folio_page(new_folio, 0);
-	list_for_each_entry(folio, &pagelist, lru) {
-		int i, nr_pages = folio_nr_pages(folio);
-
-		while (index < folio->index) {
-			clear_highpage(dst);
-			index++;
-			dst++;
-		}
-
-		for (i = 0; i < nr_pages; i++) {
-			if (copy_mc_highpage(dst, folio_page(folio, i)) > 0) {
-				result = SCAN_COPY_MC;
-				goto rollback;
-			}
-			index++;
-			dst++;
-		}
-	}
-	while (index < end) {
-		clear_highpage(dst);
-		index++;
-		dst++;
-	}
-
-	if (nr_none) {
-		struct vm_area_struct *vma;
-		int nr_none_check = 0;
-
-		i_mmap_lock_read(mapping);
-		xas_lock_irq(&xas);
-
-		xas_set(&xas, start);
-		for (index = start; index < end; index++) {
-			if (!xas_next(&xas)) {
-				xas_store(&xas, XA_RETRY_ENTRY);
-				if (xas_error(&xas)) {
-					result = SCAN_STORE_FAILED;
-					goto immap_locked;
-				}
-				nr_none_check++;
-			}
-		}
-
-		if (nr_none != nr_none_check) {
-			result = SCAN_PAGE_FILLED;
-			goto immap_locked;
-		}
-
-		/*
-		 * If userspace observed a missing page in a VMA with
-		 * a MODE_MISSING userfaultfd, then it might expect a
-		 * UFFD_EVENT_PAGEFAULT for that page. If so, we need to
-		 * roll back to avoid suppressing such an event. Since
-		 * wp/minor userfaultfds don't give userspace any
-		 * guarantees that the kernel doesn't fill a missing
-		 * page with a zero page, so they don't matter here.
-		 *
-		 * Any userfaultfds registered after this point will
-		 * not be able to observe any missing pages due to the
-		 * previously inserted retry entries.
-		 */
-		mapping_rmap_tree_foreach(vma, mapping, start, end) {
-			if (userfaultfd_missing(vma)) {
-				result = SCAN_EXCEED_NONE_PTE;
-				goto immap_locked;
-			}
-		}
-
-immap_locked:
-		i_mmap_unlock_read(mapping);
-		if (result != SCAN_SUCCEED) {
-			xas_set(&xas, start);
-			for (index = start; index < end; index++) {
-				if (xas_next(&xas) == XA_RETRY_ENTRY)
-					xas_store(&xas, NULL);
-			}
-
-			xas_unlock_irq(&xas);
-			goto rollback;
-		}
-	} else {
-		xas_lock_irq(&xas);
-	}
-
-	if (is_shmem) {
-		lruvec_stat_mod_folio(new_folio, NR_SHMEM, HPAGE_PMD_NR);
-		lruvec_stat_mod_folio(new_folio, NR_SHMEM_THPS, HPAGE_PMD_NR);
-	} else {
-		lruvec_stat_mod_folio(new_folio, NR_FILE_THPS, HPAGE_PMD_NR);
-	}
-	lruvec_stat_mod_folio(new_folio, NR_FILE_PAGES, HPAGE_PMD_NR);
-
-	/*
-	 * Mark new_folio as uptodate before inserting it into the
-	 * page cache so that it isn't mistaken for an fallocated but
-	 * unwritten page.
-	 */
-	folio_mark_uptodate(new_folio);
-	folio_ref_add(new_folio, HPAGE_PMD_NR - 1);
-
-	if (is_shmem)
-		folio_mark_dirty(new_folio);
-	folio_add_lru(new_folio);
-
-	/* Join all the small entries into a single multi-index entry. */
-	xas_set_order(&xas, start, HPAGE_PMD_ORDER);
-	xas_store(&xas, new_folio);
-	WARN_ON_ONCE(xas_error(&xas));
-	xas_unlock_irq(&xas);
-
-	/*
-	 * Remove pte page tables, so we can re-fault the page as huge.  A caller
-	 * that wants the PMD mapped now is told to go and do that.
-	 */
-	retract_page_tables(mapping, start);
-	if (cc->policy.install_pmd)
-		result = SCAN_PTE_MAPPED_HUGEPAGE;
-	folio_unlock(new_folio);
-
-	/*
-	 * The collapse has succeeded, so free the old folios.
-	 */
-	list_for_each_entry_safe(folio, tmp, &pagelist, lru) {
-		list_del(&folio->lru);
-		lruvec_stat_mod_folio(folio, NR_FILE_PAGES,
-				      -folio_nr_pages(folio));
-		if (is_shmem)
-			lruvec_stat_mod_folio(folio, NR_SHMEM,
-					      -folio_nr_pages(folio));
-		folio->mapping = NULL;
-		folio_clear_active(folio);
-		folio_clear_unevictable(folio);
-		folio_unlock(folio);
-		folio_put_refs(folio, 2 + folio_nr_pages(folio));
-	}
-
-	goto out;
-
-rollback:
-	/* Something went wrong: roll back page cache changes */
-	if (nr_none) {
-		xas_lock_irq(&xas);
-		mapping->nrpages -= nr_none;
-		xas_unlock_irq(&xas);
-		shmem_uncharge(mapping->host, nr_none);
-	}
-
-	list_for_each_entry_safe(folio, tmp, &pagelist, lru) {
-		list_del(&folio->lru);
-		folio_unlock(folio);
-		folio_putback_lru(folio);
-		folio_put(folio);
-	}
-
-	new_folio->mapping = NULL;
-
-	folio_unlock(new_folio);
-	folio_put(new_folio);
-out:
-	VM_BUG_ON(!list_empty(&pagelist));
-	trace_mm_khugepaged_collapse_file(mm, new_folio, index, addr, is_shmem, file, HPAGE_PMD_NR, result);
-	return result;
-}
-
-static enum scan_result collapse_scan_file(struct mm_struct *mm,
-		unsigned long addr, struct file *file, pgoff_t start,
-		struct collapse_control *cc)
-{
-	const unsigned int max_ptes_none = collapse_max_ptes_none(cc, NULL, HPAGE_PMD_ORDER);
-	const unsigned int max_ptes_swap = collapse_max_ptes_swap(cc, HPAGE_PMD_ORDER);
-	struct folio *folio = NULL;
-	struct address_space *mapping = file->f_mapping;
-	XA_STATE(xas, &mapping->i_pages, start);
-	int present, swap;
-	int node = NUMA_NO_NODE;
-	enum scan_result result = SCAN_SUCCEED;
-
-	present = 0;
-	swap = 0;
-	collapse_control_init_scan(cc);
-	rcu_read_lock();
-	xas_for_each(&xas, folio, start + HPAGE_PMD_NR - 1) {
-		if (xas_retry(&xas, folio))
-			continue;
-
-		if (xa_is_value(folio)) {
-			swap += 1 << xas_get_order(&xas);
-			if (swap > max_ptes_swap) {
-				result = SCAN_EXCEED_SWAP_PTE;
-				count_vm_event(THP_SCAN_EXCEED_SWAP_PTE);
-				break;
-			}
-			continue;
-		}
-
-		if (!folio_try_get(folio)) {
-			xas_reset(&xas);
-			continue;
-		}
-
-		if (unlikely(folio != xas_reload(&xas))) {
-			folio_put(folio);
-			xas_reset(&xas);
-			continue;
-		}
-
-		if (is_pmd_order(folio_order(folio))) {
-			result = SCAN_PTE_MAPPED_HUGEPAGE;
-			/*
-			 * PMD-sized THP implies that we can only try
-			 * retracting the PTE table.
-			 */
-			folio_put(folio);
-			break;
-		}
-
-		node = folio_nid(folio);
-		if (collapse_scan_abort(node, cc)) {
-			result = SCAN_SCAN_ABORT;
-			folio_put(folio);
-			break;
-		}
-		cc->node_load[node]++;
-
-		if (!folio_test_lru(folio)) {
-			result = SCAN_PAGE_LRU;
-			folio_put(folio);
-			break;
-		}
-
-		if (folio_expected_ref_count(folio) + 1 != folio_ref_count(folio)) {
-			result = SCAN_PAGE_COUNT;
-			folio_put(folio);
-			break;
-		}
-
-		/*
-		 * We probably should check if the folio is referenced
-		 * here, but nobody would transfer pte_young() to
-		 * folio_test_referenced() for us.  And rmap walk here
-		 * is just too costly...
-		 */
-
-		present += folio_nr_pages(folio);
-		folio_put(folio);
-
-		if (need_resched()) {
-			xas_pause(&xas);
-			cond_resched_rcu();
-		}
-	}
-	rcu_read_unlock();
-	if (result == SCAN_PTE_MAPPED_HUGEPAGE)
-		cc->progress++;
-	else
-		cc->progress += HPAGE_PMD_NR;
-
-	if (result == SCAN_SUCCEED) {
-		if (present < HPAGE_PMD_NR - max_ptes_none) {
-			result = SCAN_EXCEED_NONE_PTE;
-			count_vm_event(THP_SCAN_EXCEED_NONE_PTE);
-		} else {
-			result = collapse_file(mm, addr, file, start, cc);
-		}
-	}
-
-	trace_mm_khugepaged_scan_file(mm, folio, file, present, swap, result);
-	return result;
-}
-
-/*
- * Try to collapse a single PMD starting at a PMD aligned addr, and return
- * the results.
- */
-static enum scan_result collapse_single_pmd(unsigned long addr,
-		unsigned long end, struct vm_area_struct *vma,
-		bool *lock_dropped, struct collapse_control *cc)
-{
-	struct mm_struct *mm = vma->vm_mm;
-	bool triggered_wb = false;
-	enum scan_result result;
-	struct file *file;
-	pgoff_t pgoff;
-
-	mmap_assert_locked(mm);
-
-	if (vma_is_anonymous(vma)) {
-		result = collapse_scan_anon_pmd(vma, addr, end, cc);
-		if (!cc->select_orders)
-			goto end;
-
-		/* collapse_anon_pmd() takes mmap_lock itself, where it needs it */
-		mmap_read_unlock(mm);
-		*lock_dropped = true;
-
-		result = collapse_anon_pmd(mm, addr, end, cc);
-		goto end;
-	}
-
-	file = get_file(vma->vm_file);
-	pgoff = linear_page_index(vma, addr);
-
-	mmap_read_unlock(mm);
-	*lock_dropped = true;
-retry:
-	result = collapse_scan_file(mm, addr, file, pgoff, cc);
-
-	/* Dirty pages are worth a writeback and one more try, if asked for */
-	if (cc->policy.writeback_dirty && result == SCAN_PAGE_DIRTY_OR_WRITEBACK &&
-	    !triggered_wb && mapping_can_writeback(file->f_mapping)) {
-		const loff_t lstart = (loff_t)pgoff << PAGE_SHIFT;
-		const loff_t lend = lstart + HPAGE_PMD_SIZE - 1;
-
-		filemap_write_and_wait_range(file->f_mapping, lstart, lend);
-		triggered_wb = true;
-		goto retry;
-	}
-	fput(file);
-
-	if (result == SCAN_PTE_MAPPED_HUGEPAGE) {
-		mmap_read_lock(mm);
-		if (collapse_test_exit_or_disable(mm))
-			result = SCAN_ANY_PROCESS;
-		else
-			result = try_collapse_pte_mapped_thp(mm, addr,
-							cc->policy.install_pmd);
-		if (result == SCAN_PMD_MAPPED)
-			result = SCAN_SUCCEED;
-		mmap_read_unlock(mm);
-	}
-end:
-	return result;
-}
-
 static void collapse_scan_mm_slot(unsigned int progress_max,
 		enum scan_result *result, struct collapse_control *cc)
 	__releases(&khugepaged_mm_lock)
-- 
2.54.0

[RFC PATCH 33/57] mm/collapse: split collapse into a scan and a run

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

collapse_single_pmd() did both halves of a collapse behind one call, and
dropped mmap_lock somewhere in the middle.  Which of its paths dropped it
was not something a caller could see, so it was handed back a bool and had
to keep track.

Both callers did that badly: khugepaged broke out of its VMA walk after
every table, collapsed or not, and MADV_COLLAPSE carried lock state
through its loop and re-took the lock only to hand it back.

Split it in two, with the lock as the boundary:

 - collapse_scan_pmd() judges one table and returns with mmap_lock still
   held.  It only reads, and almost every table it is offered has nothing
   in it, so a caller walks a whole VMA under the one lock it took to get
   there.
 - collapse_run_pmd() is called without the lock, which the caller gives
   up first, and takes it again per round.  What it does is slow enough
   that a writer would otherwise wait behind all of it.

Whether there is anything to run is the scan's return value, so no caller
has to ask about the lock.

The file side is what makes this more than a rename.  A file collapse
works on the page cache and never sees a VMA, but the file and the offset
have to come from one: the scan takes them while it still has the VMA, in
cc->scan_file and cc->scan_pgoff, and the run is what gives the reference
back.  A scan that found file work therefore has to be run, and
collapse_control_release() warns and drops the reference rather than rest
on callers getting that right.

The orders a VMA allows now come in as an argument.  khugepaged's walk
already computes that mask once per VMA, where the old shape recomputed it
twice for every table.  MADV_COLLAPSE keeps its own copy only while it
holds the VMA, and clears it beside vma = NULL: after the lock is given
up, the next lookup may return a different VMA.

khugepaged now stays in its VMA walk across every table it refuses, and
gives the lock up only for a table it is going to collapse.

MADV_COLLAPSE does the same.  It gives up the lock the VMA walk left held,
so lru_add_drain_all() does not wait on every CPU under it, then takes it
once and walks its whole range.  It drops that lock only around a
collapse, and looks the VMA up again only then, since nothing else can
have moved it.

Holding on across refusals is not new for anonymous memory -- the old
entry already returned with the lock held when a table had nothing in it
-- but it was never true of file ranges, and never something a caller
could rely on.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 167 +++++++++++++++++++++++++++++++-----------
 mm/collapse.h   |  48 ++++++++++++-
 mm/khugepaged.c | 188 +++++++++++++++++++++++-------------------------
 mm/mremap.c     |   2 +-
 4 files changed, 262 insertions(+), 143 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 21bfbc038044..f0d80204c2bd 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -401,6 +401,10 @@ static unsigned int candidate_offset(const struct collapse_candidate *cand,
 
 void collapse_control_release(struct collapse_control *cc)
 {
+	/* Only a scan that was never run leaves this behind */
+	if (WARN_ON_ONCE(cc->scan_file))
+		fput(cc->scan_file);
+
 	kfree(cc->candidates);
 	kfree(cc->saved_ptes);
 	kfree(cc->retries);
@@ -413,6 +417,10 @@ int collapse_control_init(struct collapse_control *cc)
 {
 	cc->nr_candidates = 0;
 	cc->nr_retries = 0;
+	cc->select_orders = 0;
+	cc->scan_refusal = SCAN_FAIL;
+	cc->scan_file = NULL;
+	cc->scan_pgoff = 0;
 	cc->candidates = kmalloc_objs(*cc->candidates, COLLAPSE_MAX_CANDIDATES);
 	cc->saved_ptes = kmalloc_objs(*cc->saved_ptes, COLLAPSE_SAVED_PTES);
 	cc->retries = kmalloc_objs(*cc->retries, COLLAPSE_RETRY_STORE_SIZE);
@@ -2116,7 +2124,8 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
  */
 static enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
 					unsigned long start, unsigned long end,
-					struct collapse_control *cc)
+					struct collapse_control *cc,
+					unsigned long vma_orders)
 {
 	const unsigned long pmd_addr = start & HPAGE_PMD_MASK;
 	struct mm_struct *mm = vma->vm_mm;
@@ -2135,12 +2144,7 @@ static enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
 	/* Cleared only once a table has turned out to be there */
 	collapse_anon_scan_init(cc);
 
-	cc->select_orders = collapse_possible_orders(vma, vma->vm_flags,
-						     cc->policy.tva_type);
-	if (!cc->select_orders) {
-		cc->scan_refusal = SCAN_VMA_CHECK;
-		return cc->scan_refusal;
-	}
+	cc->select_orders = vma_orders;
 
 	/* The scan narrows select_orders to whatever is left worth trying */
 	cc->scan_refusal = collapse_scan_table(vma, pmd, start, end, cc);
@@ -2596,7 +2600,7 @@ static void count_collapse_event(unsigned int order, enum vm_event_item vm_event
 	count_mthp_stat(order, mthp_event);
 }
 
-static void collapse_control_init_scan(struct collapse_control *cc)
+static void collapse_file_scan_init(struct collapse_control *cc)
 {
 	memset(cc->node_load, 0, sizeof(cc->node_load));
 	nodes_clear(cc->alloc_nmask);
@@ -3493,7 +3497,7 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 	return result;
 }
 
-static enum scan_result collapse_scan_file(struct mm_struct *mm,
+static enum scan_result collapse_pagecache_pmd(struct mm_struct *mm,
 		unsigned long addr, struct file *file, pgoff_t start,
 		struct collapse_control *cc)
 {
@@ -3508,7 +3512,7 @@ static enum scan_result collapse_scan_file(struct mm_struct *mm,
 
 	present = 0;
 	swap = 0;
-	collapse_control_init_scan(cc);
+	collapse_file_scan_init(cc);
 	rcu_read_lock();
 	xas_for_each(&xas, folio, start + HPAGE_PMD_NR - 1) {
 		if (xas_retry(&xas, folio))
@@ -3600,46 +3604,65 @@ static enum scan_result collapse_scan_file(struct mm_struct *mm,
 }
 
 /*
- * Try to collapse a single PMD starting at a PMD aligned addr, and return
- * the results.
+ * Judge one table's worth of a file VMA.  All it needs of the VMA is the file and
+ * the offset, which it takes while it still has both; the collapse works on the
+ * page cache and never sees a VMA.
  */
-enum scan_result collapse_single_pmd(unsigned long addr,
-		unsigned long end, struct vm_area_struct *vma,
-		bool *lock_dropped, struct collapse_control *cc)
+static enum scan_result collapse_scan_file_pmd(struct vm_area_struct *vma,
+		unsigned long addr, struct collapse_control *cc)
 {
-	struct mm_struct *mm = vma->vm_mm;
+	enum scan_result result;
+	pmd_t *pmd;
+
+	/*
+	 * A file collapse only ever builds a PMD, so the whole table has to be
+	 * the VMA's -- a PMD shared with another VMA would need all of them
+	 * locked.  Not the question collapse_possible_orders() answered, which is
+	 * whether the VMA may use the order at all: this is whether the table at
+	 * @addr is wholly inside it.  While a file VMA collapses at PMD order
+	 * alone its callers hand over whole tables and this cannot fire, but the
+	 * anonymous side already hands over parts of one.
+	 */
+	if (!thp_vma_suitable_order(vma, addr, HPAGE_PMD_ORDER))
+		return SCAN_ADDRESS_RANGE;
+
+	/*
+	 * A PMD that is huge already has nothing left to collapse, and skipping
+	 * it here is what keeps mmap_lock out of a collapse that would find
+	 * nothing.  Everything else is worth the page cache scan, pmd_none()
+	 * included: a file range can be collapsed out of the cache without being
+	 * mapped first, which is why this is not the test the anonymous side
+	 * makes.
+	 */
+	result = find_pmd_or_thp_or_none(vma->vm_mm, addr & HPAGE_PMD_MASK, &pmd);
+	if (result == SCAN_PMD_MAPPED)
+		return result;
+
+	cc->scan_file = get_file(vma->vm_file);
+	cc->scan_pgoff = linear_page_index(vma, addr);
+
+	return SCAN_SUCCEED;
+}
+
+/*
+ * Build a PMD over what the page cache holds, and map it over the range if a huge
+ * folio is already there but mapped by PTEs.  Runs with no mmap_lock, which the
+ * caller gave up, and takes it again only for that last step.
+ */
+static enum scan_result collapse_file_pmd(struct mm_struct *mm,
+		unsigned long addr, struct collapse_control *cc)
+{
+	struct file *file = cc->scan_file;
 	bool triggered_wb = false;
 	enum scan_result result;
-	struct file *file;
-	pgoff_t pgoff;
 
-	mmap_assert_locked(mm);
-
-	if (vma_is_anonymous(vma)) {
-		result = collapse_scan_anon_pmd(vma, addr, end, cc);
-		if (!cc->select_orders)
-			goto end;
-
-		/* collapse_anon_pmd() takes mmap_lock itself, where it needs it */
-		mmap_read_unlock(mm);
-		*lock_dropped = true;
-
-		result = collapse_anon_pmd(mm, addr, end, cc);
-		goto end;
-	}
-
-	file = get_file(vma->vm_file);
-	pgoff = linear_page_index(vma, addr);
-
-	mmap_read_unlock(mm);
-	*lock_dropped = true;
 retry:
-	result = collapse_scan_file(mm, addr, file, pgoff, cc);
+	result = collapse_pagecache_pmd(mm, addr, file, cc->scan_pgoff, cc);
 
 	/* Dirty pages are worth a writeback and one more try, if asked for */
 	if (cc->policy.writeback_dirty && result == SCAN_PAGE_DIRTY_OR_WRITEBACK &&
 	    !triggered_wb && mapping_can_writeback(file->f_mapping)) {
-		const loff_t lstart = (loff_t)pgoff << PAGE_SHIFT;
+		const loff_t lstart = (loff_t)cc->scan_pgoff << PAGE_SHIFT;
 		const loff_t lend = lstart + HPAGE_PMD_SIZE - 1;
 
 		filemap_write_and_wait_range(file->f_mapping, lstart, lend);
@@ -3647,6 +3670,7 @@ enum scan_result collapse_single_pmd(unsigned long addr,
 		goto retry;
 	}
 	fput(file);
+	cc->scan_file = NULL;
 
 	if (result == SCAN_PTE_MAPPED_HUGEPAGE) {
 		mmap_read_lock(mm);
@@ -3659,6 +3683,67 @@ enum scan_result collapse_single_pmd(unsigned long addr,
 			result = SCAN_SUCCEED;
 		mmap_read_unlock(mm);
 	}
-end:
+
 	return result;
 }
+
+/*
+ * Scan one table's worth of @vma and decide whether there is anything to collapse
+ * in it.  The caller holds mmap_lock for reading and still holds it when this
+ * returns: what is looked at is either the VMA or a page table that the lock
+ * keeps in place.
+ *
+ * Returns whether collapse_run_pmd() has anything to do, and a scan that found
+ * something has to be run: the file side takes a reference on the file while it
+ * still has the VMA to take it from, and the run is what gives it back.  What the
+ * scan turned down is left in cc->scan_refusal either way.
+ */
+bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
+		unsigned long end, struct collapse_control *cc,
+		unsigned long vma_orders)
+{
+	struct mm_struct *mm = vma->vm_mm;
+
+	mmap_assert_locked(mm);
+
+	/*
+	 * What the scan answers with, so cleared before it runs.
+	 * collapse_anon_scan_init() clears the orders too, but only once the
+	 * table has turned out to be there.
+	 */
+	cc->select_orders = 0;
+
+	/* Ours to give back only if the last scan was never run */
+	if (WARN_ON_ONCE(cc->scan_file)) {
+		fput(cc->scan_file);
+		cc->scan_file = NULL;
+	}
+
+	if (unlikely(collapse_test_exit_or_disable(mm)))
+		cc->scan_refusal = SCAN_ANY_PROCESS;
+	else if (addr < vma->vm_start || end > vma->vm_end)
+		cc->scan_refusal = SCAN_ADDRESS_RANGE;
+	else if (!vma_orders)
+		cc->scan_refusal = SCAN_VMA_CHECK;
+	else if (vma_is_anonymous(vma))
+		collapse_scan_anon_pmd(vma, addr, end, cc, vma_orders);
+	else
+		cc->scan_refusal = collapse_scan_file_pmd(vma, addr, cc);
+
+	return cc->select_orders || cc->scan_file;
+}
+
+/*
+ * Collapse what the scan selected.  Called with no mmap_lock: the caller gives it
+ * up first, because a collapse takes it again for each round and revalidates
+ * under it, and holding it across the whole collapse would keep a writer to the
+ * address space waiting for it.
+ */
+enum scan_result collapse_run_pmd(struct mm_struct *mm, unsigned long addr,
+		unsigned long end, struct collapse_control *cc)
+{
+	if (cc->scan_file)
+		return collapse_file_pmd(mm, addr, cc);
+	else
+		return collapse_anon_pmd(mm, addr, end, cc);
+}
diff --git a/mm/collapse.h b/mm/collapse.h
index dc60806fb81e..4af7bb9c4261 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -150,6 +150,13 @@ struct collapse_control {
 	 */
 	enum scan_result scan_refusal;
 
+	/*
+	 * A reference the file side takes while it still has the VMA, since the
+	 * collapse runs without it, and the offset it decided on.
+	 */
+	struct file *scan_file;
+	pgoff_t scan_pgoff;
+
 	/* Why the last window was refused */
 	enum scan_result select_result;
 
@@ -187,12 +194,47 @@ static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
 		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
 }
 
+/*
+ * A caller states what it allows in the policy, takes a control for the arrays a
+ * round needs, and then hands over one PTE table's worth of a VMA at a time:
+ *
+ *	collapse_control_init(cc);		once per control
+ *	fill in cc->policy;			what this caller allows
+ *	collapse_scan_pmd(vma, addr, end, cc);	per table, as often as wanted
+ *	collapse_run_pmd(mm, addr, end, cc);	when the scan found work
+ *	collapse_control_release(cc);
+ *
+ * The caller holds mmap_lock for reading and passes a range within one PTE table
+ * of @vma.  A range the VMA does not cover is refused, which is also how a caller
+ * learns that its own range shrank.
+ *
+ * A scan returns with that lock still held: it only reads, and almost every table
+ * it is offered has nothing to collapse, so a caller walks a whole VMA under the
+ * one lock it took to get there.
+ *
+ * A collapse is called without it: the caller gives the lock up first, and with it
+ * @vma and anything derived under it, so a caller carrying on has to look up
+ * again.  What the collapse does -- allocate, quiesce, copy, flush -- is slow
+ * enough that a writer would wait behind it, so it takes the lock again per round
+ * instead, and revalidates rather than trusting what the scan saw.
+ *
+ * A scan that found something has to be run: the file side takes a reference on
+ * the file while it still has the VMA to take it from, and the run is what gives
+ * it back.  What it turned down is left in cc->scan_refusal, for a caller that has
+ * to report why a table was not collapsed.
+ *
+ * A control is not reentrant: it carries the arrays a round works out of, so one
+ * per collapsing thread.
+ */
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
-enum scan_result collapse_single_pmd(unsigned long addr, unsigned long end,
-		struct vm_area_struct *vma, bool *lock_dropped,
-		struct collapse_control *cc);
+bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
+		unsigned long end, struct collapse_control *cc,
+		unsigned long vma_orders);
+enum scan_result collapse_run_pmd(struct mm_struct *mm, unsigned long addr,
+		unsigned long end, struct collapse_control *cc);
 
+/* Which orders a VMA may collapse to, empty when it may not collapse at all */
 unsigned long collapse_possible_orders(struct vm_area_struct *vma,
 		vm_flags_t vm_flags, enum tva_type tva_flags);
 
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index b7fc93e11d6b..47c134cd4129 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -485,64 +485,6 @@ static void collapse_policy_khugepaged(struct collapse_policy *p)
 	p->tva_type = TVA_KHUGEPAGED;
 }
 
-/* MADV_COLLAPSE was asked for explicitly, so it is not held to those. */
-static void collapse_policy_forced(struct collapse_policy *p)
-{
-	p->max_ptes_none = HPAGE_PMD_NR;
-	p->max_ptes_swap = HPAGE_PMD_NR;
-	p->max_ptes_shared = HPAGE_PMD_NR;
-	p->strict_sub_pmd = false;
-	p->skip_lazyfree = false;
-	p->require_referenced = false;
-	p->install_pmd = true;
-	p->writeback_dirty = true;
-	p->gfp = GFP_TRANSHUGE;
-	p->tva_type = TVA_FORCED_COLLAPSE;
-}
-
-/*
- * If mmap_lock temporarily dropped, revalidate vma
- * after taking the mmap_lock again.
- * Returns enum scan_result value.
- */
-
-static enum scan_result hugepage_vma_revalidate(struct mm_struct *mm, unsigned long address,
-		bool expect_anon, struct vm_area_struct **vmap,
-		struct collapse_control *cc, unsigned int order)
-{
-	struct vm_area_struct *vma;
-	enum tva_type type = cc->policy.tva_type;
-
-	if (unlikely(collapse_test_exit_or_disable(mm)))
-		return SCAN_ANY_PROCESS;
-
-	*vmap = vma = find_vma(mm, address);
-	if (!vma)
-		return SCAN_VMA_NULL;
-
-	/*
-	 * We cannot collapse VMA regions that do not span the full PMD. This is
-	 * due to the potential of the PMD being shared by another VMA leaving
-	 * us vulnerable to a race condition. Always check the PMD order here to
-	 * ensure its not shared by another VMA. We'd need to lock all VMAs in
-	 * the PMD range to support this.
-	 */
-	if (!thp_vma_suitable_order(vma, address, PMD_ORDER))
-		return SCAN_ADDRESS_RANGE;
-	if (!thp_vma_allowable_orders(vma, vma->vm_flags, type, BIT(order)))
-		return SCAN_VMA_CHECK;
-	/*
-	 * Anon VMA expected, the address may be unmapped then
-	 * remapped to file after khugepaged reacquired the mmap_lock.
-	 *
-	 * thp_vma_allowable_orders() may return true for qualified file
-	 * vmas.
-	 */
-	if (expect_anon && (!(*vmap)->anon_vma || !vma_is_anonymous(*vmap)))
-		return SCAN_PAGE_ANON;
-	return SCAN_SUCCEED;
-}
-
 static void collect_mm_slot(struct mm_slot *slot)
 {
 	struct mm_struct *mm = slot->mm;
@@ -636,8 +578,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			khugepaged_scan.address = hstart;
 
 		while (khugepaged_scan.address < hend) {
-			unsigned long pmd_addr, range_end;
-			bool lock_dropped = false;
+			unsigned long pmd_addr, range_end, start;
 
 			/* One table's worth at most, and never past the VMA */
 			pmd_addr = khugepaged_scan.address & HPAGE_PMD_MASK;
@@ -649,24 +590,24 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 
 			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart);
 
-			*result = collapse_single_pmd(khugepaged_scan.address,
-						      range_end, vma,
-						      &lock_dropped, cc);
-			if (*result == SCAN_SUCCEED)
-				++khugepaged_pages_collapsed;
+			start = khugepaged_scan.address;
 			/* move to next address */
 			khugepaged_scan.address = range_end;
-			if (lock_dropped)
-				/*
-				 * We released mmap_lock so break loop.  Note
-				 * that we drop mmap_lock before all hugepage
-				 * allocations, so if allocation fails, we are
-				 * guaranteed to break here and report the
-				 * correct result back to caller.
-				 */
-				goto breakouterloop_mmap_lock;
-			if (cc->progress >= progress_max)
-				goto breakouterloop;
+
+			/* If nothing to collapse, the lock is still ours */
+			if (!collapse_scan_pmd(vma, start, range_end, cc, orders)) {
+				*result = cc->scan_refusal;
+				if (cc->progress >= progress_max)
+					goto breakouterloop;
+				continue;
+			}
+
+			/* collapse_run_pmd() takes its own locks, so give this up */
+			mmap_read_unlock(mm);
+			*result = collapse_run_pmd(mm, start, range_end, cc);
+			if (*result == SCAN_SUCCEED)
+				++khugepaged_pages_collapsed;
+			goto breakouterloop_mmap_lock;
 		}
 	}
 breakouterloop:
@@ -904,6 +845,21 @@ bool current_is_khugepaged(void)
 	return kthread_func(current) == khugepaged;
 }
 
+/* MADV_COLLAPSE was asked for explicitly, so it is not held to those. */
+static void collapse_policy_forced(struct collapse_policy *p)
+{
+	p->max_ptes_none = HPAGE_PMD_NR;
+	p->max_ptes_swap = HPAGE_PMD_NR;
+	p->max_ptes_shared = HPAGE_PMD_NR;
+	p->strict_sub_pmd = false;
+	p->skip_lazyfree = false;
+	p->require_referenced = false;
+	p->install_pmd = true;
+	p->writeback_dirty = true;
+	p->gfp = GFP_TRANSHUGE;
+	p->tva_type = TVA_FORCED_COLLAPSE;
+}
+
 static int madvise_collapse_errno(enum scan_result r)
 {
 	/*
@@ -942,9 +898,10 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	struct collapse_control *cc;
 	struct mm_struct *mm = vma->vm_mm;
 	unsigned long hstart, hend, addr;
+	/* What the VMA allows; valid only while its lock is held */
+	unsigned long vma_orders;
 	enum scan_result last_fail = SCAN_FAIL;
 	int thps = 0;
-	bool mmap_unlocked = false;
 	int err;
 
 	BUG_ON(vma->vm_start > start);
@@ -971,28 +928,67 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	}
 
 	mmgrab(mm);
+
+	/*
+	 * Nothing below wants the lock the VMA walk left held, and
+	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
+	 * walk carries on under mmap_lock and its own caller is what drops it,
+	 * so reporting this only tells the walk that its VMA is now stale.
+	 */
+	mmap_read_unlock(mm);
+	*lock_dropped = true;
+	vma = NULL;
+	vma_orders = 0;
 	lru_add_drain_all();
 
 	for (addr = hstart; addr < hend; addr += HPAGE_PMD_SIZE) {
-		enum scan_result result = SCAN_FAIL;
+		enum scan_result result;
 
-		if (mmap_unlocked) {
+		/*
+		 * A collapse gives the lock up, and the VMA has to be found
+		 * again after one: it can shrink while nothing is held.  A scan
+		 * that finds nothing to collapse leaves the lock alone, so a
+		 * range that is already collapsed walks it without relocking.
+		 *
+		 * Reschedule only here, where nothing is held: a preemption
+		 * point under a lock is a writer waiting longer.
+		 */
+		if (!vma) {
 			cond_resched();
 			mmap_read_lock(mm);
-			mmap_unlocked = false;
-			*lock_dropped = true;
-			result = hugepage_vma_revalidate(mm, addr, false, &vma,
-							 cc, HPAGE_PMD_ORDER);
-			if (result != SCAN_SUCCEED) {
-				last_fail = result;
-				goto out_nolock;
+			vma = vma_lookup(mm, addr);
+			if (!vma) {
+				mmap_read_unlock(mm);
+				hend = addr;
+				break;
 			}
-
-			hend = min(hend, vma->vm_end & HPAGE_PMD_MASK);
+			vma_orders = collapse_possible_orders(vma,
+					vma->vm_flags, TVA_FORCED_COLLAPSE);
 		}
 
-		result = collapse_single_pmd(addr, addr + HPAGE_PMD_SIZE, vma,
-					     &mmap_unlocked, cc);
+		/* If nothing to collapse, the lock is still ours */
+		if (!collapse_scan_pmd(vma, addr, addr + HPAGE_PMD_SIZE, cc,
+				       vma_orders)) {
+			result = cc->scan_refusal;
+		} else {
+			/* collapse_run_pmd() takes its own locks, so give this up */
+			mmap_read_unlock(mm);
+			vma = NULL;
+			/* The mask belonged to that lock, not to this range */
+			vma_orders = 0;
+
+			result = collapse_run_pmd(mm, addr,
+						  addr + HPAGE_PMD_SIZE, cc);
+		}
+
+		/*
+		 * The VMA shrank under us, so the rest of the range was never
+		 * ours to collapse: stop, and expect only what came before.
+		 */
+		if (result == SCAN_VMA_NULL || result == SCAN_ADDRESS_RANGE) {
+			hend = addr;
+			break;
+		}
 
 		switch (result) {
 		case SCAN_SUCCEED:
@@ -1015,18 +1011,14 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 		default:
 			last_fail = result;
 			/* Other error, exit */
-			goto out_maybelock;
+			goto out;
 		}
 	}
 
-out_maybelock:
-	/* Caller expects us to hold mmap_lock on return */
-	if (mmap_unlocked) {
-		*lock_dropped = true;
+out:
+	/* The VMA walk this returns to expects the lock it was holding */
+	if (!vma)
 		mmap_read_lock(mm);
-	}
-out_nolock:
-	mmap_assert_locked(mm);
 	mmdrop(mm);
 	collapse_control_release(cc);
 	kfree(cc);
diff --git a/mm/mremap.c b/mm/mremap.c
index e8df5cdb0ac9..a25ac3db787a 100644
--- a/mm/mremap.c
+++ b/mm/mremap.c
@@ -244,7 +244,7 @@ static int move_ptes(struct pagetable_move_control *pmc,
 		goto out;
 	}
 	/*
-	 * Now new_pte is none, so collapse_scan_file() path can not find
+	 * Now new_pte is none, so collapse_pagecache_pmd() path can not find
 	 * this by traversing file->f_mapping, so there is no concurrency with
 	 * retract_page_tables(). In addition, we already hold the exclusive
 	 * mmap_lock, so this new_pte page is stable, so there is no need to get
-- 
2.54.0

[RFC PATCH 34/57] mm/collapse: implement MADV_COLLAPSE in madvise.c

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

MADV_COLLAPSE is a madvise operation, but its implementation sat in
khugepaged.c.  The daemon's file therefore also held a syscall's worth of
code that has nothing to do with the daemon: the walk over the user's
range, the per-PMD loop, and the errno translation that reports back
through madvise(2).

Move it to madvise.c, among the operations it belongs with, along with the
errno map and the policy it states for itself.  It takes a struct
madvise_behavior like every one of those operations, which is where the
range, the VMA and the lock-dropped flag it used to be handed separately
already live.

It stays a caller of the same interface khugepaged uses, so nothing about
the collapse changes.  The eligibility test reads
collapse_possible_orders() rather than collapse_possible(), a static
wrapper around it that madvise.c cannot reach.

With the declaration in huge_mm.h no longer needed, the
!CONFIG_TRANSPARENT_HUGEPAGE stub moves in with it.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/linux/huge_mm.h |   9 --
 mm/khugepaged.c         | 182 -------------------------------------
 mm/madvise.c            | 195 +++++++++++++++++++++++++++++++++++++++-
 3 files changed, 193 insertions(+), 193 deletions(-)

diff --git a/include/linux/huge_mm.h b/include/linux/huge_mm.h
index c745f7ad2298..8ca0fa3be2ac 100644
--- a/include/linux/huge_mm.h
+++ b/include/linux/huge_mm.h
@@ -510,8 +510,6 @@ change_huge_pud(struct mmu_gather *tlb, struct vm_area_struct *vma,
 
 int hugepage_madvise(struct vm_area_struct *vma, vm_flags_t *vm_flags,
 		     int advice);
-int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
-		     unsigned long end, bool *lock_dropped);
 void vma_adjust_trans_huge(struct vm_area_struct *vma, unsigned long start,
 			   unsigned long end, struct vm_area_struct *next);
 spinlock_t *__pmd_trans_huge_lock(pmd_t *pmd, struct vm_area_struct *vma);
@@ -715,13 +713,6 @@ static inline int hugepage_madvise(struct vm_area_struct *vma,
 	return -EINVAL;
 }
 
-static inline int madvise_collapse(struct vm_area_struct *vma,
-				   unsigned long start,
-				   unsigned long end, bool *lock_dropped)
-{
-	return -EINVAL;
-}
-
 static inline void vma_adjust_trans_huge(struct vm_area_struct *vma,
 					 unsigned long start,
 					 unsigned long end,
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 47c134cd4129..967cc472b6dc 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -844,185 +844,3 @@ bool current_is_khugepaged(void)
 {
 	return kthread_func(current) == khugepaged;
 }
-
-/* MADV_COLLAPSE was asked for explicitly, so it is not held to those. */
-static void collapse_policy_forced(struct collapse_policy *p)
-{
-	p->max_ptes_none = HPAGE_PMD_NR;
-	p->max_ptes_swap = HPAGE_PMD_NR;
-	p->max_ptes_shared = HPAGE_PMD_NR;
-	p->strict_sub_pmd = false;
-	p->skip_lazyfree = false;
-	p->require_referenced = false;
-	p->install_pmd = true;
-	p->writeback_dirty = true;
-	p->gfp = GFP_TRANSHUGE;
-	p->tva_type = TVA_FORCED_COLLAPSE;
-}
-
-static int madvise_collapse_errno(enum scan_result r)
-{
-	/*
-	 * MADV_COLLAPSE breaks from existing madvise(2) conventions to provide
-	 * actionable feedback to caller, so they may take an appropriate
-	 * fallback measure depending on the nature of the failure.
-	 */
-	switch (r) {
-	case SCAN_ALLOC_HUGE_PAGE_FAIL:
-		return -ENOMEM;
-	case SCAN_CGROUP_CHARGE_FAIL:
-	case SCAN_EXCEED_NONE_PTE:
-		return -EBUSY;
-	/* Resource temporary unavailable - trying again might succeed */
-	case SCAN_PAGE_COUNT:
-	case SCAN_PAGE_LOCK:
-	case SCAN_PAGE_LRU:
-	case SCAN_DEL_PAGE_LRU:
-	case SCAN_PAGE_FILLED:
-	case SCAN_PAGE_HAS_PRIVATE:
-	case SCAN_PAGE_DIRTY_OR_WRITEBACK:
-		return -EAGAIN;
-	/*
-	 * Other: Trying again likely not to succeed / error intrinsic to
-	 * specified memory range. khugepaged likely won't be able to collapse
-	 * either.
-	 */
-	default:
-		return -EINVAL;
-	}
-}
-
-int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
-		     unsigned long end, bool *lock_dropped)
-{
-	struct collapse_control *cc;
-	struct mm_struct *mm = vma->vm_mm;
-	unsigned long hstart, hend, addr;
-	/* What the VMA allows; valid only while its lock is held */
-	unsigned long vma_orders;
-	enum scan_result last_fail = SCAN_FAIL;
-	int thps = 0;
-	int err;
-
-	BUG_ON(vma->vm_start > start);
-	BUG_ON(vma->vm_end < end);
-
-	if (!collapse_possible(vma, vma->vm_flags, TVA_FORCED_COLLAPSE))
-		return -EINVAL;
-
-	hstart = ALIGN(start, HPAGE_PMD_SIZE);
-	hend = ALIGN_DOWN(end, HPAGE_PMD_SIZE);
-
-	if (hstart >= hend)
-		return 0;
-
-	cc = kmalloc_obj(*cc);
-	if (!cc)
-		return -ENOMEM;
-	collapse_policy_forced(&cc->policy);
-	cc->progress = 0;
-	err = collapse_control_init(cc);
-	if (err) {
-		kfree(cc);
-		return err;
-	}
-
-	mmgrab(mm);
-
-	/*
-	 * Nothing below wants the lock the VMA walk left held, and
-	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
-	 * walk carries on under mmap_lock and its own caller is what drops it,
-	 * so reporting this only tells the walk that its VMA is now stale.
-	 */
-	mmap_read_unlock(mm);
-	*lock_dropped = true;
-	vma = NULL;
-	vma_orders = 0;
-	lru_add_drain_all();
-
-	for (addr = hstart; addr < hend; addr += HPAGE_PMD_SIZE) {
-		enum scan_result result;
-
-		/*
-		 * A collapse gives the lock up, and the VMA has to be found
-		 * again after one: it can shrink while nothing is held.  A scan
-		 * that finds nothing to collapse leaves the lock alone, so a
-		 * range that is already collapsed walks it without relocking.
-		 *
-		 * Reschedule only here, where nothing is held: a preemption
-		 * point under a lock is a writer waiting longer.
-		 */
-		if (!vma) {
-			cond_resched();
-			mmap_read_lock(mm);
-			vma = vma_lookup(mm, addr);
-			if (!vma) {
-				mmap_read_unlock(mm);
-				hend = addr;
-				break;
-			}
-			vma_orders = collapse_possible_orders(vma,
-					vma->vm_flags, TVA_FORCED_COLLAPSE);
-		}
-
-		/* If nothing to collapse, the lock is still ours */
-		if (!collapse_scan_pmd(vma, addr, addr + HPAGE_PMD_SIZE, cc,
-				       vma_orders)) {
-			result = cc->scan_refusal;
-		} else {
-			/* collapse_run_pmd() takes its own locks, so give this up */
-			mmap_read_unlock(mm);
-			vma = NULL;
-			/* The mask belonged to that lock, not to this range */
-			vma_orders = 0;
-
-			result = collapse_run_pmd(mm, addr,
-						  addr + HPAGE_PMD_SIZE, cc);
-		}
-
-		/*
-		 * The VMA shrank under us, so the rest of the range was never
-		 * ours to collapse: stop, and expect only what came before.
-		 */
-		if (result == SCAN_VMA_NULL || result == SCAN_ADDRESS_RANGE) {
-			hend = addr;
-			break;
-		}
-
-		switch (result) {
-		case SCAN_SUCCEED:
-		case SCAN_PMD_MAPPED:
-			++thps;
-			break;
-		/* Whitelisted set of results where continuing OK */
-		case SCAN_NO_PTE_TABLE:
-		case SCAN_PTE_NON_PRESENT:
-		case SCAN_PTE_UFFD:
-		case SCAN_LACK_REFERENCED_PAGE:
-		case SCAN_PAGE_NULL:
-		case SCAN_PAGE_COUNT:
-		case SCAN_PAGE_LOCK:
-		case SCAN_PAGE_COMPOUND:
-		case SCAN_PAGE_LRU:
-		case SCAN_DEL_PAGE_LRU:
-			last_fail = result;
-			break;
-		default:
-			last_fail = result;
-			/* Other error, exit */
-			goto out;
-		}
-	}
-
-out:
-	/* The VMA walk this returns to expects the lock it was holding */
-	if (!vma)
-		mmap_read_lock(mm);
-	mmdrop(mm);
-	collapse_control_release(cc);
-	kfree(cc);
-
-	return thps == ((hend - hstart) >> HPAGE_PMD_SHIFT) ? 0
-			: madvise_collapse_errno(last_fail);
-}
diff --git a/mm/madvise.c b/mm/madvise.c
index c179938097bf..76ddf61f043f 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -894,6 +894,198 @@ bool madvise_dontneed_free_valid_vma(struct madvise_behavior *madv_behavior)
 	return true;
 }
 
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+#include "collapse.h"
+
+/* MADV_COLLAPSE was asked for explicitly, so it is not held to those. */
+static void collapse_policy_forced(struct collapse_policy *p)
+{
+	p->max_ptes_none = HPAGE_PMD_NR;
+	p->max_ptes_swap = HPAGE_PMD_NR;
+	p->max_ptes_shared = HPAGE_PMD_NR;
+	p->strict_sub_pmd = false;
+	p->skip_lazyfree = false;
+	p->require_referenced = false;
+	p->install_pmd = true;
+	p->writeback_dirty = true;
+	p->gfp = GFP_TRANSHUGE;
+	p->tva_type = TVA_FORCED_COLLAPSE;
+}
+
+static int madvise_collapse_errno(enum scan_result r)
+{
+	/*
+	 * MADV_COLLAPSE breaks from existing madvise(2) conventions to provide
+	 * actionable feedback to caller, so they may take an appropriate
+	 * fallback measure depending on the nature of the failure.
+	 */
+	switch (r) {
+	case SCAN_ALLOC_HUGE_PAGE_FAIL:
+		return -ENOMEM;
+	case SCAN_CGROUP_CHARGE_FAIL:
+	case SCAN_EXCEED_NONE_PTE:
+		return -EBUSY;
+	/* Resource temporary unavailable - trying again might succeed */
+	case SCAN_PAGE_COUNT:
+	case SCAN_PAGE_LOCK:
+	case SCAN_PAGE_LRU:
+	case SCAN_DEL_PAGE_LRU:
+	case SCAN_PAGE_FILLED:
+	case SCAN_PAGE_HAS_PRIVATE:
+	case SCAN_PAGE_DIRTY_OR_WRITEBACK:
+		return -EAGAIN;
+	/*
+	 * Other: Trying again likely not to succeed / error intrinsic to
+	 * specified memory range. khugepaged likely won't be able to collapse
+	 * either.
+	 */
+	default:
+		return -EINVAL;
+	}
+}
+
+static int madvise_collapse(struct madvise_behavior *madv_behavior)
+{
+	struct madvise_behavior_range *range = &madv_behavior->range;
+	struct vm_area_struct *vma = madv_behavior->vma;
+	struct mm_struct *mm = madv_behavior->mm;
+	unsigned long hstart, hend, addr;
+	struct collapse_control *cc;
+	unsigned long vma_orders;
+	enum scan_result last_fail = SCAN_FAIL;
+	int thps = 0;
+	int err;
+
+	BUG_ON(vma->vm_start > range->start);
+	BUG_ON(vma->vm_end < range->end);
+
+	if (!collapse_possible_orders(vma, vma->vm_flags, TVA_FORCED_COLLAPSE))
+		return -EINVAL;
+
+	hstart = ALIGN(range->start, HPAGE_PMD_SIZE);
+	hend = ALIGN_DOWN(range->end, HPAGE_PMD_SIZE);
+
+	if (hstart >= hend)
+		return 0;
+
+	cc = kmalloc_obj(*cc);
+	if (!cc)
+		return -ENOMEM;
+	collapse_policy_forced(&cc->policy);
+	cc->progress = 0;
+	err = collapse_control_init(cc);
+	if (err) {
+		kfree(cc);
+		return err;
+	}
+
+	mmgrab(mm);
+
+	/*
+	 * Nothing below wants the lock the VMA walk left held, and
+	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
+	 * walk carries on under mmap_lock and its own caller is what drops it,
+	 * so reporting this only tells the walk that its VMA is now stale.
+	 */
+	mmap_read_unlock(mm);
+	mark_mmap_lock_dropped(madv_behavior);
+	vma = NULL;
+	vma_orders = 0;
+	lru_add_drain_all();
+
+	for (addr = hstart; addr < hend; addr += HPAGE_PMD_SIZE) {
+		enum scan_result result;
+
+		/*
+		 * A collapse gives the lock up, and the VMA has to be found
+		 * again after one: it can shrink while nothing is held.  A scan
+		 * that finds nothing to collapse leaves the lock alone, so a
+		 * range that is already collapsed walks it without relocking.
+		 *
+		 * Reschedule only here, where nothing is held: a preemption
+		 * point under a lock is a writer waiting longer.
+		 */
+		if (!vma) {
+			cond_resched();
+			mmap_read_lock(mm);
+			vma = vma_lookup(mm, addr);
+			if (!vma) {
+				mmap_read_unlock(mm);
+				hend = addr;
+				break;
+			}
+			vma_orders = collapse_possible_orders(vma,
+					vma->vm_flags, TVA_FORCED_COLLAPSE);
+		}
+
+		/* If nothing to collapse, the lock is still ours */
+		if (!collapse_scan_pmd(vma, addr, addr + HPAGE_PMD_SIZE, cc,
+				       vma_orders)) {
+			result = cc->scan_refusal;
+		} else {
+			/* collapse_run_pmd() takes its own locks, so give this up */
+			mmap_read_unlock(mm);
+			vma = NULL;
+			/* The mask belonged to that lock, not to this range */
+			vma_orders = 0;
+
+			result = collapse_run_pmd(mm, addr,
+						  addr + HPAGE_PMD_SIZE, cc);
+		}
+
+		/*
+		 * The VMA shrank under us, so the rest of the range was never
+		 * ours to collapse: stop, and expect only what came before.
+		 */
+		if (result == SCAN_VMA_NULL || result == SCAN_ADDRESS_RANGE) {
+			hend = addr;
+			break;
+		}
+
+		switch (result) {
+		case SCAN_SUCCEED:
+		case SCAN_PMD_MAPPED:
+			++thps;
+			break;
+		/* Whitelisted set of results where continuing OK */
+		case SCAN_NO_PTE_TABLE:
+		case SCAN_PTE_NON_PRESENT:
+		case SCAN_PTE_UFFD:
+		case SCAN_LACK_REFERENCED_PAGE:
+		case SCAN_PAGE_NULL:
+		case SCAN_PAGE_COUNT:
+		case SCAN_PAGE_LOCK:
+		case SCAN_PAGE_COMPOUND:
+		case SCAN_PAGE_LRU:
+		case SCAN_DEL_PAGE_LRU:
+			last_fail = result;
+			break;
+		default:
+			last_fail = result;
+			/* Other error, exit */
+			goto out;
+		}
+	}
+
+out:
+	/* The VMA walk this returns to expects the lock it was holding */
+	if (!vma)
+		mmap_read_lock(mm);
+	mmdrop(mm);
+	collapse_control_release(cc);
+	kfree(cc);
+
+	return thps == ((hend - hstart) >> HPAGE_PMD_SHIFT) ? 0
+			: madvise_collapse_errno(last_fail);
+}
+
+#else
+static int madvise_collapse(struct madvise_behavior *madv_behavior)
+{
+	return -EINVAL;
+}
+#endif /* CONFIG_TRANSPARENT_HUGEPAGE */
+
 static long madvise_dontneed_free(struct madvise_behavior *madv_behavior)
 {
 	struct mm_struct *mm = madv_behavior->mm;
@@ -1361,8 +1553,7 @@ static int madvise_vma_behavior(struct madvise_behavior *madv_behavior)
 	case MADV_DONTNEED_LOCKED:
 		return madvise_dontneed_free(madv_behavior);
 	case MADV_COLLAPSE:
-		return madvise_collapse(vma, range->start, range->end,
-			&madv_behavior->lock_dropped);
+		return madvise_collapse(madv_behavior);
 	case MADV_GUARD_INSTALL:
 		return madvise_guard_install(madv_behavior);
 	case MADV_GUARD_REMOVE:
-- 
2.54.0

[RFC PATCH 35/57] mm/madvise: drop MADV_COLLAPSE's redundant mm reference

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

madvise_collapse() holds an mmgrab() reference across its work, which
nothing needs.  mmgrab() pins the mm_struct alone; every caller already
holds mm_users, which keeps the address space itself alive and so implies
it:

 - madvise(2) works on current->mm, which lives as long as the task is in
   the syscall;
 - process_madvise(2) reaches a remote mm through mm_access(), which takes
   an mm_users reference and holds it until the syscall returns;
 - io_uring passes current->mm;
 - DAMON takes one with get_task_mm() and drops it after the call.

Drop the mmgrab()/mmdrop() pair.  It has been there since
commit 7d8faaf15545 ("mm/madvise: introduce MADV_COLLAPSE sync hugepage collapse").

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/madvise.c | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/mm/madvise.c b/mm/madvise.c
index 76ddf61f043f..c1bb425be3f4 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -979,8 +979,6 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 		return err;
 	}
 
-	mmgrab(mm);
-
 	/*
 	 * Nothing below wants the lock the VMA walk left held, and
 	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
@@ -1071,7 +1069,6 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 	/* The VMA walk this returns to expects the lock it was holding */
 	if (!vma)
 		mmap_read_lock(mm);
-	mmdrop(mm);
 	collapse_control_release(cc);
 	kfree(cc);
 
-- 
2.54.0

[RFC PATCH 36/57] mm/collapse: report what the scan found

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The anonymous scan reports nothing, and its verdict decides everything
after it: which orders selection may still try, and whether the table is
refused outright.  The only way to see it was to infer it from what the
candidates did afterwards, or from their absence.

Add mm_collapse_scan: where the table was scanned, how many slots were
holes or the zeropage, how many were swapped out, which orders survived
the scan, and the verdict.

The first two counts are why a table yields a smaller window than
expected.  The third is what a PMD candidate will have to read back.  The
order mask is the scan's whole output to selection in one number: a table
refused as a unit shows an empty mask, which distinguishes it at a glance
from one that merely lost the PMD order.

The file scan already reported through mm_khugepaged_scan_file; this gives
the anonymous side the same.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 34 ++++++++++++++++++++++++++++++
 mm/collapse.c                      |  2 ++
 2 files changed, 36 insertions(+)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 86131845b761..573cf5428969 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -126,6 +126,40 @@ TRACE_EVENT(mm_collapse_huge_page,
 		__entry->order)
 );
 
+TRACE_EVENT(mm_collapse_scan,
+
+	TP_PROTO(struct mm_struct *mm, unsigned long addr, int none_or_zero,
+		 int unmapped, unsigned long orders, int result),
+
+	TP_ARGS(mm, addr, none_or_zero, unmapped, orders, result),
+
+	TP_STRUCT__entry(
+		__field(struct mm_struct *, mm)
+		__field(unsigned long, addr)
+		__field(int, none_or_zero)
+		__field(int, unmapped)
+		__field(unsigned long, orders)
+		__field(int, result)
+	),
+
+	TP_fast_assign(
+		__entry->mm = mm;
+		__entry->addr = addr;
+		__entry->none_or_zero = none_or_zero;
+		__entry->unmapped = unmapped;
+		__entry->orders = orders;
+		__entry->result = result;
+	),
+
+	TP_printk("mm=%p, addr=0x%lx, none_or_zero=%d, unmapped=%d, orders=0x%lx, result=%s",
+		__entry->mm,
+		__entry->addr,
+		__entry->none_or_zero,
+		__entry->unmapped,
+		__entry->orders,
+		__print_symbolic(__entry->result, SCAN_STATUS))
+);
+
 TRACE_EVENT(mm_collapse_candidate,
 
 	TP_PROTO(struct mm_struct *mm, unsigned long addr, unsigned int order,
diff --git a/mm/collapse.c b/mm/collapse.c
index f0d80204c2bd..b750a1fc81a5 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -2097,6 +2097,8 @@ static enum scan_result collapse_scan_table(struct vm_area_struct *vma,
 		cc->select_orders &= ~BIT(HPAGE_PMD_ORDER);
 
 	cc->scan_unmapped = unmapped;
+	trace_mm_collapse_scan(vma->vm_mm, start, none_or_zero, unmapped,
+			       cc->select_orders, result);
 	return result;
 }
 
-- 
2.54.0

[RFC PATCH 37/57] mm/collapse: report what the fault-in pass paid

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The fault-in pass is the one place a collapse does work on someone else's
behalf: a swap read, or a CoW break, for every slot that needs one.  How
much of that a round pays is invisible, and it is the first thing to look
at when collapses are slow, or when a workload notices khugepaged at all.

Add mm_collapse_faultin: the faults taken across the round, with the
outcome.  A round that collapses a full table without faulting anything
and one that reads sixty-four pages back from swap are otherwise
indistinguishable.

The mm is captured before the walk, because the pass returns with
mmap_lock dropped on failure and the VMA is then unsafe to touch at the
report.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 24 ++++++++++++++++++++++++
 mm/collapse.c                      | 19 ++++++++++++++-----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 573cf5428969..c2314e26111c 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -160,6 +160,30 @@ TRACE_EVENT(mm_collapse_scan,
 		__print_symbolic(__entry->result, SCAN_STATUS))
 );
 
+TRACE_EVENT(mm_collapse_faultin,
+
+	TP_PROTO(struct mm_struct *mm, unsigned int nr_faults, int result),
+
+	TP_ARGS(mm, nr_faults, result),
+
+	TP_STRUCT__entry(
+		__field(struct mm_struct *, mm)
+		__field(unsigned int, nr_faults)
+		__field(int, result)
+	),
+
+	TP_fast_assign(
+		__entry->mm = mm;
+		__entry->nr_faults = nr_faults;
+		__entry->result = result;
+	),
+
+	TP_printk("mm=%p, nr_faults=%u, result=%s",
+		__entry->mm,
+		__entry->nr_faults,
+		__print_symbolic(__entry->result, SCAN_STATUS))
+);
+
 TRACE_EVENT(mm_collapse_candidate,
 
 	TP_PROTO(struct mm_struct *mm, unsigned long addr, unsigned int order,
diff --git a/mm/collapse.c b/mm/collapse.c
index b750a1fc81a5..1b5db42b6991 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -509,8 +509,10 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 
 /*
  * Bring one address to a state the freeze will accept: present, and exclusive if
- * it is anonymous.  Returns with mmap_lock dropped on every failure, because the
- * fault path may drop it and the caller cannot tell which case it is in.
+ * it is anonymous.  Every fault it takes to get there counts in *nr_faults, each
+ * one an allocation or a read the round is paying for.  Returns with mmap_lock
+ * dropped on every failure, because the fault path may drop it and the caller
+ * cannot tell which case it is in.
  *
  * SCAN_EXCEED_SWAP_PTE is the exception: it is a verdict on this candidate
  * rather than on the round, nothing was faulted to reach it, and it keeps the
@@ -518,7 +520,8 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
  */
 static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 					      struct collapse_candidate *cand,
-					      pmd_t *pmd, unsigned long addr)
+					      pmd_t *pmd, unsigned long addr,
+					      unsigned int *nr_faults)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	const unsigned int flags = FAULT_FLAG_ALLOW_RETRY | FAULT_FLAG_UNSHARE |
@@ -571,6 +574,7 @@ static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 
 		/* Only swap or shared PTEs reach here; the rest broke out */
 		ret = handle_mm_fault(vma, addr, flags, NULL);
+		(*nr_faults)++;
 		/*
 		 * Not a verdict on this window: the fault dropped the lock to
 		 * wait, which is what a swap-in normally does.  Distinct from
@@ -600,7 +604,9 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 					 struct collapse_control *cc,
 					 pmd_t *pmd)
 {
+	struct mm_struct *mm = vma->vm_mm;
 	enum scan_result result = SCAN_SUCCEED;
+	unsigned int nr_faults = 0;
 	unsigned int i;
 
 	for (i = 0; i < cc->nr_candidates; i++) {
@@ -616,7 +622,8 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 		     j++, addr += PAGE_SIZE) {
 			enum scan_result r;
 
-			r = collapse_faultin_addr(vma, cand, pmd, addr);
+			r = collapse_faultin_addr(vma, cand, pmd, addr,
+						  &nr_faults);
 			/*
 			 * The one failure that judges this candidate rather
 			 * than the round, and so the one that leaves the lock
@@ -627,7 +634,7 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 			if (r == SCAN_EXCEED_SWAP_PTE) {
 				cand->state = CAND_SKIPPED;
 				cand->result = r;
-				collapse_trace_candidate(vma->vm_mm, cand,
+				collapse_trace_candidate(mm, cand,
 							 COLLAPSE_PASS_FAULTIN);
 				break;
 			}
@@ -638,6 +645,8 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 		}
 	}
 out:
+	/* @vma is unsafe on the failure path: the callee dropped mmap_lock */
+	trace_mm_collapse_faultin(mm, nr_faults, result);
 	return result;
 }
 
-- 
2.54.0

[RFC PATCH 38/57] mm/collapse: report the round, and what it made faulters wait

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A round is the unit the engine actually works in, and nothing reports one.
The per-candidate events say which windows were taken and which were
refused.  They do not say how large the batch was, how much of it landed,
or the number that matters most for whether batching was the right idea:
how long a faulter on a source is held up.

That wait has a definite span.  A thread touching a source sleeps on the
folio lock the freeze took, and wakes when the putback drops it.  So the
interval from the first freeze to the last putback is what the round costs
anyone unlucky enough to touch it.

Add mm_collapse_round: that interval in microseconds, with the candidates
collected, the ones installed, and the outcome.  It is the thing to watch
if a larger batch is ever proposed.

collapse_finish() returns the number of candidates installed, and this is
its first consumer.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 31 ++++++++++++++++++++++++++++++
 mm/collapse.c                      | 17 +++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index c2314e26111c..d7c0195ace92 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -160,6 +160,37 @@ TRACE_EVENT(mm_collapse_scan,
 		__print_symbolic(__entry->result, SCAN_STATUS))
 );
 
+TRACE_EVENT(mm_collapse_round,
+
+	TP_PROTO(struct mm_struct *mm, unsigned int nr_candidates,
+		 unsigned int nr_installed, int result, u64 freeze_to_wake_us),
+
+	TP_ARGS(mm, nr_candidates, nr_installed, result, freeze_to_wake_us),
+
+	TP_STRUCT__entry(
+		__field(struct mm_struct *, mm)
+		__field(unsigned int, nr_candidates)
+		__field(unsigned int, nr_installed)
+		__field(int, result)
+		__field(u64, freeze_to_wake_us)
+	),
+
+	TP_fast_assign(
+		__entry->mm = mm;
+		__entry->nr_candidates = nr_candidates;
+		__entry->nr_installed = nr_installed;
+		__entry->result = result;
+		__entry->freeze_to_wake_us = freeze_to_wake_us;
+	),
+
+	TP_printk("mm=%p, nr_candidates=%u, nr_installed=%u, result=%s, freeze_to_wake_us=%llu",
+		__entry->mm,
+		__entry->nr_candidates,
+		__entry->nr_installed,
+		__print_symbolic(__entry->result, SCAN_STATUS),
+		__entry->freeze_to_wake_us)
+);
+
 TRACE_EVENT(mm_collapse_faultin,
 
 	TP_PROTO(struct mm_struct *mm, unsigned int nr_faults, int result),
diff --git a/mm/collapse.c b/mm/collapse.c
index 1b5db42b6991..d0d28e8dfcea 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -8,6 +8,7 @@
 #include <linux/huge_mm.h>
 #include <linux/hugetlb.h>	/* x86 flush_tlb_range() uses hstate_vma() */
 #include <linux/leafops.h>
+#include <linux/math64.h>
 #include <linux/memcontrol.h>
 #include <linux/mm.h>
 #include <linux/mmu_notifier.h>
@@ -20,6 +21,7 @@
 #include <linux/sizes.h>
 #include <linux/slab.h>
 #include <linux/swap.h>
+#include <linux/timekeeping.h>
 #include <linux/userfaultfd_k.h>
 #include <linux/vmstat.h>
 
@@ -1803,6 +1805,8 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 	struct mmu_notifier_range range;
 	struct vm_area_struct *vma;
 	enum scan_result result;
+	unsigned int nr_installed;
+	u64 latency = 0;
 	pmd_t *pmd;
 
 	collapse_reserve(mm, cc);
@@ -1838,6 +1842,13 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 				cc->batch_start, cc->batch_end);
 	mmu_notifier_invalidate_range_start(&range);
 
+	/*
+	 * What the faulters on this batch's sources are made to wait: they sleep
+	 * from the freeze that took their folio's lock to the putback that drops
+	 * it.  Measured per round rather than argued about.
+	 */
+	latency = ktime_get_ns();
+
 	/*
 	 * None of these can fail as a whole: the freeze takes the sources it
 	 * can and drops the candidates it cannot, and each pass after it works
@@ -1849,12 +1860,16 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 	collapse_install(vma, cc, pmd);
 	collapse_putback(vma, cc);
 
+	latency = ktime_get_ns() - latency;
+
 	mmu_notifier_invalidate_range_end(&range);
 
 out_unlock:
 	mmap_read_unlock(mm);
 out:
-	collapse_finish(mm, cc, result);
+	nr_installed = collapse_finish(mm, cc, result);
+	trace_mm_collapse_round(mm, cc->nr_candidates, nr_installed, result,
+				div_u64(latency, NSEC_PER_USEC));
 }
 
 /*
-- 
2.54.0

[RFC PATCH 39/57] mm/collapse: name the file collapse's tracepoints after collapse

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The two file events are named for khugepaged, from when khugepaged was the
only thing that collapsed and the code lived in its file.  Neither is true
now: MADV_COLLAPSE reaches them through the same entry point, and they are
emitted from collapse.c beside the events that do use the collapse name.

Rename mm_khugepaged_scan_file to mm_collapse_scan_file, and
mm_khugepaged_collapse_file to mm_collapse_file.  Every event a collapse
emits is then under one prefix: the anonymous and file scans, the
per-candidate verdicts, the fault-in, the round, and the file collapse
itself.

mm_khugepaged_scan keeps its name.  That one is the daemon reporting a
scan pass of its own, from khugepaged.c, and it is not something a
collapse emits.

Renaming a tracepoint breaks anything watching the old name.  In tree that
is raw_tp_null_args[], which tells the BPF verifier that the folio
argument of both may be NULL.  Without a matching entry the verifier would
let a program dereference it unchecked.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 4 ++--
 kernel/bpf/btf.c                   | 4 ++--
 mm/collapse.c                      | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index d7c0195ace92..5d0891e03bb0 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -308,7 +308,7 @@ TRACE_EVENT(mm_collapse_huge_page_swapin,
 		__entry->order)
 );
 
-TRACE_EVENT(mm_khugepaged_scan_file,
+TRACE_EVENT(mm_collapse_scan_file,
 
 	TP_PROTO(struct mm_struct *mm, struct folio *folio, struct file *file,
 		 int present, int swap, int result),
@@ -342,7 +342,7 @@ TRACE_EVENT(mm_khugepaged_scan_file,
 		__print_symbolic(__entry->result, SCAN_STATUS))
 );
 
-TRACE_EVENT(mm_khugepaged_collapse_file,
+TRACE_EVENT(mm_collapse_file,
 	TP_PROTO(struct mm_struct *mm, struct folio *new_folio, pgoff_t index,
 			unsigned long addr, bool is_shmem, struct file *file,
 			int nr, int result),
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index c4673a54c4ba..58f78274d989 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -6705,8 +6705,8 @@ static const struct bpf_raw_tp_null_args raw_tp_null_args[] = {
 	/* huge_memory */
 	{ "mm_khugepaged_scan_pmd", 0x10 },
 	{ "mm_collapse_huge_page_isolate", 0x1 },
-	{ "mm_khugepaged_scan_file", 0x10 },
-	{ "mm_khugepaged_collapse_file", 0x10 },
+	{ "mm_collapse_scan_file", 0x10 },
+	{ "mm_collapse_file", 0x10 },
 	/* kmem */
 	{ "mm_page_alloc", 0x1 },
 	{ "mm_page_pcpu_drain", 0x1 },
diff --git a/mm/collapse.c b/mm/collapse.c
index d0d28e8dfcea..6c17c83a4e21 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -3519,7 +3519,7 @@ static enum scan_result collapse_file(struct mm_struct *mm, unsigned long addr,
 	folio_put(new_folio);
 out:
 	VM_BUG_ON(!list_empty(&pagelist));
-	trace_mm_khugepaged_collapse_file(mm, new_folio, index, addr, is_shmem, file, HPAGE_PMD_NR, result);
+	trace_mm_collapse_file(mm, new_folio, index, addr, is_shmem, file, HPAGE_PMD_NR, result);
 	return result;
 }
 
@@ -3625,7 +3625,7 @@ static enum scan_result collapse_pagecache_pmd(struct mm_struct *mm,
 		}
 	}
 
-	trace_mm_khugepaged_scan_file(mm, folio, file, present, swap, result);
+	trace_mm_collapse_scan_file(mm, folio, file, present, swap, result);
 	return result;
 }
 
-- 
2.54.0

[RFC PATCH 40/57] mm/collapse: remove the tracepoints of the mechanism that is gone

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Four tracepoints lost their only emitter when the anonymous mechanism was
deleted: mm_khugepaged_scan_pmd, mm_collapse_huge_page,
mm_collapse_huge_page_isolate and mm_collapse_huge_page_swapin.  Enabling
one now does nothing.

Remove them, and the two entries they have in raw_tp_null_args[], which
can never match again.  A tool that still asks for one fails to attach
rather than sitting on an event that never fires.

The engine's own events cover the same ground: mm_collapse_scan,
mm_collapse_candidate, mm_collapse_faultin and mm_collapse_round.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/huge_memory.h | 123 -----------------------------
 kernel/bpf/btf.c                   |   2 -
 2 files changed, 125 deletions(-)

diff --git a/include/trace/events/huge_memory.h b/include/trace/events/huge_memory.h
index 5d0891e03bb0..6aabf4235648 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/huge_memory.h
@@ -65,67 +65,6 @@ COLLAPSE_PASS_STATUS
 #define EM(a, b)	{a, b},
 #define EMe(a, b)	{a, b}
 
-TRACE_EVENT(mm_khugepaged_scan_pmd,
-
-	TP_PROTO(struct mm_struct *mm, struct folio *folio,
-		 int referenced, int none_or_zero, int status, int unmapped),
-
-	TP_ARGS(mm, folio, referenced, none_or_zero, status, unmapped),
-
-	TP_STRUCT__entry(
-		__field(struct mm_struct *, mm)
-		__field(unsigned long, pfn)
-		__field(int, referenced)
-		__field(int, none_or_zero)
-		__field(int, status)
-		__field(int, unmapped)
-	),
-
-	TP_fast_assign(
-		__entry->mm = mm;
-		__entry->pfn = folio ? folio_pfn(folio) : -1;
-		__entry->referenced = referenced;
-		__entry->none_or_zero = none_or_zero;
-		__entry->status = status;
-		__entry->unmapped = unmapped;
-	),
-
-	TP_printk("mm=%p, scan_pfn=0x%lx, referenced=%d, none_or_zero=%d, status=%s, unmapped=%d",
-		__entry->mm,
-		__entry->pfn,
-		__entry->referenced,
-		__entry->none_or_zero,
-		__print_symbolic(__entry->status, SCAN_STATUS),
-		__entry->unmapped)
-);
-
-TRACE_EVENT(mm_collapse_huge_page,
-
-	TP_PROTO(struct mm_struct *mm, int isolated, int status, unsigned int order),
-
-	TP_ARGS(mm, isolated, status, order),
-
-	TP_STRUCT__entry(
-		__field(struct mm_struct *, mm)
-		__field(int, isolated)
-		__field(int, status)
-		__field(unsigned int, order)
-	),
-
-	TP_fast_assign(
-		__entry->mm = mm;
-		__entry->isolated = isolated;
-		__entry->status = status;
-		__entry->order = order;
-	),
-
-	TP_printk("mm=%p, isolated=%d, status=%s, order=%u",
-		__entry->mm,
-		__entry->isolated,
-		__print_symbolic(__entry->status, SCAN_STATUS),
-		__entry->order)
-);
-
 TRACE_EVENT(mm_collapse_scan,
 
 	TP_PROTO(struct mm_struct *mm, unsigned long addr, int none_or_zero,
@@ -246,68 +185,6 @@ TRACE_EVENT(mm_collapse_candidate,
 		__print_symbolic(__entry->result, SCAN_STATUS))
 );
 
-TRACE_EVENT(mm_collapse_huge_page_isolate,
-
-	TP_PROTO(struct folio *folio, int none_or_zero,
-		 int referenced, int status, unsigned int order),
-
-	TP_ARGS(folio, none_or_zero, referenced, status, order),
-
-	TP_STRUCT__entry(
-		__field(unsigned long, pfn)
-		__field(int, none_or_zero)
-		__field(int, referenced)
-		__field(int, status)
-		__field(unsigned int, order)
-	),
-
-	TP_fast_assign(
-		__entry->pfn = folio ? folio_pfn(folio) : -1;
-		__entry->none_or_zero = none_or_zero;
-		__entry->referenced = referenced;
-		__entry->status = status;
-		__entry->order = order;
-	),
-
-	TP_printk("scan_pfn=0x%lx, none_or_zero=%d, referenced=%d, status=%s, order=%u",
-		__entry->pfn,
-		__entry->none_or_zero,
-		__entry->referenced,
-		__print_symbolic(__entry->status, SCAN_STATUS),
-		__entry->order)
-);
-
-TRACE_EVENT(mm_collapse_huge_page_swapin,
-
-	TP_PROTO(struct mm_struct *mm, int swapped_in, int referenced, int ret,
-		 unsigned int order),
-
-	TP_ARGS(mm, swapped_in, referenced, ret, order),
-
-	TP_STRUCT__entry(
-		__field(struct mm_struct *, mm)
-		__field(int, swapped_in)
-		__field(int, referenced)
-		__field(int, ret)
-		__field(unsigned int, order)
-	),
-
-	TP_fast_assign(
-		__entry->mm = mm;
-		__entry->swapped_in = swapped_in;
-		__entry->referenced = referenced;
-		__entry->ret = ret;
-		__entry->order = order;
-	),
-
-	TP_printk("mm=%p, swapped_in=%d, referenced=%d, ret=%d, order=%u",
-		__entry->mm,
-		__entry->swapped_in,
-		__entry->referenced,
-		__entry->ret,
-		__entry->order)
-);
-
 TRACE_EVENT(mm_collapse_scan_file,
 
 	TP_PROTO(struct mm_struct *mm, struct folio *folio, struct file *file,
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index 58f78274d989..22fc8f974be2 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -6703,8 +6703,6 @@ static const struct bpf_raw_tp_null_args raw_tp_null_args[] = {
 	/* host1x */
 	{ "host1x_cdma_push_gather", 0x10000 },
 	/* huge_memory */
-	{ "mm_khugepaged_scan_pmd", 0x10 },
-	{ "mm_collapse_huge_page_isolate", 0x1 },
 	{ "mm_collapse_scan_file", 0x10 },
 	{ "mm_collapse_file", 0x10 },
 	/* kmem */
-- 
2.54.0

[RFC PATCH 41/57] mm/collapse: give collapse its own trace header

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Every event in include/trace/events/huge_memory.h is a collapse event, and
all but one is emitted from mm/collapse.c.  The header is named for
huge_memory.c, which emits none of them, and the tracepoints are built by
khugepaged.c, which emits one.

Rename it to trace/events/collapse.h, with TRACE_SYSTEM to match, and
build the tracepoints in collapse.c.

Two things follow.  The pass names a candidate event prints are needed
only where the tracepoints are built, so enum collapse_pass moves out of
mm/collapse.h into collapse.c.  And khugepaged.c becomes an ordinary
includer, for the one event it does emit.

The tracefs directory moves with the trace system: events/huge_memory
becomes events/collapse, so anything enabling those events by system name
has to follow.

khugepaged's own selftest does that, and is updated here.  So are the two
other in-tree references to the old name: the MAINTAINERS entry for the
header, and the group raw_tp_null_args[] keeps its collapse entries under.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 MAINTAINERS                                     |  2 +-
 .../trace/events/{huge_memory.h => collapse.h}  |  8 ++++----
 kernel/bpf/btf.c                                |  6 +++---
 mm/collapse.c                                   | 17 ++++++++++++++++-
 mm/collapse.h                                   | 13 -------------
 mm/khugepaged.c                                 |  3 +--
 .../selftests/mm/khugepaged_sync_check.c        | 10 +++++-----
 tools/testing/selftests/mm/vm_util.c            |  2 +-
 8 files changed, 31 insertions(+), 30 deletions(-)
 rename include/trace/events/{huge_memory.h => collapse.h} (98%)

diff --git a/MAINTAINERS b/MAINTAINERS
index 0f513b42bc18..7c179b333e4e 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -17273,7 +17273,7 @@ F:	Documentation/ABI/testing/sysfs-kernel-mm-transparent-hugepage
 F:	Documentation/admin-guide/mm/transhuge.rst
 F:	include/linux/huge_mm.h
 F:	include/linux/khugepaged.h
-F:	include/trace/events/huge_memory.h
+F:	include/trace/events/collapse.h
 F:	mm/huge_memory.c
 F:	mm/khugepaged.c
 F:	mm/mm_slot.h
diff --git a/include/trace/events/huge_memory.h b/include/trace/events/collapse.h
similarity index 98%
rename from include/trace/events/huge_memory.h
rename to include/trace/events/collapse.h
index 6aabf4235648..a3af5d8cc9aa 100644
--- a/include/trace/events/huge_memory.h
+++ b/include/trace/events/collapse.h
@@ -1,9 +1,9 @@
 /* SPDX-License-Identifier: GPL-2.0 */
 #undef TRACE_SYSTEM
-#define TRACE_SYSTEM huge_memory
+#define TRACE_SYSTEM collapse
 
-#if !defined(__HUGE_MEMORY_H) || defined(TRACE_HEADER_MULTI_READ)
-#define __HUGE_MEMORY_H
+#if !defined(_TRACE_COLLAPSE_H) || defined(TRACE_HEADER_MULTI_READ)
+#define _TRACE_COLLAPSE_H
 
 #include  <linux/tracepoint.h>
 
@@ -282,5 +282,5 @@ TRACE_EVENT(mm_khugepaged_scan,
 		__entry->full_scan_finished)
 );
 
-#endif /* __HUGE_MEMORY_H */
+#endif /* _TRACE_COLLAPSE_H */
 #include <trace/define_trace.h>
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index 22fc8f974be2..505bdbb534e4 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -6683,6 +6683,9 @@ static const struct bpf_raw_tp_null_args raw_tp_null_args[] = {
 	{ "cachefiles_ondemand_cread", 0x1 },
 	{ "cachefiles_ondemand_fd_write", 0x1 },
 	{ "cachefiles_ondemand_fd_release", 0x1 },
+	/* collapse */
+	{ "mm_collapse_scan_file", 0x10 },
+	{ "mm_collapse_file", 0x10 },
 	/* ext4, from ext4__mballoc event class */
 	{ "ext4_mballoc_discard", 0x10 },
 	{ "ext4_mballoc_free", 0x10 },
@@ -6702,9 +6705,6 @@ static const struct bpf_raw_tp_null_args raw_tp_null_args[] = {
 	{ "time_out_leases", 0x10 },
 	/* host1x */
 	{ "host1x_cdma_push_gather", 0x10000 },
-	/* huge_memory */
-	{ "mm_collapse_scan_file", 0x10 },
-	{ "mm_collapse_file", 0x10 },
 	/* kmem */
 	{ "mm_page_alloc", 0x1 },
 	{ "mm_page_pcpu_drain", 0x1 },
diff --git a/mm/collapse.c b/mm/collapse.c
index 6c17c83a4e21..952e3f62d920 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -26,10 +26,25 @@
 #include <linux/vmstat.h>
 
 #include <asm/tlb.h>
-#include <trace/events/huge_memory.h>
 #include "collapse.h"
 #include "internal.h"
 
+/*
+ * Which pass of a round reached a verdict on a candidate.  Named by the trace
+ * header, which only this file builds, so it need not be shared.
+ */
+enum collapse_pass {
+	COLLAPSE_PASS_ALLOC,
+	COLLAPSE_PASS_REVALIDATE,
+	COLLAPSE_PASS_FAULTIN,
+	COLLAPSE_PASS_FREEZE,
+	COLLAPSE_PASS_COPY,
+	COLLAPSE_PASS_INSTALL,
+};
+
+#define CREATE_TRACE_POINTS
+#include <trace/events/collapse.h>
+
 /*
  * Anonymous collapse, in rounds.
  *
diff --git a/mm/collapse.h b/mm/collapse.h
index 4af7bb9c4261..9e2cec1f250b 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -16,19 +16,6 @@
 struct collapse_candidate;
 struct collapse_retry;
 
-/*
- * Which pass of a round reached a verdict on a candidate.  Only collapse.c
- * produces these; the trace header khugepaged.c builds names them.
- */
-enum collapse_pass {
-	COLLAPSE_PASS_ALLOC,
-	COLLAPSE_PASS_REVALIDATE,
-	COLLAPSE_PASS_FAULTIN,
-	COLLAPSE_PASS_FREEZE,
-	COLLAPSE_PASS_COPY,
-	COLLAPSE_PASS_INSTALL,
-};
-
 enum scan_result {
 	SCAN_FAIL,
 	SCAN_SUCCEED,
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 967cc472b6dc..f3a7aad5e8f2 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -31,8 +31,7 @@
 #include "page_alloc.h"
 #include "mm_slot.h"
 
-#define CREATE_TRACE_POINTS
-#include <trace/events/huge_memory.h>
+#include <trace/events/collapse.h>
 
 static struct task_struct *khugepaged_thread __read_mostly;
 static DEFINE_MUTEX(khugepaged_mutex);
diff --git a/tools/testing/selftests/mm/khugepaged_sync_check.c b/tools/testing/selftests/mm/khugepaged_sync_check.c
index 45001996b57a..4c37b697d3dd 100644
--- a/tools/testing/selftests/mm/khugepaged_sync_check.c
+++ b/tools/testing/selftests/mm/khugepaged_sync_check.c
@@ -41,7 +41,7 @@ static unsigned long hpage_pmd_size;
 /*
  * Each step switches the events off again, but a helper can still give up
  * on us in between (a failing sysfs write ends the test from inside
- * thp_write_num()), and huge_memory events left on are the whole machine's
+ * thp_write_num()), and collapse events left on are the whole machine's
  * problem, not this test's.
  */
 static void trace_events_off(void)
@@ -118,7 +118,7 @@ static void one_step(int iteration)
 	if (tracing_clear_trace())
 		ksft_exit_fail_msg("Cannot clear the trace buffer\n");
 	if (tracing_events_enable(trace_events_fd, true))
-		ksft_exit_fail_msg("Cannot enable huge_memory events\n");
+		ksft_exit_fail_msg("Cannot enable collapse events\n");
 
 	if (madvise(p, hpage_pmd_size, MADV_HUGEPAGE))
 		ksft_exit_fail_perror("madvise(MADV_HUGEPAGE)");
@@ -127,7 +127,7 @@ static void one_step(int iteration)
 
 	/* Off before anything that can give up: the events are system-wide. */
 	if (tracing_events_enable(trace_events_fd, false))
-		ksft_exit_fail_msg("Cannot disable huge_memory events\n");
+		ksft_exit_fail_msg("Cannot disable collapse events\n");
 	if (!passed)
 		ksft_exit_fail_msg("khugepaged did not complete a full pass\n");
 
@@ -164,9 +164,9 @@ int main(void)
 	kpageflags_fd = open("/proc/kpageflags", O_RDONLY);
 	if (kpageflags_fd < 0)
 		ksft_exit_skip("open(\"/proc/kpageflags\") requires root\n");
-	trace_events_fd = tracing_events_open("huge_memory");
+	trace_events_fd = tracing_events_open("collapse");
 	if (trace_events_fd < 0)
-		ksft_exit_skip("huge_memory events require tracefs and root\n");
+		ksft_exit_skip("collapse events require tracefs and root\n");
 	atexit(trace_events_off);
 
 	ksft_set_plan(NR_ITERATIONS);
diff --git a/tools/testing/selftests/mm/vm_util.c b/tools/testing/selftests/mm/vm_util.c
index ee1334778391..32cba59ae49c 100644
--- a/tools/testing/selftests/mm/vm_util.c
+++ b/tools/testing/selftests/mm/vm_util.c
@@ -601,7 +601,7 @@ bool is_range_backed_by_folio_orders(char *start, size_t len, int order,
 #define TRACEFS_ROOT "/sys/kernel/tracing"
 
 /*
- * Open the enable file of one ftrace event subsystem (e.g. "huge_memory").
+ * Open the enable file of one ftrace event subsystem (e.g. "collapse").
  * Returns a descriptor for tracing_events_enable(), or -1 if tracefs or the
  * subsystem is not there.  The events are system-wide state: whoever
  * switches them on owns them until it switches them off, including on the
-- 
2.54.0

[RFC PATCH 42/57] mm/collapse: allow error injection into the freeze

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The freeze refuses a candidate for reasons a test cannot arrange: a
co-mapper appearing between the two sweeps, a GUP-fast pin, writeback.
What the round does then -- the candidate drops out, its neighbours carry
on, the region goes back to selection at a lower order -- is worth
exercising on purpose.

Tag collapse_freeze_candidate(), noinline so fail_function can find the
symbol.  ERRNO is the only injection type for a function that returns a
value rather than a pointer or a bool, so an injected result is a negative
errno.  It only ever reaches cand->result, where every caller compares it
against SCAN_SUCCEED and finds it unequal, which is the refusal the test
wants.  Tracing prints it as a number, having no symbol for it.

A forced return happens at function entry, so this reaches the refusal,
not the freeze's own unwind path.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 952e3f62d920..76a53616b240 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -4,6 +4,7 @@
 #include <linux/backing-dev.h>
 #include <linux/bitops.h>
 #include <linux/dax.h>
+#include <linux/error-injection.h>
 #include <linux/highmem.h>
 #include <linux/huge_mm.h>
 #include <linux/hugetlb.h>	/* x86 flush_tlb_range() uses hstate_vma() */
@@ -937,7 +938,7 @@ static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
  * collapse_freeze() issues the ranged TLB flush over everything that froze
  * before dropping the ptl.  No copy may run before it completes.
  */
-static enum scan_result collapse_freeze_candidate(struct mm_struct *mm,
+static noinline enum scan_result collapse_freeze_candidate(struct mm_struct *mm,
 		struct collapse_candidate *cand, pte_t *pte)
 {
 	const unsigned int nr_pages = candidate_nr_pages(cand);
@@ -1085,6 +1086,7 @@ static enum scan_result collapse_freeze_candidate(struct mm_struct *mm,
 	collapse_unfreeze_candidate(mm, cand, pte, nr_saved, nr_frozen);
 	return result;
 }
+ALLOW_ERROR_INJECTION(collapse_freeze_candidate, ERRNO);
 
 /*
  * Raise the two barriers on the sources of every candidate: migration entries in
-- 
2.54.0

[RFC PATCH 43/57] mm/khugepaged: check the scan budget before the work, not after

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

pages_to_scan is meant to bound what one khugepaged pass does, but
collapse_scan_mm_slot() tested it in only one place: after a table had
been scanned and turned out to hold nothing.

Neither of the other two ways of spending the budget reached that test.  A
VMA the pass skips is charged for and walked past without asking -- one no
order can be collapsed at, or one the cursor is already past the end of.
A table that does hold a candidate leaves through the collapse.

So a pass over an address space of thousands of VMAs khugepaged cannot use
walks every one of them, however low pages_to_scan is set.

Ask at the top of both loops instead, where the other reasons to stop a
pass are already asked.  The outer loop asks before it judges a VMA, the
inner one before it scans a table.

Stopping the outer loop only works if the cursor moves, and it did not for
a skipped VMA.  Advance khugepaged_scan.address past one, so a pass that
runs out of budget resumes after the VMAs it has already judged.  Without
that, an address space with more skippable VMAs than the budget would be
walked from the same place every pass and never scanned at all.

pages_to_scan now bounds a pass that finds nothing to collapse, where
before it did not.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/khugepaged.c | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index f3a7aad5e8f2..cc5ff429d811 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -553,9 +553,19 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			cc->progress++;
 			break;
 		}
+
+		/*
+		 * Before the VMA is judged, so that a pass over an address space
+		 * of VMAs it skips is bounded by the budget too: each one is
+		 * charged for, and none of them was being asked to be scanned.
+		 */
+		if (cc->progress >= progress_max)
+			break;
+
 		orders = collapse_possible_orders(vma, vma->vm_flags,
 						  TVA_KHUGEPAGED);
 		if (!orders) {
+			khugepaged_scan.address = vma->vm_end;
 			cc->progress++;
 			continue;
 		}
@@ -570,6 +580,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 		hstart = ALIGN(vma->vm_start, window);
 		hend = ALIGN_DOWN(vma->vm_end, window);
 		if (khugepaged_scan.address > hend) {
+			khugepaged_scan.address = vma->vm_end;
 			cc->progress++;
 			continue;
 		}
@@ -584,7 +595,8 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			range_end = min(hend, pmd_addr + HPAGE_PMD_SIZE);
 
 			cond_resched();
-			if (unlikely(collapse_test_exit_or_disable(mm)))
+			if (unlikely(collapse_test_exit_or_disable(mm)) ||
+			    cc->progress >= progress_max)
 				goto breakouterloop;
 
 			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart);
@@ -596,8 +608,6 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			/* If nothing to collapse, the lock is still ours */
 			if (!collapse_scan_pmd(vma, start, range_end, cc, orders)) {
 				*result = cc->scan_refusal;
-				if (cc->progress >= progress_max)
-					goto breakouterloop;
 				continue;
 			}
 
-- 
2.54.0

[RFC PATCH 44/57] mm/khugepaged: hold the address space open across a scan

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Preparation for taking a per-VMA read lock instead of mmap_lock.

What tells khugepaged an address space is going away is the barrier in
__khugepaged_exit(): it runs before exit_mmap() and takes mmap_lock for
writing, which waits for a scan holding it for reading.  A scan under a
per-VMA lock holds no mmap_lock, so nothing waits for it and exit_mmap()
frees the page tables it is walking.

Take a reference on mm_users for the pass instead.  __mmput() cannot start
while one is held, so neither can exit_mmap(), whatever lock the pass
uses.

Drop it with mmput_async(), so the last reference does not tear an address
space down inside khugepaged.  Drop it before the exiting mm is judged,
too: that judgement needs the true count to release the slot.

The reference is also what the exiting-mm checks were reading, so an
address space whose owner has gone now shows as one reference rather than
none.  The three checks inside the pass ask collapse_test_exit_mmref()
instead; the slot-release judgement keeps the old test, running after the
reference is dropped.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.h   | 19 +++++++++++++++++--
 mm/khugepaged.c | 31 +++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 9e2cec1f250b..74e513c5c76c 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -170,6 +170,11 @@ struct collapse_control {
 	pte_t *saved_ptes;
 };
 
+static inline int collapse_disabled(struct mm_struct *mm)
+{
+	return mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
+}
+
 static inline int collapse_test_exit(struct mm_struct *mm)
 {
 	return atomic_read(&mm->mm_users) == 0;
@@ -177,8 +182,18 @@ static inline int collapse_test_exit(struct mm_struct *mm)
 
 static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
 {
-	return collapse_test_exit(mm) ||
-		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
+	return collapse_test_exit(mm) || collapse_disabled(mm);
+}
+
+/* The owner has gone: the caller's own reference is the only one left */
+static inline int collapse_test_exit_mmref(struct mm_struct *mm)
+{
+	return atomic_read(&mm->mm_users) == 1;
+}
+
+static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
+{
+	return collapse_test_exit_mmref(mm) || collapse_disabled(mm);
 }
 
 /*
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index cc5ff429d811..f3ea1846990e 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -531,16 +531,31 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	spin_unlock(&khugepaged_mm_lock);
 
 	mm = slot->mm;
+	vma = NULL;
+
+	/*
+	 * A reference on mm_users for as long as the pass works on this address
+	 * space.  __mmput() cannot start while one is held, so neither can
+	 * exit_mmap(), and the VMAs and page tables stay where they are.
+	 *
+	 * Once per pass, not once per table: the reference is what makes the
+	 * address space safe to work on, and a pass is how long that is wanted
+	 * for.  Nothing else in mm takes it per unit of work -- DAMON takes one
+	 * per target and walks every region under it, swapoff one per mm across
+	 * the whole address space, userfaultfd one per call.
+	 */
+	if (!mmget_not_zero(mm))
+		goto breakouterloop_no_mmput;
+
 	/*
 	 * Don't wait for semaphore (to avoid long wait times).  Just move to
 	 * the next mm on the list.
 	 */
-	vma = NULL;
 	if (unlikely(!mmap_read_trylock(mm)))
 		goto breakouterloop_mmap_lock;
 
 	cc->progress++;
-	if (unlikely(collapse_test_exit_or_disable(mm)))
+	if (unlikely(collapse_test_exit_or_disable_mmref(mm)))
 		goto breakouterloop;
 
 	vma_iter_init(&vmi, mm, khugepaged_scan.address);
@@ -549,7 +564,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 		unsigned long orders;
 
 		cond_resched();
-		if (unlikely(collapse_test_exit_or_disable(mm))) {
+		if (unlikely(collapse_test_exit_or_disable_mmref(mm))) {
 			cc->progress++;
 			break;
 		}
@@ -595,7 +610,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			range_end = min(hend, pmd_addr + HPAGE_PMD_SIZE);
 
 			cond_resched();
-			if (unlikely(collapse_test_exit_or_disable(mm)) ||
+			if (unlikely(collapse_test_exit_or_disable_mmref(mm)) ||
 			    cc->progress >= progress_max)
 				goto breakouterloop;
 
@@ -622,6 +637,14 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 breakouterloop:
 	mmap_read_unlock(mm); /* exit_mmap will destroy ptes after this */
 breakouterloop_mmap_lock:
+	/*
+	 * Not mmput(): the last reference would run exit_mmap() here, and
+	 * khugepaged is not the thread that should tear an address space down.
+	 * Dropped before the exiting mm is judged below, so that judgement still
+	 * sees the true count.
+	 */
+	mmput_async(mm);
+breakouterloop_no_mmput:
 
 	spin_lock(&khugepaged_mm_lock);
 	VM_BUG_ON(khugepaged_scan.mm_slot != slot);
-- 
2.54.0

[RFC PATCH 45/57] mm/collapse: take a per-VMA read lock for the round

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A round takes mmap_lock for reading and holds it across the freeze, the
copy and the install.  It is a read lock, so it blocks no faults -- but
every writer to the address space waits behind it, wherever in the mm that
writer is working.

A round works inside one VMA.  Everything it has to be excluded from --
fork, split, merge, unmap, and the free_pgtables() that follows them --
takes vma_start_write() on the VMA it touches first.  So a read lock on
that VMA excludes exactly what an mmap_read excluded, and an mmap_write
elsewhere in the mm stops waiting for a collapse it has nothing to do
with.

Look the VMA up with lock_vma_under_rcu() per round.  The round therefore
no longer works on the VMA the scan was given: it may be a different one
at that address, or the same one shrunk.  collapse_revalidate() already
re-asks every question the scan asked, and now also checks that the VMA
still covers what the round collected.

That lookup also fails on a VMA that is merely being written to, because
vma_start_read() fails while a writer holds vma_start_write().  A caller
cannot tell that from a VMA that has gone, so report SCAN_VMA_LOCK for
both, which MADV_COLLAPSE turns into -EAGAIN.  Treating it as a range that
shrank would report success for a collapse that never happened.

The fault-in pass gives up the same lock it was called under, so its
unlocks move with it, and the comments that carry the exclusion argument
name the lock they argue from.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 include/trace/events/collapse.h |  1 +
 mm/collapse.c                   | 77 +++++++++++++++++++--------------
 mm/collapse.h                   |  1 +
 mm/madvise.c                    |  1 +
 4 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/include/trace/events/collapse.h b/include/trace/events/collapse.h
index a3af5d8cc9aa..591030973d66 100644
--- a/include/trace/events/collapse.h
+++ b/include/trace/events/collapse.h
@@ -30,6 +30,7 @@
 	EM( SCAN_PAGE_COMPOUND,		"page_compound")		\
 	EM( SCAN_ANY_PROCESS,		"no_process_for_page")		\
 	EM( SCAN_VMA_NULL,		"vma_null")			\
+	EM( SCAN_VMA_LOCK,		"vma_not_lockable")		\
 	EM( SCAN_VMA_CHECK,		"vma_check_failed")		\
 	EM( SCAN_ADDRESS_RANGE,		"not_suitable_address_range")	\
 	EM( SCAN_DEL_PAGE_LRU,		"could_not_delete_page_from_lru")\
diff --git a/mm/collapse.c b/mm/collapse.c
index 76a53616b240..b6e92e24a3c5 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -52,7 +52,8 @@ enum collapse_pass {
  * The folios mapped across a window of PTEs become one folio of that window's
  * order, with the sources quiesced by the two barriers migration uses --
  * migration entries in their PTEs, then a frozen refcount -- so the copy itself
- * needs no lock.  The engine runs under mmap_read throughout.
+ * needs no lock.  A collapse takes a read lock on the VMA for each round; a
+ * scan still runs under the mmap_lock its caller holds.
  *
  * A round carries a batch of candidate windows through the passes together,
  * rather than carrying one window through the whole collapse.  [ptl] and
@@ -450,11 +451,11 @@ int collapse_control_init(struct collapse_control *cc)
 }
 
 /*
- * The scan and the allocation both dropped mmap_lock, so nothing seen before it
- * can be trusted: check the VMA the round just looked up and the PTE table
- * again, and that they still allow every provisioned candidate.
+ * The scan and the allocation both ran unlocked, so nothing seen before can be
+ * trusted: check the VMA the round just locked and the PTE table again, and
+ * that they still allow every provisioned candidate.
  *
- * The VMA was found by address, so it need not be the one the scan saw, nor
+ * The VMA was locked by address, so it need not be the one the scan saw, nor
  * still cover everything the round collected -- thp_vma_suitable_order() asks
  * that of each candidate, since a window is aligned to its own order.  A VMA
  * that shrank under a candidate therefore refuses that candidate and no more,
@@ -528,9 +529,9 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 /*
  * Bring one address to a state the freeze will accept: present, and exclusive if
  * it is anonymous.  Every fault it takes to get there counts in *nr_faults, each
- * one an allocation or a read the round is paying for.  Returns with mmap_lock
- * dropped on every failure, because the fault path may drop it and the caller
- * cannot tell which case it is in.
+ * one an allocation or a read the round is paying for.  Returns with the VMA
+ * read lock dropped on every failure, because the fault path may drop it and
+ * the caller cannot tell which case it is in.
  *
  * SCAN_EXCEED_SWAP_PTE is the exception: it is a verdict on this candidate
  * rather than on the round, nothing was faulted to reach it, and it keeps the
@@ -543,6 +544,7 @@ static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 {
 	struct mm_struct *mm = vma->vm_mm;
 	const unsigned int flags = FAULT_FLAG_ALLOW_RETRY | FAULT_FLAG_UNSHARE |
+		FAULT_FLAG_VMA_LOCK |
 		(mm != current->mm ? FAULT_FLAG_REMOTE : 0);
 	unsigned int tries;
 
@@ -553,7 +555,7 @@ static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 
 		pte = pte_offset_map(pmd, addr);
 		if (!pte) {
-			mmap_read_unlock(mm);
+			vma_end_read(vma);
 			return SCAN_NO_PTE_TABLE;
 		}
 		ptent = ptep_get_lockless(pte);
@@ -594,14 +596,14 @@ static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 		ret = handle_mm_fault(vma, addr, flags, NULL);
 		(*nr_faults)++;
 		/*
-		 * Not a verdict on this window: the fault dropped the lock to
-		 * wait, which is what a swap-in normally does.  Distinct from
+		 * Not a verdict on this window: the fault dropped the VMA lock
+		 * to wait, which is what a swap-in normally does.  Distinct from
 		 * SCAN_PAGE_LOCK, a folio someone else holds locked.
 		 */
 		if (ret & VM_FAULT_RETRY)
 			return SCAN_LOCK_DROPPED;
 		if (ret & VM_FAULT_ERROR) {
-			mmap_read_unlock(mm);
+			vma_end_read(vma);
 			return SCAN_FAIL;
 		}
 	}
@@ -611,7 +613,7 @@ static enum scan_result collapse_faultin_addr(struct vm_area_struct *vma,
 
 /*
  * Make every source the round needs present and exclusively owned by this mm,
- * by faulting it in as an ordinary access would.  Sleeps, and drops mmap_lock on
+ * by faulting it in as an ordinary access would.  Sleeps, and drops the VMA read
  * failure, since a fault may have to be retried with it released.
  *
  * Anything faulted in lands on a per-CPU LRU batch, holding a reference the
@@ -663,7 +665,7 @@ static enum scan_result collapse_faultin(struct vm_area_struct *vma,
 		}
 	}
 out:
-	/* @vma is unsafe on the failure path: the callee dropped mmap_lock */
+	/* @vma is unsafe on the failure path: the callee dropped its read lock */
 	trace_mm_collapse_faultin(mm, nr_faults, result);
 	return result;
 }
@@ -785,7 +787,7 @@ static void collapse_unfreeze_candidate(struct mm_struct *mm,
  * next page.  No layout is refused for its shape -- the next slot simply starts
  * its own span -- so partially mapped and compound sources collapse too.
  *
- * Caller holds mmap_read and the table's ptl.
+ * Caller holds the VMA read lock and the table's ptl.
  */
 static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
 						 struct collapse_control *cc,
@@ -929,7 +931,7 @@ static enum scan_result collapse_check_candidate(struct vm_area_struct *vma,
  *
  * On entry:
  *
- *  - mmap_read is held, and the table's ptl for the whole freeze;
+ *  - the VMA read lock is held, and the table's ptl for the whole freeze;
  *  - collapse_check_candidate() has accepted the candidate under that same ptl
  *    hold;
  *  - the round is covered by an mmu_notifier_invalidate_range_start() issued
@@ -1382,11 +1384,12 @@ static bool collapse_abort_slot(struct vm_area_struct *vma, struct folio *folio,
 
 /*
  * Abort one frozen candidate at install time: it took a machine check during the
- * copy, or some of its slots no longer hold our migration entries.  mmap_read
- * (held freeze..putback) blocks fork, mremap and munmap, and faults wait on the
- * migration entries -- but madvise-class operations run under mmap_read too, so a
- * concurrent MADV_DONTNEED may have zapped frozen slots, and a fault may have
- * refilled a zapped one.
+ * copy, or some of its slots no longer hold our migration entries.  The VMA read
+ * lock (held freeze..putback) blocks fork, mremap and munmap, each of which
+ * takes vma_start_write() on the VMA it touches, and faults wait on the
+ * migration entries -- but MADV_DONTNEED and MADV_FREE take a VMA read lock of
+ * their own, which ours does not exclude, so a concurrent zap may have taken
+ * frozen slots, and a fault may have refilled a zapped one.
  *
  * Slots still holding our entries are restored from the saved values (no TLB
  * flush: identical translation).  Foreign slots are left exactly as found --
@@ -1500,8 +1503,8 @@ static bool collapse_verify_candidate(struct collapse_candidate *cand,
  * The PMD terminal layer: verify, detach the table, deposit a fresh one and
  * install the leaf, as one atomic section under the pmd lock.  A pmd_none window
  * never exists -- faults stay held at pte level by the migration entries
- * throughout -- which is what lets PMD collapse run under mmap_read like the rest
- * of the engine.
+ * throughout -- which is what lets PMD collapse run under a VMA read lock like
+ * the rest of the engine.
  */
 static void collapse_install_pmd(struct vm_area_struct *vma,
 				 struct collapse_control *cc, pmd_t *pmd)
@@ -1571,8 +1574,8 @@ static void collapse_install_pmd(struct vm_area_struct *vma,
 	 * walks on the sources are unreachable -- refcounts frozen, folio locks
 	 * held from freeze to putback -- non-rmap pte walkers see migration
 	 * entries, pmd-level observers see the old table or the leaf and never an
-	 * intermediate, and fork, mremap and munmap take mmap_write, which our
-	 * mmap_read excludes.
+	 * intermediate, and fork, mremap and munmap take vma_start_write() on the
+	 * VMA they touch, which our VMA read lock excludes.
 	 *
 	 * The flush inside pmdp_collapse_flush() is the round's second over this
 	 * range: the freeze displaced every leaf here and flushed before dropping
@@ -1829,13 +1832,23 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 	collapse_reserve(mm, cc);
 	collapse_deposit(mm, cc);
 
+	/*
+	 * A read lock on the VMA rather than on the mm.  A round works inside one
+	 * VMA, and everything else it has to be excluded from -- fork, split,
+	 * merge, unmap, and the free_pgtables() that follows them -- takes
+	 * vma_start_write() on the VMA it touches first.  An mmap_write elsewhere
+	 * in the mm no longer waits behind a collapse.
+	 */
 retry:
-	mmap_read_lock(mm);
-
-	vma = find_vma(mm, pmd_addr);
+	vma = lock_vma_under_rcu(mm, pmd_addr);
 	if (!vma) {
-		result = SCAN_VMA_NULL;
-		goto out_unlock;
+		/*
+		 * Not only a VMA that has gone: lock_vma_under_rcu() also fails
+		 * on one being written to right now.  A caller cannot tell the
+		 * two apart, so say what is true of both -- try again.
+		 */
+		result = SCAN_VMA_LOCK;
+		goto out;
 	}
 
 	result = collapse_revalidate(vma, pmd_addr, cc, &pmd);
@@ -1852,7 +1865,7 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 	if (result == SCAN_LOCK_DROPPED && --passes)
 		goto retry;
 	if (result != SCAN_SUCCEED)
-		goto out;	/* the callee released mmap_lock */
+		goto out;	/* the callee released the VMA read lock */
 
 	/* One invalidate window spans the batch, as collapse_revalidate() left it */
 	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm,
@@ -1882,7 +1895,7 @@ static void collapse_round(struct mm_struct *mm, unsigned long pmd_addr,
 	mmu_notifier_invalidate_range_end(&range);
 
 out_unlock:
-	mmap_read_unlock(mm);
+	vma_end_read(vma);
 out:
 	nr_installed = collapse_finish(mm, cc, result);
 	trace_mm_collapse_round(mm, cc->nr_candidates, nr_installed, result,
diff --git a/mm/collapse.h b/mm/collapse.h
index 74e513c5c76c..e5a0ffab049c 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -39,6 +39,7 @@ enum scan_result {
 	SCAN_PAGE_COMPOUND,
 	SCAN_ANY_PROCESS,
 	SCAN_VMA_NULL,
+	SCAN_VMA_LOCK,
 	SCAN_VMA_CHECK,
 	SCAN_ADDRESS_RANGE,
 	SCAN_DEL_PAGE_LRU,
diff --git a/mm/madvise.c b/mm/madvise.c
index c1bb425be3f4..bd9123ee3cb1 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -933,6 +933,7 @@ static int madvise_collapse_errno(enum scan_result r)
 	case SCAN_PAGE_FILLED:
 	case SCAN_PAGE_HAS_PRIVATE:
 	case SCAN_PAGE_DIRTY_OR_WRITEBACK:
+	case SCAN_VMA_LOCK:
 		return -EAGAIN;
 	/*
 	 * Other: Trying again likely not to succeed / error intrinsic to
-- 
2.54.0

[RFC PATCH 46/57] mm/khugepaged: scan under a per-VMA read lock

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A pass took mmap_lock for the whole walk: every VMA of the address space
judged, and every table of every VMA scanned, under one lock.  A writer
anywhere in the mm waits for all of it.  And a pass that finds nothing to
collapse -- which is what a pass over an already-collapsed address space
is -- holds the lock for the whole sweep to say so.

Take a read lock on one VMA at a time instead.  lock_next_vma() finds the
next VMA at or after the cursor and locks it, falling back to mmap_lock
only where it cannot.  A scan holds that lock across every table of that
VMA and no longer.

What a round has to be excluded from already takes vma_start_write() on
the VMA it touches, so the exclusion is the same.  What changes is that it
is scoped to the VMA being scanned.

Two things follow from the iterator no longer being carried by mmap_lock.

The cursor moves by hand, because nothing else advances it now: past a VMA
that was walked, past one skipped without being looked at, and past each
table a scan was offered.

And the end of the address space has to be recognised rather than fallen
out of.  scan_complete says whether lock_next_vma() ran out of VMAs: an
error is not the end, and treating it as one would release the slot with
the address space half scanned.

The scan asserted mmap_assert_locked() on the way in.  Its two callers no
longer agree on what they hold -- khugepaged a VMA read lock from here,
MADV_COLLAPSE still mmap_lock -- so there is no single lock to assert, and
the assert goes.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   |  25 +++++-----
 mm/khugepaged.c | 128 +++++++++++++++++++++++++++++++-----------------
 2 files changed, 96 insertions(+), 57 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index b6e92e24a3c5..b3343595bdf2 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -52,8 +52,8 @@ enum collapse_pass {
  * The folios mapped across a window of PTEs become one folio of that window's
  * order, with the sources quiesced by the two barriers migration uses --
  * migration entries in their PTEs, then a frozen refcount -- so the copy itself
- * needs no lock.  A collapse takes a read lock on the VMA for each round; a
- * scan still runs under the mmap_lock its caller holds.
+ * needs no lock.  A scan runs under a read lock on the VMA it was handed; a
+ * collapse is given none and takes its own, for one round at a time.
  *
  * A round carries a batch of candidate windows through the passes together,
  * rather than carrying one window through the whole collapse.  [ptl] and
@@ -1949,9 +1949,10 @@ static enum scan_result collapse_scan_table(struct vm_area_struct *vma,
 	 * wait for a scan of the whole table.
 	 *
 	 * pte_offset_map() holds rcu_read_lock() until pte_unmap(), which is
-	 * what keeps the table itself from being freed underneath the walk;
-	 * mmap_lock keeps the VMA attached, without which free_pgtables() could
-	 * free it without waiting for RCU at all.  Nothing below here sleeps.
+	 * what keeps the table itself from being freed underneath the walk; the
+	 * VMA read lock keeps the VMA attached, without which free_pgtables()
+	 * could free it without waiting for RCU at all.  Nothing below here
+	 * sleeps.
 	 */
 	pte = pte_offset_map(pmd, start);
 	if (!pte) {
@@ -2174,9 +2175,9 @@ static void collapse_anon_scan_init(struct collapse_control *cc)
 /*
  * Judge one table's worth of @vma, leaving in @cc what a collapse could use:
  * which orders are still worth attempting, and why the table was turned down if
- * some order was.  Holds mmap_lock throughout -- it only reads -- and a caller
- * that acts on what it found hands the range to collapse_anon_pmd() afterwards,
- * without the lock.
+ * some order was.  Holds the read lock it was called under throughout -- it only
+ * reads -- and a caller that acts on what it found hands the range to
+ * collapse_anon_pmd() afterwards, without any lock.
  */
 static enum scan_result collapse_scan_anon_pmd(struct vm_area_struct *vma,
 					unsigned long start, unsigned long end,
@@ -3745,9 +3746,9 @@ static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 
 /*
  * Scan one table's worth of @vma and decide whether there is anything to collapse
- * in it.  The caller holds mmap_lock for reading and still holds it when this
- * returns: what is looked at is either the VMA or a page table that the lock
- * keeps in place.
+ * in it.  The caller holds a read lock and still holds it when this returns:
+ * what is looked at is either the VMA or a page table that the lock keeps in
+ * place.
  *
  * Returns whether collapse_run_pmd() has anything to do, and a scan that found
  * something has to be run: the file side takes a reference on the file while it
@@ -3760,8 +3761,6 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 {
 	struct mm_struct *mm = vma->vm_mm;
 
-	mmap_assert_locked(mm);
-
 	/*
 	 * What the scan answers with, so cleared before it runs.
 	 * collapse_anon_scan_init() clears the orders too, but only once the
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index f3ea1846990e..1d77d9a8046d 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -511,10 +511,10 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	__releases(&khugepaged_mm_lock)
 	__acquires(&khugepaged_mm_lock)
 {
-	struct vma_iterator vmi;
 	struct mm_slot *slot;
 	struct mm_struct *mm;
 	struct vm_area_struct *vma;
+	bool scan_complete = false;
 	unsigned int progress_prev = cc->progress;
 
 	lockdep_assert_held(&khugepaged_mm_lock);
@@ -534,55 +534,82 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	vma = NULL;
 
 	/*
-	 * A reference on mm_users for as long as the pass works on this address
-	 * space.  __mmput() cannot start while one is held, so neither can
-	 * exit_mmap(), and the VMAs and page tables stay where they are.
+	 * Hold the address space open for the pass.  A collapse works under a
+	 * per-VMA read lock, and the barrier __khugepaged_exit() puts in front
+	 * of exit_mmap() -- mmap_write_lock() -- waits for a reader of
+	 * mmap_lock, not for a reader of one VMA.  A reference on mm_users
+	 * stops __mmput(), and so both of those, from starting at all.
 	 *
-	 * Once per pass, not once per table: the reference is what makes the
-	 * address space safe to work on, and a pass is how long that is wanted
-	 * for.  Nothing else in mm takes it per unit of work -- DAMON takes one
-	 * per target and walks every region under it, swapoff one per mm across
-	 * the whole address space, userfaultfd one per call.
+	 * Once per pass rather than once per table: the reference is what makes
+	 * the address space safe to work on, and the pass is how long that is
+	 * wanted for.  Nothing else in mm takes it per unit of work -- DAMON
+	 * takes one per target and walks every region under it, swapoff one per
+	 * mm across the whole address space, userfaultfd one per call.  It is
+	 * dropped below before the exiting mm is judged, so that judgement still
+	 * sees the true count.
 	 */
 	if (!mmget_not_zero(mm))
 		goto breakouterloop_no_mmput;
 
-	/*
-	 * Don't wait for semaphore (to avoid long wait times).  Just move to
-	 * the next mm on the list.
-	 */
-	if (unlikely(!mmap_read_trylock(mm)))
-		goto breakouterloop_mmap_lock;
-
 	cc->progress++;
-	if (unlikely(collapse_test_exit_or_disable_mmref(mm)))
-		goto breakouterloop;
 
-	vma_iter_init(&vmi, mm, khugepaged_scan.address);
-	for_each_vma(vmi, vma) {
+	/*
+	 * One VMA at a time, each held by its own read lock rather than by
+	 * mmap_lock over the whole address space.  lock_next_vma() locks what it
+	 * finds, falling back to mmap_lock only where it cannot.
+	 *
+	 * Whether this mm still wants collapsing is asked once, at the top of
+	 * each round of the loop.  Asking again before entering it only repeats
+	 * the same question: nothing between the two can answer it differently.
+	 */
+	for (;;) {
 		unsigned long hstart, hend, window;
+		struct vma_iterator vmi;
 		unsigned long orders;
 
 		cond_resched();
+		/*
+		 * Our reference is the reason the count cannot fall to zero, so
+		 * it is also what an address space whose owner has gone looks
+		 * like.  Stopping is what frees it: nothing else here would.
+		 */
 		if (unlikely(collapse_test_exit_or_disable_mmref(mm))) {
 			cc->progress++;
-			break;
+			goto breakouterloop;
 		}
 
 		/*
-		 * Before the VMA is judged, so that a pass over an address space
-		 * of VMAs it skips is bounded by the budget too: each one is
-		 * charged for, and none of them was being asked to be scanned.
+		 * Before a VMA is locked, so that a pass over an address space
+		 * of VMAs it skips is bounded by the budget too, and so that a
+		 * collapse returning here does not lock one to be told it is
+		 * out of budget.
 		 */
 		if (cc->progress >= progress_max)
-			break;
+			goto breakouterloop;
+
+		/* The first VMA at or after the cursor, which often sits in a gap */
+		rcu_read_lock();
+		vma_iter_init(&vmi, mm, khugepaged_scan.address);
+		vma = lock_next_vma(mm, &vmi, khugepaged_scan.address);
+		rcu_read_unlock();
+
+		/*
+		 * NULL is the end of the address space, and the only thing that
+		 * finishes this mm.  An error is a fatal signal or the unlikely
+		 * reference count overflow: leave the mm for the next pass
+		 * rather than treat it as walked.
+		 */
+		if (IS_ERR_OR_NULL(vma)) {
+			scan_complete = !IS_ERR(vma);
+			vma = NULL;
+			goto breakouterloop;
+		}
 
 		orders = collapse_possible_orders(vma, vma->vm_flags,
 						  TVA_KHUGEPAGED);
 		if (!orders) {
-			khugepaged_scan.address = vma->vm_end;
 			cc->progress++;
-			continue;
+			goto next_vma;
 		}
 
 		/*
@@ -595,9 +622,8 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 		hstart = ALIGN(vma->vm_start, window);
 		hend = ALIGN_DOWN(vma->vm_end, window);
 		if (khugepaged_scan.address > hend) {
-			khugepaged_scan.address = vma->vm_end;
 			cc->progress++;
-			continue;
+			goto next_vma;
 		}
 		if (khugepaged_scan.address < hstart)
 			khugepaged_scan.address = hstart;
@@ -605,19 +631,24 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 		while (khugepaged_scan.address < hend) {
 			unsigned long pmd_addr, range_end, start;
 
+			cond_resched();
+
+			if (unlikely(collapse_test_exit_or_disable_mmref(mm)) ||
+			    cc->progress >= progress_max) {
+				vma_end_read(vma);
+				vma = NULL;
+				goto breakouterloop;
+			}
+
 			/* One table's worth at most, and never past the VMA */
 			pmd_addr = khugepaged_scan.address & HPAGE_PMD_MASK;
 			range_end = min(hend, pmd_addr + HPAGE_PMD_SIZE);
-
-			cond_resched();
-			if (unlikely(collapse_test_exit_or_disable_mmref(mm)) ||
-			    cc->progress >= progress_max)
-				goto breakouterloop;
+			start = khugepaged_scan.address;
 
 			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart);
+			VM_WARN_ON_ONCE(range_end > hend);
 
-			start = khugepaged_scan.address;
-			/* move to next address */
+			/* Move the cursor on regardless of what the scan says */
 			khugepaged_scan.address = range_end;
 
 			/* If nothing to collapse, the lock is still ours */
@@ -627,21 +658,30 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			}
 
 			/* collapse_run_pmd() takes its own locks, so give this up */
-			mmap_read_unlock(mm);
+			vma_end_read(vma);
+			vma = NULL;
+
 			*result = collapse_run_pmd(mm, start, range_end, cc);
 			if (*result == SCAN_SUCCEED)
-				++khugepaged_pages_collapsed;
-			goto breakouterloop_mmap_lock;
+				khugepaged_pages_collapsed++;
+			goto breakouterloop;
 		}
+next_vma:
+		/*
+		 * Past this VMA: the cursor has to move by hand, where the
+		 * mmap_lock iterator used to carry it.  A VMA that was walked
+		 * is charged by the scan itself, one table at a time; only one
+		 * passed over without being looked at is charged here.
+		 */
+		khugepaged_scan.address = vma->vm_end;
+		vma_end_read(vma);
+		vma = NULL;
 	}
+
 breakouterloop:
-	mmap_read_unlock(mm); /* exit_mmap will destroy ptes after this */
-breakouterloop_mmap_lock:
 	/*
 	 * Not mmput(): the last reference would run exit_mmap() here, and
 	 * khugepaged is not the thread that should tear an address space down.
-	 * Dropped before the exiting mm is judged below, so that judgement still
-	 * sees the true count.
 	 */
 	mmput_async(mm);
 breakouterloop_no_mmput:
@@ -652,7 +692,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	 * Release the current mm_slot if this mm is about to die, or
 	 * if we scanned all vmas of this mm, or THP got disabled.
 	 */
-	if (collapse_test_exit_or_disable(mm) || !vma) {
+	if (collapse_test_exit_or_disable(mm) || scan_complete) {
 		/*
 		 * Make sure that if mm_users is reaching zero while
 		 * khugepaged runs here, khugepaged_exit will find
-- 
2.54.0

[RFC PATCH 47/57] mm/madvise: collapse under a per-VMA read lock

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

MADV_COLLAPSE arrives under mmap_lock, gives it up, and takes it again per
PMD to scan and hand a table to the collapse.  Every writer to the address
space waits behind each of those, and the range can be as large as the
caller asked for.

Take a read lock on the VMA instead: MADVISE_VMA_READ_LOCK in the lock
mode, lock_vma_under_rcu() per PMD, and the walk told not to release what
the behaviour already let go.

A scan that finds nothing keeps the lock, so a range that is already
collapsed walks it without relocking.  A collapse gives the lock up and
looks the VMA up again afterwards: it can shrink while nothing is held,
which the scan reports as a refused range like any other.

Remote madvise is the exception.  process_madvise() has to untag the range
with untagged_addr_remote() before any VMA is looked at.  That reads mm
state mmap_lock protects, so remote MADV_COLLAPSE keeps the mmap_read it
has.

A VMA that cannot be locked is reported as SCAN_VMA_LOCK, which reaches
the caller as -EAGAIN.  lock_vma_under_rcu() also fails on a VMA being
written to, and reporting that as a range which shrank would tell the
caller a collapse succeeded where none was attempted.  With that, nothing
produces SCAN_VMA_NULL any more, and the test for it goes.

Both callers hold a read lock on the VMA now.  That is what the changes
outside madvise.c are for: the engine's interface documented mmap_lock as
its precondition, and only here does that stop being true of every caller.
The scan asserts the lock again too, which was not possible while the
callers disagreed.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 24 +++++++------
 mm/collapse.h |  8 ++---
 mm/madvise.c  | 98 +++++++++++++++++++++++++++++++++++++++------------
 3 files changed, 93 insertions(+), 37 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index b3343595bdf2..1e3b2d202ffe 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -3685,8 +3685,8 @@ static enum scan_result collapse_scan_file_pmd(struct vm_area_struct *vma,
 
 	/*
 	 * A PMD that is huge already has nothing left to collapse, and skipping
-	 * it here is what keeps mmap_lock out of a collapse that would find
-	 * nothing.  Everything else is worth the page cache scan, pmd_none()
+	 * it here is what keeps a collapse that would find nothing from being
+	 * run at all.  Everything else is worth the page cache scan, pmd_none()
 	 * included: a file range can be collapsed out of the cache without being
 	 * mapped first, which is why this is not the test the anonymous side
 	 * makes.
@@ -3703,8 +3703,8 @@ static enum scan_result collapse_scan_file_pmd(struct vm_area_struct *vma,
 
 /*
  * Build a PMD over what the page cache holds, and map it over the range if a huge
- * folio is already there but mapped by PTEs.  Runs with no mmap_lock, which the
- * caller gave up, and takes it again only for that last step.
+ * folio is already there but mapped by PTEs.  Runs with no lock on the VMA,
+ * which the caller gave up, and takes mmap_lock only for that last step.
  */
 static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 		unsigned long addr, struct collapse_control *cc)
@@ -3746,9 +3746,9 @@ static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 
 /*
  * Scan one table's worth of @vma and decide whether there is anything to collapse
- * in it.  The caller holds a read lock and still holds it when this returns:
- * what is looked at is either the VMA or a page table that the lock keeps in
- * place.
+ * in it.  The caller holds a read lock on @vma and still holds it when this
+ * returns: what is looked at is either the VMA or a page table that the lock
+ * keeps in place.
  *
  * Returns whether collapse_run_pmd() has anything to do, and a scan that found
  * something has to be run: the file side takes a reference on the file while it
@@ -3761,6 +3761,8 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 {
 	struct mm_struct *mm = vma->vm_mm;
 
+	vma_assert_locked(vma);
+
 	/*
 	 * What the scan answers with, so cleared before it runs.
 	 * collapse_anon_scan_init() clears the orders too, but only once the
@@ -3789,10 +3791,10 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 }
 
 /*
- * Collapse what the scan selected.  Called with no mmap_lock: the caller gives it
- * up first, because a collapse takes it again for each round and revalidates
- * under it, and holding it across the whole collapse would keep a writer to the
- * address space waiting for it.
+ * Collapse what the scan selected.  Called with no lock on the VMA the scan
+ * looked at: the caller gives that up first, because a collapse takes its own
+ * for each round and revalidates under it, and holding one across the whole
+ * collapse would keep a writer to the VMA waiting for it.
  */
 enum scan_result collapse_run_pmd(struct mm_struct *mm, unsigned long addr,
 		unsigned long end, struct collapse_control *cc)
diff --git a/mm/collapse.h b/mm/collapse.h
index e5a0ffab049c..aeaca305e71a 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -207,8 +207,8 @@ static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
  *	collapse_run_pmd(mm, addr, end, cc);	when the scan found work
  *	collapse_control_release(cc);
  *
- * The caller holds mmap_lock for reading and passes a range within one PTE table
- * of @vma.  A range the VMA does not cover is refused, which is also how a caller
+ * The caller holds a read lock on @vma and passes a range within one PTE table
+ * of it.  A range the VMA does not cover is refused, which is also how a caller
  * learns that its own range shrank.
  *
  * A scan returns with that lock still held: it only reads, and almost every table
@@ -218,8 +218,8 @@ static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
  * A collapse is called without it: the caller gives the lock up first, and with it
  * @vma and anything derived under it, so a caller carrying on has to look up
  * again.  What the collapse does -- allocate, quiesce, copy, flush -- is slow
- * enough that a writer would wait behind it, so it takes the lock again per round
- * instead, and revalidates rather than trusting what the scan saw.
+ * enough that a writer to the VMA would wait behind it, so it takes its own lock
+ * per round instead, and revalidates rather than trusting what the scan saw.
  *
  * A scan that found something has to be run: the file side takes a reference on
  * the file while it still has the VMA to take it from, and the run is what gives
diff --git a/mm/madvise.c b/mm/madvise.c
index bd9123ee3cb1..5f6d815d70ad 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -276,6 +276,18 @@ static void mark_mmap_lock_dropped(struct madvise_behavior *madv_behavior)
 	madv_behavior->lock_dropped = true;
 }
 
+/*
+ * The VMA-lock counterpart, for a behaviour that releases the VMA it was handed
+ * and locks what it needs for itself.  The walk has nothing left to release,
+ * and unlike the mmap_lock case it has nothing to carry on with either: the VMA
+ * fast path applies to one VMA and returns.
+ */
+static void mark_vma_lock_dropped(struct madvise_behavior *madv_behavior)
+{
+	VM_WARN_ON_ONCE(madv_behavior->lock_mode != MADVISE_VMA_READ_LOCK);
+	madv_behavior->lock_dropped = true;
+}
+
 /*
  * Schedule all required I/O operations.  Do not wait for completion.
  */
@@ -948,6 +960,8 @@ static int madvise_collapse_errno(enum scan_result r)
 static int madvise_collapse(struct madvise_behavior *madv_behavior)
 {
 	struct madvise_behavior_range *range = &madv_behavior->range;
+	const bool vma_locked =
+		madv_behavior->lock_mode == MADVISE_VMA_READ_LOCK;
 	struct vm_area_struct *vma = madv_behavior->vma;
 	struct mm_struct *mm = madv_behavior->mm;
 	unsigned long hstart, hend, addr;
@@ -981,13 +995,22 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 	}
 
 	/*
-	 * Nothing below wants the lock the VMA walk left held, and
-	 * lru_add_drain_all() waits on every CPU, so give it up first.  The
-	 * walk carries on under mmap_lock and its own caller is what drops it,
-	 * so reporting this only tells the walk that its VMA is now stale.
+	 * Give up whatever the caller locked for us.  lru_add_drain_all() below
+	 * must not run under a lock, and the loop locks what it works on for
+	 * itself, one VMA at a time, so the caller's VMA is of no use past here.
+	 *
+	 * Which lock that is depends on how we were reached.  A range inside one
+	 * VMA arrives with that VMA read-locked and nothing else; a range that
+	 * spans VMAs arrives under mmap_lock, because try_vma_read_lock() took
+	 * it and turned the walk generic.
 	 */
-	mmap_read_unlock(mm);
-	mark_mmap_lock_dropped(madv_behavior);
+	if (vma_locked) {
+		vma_end_read(vma);
+		mark_vma_lock_dropped(madv_behavior);
+	} else {
+		mmap_read_unlock(mm);
+		mark_mmap_lock_dropped(madv_behavior);
+	}
 	vma = NULL;
 	vma_orders = 0;
 	lru_add_drain_all();
@@ -996,22 +1019,36 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 		enum scan_result result;
 
 		/*
-		 * A collapse gives the lock up, and the VMA has to be found
-		 * again after one: it can shrink while nothing is held.  A scan
-		 * that finds nothing to collapse leaves the lock alone, so a
-		 * range that is already collapsed walks it without relocking.
+		 * On another process, the reference this call holds is what
+		 * keeps the address space from being torn down -- so if it is
+		 * the only one left, the owner has gone and every page of it is
+		 * waiting on us to stop.  Nothing else here would notice: the
+		 * range is the caller's, and it can be enormous.
+		 */
+		if (mm != current->mm && collapse_test_exit_mmref(mm)) {
+			hend = addr;
+			break;
+		}
+
+		/*
+		 * A collapse gives the VMA read lock up, and the VMA has to be
+		 * found again after one: it can shrink while nothing is held.
 		 *
 		 * Reschedule only here, where nothing is held: a preemption
 		 * point under a lock is a writer waiting longer.
 		 */
 		if (!vma) {
 			cond_resched();
-			mmap_read_lock(mm);
-			vma = vma_lookup(mm, addr);
-			if (!vma) {
-				mmap_read_unlock(mm);
-				hend = addr;
-				break;
+			vma = lock_vma_under_rcu(mm, addr);
+			if (IS_ERR_OR_NULL(vma)) {
+				/*
+				 * Not only a VMA that has gone: this also fails
+				 * on one being written to right now.  Say what
+				 * is true of both -- try again.
+				 */
+				vma = NULL;
+				last_fail = SCAN_VMA_LOCK;
+				goto out;
 			}
 			vma_orders = collapse_possible_orders(vma,
 					vma->vm_flags, TVA_FORCED_COLLAPSE);
@@ -1023,7 +1060,7 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 			result = cc->scan_refusal;
 		} else {
 			/* collapse_run_pmd() takes its own locks, so give this up */
-			mmap_read_unlock(mm);
+			vma_end_read(vma);
 			vma = NULL;
 			/* The mask belonged to that lock, not to this range */
 			vma_orders = 0;
@@ -1036,7 +1073,7 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 		 * The VMA shrank under us, so the rest of the range was never
 		 * ours to collapse: stop, and expect only what came before.
 		 */
-		if (result == SCAN_VMA_NULL || result == SCAN_ADDRESS_RANGE) {
+		if (result == SCAN_ADDRESS_RANGE) {
 			hend = addr;
 			break;
 		}
@@ -1067,8 +1104,14 @@ static int madvise_collapse(struct madvise_behavior *madv_behavior)
 	}
 
 out:
-	/* The VMA walk this returns to expects the lock it was holding */
-	if (!vma)
+	if (vma)
+		vma_end_read(vma);
+	/*
+	 * Hand mmap_lock back only to a caller that is going to carry on with
+	 * it: the generic walk finds the next VMA under it.  The VMA fast path
+	 * applies to one VMA and returns, so it wants nothing back.
+	 */
+	if (!vma_locked)
 		mmap_read_lock(mm);
 	collapse_control_release(cc);
 	kfree(cc);
@@ -1866,7 +1909,10 @@ int madvise_walk_vmas(struct madvise_behavior *madv_behavior)
 	if (madv_behavior->lock_mode == MADVISE_VMA_READ_LOCK &&
 	    try_vma_read_lock(madv_behavior)) {
 		error = madvise_vma_behavior(madv_behavior);
-		vma_end_read(madv_behavior->vma);
+		/* A behaviour that let the VMA go has nothing left to release */
+		if (!madv_behavior->lock_dropped)
+			vma_end_read(madv_behavior->vma);
+		madv_behavior->lock_dropped = false;
 		return error;
 	}
 
@@ -1941,8 +1987,16 @@ static enum madvise_lock_mode get_lock_mode(struct madvise_behavior *madv_behavi
 	case MADV_PAGEOUT:
 	case MADV_POPULATE_READ:
 	case MADV_POPULATE_WRITE:
-	case MADV_COLLAPSE:
 		return MADVISE_MMAP_READ_LOCK;
+	case MADV_COLLAPSE:
+		/*
+		 * Only for this process.  On another one the range has to be
+		 * untagged with untagged_addr_remote(), which reads mm state
+		 * that mmap_lock protects, before any VMA is looked at.
+		 */
+		if (madv_behavior->mm != current->mm)
+			return MADVISE_MMAP_READ_LOCK;
+		fallthrough;
 	case MADV_GUARD_INSTALL:
 	case MADV_GUARD_REMOVE:
 	case MADV_DONTNEED:
-- 
2.54.0

[RFC PATCH 48/57] mm/collapse: assert the mm reference the engine relies on

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Every caller of the engine holds a reference on mm_users for as long as it
works: khugepaged for a pass, MADV_COLLAPSE for a call.  So mm_users
cannot reach zero underneath one, yet collapse.c still asked whether it
had, in three places where the answer can only be no.

Ask only whether collapsing was turned off, which prctl() can do at any
point, and assert the rest where the interface begins.

A caller that arrives without a reference is a bug in the caller.  A debug
build says so there, rather than leaving it to be found when an address
space is freed under a collapse.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 1e3b2d202ffe..28631c734fcb 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -473,7 +473,7 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 	enum scan_result result;
 	unsigned int i, nr_live = 0;
 
-	if (unlikely(collapse_test_exit_or_disable(mm)))
+	if (unlikely(collapse_disabled(mm)))
 		return SCAN_ANY_PROCESS;
 
 	if (!vma->anon_vma || !vma_is_anonymous(vma))
@@ -3731,7 +3731,7 @@ static enum scan_result collapse_file_pmd(struct mm_struct *mm,
 
 	if (result == SCAN_PTE_MAPPED_HUGEPAGE) {
 		mmap_read_lock(mm);
-		if (collapse_test_exit_or_disable(mm))
+		if (collapse_disabled(mm))
 			result = SCAN_ANY_PROCESS;
 		else
 			result = try_collapse_pte_mapped_thp(mm, addr,
@@ -3762,6 +3762,8 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 	struct mm_struct *mm = vma->vm_mm;
 
 	vma_assert_locked(vma);
+	/* The caller holds a reference on it, so it cannot have gone away */
+	VM_WARN_ON_ONCE(collapse_test_exit(mm));
 
 	/*
 	 * What the scan answers with, so cleared before it runs.
@@ -3776,7 +3778,7 @@ bool collapse_scan_pmd(struct vm_area_struct *vma, unsigned long addr,
 		cc->scan_file = NULL;
 	}
 
-	if (unlikely(collapse_test_exit_or_disable(mm)))
+	if (unlikely(collapse_disabled(mm)))
 		cc->scan_refusal = SCAN_ANY_PROCESS;
 	else if (addr < vma->vm_start || end > vma->vm_end)
 		cc->scan_refusal = SCAN_ADDRESS_RANGE;
-- 
2.54.0

[RFC PATCH 49/57] mm/khugepaged: drop the mmap_lock barrier from __khugepaged_exit()

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

__khugepaged_exit() took mmap_lock for writing and dropped it again, to
wait for a scan holding it for reading before exit_mmap() destroyed the
page tables underneath it.

A scan now holds a reference on mm_users for as long as it runs, so it
cannot be in flight here at all: __mmput() runs only once that count has
reached zero, and asserts so on entry.  The barrier is redundant.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/khugepaged.c | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 1d77d9a8046d..2b3364bdb100 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -434,16 +434,6 @@ void __khugepaged_exit(struct mm_struct *mm)
 		mm_flags_clear(MMF_VM_HUGEPAGE, mm);
 		mm_slot_free(mm_slot_cache, slot);
 		mmdrop(mm);
-	} else if (slot) {
-		/*
-		 * This is required to serialize against
-		 * collapse_test_exit() (which is guaranteed to run
-		 * under mmap_lock read mode). Stop here (after we return all
-		 * pagetables will be destroyed) until khugepaged has finished
-		 * working on the pagetables under the mmap_lock.
-		 */
-		mmap_write_lock(mm);
-		mmap_write_unlock(mm);
 	}
 }
 
-- 
2.54.0

[RFC PATCH 50/57] selftests/mm: attribute collapses by candidate event alone

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The check counts collapses of its window by matching isolate events
against the source PFNs it recorded beforehand.  That tracepoint went with
the mechanism that emitted it, so nothing would match.

Count the engine's per-candidate events instead: an install that
succeeded, at the window's address and order, is one collapse of that
window.

Preparing the window still checks that the sources are present, which is
the other thing those PFN lookups were doing.  The recorded PFNs
themselves are no longer needed, so the array goes.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 .../selftests/mm/khugepaged_sync_check.c      | 55 +++++++++----------
 1 file changed, 26 insertions(+), 29 deletions(-)

diff --git a/tools/testing/selftests/mm/khugepaged_sync_check.c b/tools/testing/selftests/mm/khugepaged_sync_check.c
index 4c37b697d3dd..2a0c247aeec5 100644
--- a/tools/testing/selftests/mm/khugepaged_sync_check.c
+++ b/tools/testing/selftests/mm/khugepaged_sync_check.c
@@ -7,10 +7,9 @@
  * advancing by two is a completion barrier for one full pass that
  * started after setup (khugepaged_full_pass()). Verify the pair gives
  * deterministic, attributable results: one barrier step over one
- * prepared window produces exactly one collapse attempt on that
- * window's source pages (mm_collapse_huge_page_isolate events filtered
- * by source PFN and order) and the window is collapsed
- * afterwards, repeatably.
+ * prepared window produces exactly one collapse of that window --
+ * mm_collapse_candidate install events at its address and order -- and
+ * the window is collapsed afterwards, repeatably.
  *
  * scan_sleep_millisecs is set to 60s to prove the wake path: without
  * the wake, one barrier step would sleep multiples of that and blow
@@ -51,11 +50,10 @@ static void trace_events_off(void)
 }
 
 /*
- * Count collapse attempts attributable to our window: isolate events whose
- * scan_pfn is one of the window's source PFNs, reported once per attempt.
+ * Count the collapses attributable to our window: per-candidate install
+ * events at the window's address and order, one per collapse.
  */
-static int count_attributed(unsigned long *pfns, int nr_pfns,
-			    unsigned int order)
+static int count_attributed(unsigned long addr, unsigned int order)
 {
 	char line[1024];
 	int count = 0;
@@ -66,25 +64,25 @@ static int count_attributed(unsigned long *pfns, int nr_pfns,
 		ksft_exit_fail_msg("Cannot open trace buffer\n");
 
 	while (fgets(line, sizeof(line), fp)) {
+		char *s;
 		unsigned long val;
 		unsigned int ord;
-		char *s, *o;
-		int i;
+		char *o;
 
-		s = strstr(line, "mm_collapse_huge_page_isolate:");
-		if (!s)
-			continue;
-		if (sscanf(s, "mm_collapse_huge_page_isolate: scan_pfn=0x%lx",
-			   &val) != 1)
-			continue;
-		o = strstr(s, "order=");
-		if (!o || sscanf(o, "order=%u", &ord) != 1 || ord != order)
-			continue;
-		for (i = 0; i < nr_pfns; i++) {
-			if (val == pfns[i]) {
-				count++;
-				break;
-			}
+		s = strstr(line, "mm_collapse_candidate:");
+		if (s) {
+			if (!strstr(s, "pass=install") ||
+			    !strstr(s, "result=succeeded"))
+				continue;
+			o = strstr(s, "addr=");
+			if (!o || sscanf(o, "addr=0x%lx", &val) != 1 ||
+			    val != addr)
+				continue;
+			o = strstr(s, "order=");
+			if (!o || sscanf(o, "order=%u", &ord) != 1 ||
+			    ord != order)
+				continue;
+			count++;
 		}
 	}
 	fclose(fp);
@@ -95,7 +93,6 @@ static void one_step(int iteration)
 {
 	const size_t window = getpagesize() << TARGET_ORDER;
 	const int nr_pages = 1 << TARGET_ORDER;
-	unsigned long pfns[1 << TARGET_ORDER];
 	bool collapsed, passed;
 	int attributed;
 	char *p;
@@ -106,11 +103,11 @@ static void one_step(int iteration)
 	if (p != BASE_ADDR)
 		ksft_exit_fail_perror("mmap() window");
 
-	/* Prepare one window; record its source PFNs. */
+	/* Prepare one window, and check the sources really are present. */
 	for (i = 0; i < nr_pages; i++) {
 		p[i * getpagesize()] = i + 1;
-		pfns[i] = pagemap_get_pfn(pagemap_fd, p + i * getpagesize());
-		if (pfns[i] == -1UL)
+		if (pagemap_get_pfn(pagemap_fd,
+				    p + i * getpagesize()) == -1UL)
 			ksft_exit_fail_msg("Source page not present\n");
 	}
 
@@ -133,7 +130,7 @@ static void one_step(int iteration)
 
 	collapsed = is_range_backed_by_folio_orders(p, window, TARGET_ORDER,
 						    pagemap_fd, kpageflags_fd);
-	attributed = count_attributed(pfns, nr_pages, TARGET_ORDER);
+	attributed = count_attributed((unsigned long)p, TARGET_ORDER);
 
 	ksft_test_result(collapsed && attributed == 1,
 			 "step %d: window collapsed, %d attributed result(s)\n",
-- 
2.54.0

[RFC PATCH 51/57] selftests/mm: cover collapse inside a sub-PMD VMA

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

The new engine can collapse an mTHP inside a VMA smaller than a PMD.  The
old mechanism could not: its coverage was rooted at whole PMD-aligned
spans, so a VMA that could not hold one was passed over however many
mTHP-sized windows it held.  On arm64 with 64K pages, where a PMD is 512M,
that was every VMA below 512M.

Nothing covers that, so add two cases:

 - a VMA of exactly one window, the smallest thing that can be collapsed
   at all;
 - a VMA of several windows, which also walks from one window to the next
   inside a single VMA.

The second takes as many windows as fit in half a table.  At the order
just below the PMD order two windows are already a whole table, so there
is no room for several and the case skips.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c | 65 +++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index 172e7307eeee..85f138cfb9b2 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1536,6 +1536,69 @@ static void collapse_order_mixed_sources(struct collapse_context *c,
 	ksft_test_result_report(exit_status, "%s\n", __func__);
 }
 
+/*
+ * A VMA smaller than a PMD is a valid collapse target, so long as it holds a
+ * naturally aligned window of the target order.  This is the case khugepaged
+ * used to pass over entirely, its coverage being rooted at whole PMD-aligned
+ * spans, which on arm64 with 64K pages meant every VMA below 512M.
+ */
+static void __collapse_order_sub_pmd_vma(struct collapse_context *c,
+					 struct mem_ops *ops, int nr_windows,
+					 const char *name)
+{
+	size_t size = nr_windows * mthp_window_size();
+	void *p;
+
+	mthp_push_target_order();
+
+	p = mmap(BASE_ADDR, size, PROT_READ | PROT_WRITE,
+		 MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+	if (p != BASE_ADDR)
+		ksft_exit_fail_msg("Failed to allocate VMA at %p\n", BASE_ADDR);
+
+	fill_memory(p, 0, size);
+	if (!window_not_collapsed(p, size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	madvise(p, size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse inside a sub-PMD VMA (%d windows)...",
+		       nr_windows);
+	if (!khugepaged_wait_full_pass())
+		fail("Timeout");
+	else if (window_collapsed(p, size))
+		success("OK");
+	else
+		fail("Fail");
+
+	validate_memory(p, 0, size);
+	munmap(p, size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", name);
+}
+
+static void collapse_order_sub_pmd_vma(struct collapse_context *c,
+				       struct mem_ops *ops)
+{
+	__collapse_order_sub_pmd_vma(c, ops, 1, __func__);
+}
+
+static void collapse_order_sub_pmd_range(struct collapse_context *c,
+					 struct mem_ops *ops)
+{
+	size_t window = mthp_window_size();
+	int nr_windows = 16;
+
+	while (nr_windows > 1 && nr_windows * window > hpage_pmd_size / 2)
+		nr_windows /= 2;
+
+	if (nr_windows == 1) {
+		ksft_test_result_skip("%s: no room for multiple windows below the PMD\n",
+				      __func__);
+		return;
+	}
+	__collapse_order_sub_pmd_vma(c, ops, nr_windows, __func__);
+}
+
 static void usage(void)
 {
 	fprintf(stderr, "\nUsage: ./khugepaged [OPTIONS] <test type> [dir]\n\n");
@@ -1823,6 +1886,8 @@ int main(int argc, char **argv)
 		TEST(collapse_order_partial_window, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_max_ptes_none, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_mixed_sources, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_sub_pmd_vma, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_sub_pmd_range, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
-- 
2.54.0

[RFC PATCH 52/57] selftests/mm: cover a hole-y window in a sub-PMD VMA

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A collapse over a partially populated window has to zero the slots it
found empty.  The destination comes from the allocator holding whatever
was last written to it, and a slot the process never touched must read
back zero.

Nothing in the suite checks that.  validate_memory() only re-reads the
pattern fill_memory() wrote, and every caller passes it the faulted extent
alone, so a collapse that left stale bytes behind the holes would pass
every case here.

Check it where population and sub-PMD eligibility meet, which no existing
case covers either.  One page is faulted in a window-sized VMA; the case
expects the window collapsed, the faulted page unchanged, and every byte
behind the holes zero.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c | 50 +++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index 85f138cfb9b2..b61e32566d47 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1599,6 +1599,55 @@ static void collapse_order_sub_pmd_range(struct collapse_context *c,
 	__collapse_order_sub_pmd_vma(c, ops, nr_windows, __func__);
 }
 
+/*
+ * A partially populated window in a sub-PMD VMA: population and
+ * sub-PMD eligibility at once. The unfaulted slots must come back
+ * zero-filled in the collapsed folio, and the faulted ones unchanged.
+ */
+static void collapse_order_sub_pmd_holes(struct collapse_context *c,
+					 struct mem_ops *ops)
+{
+	size_t size = mthp_window_size();
+	char *bytes;
+	void *p;
+	size_t i;
+
+	mthp_push_target_order();
+
+	p = mmap(BASE_ADDR, size, PROT_READ | PROT_WRITE,
+		 MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+	bytes = p;
+	if (p != BASE_ADDR)
+		ksft_exit_fail_msg("Failed to allocate VMA at %p\n", BASE_ADDR);
+
+	fill_memory(p, 0, page_size);
+	if (!window_not_collapsed(p, size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	madvise(p, size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse hole-y window inside a sub-PMD VMA...");
+	if (!khugepaged_wait_full_pass()) {
+		fail("Timeout");
+	} else if (window_collapsed(p, size)) {
+		/* The unfaulted tail must be zero-filled. */
+		for (i = page_size; i < size; i++) {
+			if (bytes[i])
+				break;
+		}
+		if (i == size)
+			success("OK");
+		else
+			fail("Fail");
+	} else {
+		fail("Fail");
+	}
+
+	validate_memory(p, 0, page_size);
+	munmap(p, size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 static void usage(void)
 {
 	fprintf(stderr, "\nUsage: ./khugepaged [OPTIONS] <test type> [dir]\n\n");
@@ -1888,6 +1937,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_mixed_sources, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_vma, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_range, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_sub_pmd_holes, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
-- 
2.54.0

[RFC PATCH 53/57] selftests/mm: cover collapse of mlocked ranges

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Collapsing an mlocked range makes the teardown do something it does
nowhere else: the sources have to be munlocked while the destination
arrives already mlocked, and munlocking takes a reference.  So a teardown
that reaches a source before it is unfrozen fails on a refcount that is
not allowed to move.  That is the one ordering constraint in the putback
with no other way to be caught.

An mlocked range was collapsible before, as long as the whole VMA was
locked.  This case is the other shape.  mlock() over part of a VMA splits
it, leaving the locked part smaller than a PMD, which khugepaged passed
over for as long as its coverage was rooted at PMD-aligned spans.  A
partially mlocked region therefore went uncollapsed however long it lived.

Cover it deterministically: mlock a window, collapse it, check the
contents survive.  Drive it under contention too, with a thread mlocking
and munlocking random spans across the race harness's region, since the
ordering only breaks when a teardown and an mlock overlap.  The plain
racers never touch VM_LOCKED at all.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c      | 38 ++++++++++++++++++++
 tools/testing/selftests/mm/khugepaged_race.c | 29 +++++++++++++--
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index b61e32566d47..208300ecb344 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1599,6 +1599,43 @@ static void collapse_order_sub_pmd_range(struct collapse_context *c,
 	__collapse_order_sub_pmd_vma(c, ops, nr_windows, __func__);
 }
 
+/*
+ * Collapse of an mlocked window: source teardown munlocks the old
+ * pages while the new folio arrives mlocked via folio_add_lru_vma().
+ * A teardown that touches the sources while they are still frozen
+ * blows up exactly here (munlock_folio() takes a reference).
+ */
+static void collapse_order_mlocked(struct collapse_context *c,
+				   struct mem_ops *ops)
+{
+	size_t window = mthp_window_size();
+	void *p;
+
+	mthp_push_target_order();
+
+	p = ops->setup_area(1);
+	ops->fault(p, 0, window);
+	if (mlock(p, window))
+		ksft_exit_fail_perror("mlock()");
+	if (!window_not_collapsed(p, hpage_pmd_size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	madvise(p, hpage_pmd_size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse fully populated mlocked window...");
+	if (!khugepaged_wait_full_pass())
+		fail("Timeout");
+	else if (window_collapsed(p, window))
+		success("OK");
+	else
+		fail("Fail");
+
+	validate_memory(p, 0, window);
+	munlock(p, window);
+	ops->cleanup_area(p, hpage_pmd_size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 /*
  * A partially populated window in a sub-PMD VMA: population and
  * sub-PMD eligibility at once. The unfaulted slots must come back
@@ -1938,6 +1975,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_sub_pmd_vma, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_range, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_holes, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_mlocked, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
diff --git a/tools/testing/selftests/mm/khugepaged_race.c b/tools/testing/selftests/mm/khugepaged_race.c
index 6682bbae0a8f..a4710130aabf 100644
--- a/tools/testing/selftests/mm/khugepaged_race.c
+++ b/tools/testing/selftests/mm/khugepaged_race.c
@@ -219,6 +219,29 @@ static void *forker_fn(void *arg)
 	return NULL;
 }
 
+/*
+ * mlock/munlock cycling over the shared areas: collapse of an mlocked
+ * range munlocks the sources at teardown and mlocks the new folio --
+ * the interaction the fuzzer caught (munlock on a frozen source) and
+ * the plain racers never drove.
+ */
+static void *mlocker_fn(void *arg)
+{
+	unsigned int seed = (unsigned long)arg;
+
+	while (!stop) {
+		unsigned long page_idx = rand_page(&seed);
+		unsigned long nr = 1UL << (rand_r(&seed) % 8);	/* 1..128 pages */
+
+		if (rand_r(&seed) & 1)
+			mlock(region + page_idx * page_size, nr * page_size);
+		else
+			munlock(region + page_idx * page_size, nr * page_size);
+		usleep(rand_r(&seed) % 1000);
+	}
+	return NULL;
+}
+
 static void *mremapper_fn(void *arg)
 {
 	unsigned int seed = (unsigned long)arg;
@@ -328,14 +351,14 @@ int main(int argc, char **argv)
 {
 	static const char * const thread_names[] = {
 		"faulter", "faulter2", "dontneed", "pinner", "forker",
-		"mremapper", "pageout", "compactor",
+		"mremapper", "mlocker", "pageout", "compactor",
 	};
 	void *(*const thread_fns[])(void *) = {
 		faulter_fn, faulter_fn, dontneed_fn, pinner_fn, forker_fn,
-		mremapper_fn, pageout_fn, compactor_fn,
+		mremapper_fn, mlocker_fn, pageout_fn, compactor_fn,
 	};
 	enum { T_FAULTER, T_FAULTER2, T_DONTNEED, T_PINNER, T_FORKER,
-	       T_MREMAPPER, T_PAGEOUT, T_COMPACTOR };
+	       T_MREMAPPER, T_MLOCKER, T_PAGEOUT, T_COMPACTOR };
 	const unsigned long pageout_bit = 1UL << T_PAGEOUT;
 	const unsigned long compactor_bit = 1UL << T_COMPACTOR;
 	const int nr_threads = ARRAY_SIZE(thread_names);
-- 
2.54.0

[RFC PATCH 54/57] selftests/mm: cover collapse beside a MADV_FREE'd page

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A clean lazyfree page must not be collapsed.  Copying it into a folio that
is not lazyfree would quietly take back memory the process offered to the
kernel, and reclaim would no longer be free to drop it.  khugepaged
refuses the window that holds one.

What it should not do is give up on the rest of the table.  One page a
process no longer needs is a poor reason to leave a whole PMD's worth of
memory without large folios.  Yet refusing per table is what khugepaged
did: any single disqualified PTE ended the scan.

So collapse a table with one MADV_FREE'd page in it, and expect three
things: the windows beside it collapsed, the window holding it not, and
the page itself still backed by an order-0 folio.

That last one matters because selection descends orders on a refusal, so
checking the target order alone would not notice a smaller window
swallowing it.

The freed page's contents are not checked -- reclaim is entitled to have
dropped them -- but everything else must still read back.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c | 44 +++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index 208300ecb344..58cb7364652e 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1685,6 +1685,49 @@ static void collapse_order_sub_pmd_holes(struct collapse_context *c,
 	ksft_test_result_report(exit_status, "%s\n", __func__);
 }
 
+/*
+ * One MADV_FREE'd page must not stop the windows beside it from collapsing.
+ * khugepaged still refuses the window holding it: collapsing would copy the
+ * page into a folio that is not lazyfree, quietly making memory the process
+ * offered up undroppable again.
+ */
+static void collapse_order_lazyfree_window(struct collapse_context *c,
+					   struct mem_ops *ops)
+{
+	size_t window = mthp_window_size();
+	void *p;
+
+	mthp_push_target_order();
+
+	p = ops->setup_area(1);
+	ops->fault(p, 0, hpage_pmd_size);
+	if (!window_not_collapsed(p, hpage_pmd_size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	/* Clean and lazyfree: do not touch this page again. */
+	if (madvise(p, page_size, MADV_FREE))
+		ksft_exit_fail_perror("MADV_FREE");
+
+	madvise(p, hpage_pmd_size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse the windows beside a MADV_FREE'd page...");
+	if (!khugepaged_wait_full_pass())
+		fail("Timeout");
+	else if (window_collapsed(p + window, hpage_pmd_size - window) &&
+		 window_not_collapsed(p, window) &&
+		 /* Left alone at every order, not just the target one */
+		 is_range_backed_by_folio_orders(p, page_size, 0,
+						 pagemap_fd, kpageflags_fd))
+		success("OK");
+	else
+		fail("Fail");
+
+	/* Everything but the freed page, whose contents may be gone. */
+	validate_memory(p, page_size, hpage_pmd_size);
+	ops->cleanup_area(p, hpage_pmd_size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 static void usage(void)
 {
 	fprintf(stderr, "\nUsage: ./khugepaged [OPTIONS] <test type> [dir]\n\n");
@@ -1976,6 +2019,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_sub_pmd_range, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_holes, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_mlocked, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_lazyfree_window, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
-- 
2.54.0

[RFC PATCH 55/57] selftests/mm: cover collapse beside a pinned page

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

A GUP pin holds a reference nobody else can account for, so a pinned page
cannot be collapsed.  The freeze needs the folio's reference count to
match what its mappings and cache membership explain, and a pin makes it
not.

The window holding one has to be refused.  The windows beside it should
not be, as with a lazyfree page.

Pin one page for the duration of a khugepaged pass through gup_test, which
makes the refusal deterministic instead of a race.  Check the same three
things the lazyfree case does: the rest of the table collapsed, that
window not, and the pinned page still backed by an order-0 folio, since
selection descends orders on a refusal.

Skipped without CONFIG_GUP_TEST or the privilege to use it.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c | 64 +++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index 58cb7364652e..bd684cd25fed 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -21,7 +21,9 @@
 
 #include "linux/magic.h"
 
+#include <sys/ioctl.h>
 #include "vm_util.h"
+#include "../../../../mm/gup_test.h"
 #include "hugepage_settings.h"
 
 #define BASE_ADDR ((void *)(1UL << 30))
@@ -1728,6 +1730,67 @@ static void collapse_order_lazyfree_window(struct collapse_context *c,
 	ksft_test_result_report(exit_status, "%s\n", __func__);
 }
 
+/*
+ * A GUP pin on one page holds a reference the collapse cannot account for, so
+ * khugepaged must leave the window holding it alone -- and collapse the rest.
+ * The pin is held across the whole khugepaged pass, which is what makes the
+ * refusal deterministic rather than a race.
+ */
+static void collapse_order_pinned_window(struct collapse_context *c,
+					 struct mem_ops *ops)
+{
+	struct pin_longterm_test pin = {};
+	size_t window = mthp_window_size();
+	void *p;
+	int fd;
+
+	fd = open("/sys/kernel/debug/gup_test", O_RDWR);
+	if (fd < 0) {
+		ksft_test_result_skip("%s: gup_test needs CONFIG_GUP_TEST and root\n",
+				      __func__);
+		return;
+	}
+
+	mthp_push_target_order();
+
+	p = ops->setup_area(1);
+	ops->fault(p, 0, hpage_pmd_size);
+	if (!window_not_collapsed(p, hpage_pmd_size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	pin.addr = (__u64)(unsigned long)p;
+	pin.size = page_size;
+	pin.flags = PIN_LONGTERM_TEST_FLAG_USE_WRITE;
+	if (ioctl(fd, PIN_LONGTERM_TEST_START, &pin)) {
+		ops->cleanup_area(p, hpage_pmd_size);
+		close(fd);
+		thp_pop_settings();
+		ksft_test_result_skip("%s: cannot pin\n", __func__);
+		return;
+	}
+
+	madvise(p, hpage_pmd_size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse the windows beside a pinned page...");
+	if (!khugepaged_wait_full_pass())
+		fail("Timeout");
+	else if (window_collapsed(p + window, hpage_pmd_size - window) &&
+		 window_not_collapsed(p, window) &&
+		 /* Left alone at every order, not just the target one */
+		 is_range_backed_by_folio_orders(p, page_size, 0,
+						 pagemap_fd, kpageflags_fd))
+		success("OK");
+	else
+		fail("Fail");
+
+	ioctl(fd, PIN_LONGTERM_TEST_STOP);
+	close(fd);
+
+	validate_memory(p, 0, hpage_pmd_size);
+	ops->cleanup_area(p, hpage_pmd_size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 static void usage(void)
 {
 	fprintf(stderr, "\nUsage: ./khugepaged [OPTIONS] <test type> [dir]\n\n");
@@ -2020,6 +2083,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_sub_pmd_holes, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_mlocked, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_lazyfree_window, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_pinned_window, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
-- 
2.54.0

[RFC PATCH 56/57] selftests/mm: cover the scaled max_ptes_shared limit

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

max_ptes_shared is written in PTEs of a whole PMD, but a range smaller
than a PMD is scanned in full and judged by the same setting.  Compared
raw the limit is then unreachable: on arm64 with 64K base pages a 2M range
holds 32 PTEs and can never exceed a 4096-PTE budget.  So it is scaled to
what was actually scanned, and what decides is the shared fraction rather
than the count.

Cover both directions in one sub-PMD range, with a forked child holding
the sources shared.  One PTE past the scaled budget: nothing in the range
collapses.  Break CoW on one more page, bringing it inside: it collapses.

The counts come from the current max_ptes_shared rather than being
hardcoded, so the case follows the setting and the order it runs at.

Without the scaling the first half passes wrongly: a few dozen shared PTEs
never reach a limit expressed in PMD units, so the range collapses when it
should not.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c | 85 +++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index bd684cd25fed..dde24756a3a2 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1791,6 +1791,90 @@ static void collapse_order_pinned_window(struct collapse_context *c,
 	ksft_test_result_report(exit_status, "%s\n", __func__);
 }
 
+/*
+ * max_ptes_shared counts PTEs of a whole PMD, but a VMA smaller than one is
+ * scanned in full and judged against the same setting, so the limit is scaled
+ * to the range actually scanned: what decides is the shared *fraction*, not
+ * the raw count. An unscaled comparison against HPAGE_PMD_NR/2 could never
+ * refuse a range this small, so both directions are checked here.
+ */
+static void collapse_order_sub_pmd_shared(struct collapse_context *c,
+					  struct mem_ops *ops)
+{
+	size_t window = mthp_window_size();
+	size_t size = 4 * window;
+	unsigned long nr_ptes = size / page_size;
+	unsigned long budget, cow;
+	int max_shared, wstatus;
+	void *p;
+
+	/*
+	 * The range has to stay strictly below a PMD to say anything about the
+	 * scaling: at exactly one PMD the scaled limit is the raw one, and the
+	 * case would pass without testing what it is here for.
+	 */
+	if (size >= hpage_pmd_size) {
+		ksft_test_result_skip("%s: four windows do not fit below the PMD\n",
+				      __func__);
+		return;
+	}
+
+	max_shared = thp_read_num("khugepaged/max_ptes_shared");
+	/* The same fraction of this range as max_shared is of a PMD. */
+	budget = (unsigned long)max_shared * nr_ptes / hpage_pmd_nr;
+	if (budget + 1 > nr_ptes) {
+		ksft_test_result_skip("%s: max_ptes_shared leaves nothing to exceed\n",
+				      __func__);
+		return;
+	}
+
+	mthp_push_target_order();
+
+	p = mmap(BASE_ADDR, size, PROT_READ | PROT_WRITE,
+		 MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+	if (p != BASE_ADDR)
+		ksft_exit_fail_msg("Failed to allocate VMA at %p\n", BASE_ADDR);
+	fill_memory(p, 0, size);
+	madvise(p, size, MADV_HUGEPAGE);
+
+	if (!fork()) {
+		/*
+		 * Everything is shared with the parent now. Break CoW on all
+		 * but budget + 1 PTEs: one PTE over the scaled limit, and far
+		 * below the unscaled one.
+		 */
+		cow = nr_ptes - budget - 1;
+		fill_memory(p, 0, cow * page_size);
+		ksft_print_msg("Refuse a sub-PMD range over the scaled max_ptes_shared...");
+		if (!khugepaged_wait_full_pass())
+			fail("Timeout");
+		else if (window_not_collapsed(p, size))
+			success("OK");
+		else
+			fail("Fail");
+
+		/* One fewer shared PTE brings it back within the limit. */
+		fill_memory(p, cow * page_size, (cow + 1) * page_size);
+		ksft_print_msg("Collapse once inside it...");
+		if (!khugepaged_wait_full_pass())
+			fail("Timeout");
+		else if (window_collapsed(p, size))
+			success("OK");
+		else
+			fail("Fail");
+
+		validate_memory(p, 0, size);
+		_exit(exit_status);
+	}
+	wait(&wstatus);
+	if (WEXITSTATUS(wstatus))
+		exit_status = WEXITSTATUS(wstatus);
+
+	munmap(p, size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 static void usage(void)
 {
 	fprintf(stderr, "\nUsage: ./khugepaged [OPTIONS] <test type> [dir]\n\n");
@@ -2084,6 +2168,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_mlocked, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_lazyfree_window, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_pinned_window, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_sub_pmd_shared, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
-- 
2.54.0

[RFC PATCH 57/57] MAINTAINERS: add an entry for collapse

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Collapse is now its own body of code: mm/collapse.c, a header stating what
it offers a caller, and its own tracepoint family.  It has two callers,
the khugepaged daemon and MADV_COLLAPSE, rather than living inside one of
them.

Add a section for it, and move what belongs to it out of the THP entry:
khugepaged.c and its header, the trace header, the selftests, and
mm_slot.h.

The selftests are matched as khugepaged*.c, which also covers the race and
sync-check harnesses the THP entry did not list.  mm_slot.h stays listed
under KSM as well, which is the other user of it.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 MAINTAINERS | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/MAINTAINERS b/MAINTAINERS
index 7c179b333e4e..4e4e5030b979 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -16955,6 +16955,20 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
 F:	include/linux/balloon.h
 F:	mm/balloon.c
 
+MEMORY MANAGEMENT - COLLAPSE
+M:	Kiryl Shutsemau <kas@kernel.org>
+L:	linux-mm@kvack.org
+S:	Maintained
+W:	http://www.linux-mm.org
+T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
+F:	include/linux/khugepaged.h
+F:	include/trace/events/collapse.h
+F:	mm/collapse.c
+F:	mm/collapse.h
+F:	mm/khugepaged.c
+F:	mm/mm_slot.h
+F:	tools/testing/selftests/mm/khugepaged*.c
+
 MEMORY MANAGEMENT - CORE
 M:	Andrew Morton <akpm@linux-foundation.org>
 M:	David Hildenbrand <david@kernel.org>
@@ -17272,12 +17286,7 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
 F:	Documentation/ABI/testing/sysfs-kernel-mm-transparent-hugepage
 F:	Documentation/admin-guide/mm/transhuge.rst
 F:	include/linux/huge_mm.h
-F:	include/linux/khugepaged.h
-F:	include/trace/events/collapse.h
 F:	mm/huge_memory.c
-F:	mm/khugepaged.c
-F:	mm/mm_slot.h
-F:	tools/testing/selftests/mm/khugepaged.c
 F:	tools/testing/selftests/mm/split_huge_page_test.c
 F:	tools/testing/selftests/mm/transhuge-stress.c
 
-- 
2.54.0

Re: [RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[Zi Yan]

On Sun Aug 16, 2026 at 6:45 PM EDT, Kiryl Shutsemau wrote:
> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
>
> Yes, I know, this is a lot of changes. But I'm happy with the overall state
> of the patchset and the only reason I tag it as RFC is that it is tricky
> to get 57 patches upstream.
>
> I wanted to give a view of the end state first. I will suggest a possible
> way to split it below.
>
> I would appreciate any feedback.
>
> TL;DR
> =====
>
> This replaces khugepaged's anonymous collapse with an engine that
> can collapse sub-PMD ranges. It is built around migration entries and
> frozen folios instead of heavy locking and isolation, aiming for better
> scalability and less disruption to the workload being collapsed.
>
> Why
> ===
>
> mTHP collapse landed in khugepaged in 7.2 and I was glad to see it.  We
> at Meta run arm64 with 64K base pages, where a PMD is 512M: PMD-order THP
> is of limited use at that size, and mTHP is exactly what we want.
>
> It turned out not to help us.
>
> khugepaged only ever looks at PMD-aligned windows, and it is not an easy
> limitation to lift.
>
> Fixing the alignment is a one-line change, but what it feeds assumes the
> PMD everywhere that matters: collapse_huge_page() clears and flushes the
> whole PMD whatever order it is collapsing, installs a PMD leaf because
> that is the only thing it can produce, and keeps everyone out with
> mmap_write_lock, anon_vma_lock_write() and an IPI broadcast while it
> does.
>
> Which is why hugepage_vma_revalidate() demands that the VMA span the
> whole PMD even for an mTHP order -- "we'd need to lock all VMAs in the
> PMD range to support this", as the comment there puts it.  A PMD-granular
> operation is only safe when one VMA owns the PMD, and that is exactly the
> restriction in the way.  The alignment is the symptom; the PMD is the
> design.
>
> So both roots have to go.

I agree that khugepaged is designed for PMD-aligned collapse and this is
a limitation we want to get rid of. It is great you are looking at them.

>
> Design
> ======
>
> The old mechanism holds the address space still because it has nothing
> else stopping the sources from moving under the copy.  The new engine
> makes the sources themselves inert instead, with the two barriers
> migration already uses, raised in that order:
>
>   1. migration entries replace the source PTEs.  Faults and GUP-slow
>      now wait on the source folio's lock, which is taken before the
>      first entry becomes visible.
>   2. the source folio's refcount is frozen to its expected value.
>      GUP-fast, pfn walkers, reclaim, compaction and memory-failure all
>      fail folio_try_get() and back off.
>
> Between the two, nothing can reach a source, so the copy runs with no
> lock held at all -- and the address space is left alone while it does.
>
> What that removes from every collapse path:
>
>   mmap_write_lock              -> mmap_read
>   anon_vma_lock_write()        -> nothing: an rmap walk needs the folio
>                                   locked, and the engine holds that lock
>                                   from freeze to putback
>   tlb_remove_table_sync_one()  -> nothing: one ranged flush per round
>   LRU isolation                -> nothing: sources are inert in place

I remember we were discussing using migration entry and the issue with
mmap_write_lock() in the context of in-place THP promotion and the
conclusion was that because MADV_DONTNEED (maybe MADV_REMOVE or
MADV_PAGEOUT) works on page table and does not change VMAs,
mmap_write_lock() is needed to prevent things being changed under
khugepaged. Anything different in normal khugepaged collapse process, so
that it is OK to use mmap_read_lock? Let me know if I misremember it.

Thanks.

-- 
Best Regards,
Yan, Zi

Re: [RFC PATCH 57/57] MAINTAINERS: add an entry for collapse

[Lorenzo Stoakes (ARM)]

On Sun, Aug 16, 2026 at 11:46:09PM +0100, Kiryl Shutsemau wrote:
> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
>
> Collapse is now its own body of code: mm/collapse.c, a header stating what
> it offers a caller, and its own tracepoint family.  It has two callers,
> the khugepaged daemon and MADV_COLLAPSE, rather than living inside one of
> them.
>
> Add a section for it, and move what belongs to it out of the THP entry:
> khugepaged.c and its header, the trace header, the selftests, and
> mm_slot.h.
>
> The selftests are matched as khugepaged*.c, which also covers the race and
> sync-check harnesses the THP entry did not list.  mm_slot.h stays listed
> under KSM as well, which is the other user of it.
>
> Assisted-by: Claude-Code:claude-opus-5
> Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
> ---
>  MAINTAINERS | 19 ++++++++++++++-----
>  1 file changed, 14 insertions(+), 5 deletions(-)
>
> diff --git a/MAINTAINERS b/MAINTAINERS
> index 7c179b333e4e..4e4e5030b979 100644
> --- a/MAINTAINERS
> +++ b/MAINTAINERS
> @@ -16955,6 +16955,20 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
>  F:	include/linux/balloon.h
>  F:	mm/balloon.c
>
> +MEMORY MANAGEMENT - COLLAPSE
> +M:	Kiryl Shutsemau <kas@kernel.org>
> +L:	linux-mm@kvack.org
> +S:	Maintained
> +W:	http://www.linux-mm.org
> +T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
> +F:	include/linux/khugepaged.h
> +F:	include/trace/events/collapse.h
> +F:	mm/collapse.c
> +F:	mm/collapse.h
> +F:	mm/khugepaged.c
> +F:	mm/mm_slot.h
> +F:	tools/testing/selftests/mm/khugepaged*.c
> +

I'll look through the series properly, but presumably this is a separation of
the THP collapse logic, and therefore should have the same maintainers/reviewers
as the rest of THP.

>  MEMORY MANAGEMENT - CORE
>  M:	Andrew Morton <akpm@linux-foundation.org>
>  M:	David Hildenbrand <david@kernel.org>
> @@ -17272,12 +17286,7 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
>  F:	Documentation/ABI/testing/sysfs-kernel-mm-transparent-hugepage
>  F:	Documentation/admin-guide/mm/transhuge.rst
>  F:	include/linux/huge_mm.h
> -F:	include/linux/khugepaged.h
> -F:	include/trace/events/collapse.h
>  F:	mm/huge_memory.c
> -F:	mm/khugepaged.c
> -F:	mm/mm_slot.h
> -F:	tools/testing/selftests/mm/khugepaged.c
>  F:	tools/testing/selftests/mm/split_huge_page_test.c
>  F:	tools/testing/selftests/mm/transhuge-stress.c
>
> --
> 2.54.0
>

--
Cheers, Lorenzo

Re: [RFC PATCH 57/57] MAINTAINERS: add an entry for collapse

[David Hildenbrand (Arm)]

On 8/17/26 10:04, Lorenzo Stoakes (ARM) wrote:
> On Sun, Aug 16, 2026 at 11:46:09PM +0100, Kiryl Shutsemau wrote:
>> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
>>
>> Collapse is now its own body of code: mm/collapse.c, a header stating what
>> it offers a caller, and its own tracepoint family.  It has two callers,
>> the khugepaged daemon and MADV_COLLAPSE, rather than living inside one of
>> them.
>>
>> Add a section for it, and move what belongs to it out of the THP entry:
>> khugepaged.c and its header, the trace header, the selftests, and
>> mm_slot.h.
>>
>> The selftests are matched as khugepaged*.c, which also covers the race and
>> sync-check harnesses the THP entry did not list.  mm_slot.h stays listed
>> under KSM as well, which is the other user of it.
>>
>> Assisted-by: Claude-Code:claude-opus-5
>> Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
>> ---
>>  MAINTAINERS | 19 ++++++++++++++-----
>>  1 file changed, 14 insertions(+), 5 deletions(-)
>>
>> diff --git a/MAINTAINERS b/MAINTAINERS
>> index 7c179b333e4e..4e4e5030b979 100644
>> --- a/MAINTAINERS
>> +++ b/MAINTAINERS
>> @@ -16955,6 +16955,20 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
>>  F:	include/linux/balloon.h
>>  F:	mm/balloon.c
>>
>> +MEMORY MANAGEMENT - COLLAPSE
>> +M:	Kiryl Shutsemau <kas@kernel.org>
>> +L:	linux-mm@kvack.org
>> +S:	Maintained
>> +W:	http://www.linux-mm.org
>> +T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
>> +F:	include/linux/khugepaged.h
>> +F:	include/trace/events/collapse.h
>> +F:	mm/collapse.c
>> +F:	mm/collapse.h
>> +F:	mm/khugepaged.c
>> +F:	mm/mm_slot.h
>> +F:	tools/testing/selftests/mm/khugepaged*.c
>> +
> 
> I'll look through the series properly, but presumably this is a separation of
> the THP collapse logic, and therefore should have the same maintainers/reviewers
> as the rest of THP.

Jup.

-- 
Cheers,

David

Re: [RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[Lorenzo Stoakes (ARM)]

+cc Pedro for discussion about perf numbers

On Sun, Aug 16, 2026 at 11:45:12PM +0100, Kiryl Shutsemau wrote:
> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
>
> Yes, I know, this is a lot of changes. But I'm happy with the overall state
> of the patchset and the only reason I tag it as RFC is that it is tricky
> to get 57 patches upstream.
>
> I wanted to give a view of the end state first. I will suggest a possible
> way to split it below.
>
> I would appreciate any feedback.

Oh Kiryl :)

We have a THP cabal meeting every couple of weeks where it would have been
useful for you to raise this first.

In any case - this series is not something we'd consider at the moment,
even broken into parts.

David and I have put THP into feature freeze - until the codebase is
subtantially improved we're not really interested in seeing significant
development.

The technical debt is substantial and has to be paid down first.

See [0] for a rough list of TODOs in this regard.

(I also intend to do at least 1 series that improves things shortly
myself.)

mTHP khugepaged was allowed in _despite_ my being reticent given this
technical debt, but with a proviso of a feature freeze afterwards.

It's still fine obviously to send a theoretical RFC for early feedback. But
there can be no un-RFC'ing any time soon.

I'd suggest looking at cleanups that could act as foundations for your
changes.

I'll try to find some time to look through the actual patches to make
concrete suggestions/see if we could break some of this out like that.

Some more comments below.

>
> TL;DR
> =====
>
> This replaces khugepaged's anonymous collapse with an engine that
> can collapse sub-PMD ranges. It is built around migration entries and
> frozen folios instead of heavy locking and isolation, aiming for better
> scalability and less disruption to the workload being collapsed.

OK interesting :)

>
> Why
> ===
>
> mTHP collapse landed in khugepaged in 7.2 and I was glad to see it.  We
> at Meta run arm64 with 64K base pages, where a PMD is 512M: PMD-order THP
> is of limited use at that size, and mTHP is exactly what we want.

Do you have some numbers that indicate to what degree mTHP khugepaged is
beneficial?

I was discussing this with Pedro recently and he pointed out that we have
very poor coverage of the actual observed benefits of mTHP khugepaged
(there are reasons to have merged it anyway but still).

Some detailed analysis on how it benefits Meta would be really beneficial,
with specifics and what settings you find useful.

Is it only useful for larger page size? What orders have you found to be
most effective?

Is it useful in itself or do you view it as a stepping stone to further
work, like this?

>
> It turned out not to help us.

I mean maybe this answers my question :)

>
> khugepaged only ever looks at PMD-aligned windows, and it is not an easy
> limitation to lift.

Yes. This assumption is very much baked in.

I guess this is coming from the perspective of having ranges that are
neither PMD-aligned nor sized (far harder to achieve with 512 MiB PMD size
obviously).

>
> Fixing the alignment is a one-line change, but what it feeds assumes the

Hmm not so sure about that... especially given how baked in these
assumptions are.

Which by the way, all speaks to the need for rework.

The first stage in my view would be to improve the code to the point that
these kinds of assumptions fall out of it, which then lays the foundations
for future changes to eliminate the assumptions.

> PMD everywhere that matters: collapse_huge_page() clears and flushes the
> whole PMD whatever order it is collapsing, installs a PMD leaf because
> that is the only thing it can produce, and keeps everyone out with
> mmap_write_lock, anon_vma_lock_write() and an IPI broadcast while it
> does.

To be clear - the anon path. I think important to clarify :)

And yeah it does IPI for any sensible arch (with
CONFIG_MMU_GATHER_RCU_TABLE_FREE) via tlb_remove_table_sync_one(). The
other arches IPI anyway on TLB invalidation.

[Though I intend to make all page table freeing RCU relatively soon which
should? Eliminate the need for this, possibly?]

It doesn't insert a PMD leaf entry for mTHP collapse though, that's
incorrect.

It does, as you say, clear down the PMD even on mTHP PTE range
installation, and has to hold the rmap lock for a longer period to account
for adjacent ranges that share a PMD entry.

All in all it's a mess yes and clearly can be improved.

>
> Which is why hugepage_vma_revalidate() demands that the VMA span the
> whole PMD even for an mTHP order -- "we'd need to lock all VMAs in the
> PMD range to support this", as the comment there puts it.  A PMD-granular
> operation is only safe when one VMA owns the PMD, and that is exactly the
> restriction in the way.  The alignment is the symptom; the PMD is the
> design.
>
> So both roots have to go.

Yeah this is pretty radical stuff so it's just an absolute no until
technical debt is paid down properly.

That means signficant rework to make this code something far from the
overly-coupled spaghetti mess that it is now.

Efforts have been ongoing with this, Nico and others have submitted
cleanups and obviously more is required here.

But yes I think decoupling the ossified PMD assumption is important.

However we have to remember that a lot of the user-visible API assumes PMD
sizing and so the code has to clearly reflect this and make it clear that

>
> Design
> ======
>
> The old mechanism holds the address space still because it has nothing
> else stopping the sources from moving under the copy.  The new engine
> makes the sources themselves inert instead, with the two barriers
> migration already uses, raised in that order:
>
>   1. migration entries replace the source PTEs.  Faults and GUP-slow
>      now wait on the source folio's lock, which is taken before the
>      first entry becomes visible.
>   2. the source folio's refcount is frozen to its expected value.
>      GUP-fast, pfn walkers, reclaim, compaction and memory-failure all
>      fail folio_try_get() and back off.
>
> Between the two, nothing can reach a source, so the copy runs with no
> lock held at all -- and the address space is left alone while it does.

Hmm, are migration entries the right mechnanism here? Are you actually
migrating the pages to a large folio here, or using them to get the
behaviour you want on fault/GUP?

Same question in general for the freezing.

I'm not suggesting quite that these are the wrong mechanisms, just querying
how they are being used here.

I do like the idea of using the folio lock though. All rmap operations now
take the folio lock when traversing the rmap.

>
> What that removes from every collapse path:
>
>   mmap_write_lock              -> mmap_read
>   anon_vma_lock_write()        -> nothing: an rmap walk needs the folio
>                                   locked, and the engine holds that lock
>                                   from freeze to putback

I do like the idea of eliminating uses of the rmap lock like this, not only
for contention's sake but also for scalable CoW purposes which introduces
challenges with regards to holding these.

In fact, migration and huge memory collapse are the really problematic
areas.

>   tlb_remove_table_sync_one()  -> nothing: one ranged flush per round

Aren't we reliant upon this synchronisation for correctness?

>   LRU isolation                -> nothing: sources are inert in place
>
> Working in windows rather than whole PMDs takes care of the other root.
> A sub-PMD window is collapsed under the page table lock, so a collapse
> disturbs only the window it collapses, and each candidate is validated
> at its own order -- a window need only fit its own VMA.  A PMD-order
> candidate still has to own the whole PMD, which is the old rule kept
> where it is still needed.
>
> Candidates are carried through the passes a batch at a time rather than
> one window at a time, so a round pays for its flush and its lock
> acquisitions once.
>
> With the barriers holding the sources still, which read lock the engine
> takes stops being part of the design.  A round works inside a single
> VMA, so patches 43-49 switch it from mmap_read to per-VMA locking: an
> mmap_write elsewhere in the mm then stops waiting for a collapse that
> has nothing to do with it.  That block is the only part of the series
> that needs per-VMA locking to be unconditional, and it is a separate
> dependency (see below); everything before it runs under mmap_read and
> does not care.
>
> Patch 7 sketches the engine as a comment naming every pass, what lock it
> takes and what it may sleep on; the details are there rather than here.
>
> What falls out beyond the lock diet:
>
>  - mTHP collapse in VMAs smaller than a PMD, which is the arm64 case
>    above: a 2M VMA on an arm64/64K machine collapses nothing today at
>    any order, and collapses to mTHP here.

I do worry about knock-ons from this. Has to be carefully checked.

>  - Hole and zeropage population at every order, so partially populated
>    windows collapse to mTHP under the same max_ptes_none policy as PMD.
>  - Sources come in spans -- any stretch of consecutive PTEs mapping
>    consecutive pages of one folio -- so partially mapped and scrambled
>    compound sources (the PTE-mapped-THP re-collapse class) work at
>    every order.
>  - A table that cannot become one huge page still yields the largest
>    windows inside it, where before a single disqualified PTE gave up
>    the whole table.

Are you permitting collapse of ranges that straddle PTEs?

Though in general I'm confused by the single disqualified PTE here -

>
> Reading the series
> ==================
>
> 57 patches is a lot to land on a list.  They go in blocks:
>
>   1-6    helpers and shared state: pte_folio(), pte_none_or_zero(),
>          mm/collapse.h, and the policy that replaces asking whether
>          khugepaged started a collapse
>   7-8    the engine's shape: entry points, a call-tree comment naming
>          every pass, and the scan filled in
>   9-23   the collapse half, top down: the round frame, then each pass
>          in turn, then selection and the retry store
>   24     per-candidate tracing, before the switch takes the old
>          tracepoints away
>   25-28  the switch: point the anon path at the engine, widen coverage
>          to sub-PMD VMAs, delete the mechanism it replaces
>   29-35  move what is left of collapse out of khugepaged.c, and
>          MADV_COLLAPSE into madvise.c
>   36-42  tracing: the engine's own events and trace header
>   43-49  per-VMA locking, and the mm reference that makes it safe
>   50-56  selftests for what the engine can now do
>   57     MAINTAINERS

Yeah as I replied there the MAINTAINERS change is totally unacceptable :)

This is THP code and belongs to the THP maintainers/reviewers.

No coup d'etat please :)

Also reviewership/maintainership in this area is predicated on significant
review contributions which are badly needed in THP. So effort in this
respect is appreciated too :)

And come to a THP cabal meeting to discuss your changes with us please!
We're friendly :)

>
> The two patches worth reading first if you read nothing else are 7 (the
> design, as a comment naming the whole call tree) and 16 (the freeze,
> which is where the safety argument lives).
>
> A possible split, if that helps:
>
>   1-2    two mm helpers, pte_folio() and pte_none_or_zero().  Both
>          convert callers outside collapse and are useful on their own
>   3-27   the engine and the switch-over.  This is the smallest unit
>          that does anything: stop earlier and the tree carries an
>          engine nothing calls
>   28     remove the mechanism the engine replaces
>   29-42  moving what is left of collapse out of khugepaged.c, and the
>          engine's own tracepoints
>   43-49  per-VMA locking
>   50-57  selftests and MAINTAINERS
>
> Keeping the removal separate leaves both engines in the tree with only
> the new one reachable, so the switch can be reverted on its own if
> something turns up.  The old mechanism is already carried that way for
> three patches inside the series, so this costs nothing but 975 lines of
> unreferenced code until 28 lands.  That safety net only lasts until the
> blocks after it land, though: once collapse has moved out of
> khugepaged.c and the locking has changed, reverting the switch no longer
> gives back a working old engine.

Again as above, no to these changes in any form at the moment.

The techical debt must be paid down first.

>
> Base and dependencies
> =====================
>
> This applies on the selftests series, not on plain mm-new:
>
>   [PATCH v4 00/19] selftests/mm: improve khugepaged coverage
>   https://lore.kernel.org/all/20260815015901.1236937-1-kirill@shutemov.name/
>
> which is on mm-new 33f61b12d297.
>
> Patches 43-49 depend on Suren's unconditional per-VMA locks:
>
>   [PATCH v6 0/5] mm: Unconditional per-VMA locks and cleanups
>   https://lore.kernel.org/all/20260813193433.3318288-1-surenb@google.com/
>
> That series is not in mm-new yet, and with patch 46 applied SMP=n does
> not build without it: lock_next_vma() is behind CONFIG_PER_VMA_LOCK in
> mmap_lock.h.  Everything up to patch 42 builds and runs on mm-new as it
> stands.  There is no fallback path by choice -- adding one would mean
> carrying two locking models through every pass.
>
> Both branches are available at
>
>   git://git.kernel.org/pub/scm/linux/kernel/git/kas/linux.git collapse/rfc-v1
>
> and the benchmark used for the numbers below, which is unposted and not a
> dependency, at
>
>   git://git.kernel.org/pub/scm/linux/kernel/git/kas/linux.git perf/bench-usemem

As above.

>
> Performance
> ===========
>
> Measuring khugepaged is awkward.  It is a background daemon, so what
> matters is what a workload feels while it runs, not what the daemon
> reports about itself -- and the usual coverage instrument is no help
> below the PMD: smaps AnonHugePages only counts PMD-order folios, so it
> reads zero however much mTHP has been collapsed.
>
> So I wrote "perf bench mem usemem" for this.  It touches a region while
> khugepaged works on it and reports the workload's own latency
> percentiles and throughput, against per-size counters that can see
> sub-PMD folios.  The branch is above; it is unposted and not a
> dependency.
>
> x86-64, production configs (no KASAN, no lockdep, no DEBUG_VM, no
> PAGE_TABLE_CHECK), interleaved rounds on an idle host, equal work on
> every arm.
>
> I measured three kernels, so the two halves of the series can be told
> apart in the numbers below:
>
>   A   the base
>   B   the new engine, still under mmap_read
>   C   B plus per-VMA locking -- what this series ends up with
>
> The engine: a sub-PMD collapse stops blanking the surrounding 2M
> -----------------------------------------------------------------
>
> base routes sub-PMD collapse through collapse_huge_page(), whose
> pmdp_collapse_flush() and tlb_remove_table_sync_one() are not gated on
> order: to collapse an order-4 window of 16 pages it clears and flushes
> the whole 512-page PMD and IPIs, then repopulates.  The engine does the
> window under the PTL.
>
> A thread reading and writing a 32G region while it is collapsed at
> order-4, 16384 collapses on every arm:
>
>                          A        B        C
>   read p99 (ns)       3071     1023     1023   -66.7%
>   write p99 (ns)      3071      927      927   -69.8%
>
> and the workload's read rate rises by 68% on both engine kernels.
>
> B == C, so this is the engine, not the locking.
>
> At PMD order the same workload is flat, and that is expected rather
> than disappointing: it is the one configuration where both mechanisms
> disturb exactly the same 2M.  Read it as no regression at PMD order.
>
> Per-VMA locking: address-space operations stop waiting on the scan
> -------------------------------------------------------------------
>
> MADV_HUGEPAGE/MADV_NOHUGEPAGE toggling against a scanning mm, which is
> what jemalloc does with its arenas.  4096 collapses on every arm:
>
>                          A        B        C
>   ops/sec           340656   340820   603305   +77%
>   p99 (ns)           77823    86015     3327   -96%
>   p99.9 (ns)         86015    94207     9215   -89%
>
> B is about 11% worse than base at p99 here, consistently across runs:
> the engine alone slightly worsens hint-toggle latency, and per-VMA
> locking is what turns it into a win.  Both halves are in this series, so
> C is what a reviewer gets, but the middle column is the honest one.
>
> The trade is real in the other direction too.  On settled memory with
> nothing to collapse and scan_sleep_millisecs=0, per-VMA locking costs
> about 47% of scan throughput against one mmap_read for the whole walk.
> That is a synthetic worst case -- the daemon wraps 8000 times a second
> there, where production defaults to 10s between passes -- and it buys
> mmap/munmap p99 of 56us against 1.4us.
>
> Collapse itself is not slower
> -----------------------------
>
> One complete pass over a 32G region, 16384 collapses, khugepaged CPU
> from /proc/<pid>/stat, 7 repeats:
>
>   A base    median 5760 ms   spread 12.7%
>   B engine  median 4990 ms   spread  3.8%
>   C pervma  median 5020 ms   spread 12.4%
>
> The base arm is bimodal, so its median moves with sampling and the
> percentage is soft.  The distribution-free statement is better: every
> engine run used less CPU than every base run.
>
> The engine also allocates one destination per folio installed, where
> base allocates 5.25 and frees the rest again: nothing is allocated until
> the sources are frozen and the collapse can no longer be refused.
>
> A measurement note, since an earlier version of this series quoted worse
> figures.  khugepaged CPU has to be measured per collapse or per
> completed pass, never over a fixed window with scan_sleep_millisecs=0:
> the daemon never sleeps, so whichever kernel finishes the work sooner
> spends the rest of the window scanning settled memory and is charged for
> it.  Measured that way the engine appeared to cost 10% more CPU;
> measured per unit of work it costs less.
>
> Costs
> =====
>
> At PMD order the engine issues two TLB flushes per collapse where the
> old mechanism issues one: the freeze's ranged flush plus the terminal
> layer's pmdp_collapse_flush().  A PMD candidate is alone in its round,
> so nothing amortizes the first.  Dropping the old per-collapse
> tlb_remove_table_sync_one() IPI presumably pays for it, but that was not
> measured and is not claimed here.
>
> There may be a way out -- a PMD migration entry over the table during
> the window, so the CPU never caches a walk to shoot down -- but that
> means teaching every pmd-level walker a new kind of entry, and I have
> not tried it.

Hmm this seems like complexity on top of complexity...

>
> PMD collapse deposits a freshly allocated page table instead of
> redepositing the detached one.  Whoever withdraws a deposited table
> frees it immediately, with nothing to hold a lockless walker off first,
> and under a read lock the detached table may still be traversed by
> GUP-fast or an RCU pte walk.  It goes to pte_free_defer() instead,
> exactly as retract_page_tables() does.  One transient table page per PMD
> collapse buys the IPI's absence.
>
> That cost goes away if zap_deposited_table() -- the only site that frees
> a deposited table outright, the others redeposit it or repopulate the
> PMD with it -- used pte_free_defer().  The deposit would no longer have
> to be quiescent and the detached table could go straight back.  It would
> defer every THP zap's table free, and I have not tried it.
>
> A shared source now costs an extra copy.  The freeze needs every page
> exclusive to this mm, so the fault-in pass breaks CoW first -- an
> allocation and a copy -- and the collapse then copies that page into the
> destination; the old mechanism copied a shared page straight into the
> new folio and broke the sharing that way.  It is bounded by
> max_ptes_shared, which khugepaged holds at zero below the PMD order, so
> in practice this is PMD-order collapse and MADV_COLLAPSE.

Hmm I guess better to be explicit.

>
> Size
> ====
>
> mm/ grows by 1915 lines net: 4475 added against 2560 deleted.
>
> That is not a claim that this is less code, but it is less than it
> looks.  khugepaged.c goes from 3283 lines to 908.  The new engine is
> 4052 lines across mm/collapse.c and mm/collapse.h, of which 1344 --
> about a third -- are comments, which is where the pipeline's invariants
> are written down.  What replaces three install paths with their own
> isolate/copy/rollback is one engine and one contract.

As above. The existing codebase is a mess and must be cleaned up first,
this isn't optional.

>
> Testing
> =======
>
> Both matrices run the mm selftests plus a race harness, on the
> validation config: KASAN, lockdep, PROVE_LOCKING, DEBUG_VM and
> PAGE_TABLE_CHECK, 16G of guest memory, swap active so the swap-in
> prepass is exercised rather than skipped.
>
>   x86-64        433 pass, 0 fail, 12 skip
>   arm64/64K     581 pass, 0 fail, 18 skip
>
> dmesg clean on both.  The arm64 skips are a pre-existing shmem
> MADV_COLLAPSE -EINVAL on 64K pages, confirmed against the base by A/B.
>
> Every one of the 57 patches builds with no new warnings; !NUMA and !MMU
> (arm nommu) build clean.  SMP=n does not build, for the reason in the
> dependencies section above.
>
> The race harness also gets longer soaks -- 1800s per driver mode, with
> memory pressure and swap -- and the engine is fuzzed with syzkaller on a
> KCOV+KASAN build.  That found two bugs the selftests could not reach: a
> teardown that dropped rmap while the source was still frozen, where
> removing an mlocked mapping munlocks and munlock_folio() takes a
> reference a frozen folio forbids; and a whole-table MADV_DONTNEED racing
> the copy window under CONFIG_PT_RECLAIM, which freed the table and left
> the sources frozen and locked.  Both are fixed, and both gained coverage
> -- the mlocked case is patch 53.

Thanks for the detailed explanation.

>
> Kiryl Shutsemau (Meta) (57):
>   mm: add pte_folio()
>   mm: add pte_none_or_zero()
>   mm/collapse: add collapse.h for the shared collapse state
>   mm/collapse: rename mthp_present_ptes to eligible_ptes
>   mm/collapse: state what a collapse may do in the policy
>   mm/collapse: move the smallest collapse order to collapse.h
>   mm/collapse: sketch the new anonymous collapse engine
>   mm/collapse: scan a table for what a collapse could use
>   mm/collapse: collect candidate windows into a round
>   mm/collapse: run a round and feed the outcomes back
>   mm/collapse: sketch the passes of a round
>   mm/collapse: allocate a destination per candidate
>   mm/collapse: revalidate a round against the VMA
>   mm/collapse: fault the sources in before the freeze
>   mm/collapse: check what a candidate would freeze
>   mm/collapse: freeze the sources behind migration entries
>   mm/collapse: copy the sources into the destinations
>   mm/collapse: install the destinations at PTE level
>   mm/collapse: install a PMD leaf as the terminal layer
>   mm/collapse: put the sources back
>   mm/collapse: settle whatever the round reached
>   mm/collapse: walk a table with a selection cursor
>   mm/collapse: give a refused region a second chance
>   mm/collapse: report each candidate's outcome to tracing
>   mm/collapse: collapse anonymous memory with the new engine
>   mm/collapse: give collapse_single_pmd() the range to work on
>   mm/collapse: scan the windows a VMA can hold
>   mm/collapse: remove the mechanism the engine replaces
>   mm/collapse: move what a collapse is judged on into collapse.c
>   mm/collapse: name the max_ptes ceiling after collapse
>   mm/khugepaged: count collapses where khugepaged makes them
>   mm/collapse: move the file collapse into collapse.c
>   mm/collapse: split collapse into a scan and a run
>   mm/collapse: implement MADV_COLLAPSE in madvise.c
>   mm/madvise: drop MADV_COLLAPSE's redundant mm reference
>   mm/collapse: report what the scan found
>   mm/collapse: report what the fault-in pass paid
>   mm/collapse: report the round, and what it made faulters wait
>   mm/collapse: name the file collapse's tracepoints after collapse
>   mm/collapse: remove the tracepoints of the mechanism that is gone
>   mm/collapse: give collapse its own trace header
>   mm/collapse: allow error injection into the freeze
>   mm/khugepaged: check the scan budget before the work, not after
>   mm/khugepaged: hold the address space open across a scan
>   mm/collapse: take a per-VMA read lock for the round
>   mm/khugepaged: scan under a per-VMA read lock
>   mm/madvise: collapse under a per-VMA read lock
>   mm/collapse: assert the mm reference the engine relies on
>   mm/khugepaged: drop the mmap_lock barrier from __khugepaged_exit()
>   selftests/mm: attribute collapses by candidate event alone
>   selftests/mm: cover collapse inside a sub-PMD VMA
>   selftests/mm: cover a hole-y window in a sub-PMD VMA
>   selftests/mm: cover collapse of mlocked ranges
>   selftests/mm: cover collapse beside a MADV_FREE'd page
>   selftests/mm: cover collapse beside a pinned page
>   selftests/mm: cover the scaled max_ptes_shared limit
>   MAINTAINERS: add an entry for collapse
>
>  MAINTAINERS                                   |   19 +-
>  fs/proc/task_mmu.c                            |    4 +-
>  include/linux/huge_mm.h                       |    9 -
>  include/linux/mm.h                            |   14 +
>  include/linux/pgtable.h                       |   17 +
>  .../events/{huge_memory.h => collapse.h}      |  176 +-
>  kernel/bpf/btf.c                              |    8 +-
>  mm/Makefile                                   |    2 +-
>  mm/collapse.c                                 | 3808 +++++++++++++++++
>  mm/collapse.h                                 |  244 ++
>  mm/hugetlb.c                                  |    8 +-
>  mm/khugepaged.c                               | 2711 +-----------
>  mm/madvise.c                                  |  251 +-
>  mm/migrate_device.c                           |    9 +-
>  mm/mremap.c                                   |    2 +-
>  tools/testing/selftests/mm/khugepaged.c       |  346 ++
>  tools/testing/selftests/mm/khugepaged_race.c  |   29 +-
>  .../selftests/mm/khugepaged_sync_check.c      |   65 +-
>  tools/testing/selftests/mm/vm_util.c          |    2 +-
>  19 files changed, 5022 insertions(+), 2702 deletions(-)
>  rename include/trace/events/{huge_memory.h => collapse.h} (60%)
>  create mode 100644 mm/collapse.c
>  create mode 100644 mm/collapse.h
>
>
> base-commit: 8b76faf42c5d342b3bf0b1fd97bdaf603ee57354
> --
> 2.54.0
>

--
Cheers, Lorenzo

[0]:https://docs.google.com/document/d/1dsAXvVtioR2Hj98E_Um6220jamdznWxYsyx8jO0aUGU/edit?usp=sharing

Re: [RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[Kiryl Shutsemau]

On Sun, Aug 16, 2026 at 10:02:53PM -0400, Zi Yan wrote:
> > What that removes from every collapse path:
> >
> >   mmap_write_lock              -> mmap_read
> >   anon_vma_lock_write()        -> nothing: an rmap walk needs the folio
> >                                   locked, and the engine holds that lock
> >                                   from freeze to putback
> >   tlb_remove_table_sync_one()  -> nothing: one ranged flush per round
> >   LRU isolation                -> nothing: sources are inert in place
> 
> I remember we were discussing using migration entry and the issue with
> mmap_write_lock() in the context of in-place THP promotion and the
> conclusion was that because MADV_DONTNEED (maybe MADV_REMOVE or
> MADV_PAGEOUT) works on page table and does not change VMAs,
> mmap_write_lock() is needed to prevent things being changed under
> khugepaged. Anything different in normal khugepaged collapse process, so
> that it is OK to use mmap_read_lock? Let me know if I misremember it.

IIRC that discussion predates Hugh's pte_offset_map() rework -- since
0d940a9b270b the helper takes rcu_read_lock() and fails if the pmd is
none, !present or huge, so mmap_write is no longer what keeps a pte
walker out.  MADV_DONTNEED is still not excluded, and the engine does not
try to: the install re-reads every slot under the ptl and publishes only
if it still holds the migration entry this round put there, leaving a
zapped slot alone and dropping the rmap the frozen source still held.

-- 
  Kiryl Shutsemau / Kirill A. Shutemov

Re: [RFC PATCH 57/57] MAINTAINERS: add an entry for collapse

[Kiryl Shutsemau]

On Mon, Aug 17, 2026 at 10:08:32AM +0200, David Hildenbrand (Arm) wrote:
> On 8/17/26 10:04, Lorenzo Stoakes (ARM) wrote:
> >> diff --git a/MAINTAINERS b/MAINTAINERS
> >> index 7c179b333e4e..4e4e5030b979 100644
> >> --- a/MAINTAINERS
> >> +++ b/MAINTAINERS
> >> @@ -16955,6 +16955,20 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
> >>  F:	include/linux/balloon.h
> >>  F:	mm/balloon.c
> >>
> >> +MEMORY MANAGEMENT - COLLAPSE
> >> +M:	Kiryl Shutsemau <kas@kernel.org>
> >> +L:	linux-mm@kvack.org
> >> +S:	Maintained
> >> +W:	http://www.linux-mm.org
> >> +T:	git git://git.kernel.org/pub/scm/linux/kernel/git/akpm/mm
> >> +F:	include/linux/khugepaged.h
> >> +F:	include/trace/events/collapse.h
> >> +F:	mm/collapse.c
> >> +F:	mm/collapse.h
> >> +F:	mm/khugepaged.c
> >> +F:	mm/mm_slot.h
> >> +F:	tools/testing/selftests/mm/khugepaged*.c
> >> +
> > 
> > I'll look through the series properly, but presumably this is a separation of
> > the THP collapse logic, and therefore should have the same maintainers/reviewers
> > as the rest of THP.
> 
> Jup.

Sure.  I just didn't want to volunteer anybody here.

-- 
  Kiryl Shutsemau / Kirill A. Shutemov

Re: [RFC PATCH 00/57] mm/collapse: rebuild collapse on migration primitives

[Kiryl Shutsemau]

On Mon, Aug 17, 2026 at 09:52:13AM +0100, Lorenzo Stoakes (ARM) wrote:
> We have a THP cabal meeting every couple of weeks where it would have been
> useful for you to raise this first.

Fair -- though my invite is on my old @linux.intel.com address.  Could
you forward it to kas@kernel.org?

> In any case - this series is not something we'd consider at the moment,
> even broken into parts.
> 
> David and I have put THP into feature freeze - until the codebase is
> subtantially improved we're not really interested in seeing significant
> development.
> 
> The technical debt is substantial and has to be paid down first.
> 
> See [0] for a rough list of TODOs in this regard.

I read the TODO list and I'll pick from it -- though I notice the
technical debt section includes "Literally all of the code in
mm/huge_memory.c and mm/khugepaged.c", which I'd argue this series is a
fairly committed attempt at :)

One clean up I wanted to do is consolidate code by functionality, not by
the THP/non-THP split.  Move all page fault handler code into mm/fault.c,
unmap code into mm/zap.c, fork's copying into mm/fork.c -- mirroring
kernel/fork.c, so the mm half of a subsystem sits under the same name.
Large folios are an integral part of mm nowadays and I don't think we
benefit from keeping THP in a separate file.  It is also an opportunity to
shift away from mm/memory.c being a kitchen sink.

David and I talked about this at LSF/MM.  I can give it a try if it fits
your idea of "feature freeze" -- and if it doesn't collide with the series
you have in flight, in which case I'd rather go after yours than around it.

> > Why
> > ===
> >
> > mTHP collapse landed in khugepaged in 7.2 and I was glad to see it.  We
> > at Meta run arm64 with 64K base pages, where a PMD is 512M: PMD-order THP
> > is of limited use at that size, and mTHP is exactly what we want.
> 
> Do you have some numbers that indicate to what degree mTHP khugepaged is
> beneficial?

Not from the fleet yet -- that experiment is still ahead of me, so I can't
give you order-by-order numbers.

What I can say is that on x86 we lean on khugepaged heavily to get THPs in
place; it is not a marginal contributor for us.  On arm64 with 64K base
pages we get nowhere near the x86 numbers without khugepaged being able to
produce mTHP at all.

I'll grant the other half of it, and more strongly than you put it: for a
64K mTHP on a 4K base page the TLB win is modest, and with today's
mechanism -- which clears and flushes the whole 2M PMD to install it -- I
can believe the disruption exceeds the gain and the net effect on a
workload is negative.  That is what I measured: a thread reading and
writing a region while khugepaged collapses it at order-4, read p99 3071ns
against 1023ns.  It shows the disruption is real and that it comes down;
whether the collapse pays for itself at that order is a separate question.

> > khugepaged only ever looks at PMD-aligned windows, and it is not an easy
> > limitation to lift.
> 
> Yes. This assumption is very much baked in.
> 
> I guess this is coming from the perspective of having ranges that are
> neither PMD-aligned nor sized (far harder to achieve with 512 MiB PMD size
> obviously).

Right, and it is the size rather than the alignment that bites.  The real
requirement is that a VMA contain a whole PMD-aligned, PMD-sized range,
which at 2M most anonymous mappings of any size manage and at 512M almost
none do.  That is why the limitation was easy to miss until the PMD got
big.

> > Fixing the alignment is a one-line change, but what it feeds assumes the
> 
> Hmm not so sure about that... especially given how baked in these
> assumptions are.

I think we agree -- that was the setup, not the claim.  Dropping the
ALIGN() is the one line; the point of the sentence is that it buys nothing
on its own, because what you then hand a sub-PMD range to still clears the
whole PMD and still demands the VMA span it.

> Which by the way, all speaks to the need for rework.
> 
> The first stage in my view would be to improve the code to the point that
> these kinds of assumptions fall out of it, which then lays the foundations
> for future changes to eliminate the assumptions.

For the plumbing, yes -- policy, file layout, the scan/run split all
improve by refactoring in place.

I don't think the locking model gets there that way, though, which is why
I built a second engine rather than morphing the first.  The old safety
argument is "hold the address space still"; the new one is "make the
sources inert".  They are not two points on a line -- mmap_write cannot go
before something else holds the sources still, and doing that inside
collapse_huge_page() means replacing the copy, the install and the rollback
at the same time, which is the whole function.  Every halfway state has
neither argument in full.

What is gradual here is the review rather than the mechanism: the engine
arrives one pass at a time, each reviewable alone, with the old one live
until one patch switches over.

> > PMD everywhere that matters: collapse_huge_page() clears and flushes the
> > whole PMD whatever order it is collapsing, installs a PMD leaf because
> > that is the only thing it can produce, and keeps everyone out with
> > mmap_write_lock, anon_vma_lock_write() and an IPI broadcast while it
> > does.
> 
> To be clear - the anon path. I think important to clarify :)

Yes, anon only.  The file path moves into collapse.c and picks up the
scan/run split, but its mechanism is untouched.

And you are right that "installs a PMD leaf" is wrong above: only a
PMD-order collapse installs one, a sub-PMD collapse repopulates the table
with PTEs.  What is order-blind is everything around it -- the PMD is
still cleared and flushed first.

I would like to bring it into the engine as well, and with it the private
copies in MAP_PRIVATE file mappings, which neither path collapses today --
the anon side requires vma_is_anonymous() and the file side works on the
page cache.

I stopped because I wanted to keep the patch count in double digits. :P

> And yeah it does IPI for any sensible arch (with
> CONFIG_MMU_GATHER_RCU_TABLE_FREE) via tlb_remove_table_sync_one(). The
> other arches IPI anyway on TLB invalidation.
> 
> [Though I intend to make all page table freeing RCU relatively soon which
> should? Eliminate the need for this, possibly?]

That would suit this well, and it is worth covering deposited page tables
in it if they are not already in scope.  Today PMD collapse has to deposit
a freshly allocated table rather than redepositing the one it detached,
because a deposited table must be safe for zap_deposited_table() to free
immediately.  Make that free RCU-deferred and the detached table can go
straight back -- one allocation less per PMD collapse.

> However we have to remember that a lot of the user-visible API assumes PMD
> sizing and so the code has to clearly reflect this and make it clear that

Your sentence got cut off, but if this is about the tunables then let me
flag what the series does with them.  max_ptes_none, _swap and _shared are
counts out of a PMD, and a range smaller than one has fewer PTEs than the
budget, so a raw comparison can never refuse it -- a 64K range on 4K pages
is 16 PTEs against a max_ptes_shared default of 256.  For swap and shared
the engine therefore compares fractions: count * HPAGE_PMD_NR against
max * nr_scanned.

max_ptes_none stays as mTHP collapse has it, all-or-nothing: 0, or
everything at that order.  That is deliberate -- allowing holes at an order
below the largest enabled one lets khugepaged fill them and collapse the
result at the next order up, which is the ratchet max_ptes_none exists to
bound.  There is room to scale it at the terminal order, where there is no
larger order to creep into, but I have not done that here.

> > Design
> > ======
> >
> > The old mechanism holds the address space still because it has nothing
> > else stopping the sources from moving under the copy.  The new engine
> > makes the sources themselves inert instead, with the two barriers
> > migration already uses, raised in that order:
> >
> >   1. migration entries replace the source PTEs.  Faults and GUP-slow
> >      now wait on the source folio's lock, which is taken before the
> >      first entry becomes visible.
> >   2. the source folio's refcount is frozen to its expected value.
> >      GUP-fast, pfn walkers, reclaim, compaction and memory-failure all
> >      fail folio_try_get() and back off.
> >
> > Between the two, nothing can reach a source, so the copy runs with no
> > lock held at all -- and the address space is left alone while it does.
> 
> Hmm, are migration entries the right mechnanism here? Are you actually
> migrating the pages to a large folio here, or using them to get the
> behaviour you want on fault/GUP?

Both, and I would argue the behaviour is not a side effect: what a
migration entry means to a waiter -- this page is going away, sleep on its
folio lock and look again -- is exactly true of a source under collapse.
Fault, GUP-slow and rmap then all do the right thing with no new code,
which is the case for reusing the entry rather than inventing a marker
every waiter would have to learn.

What is not reused is mm/migrate.c.  A migration entry encodes one PFN, so
it cannot name an N:1 destination of a different order; the engine takes
the hold-still half and does the remap itself at install.  It is a
migration in substance -- contents move to another folio, the old mappings
are replaced -- but not one migrate.c could drive.

> Same question in general for the freezing.

folio_ref_freeze() means nobody may take a new reference, which is the
property the copy needs, and is why migration and split use it too.

> > What that removes from every collapse path:
> >
> >   mmap_write_lock              -> mmap_read
> >   anon_vma_lock_write()        -> nothing: an rmap walk needs the folio
> >                                   locked, and the engine holds that lock
> >                                   from freeze to putback
> 
> I do like the idea of eliminating uses of the rmap lock like this, not only
> for contention's sake but also for scalable CoW purposes which introduces
> challenges with regards to holding these.
> 
> In fact, migration and huge memory collapse are the really problematic
> areas.

If that is about their rmap complexity, collapse gets easier here rather
than harder.

The engine takes no rmap lock and walks no rmap.  A page shared with
another process is unshared before anything is frozen -- the fault-in
pass breaks CoW, exactly as a write would -- so by freeze time a source
is exclusive to this mm and its only live mappings are the ones being
replaced.  Migration has to cope with a folio mapped from many mms; this
never sees one.

If what scalable CoW needs is that collapse stops messing with rmap,
that is what this does.

> 
> >   tlb_remove_table_sync_one()  -> nothing: one ranged flush per round
> 
> Aren't we reliant upon this synchronisation for correctness?

Not in the new engine.

In collapse_huge_page() the IPI is what makes the *detached* table safe
to use: it pmdp_collapse_flush()es the PMD and then copies out of the
table it just detached, so it has to know no lockless walker is still
inside it.

The new engine never does that -- a sub-PMD window is collapsed in place
under the page table lock, so nothing is detached, and at PMD order the
table is detached, never touched again, and freed with pte_free_defer().
The synchronisation is still there, it is RCU rather than an IPI; and on
the arches without RCU table free the flush itself IPIs, as you say.

> >  - A table that cannot become one huge page still yields the largest
> >    windows inside it, where before a single disqualified PTE gave up
> >    the whole table.
> 
> Are you permitting collapse of ranges that straddle PTEs?

No -- a candidate never crosses a page table or a VMA.  A round works
within one table, and each candidate is validated against the VMA at its
own order.

> Though in general I'm confused by the single disqualified PTE here -

Taking that literally is how I meant it: collapse_scan_pmd() goto
out_unmap's on the first PTE that fails any of its checks -- uffd, non-anon,
clean lazyfree, off the LRU, unexpected refcount -- and mthp_collapse() only
runs if the verdict came back SCAN_SUCCEED.  So one
such PTE anywhere in the table means nothing in that table collapses, at
any order, even at an order whose windows avoid it entirely.

> > There may be a way out -- a PMD migration entry over the table during
> > the window, so the CPU never caches a walk to shoot down -- but that
> > means teaching every pmd-level walker a new kind of entry, and I have
> > not tried it.
> 
> Hmm this seems like complexity on top of complexity...

I find it rather elegant, and expect it to be minimally intrusive: let such
a PMD be walkable exactly as a present one is, so software descends through
it as usual while the CPU sees a non-present entry and caches nothing.
Transparent to software, opaque to the CPU.

Out of scope for this patchset either way.

-- 
  Kiryl Shutsemau / Kirill A. Shutemov

Re: [RFC PATCH 02/57] mm: add pte_none_or_zero()

[David Hildenbrand (Arm)]

On 8/17/26 00:45, Kiryl Shutsemau wrote:
> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
> 
> A PTE that is none and one that maps the shared zeropage both stand for
> a page of zeroes the mapping does not own.  Code that cares only about
> the contents can treat the two alike.
> 
> Move khugepaged's local helper for that test to pgtable.h, below the
> is_zero_pfn() it is built on.
> 
> migrate_vma_insert_page() open-codes the same test on the slot it is
> about to fill.  Convert it.  It still tells none from the zeropage, but
> only to decide whether there is an old mapping to flush.
> 
> No functional change intended.
> 
> Assisted-by: Claude-Code:claude-opus-5
> Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
> ---
>  include/linux/pgtable.h | 17 +++++++++++++++++
>  mm/khugepaged.c         |  7 -------
>  mm/migrate_device.c     |  9 ++-------
>  3 files changed, 19 insertions(+), 14 deletions(-)
> 
> diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
> index 8c093c119e5a..bbee6d31f015 100644
> --- a/include/linux/pgtable.h
> +++ b/include/linux/pgtable.h
> @@ -2064,6 +2064,23 @@ static inline struct page *_zero_page(unsigned long addr)
>  
>  #ifdef CONFIG_MMU
>  
> +/**
> + * pte_none_or_zero - Does this PTE map nothing, or the shared zeropage?
> + * @pte: The page table entry to test.
> + *
> + * A PTE that is none and one that maps the shared zeropage both stand for a
> + * page of zeroes the mapping does not own, so code that only cares about the
> + * contents can treat them alike.
> + *
> + * Return: %true if @pte is none or maps the shared zeropage.
> + */
> +static inline bool pte_none_or_zero(pte_t pte)
> +{
> +	if (pte_none(pte))
> +		return true;
> +	return pte_present(pte) && is_zero_pfn(pte_pfn(pte));
> +}

(casually skimming over some patches)

That's just a horrible function. :)

If there is no pte_zero() then there also shouldn't be a pte_none_or_zero().

And in the code-base we have "pte_t pte_zero = {0}" which actually makes sense,
but is not what we care about here.


Just have an additional helper like:

	pte_maps_zero_page()
	pte_is_zero_page()
	pte_maps_zero_folio()
	pte_is_zero_folio()
	pte_zero_page()
	pte_zero_folio()

And let the callers still spell both cases out.

if (pte_none(pte) || pte_is_zero_page(pte))
	/* Do something amazing */

(we have both is_zero_page() and is_zero_folio() I assume "zero page" is cleaner
as this thing might soon no longer be a folio)

-- 
Cheers,

David