[PATCH v2 00/13] mm: batch of (urgentish?) fixes in oomk, cachestat, dev-migrate, percpu

[PATCH v2 00/13] mm: batch of (urgentish?) fixes in oomk, cachestat, dev-migrate, percpu

[Lorenzo Stoakes (ARM)]

A bunch of fixes for oomk, cachestat, device migration and the per-CPU
allocator.

These patches were originally AI-generated, tamed into upstreamable state.

Some are more reworked than others. Those which needed a significant
rethink are submitted as mine and marked as Reported-by the relevant
original author, others are smaller edits which are annotated as such in
the usual fashion.

Please check for correctness, sanity, and whether the cc: stable is
warranted (I have not included it for 3 where it seemed unwarranted.)

Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
v2:
* Propagated tags, thanks everybody!
* Clarified commit message in 1/13 as per Michal.
* Added a description of the fix and its performance cost to the commit
  message of 3/13, as per Usama.
* Updated 8/13 to also refuse to migrate KSM pages to device coherent
  memory, as per Sashiko.
* Added Gregory's benchmark numbers to the commit message of 13/13.

v1:
https://lore.kernel.org/r/20261111-mason-migrate-et-al-fixes-v1-0-3d2454362e8d@kernel.org

---
Chris Mason (5):
      mm: cachestat: count each page only once
      mm: cachestat: restart the walk after each entry
      mm/migrate_device: fix ksmd oops on a migrated KSM page
      mm/migrate_device: refuse to migrate shadow stacks
      percpu: queue the balance work on an unbound workqueue

Josef Bacik (4):
      mm/migrate_device: allocate THP deferred split heads
      mm/migrate_device: don't collect a range twice
      mm/migrate_device: don't let go of the fault page
      mm/migrate_device: look up fault folio after a split

Lorenzo Stoakes (ARM) (4):
      mm/mmap: free gathered pages before MMF_OOM_SKIP
      mm/migrate_device: fix use-after-free of the anon rmap
      percpu: fix contig hints that are too low
      percpu: search the whole reserved chunk

 include/linux/migrate.h  |   1 +
 include/linux/mm_types.h |   2 +
 include/linux/oom.h      |   9 ++++
 mm/filemap.c             |  21 +++++++---
 mm/migrate_device.c      | 107 ++++++++++++++++++++++++++++++++---------------
 mm/mmap.c                |   9 +++-
 mm/oom_kill.c            |   4 +-
 mm/percpu.c              |  69 +++++++++---------------------
 8 files changed, 132 insertions(+), 90 deletions(-)
---
base-commit: 7f04048247a69019ccaa9c152002d3601b19453c
change-id: 20261107-mason-migrate-et-al-fixes-0494d34884c4

Cheers,
-- 
Lorenzo Stoakes (ARM) <ljs@kernel.org>

[PATCH v2 01/13] mm/mmap: free gathered pages before MMF_OOM_SKIP

[Lorenzo Stoakes (ARM)]

The OOM killer can quickly trigger again after having selected a
victim. That victim might not have freed its memory yet, but its mm will
be marked MMF_OOM_SKIP by this point, so the OOM killer will ignore it.

This can cause the OOM killer to kill an innocent process when the first
might have freed enough memory to relieve the memory pressure.

exit_mmap() marks the mm MMF_OOM_SKIP once unmap_vmas() returns.
However unmap_vmas() does not free the last batch of pages it
gathered. That is left to tlb_finish_mmu(), which runs after
MMF_OOM_SKIP is set, and the batch can be large.

Fix the issue by freeing the gathered pages of an OOM victim before
marking the mm MMF_OOM_SKIP.

Note that the task which calls mmput() might not be the victim - if it
took a reference on the mm (e.g. ptrace), it may hold the last reference
and release it when it calls mmput().

Therefore we cannot rely on current to tell whether the mm belongs to an
OOM victim. Instead, partially revert commit b3541d912a84 ("mm: delete
unused MMF_OOM_VICTIM flag") to reinstate MMF_OOM_VICTIM, which
mark_oom_victim() sets on the victim's mm, and check that.

Reported-by: Chris Mason <mason@kernel.org>
Fixes: bf3980c85212 ("mm: drop oom code from exit_mmap")
Cc: stable@vger.kernel.org
Reviewed-by: Gregory Price (Meta) <gourry@gourry.net>
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Acked-by: Michal Hocko <mhocko@suse.com>
Reviewed-by: Pedro Falcato <pfalcato@suse.de>
Acked-by: Usama Arif <usama.arif@linux.dev>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 include/linux/mm_types.h | 2 ++
 include/linux/oom.h      | 9 +++++++++
 mm/mmap.c                | 9 ++++++++-
 mm/oom_kill.c            | 4 +++-
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 1ed35d258d2e..09e63ff12b04 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -2030,6 +2030,8 @@ enum {
 #define MMF_TOPDOWN		31	/* mm searches top down by default */
 #define MMF_TOPDOWN_MASK	BIT(MMF_TOPDOWN)
 
+#define MMF_OOM_VICTIM		32	/* mm is the oom victim */
+
 #define MMF_INIT_LEGACY_MASK	(MMF_DUMP_FILTER_MASK |\
 				 MMF_DISABLE_THP_MASK | MMF_HAS_MDWE_MASK |\
 				 MMF_VM_MERGE_ANY_MASK | MMF_TOPDOWN_MASK)
diff --git a/include/linux/oom.h b/include/linux/oom.h
index 00da05d227e6..a9390cb70541 100644
--- a/include/linux/oom.h
+++ b/include/linux/oom.h
@@ -76,6 +76,15 @@ static inline bool tsk_is_oom_victim(struct task_struct * tsk)
 	return tsk->signal->oom_mm;
 }
 
+/*
+ * Use this helper if tsk->mm != mm and the victim mm needs a special
+ * handling. This is guaranteed to stay true after once set.
+ */
+static inline bool mm_is_oom_victim(struct mm_struct *mm)
+{
+	return mm_flags_test(MMF_OOM_VICTIM, mm);
+}
+
 /*
  * Checks whether a page fault on the given mm is still reliable.
  * This is no longer true if the oom reaper started to reap the
diff --git a/mm/mmap.c b/mm/mmap.c
index 7618e72766a5..4cb18b217048 100644
--- a/mm/mmap.c
+++ b/mm/mmap.c
@@ -1313,11 +1313,18 @@ void exit_mmap(struct mm_struct *mm)
 	/* update_hiwater_rss(mm) here? but nobody should be looking */
 	/* Use ULONG_MAX here to ensure all VMAs in the mm are unmapped */
 	unmap_vmas(&tlb, &unmap);
+
+	/*
+	 * Free the memory before marking MMF_OOM_SKIP so a re-triggered OOM
+	 * killer will only skip this mm if it's already freed all it can.
+	 */
+	if (mm_is_oom_victim(mm))
+		tlb_flush_mmu(&tlb);
 	mmap_read_unlock(mm);
 
 	/*
 	 * Set MMF_OOM_SKIP to hide this task from the oom killer/reaper
-	 * because the memory has been already freed.
+	 * because the memory has been already unmapped.
 	 */
 	mm_flags_set(MMF_OOM_SKIP, mm);
 	mmap_write_lock(mm);
diff --git a/mm/oom_kill.c b/mm/oom_kill.c
index 5d48bd862c27..9de57a23893c 100644
--- a/mm/oom_kill.c
+++ b/mm/oom_kill.c
@@ -752,8 +752,10 @@ static void mark_oom_victim(struct task_struct *tsk)
 		return;
 
 	/* oom_mm is bound to the signal struct life time. */
-	if (!cmpxchg(&tsk->signal->oom_mm, NULL, mm))
+	if (!cmpxchg(&tsk->signal->oom_mm, NULL, mm)) {
 		mmgrab(tsk->signal->oom_mm);
+		mm_flags_set(MMF_OOM_VICTIM, mm);
+	}
 
 	/*
 	 * Make sure that the process is woken up from uninterruptible sleep

-- 
2.56.0

[PATCH v2 02/13] mm: cachestat: count each page only once

[Lorenzo Stoakes (ARM)]

From: Chris Mason <mason@kernel.org>

It turns out that filemap_cachestat() can unwittingly return incorrect
page counts, if large folios are present.

This can occur due to underflow - nr_cache is nearly 2^64 on a 64-bit
kernel if the input range starts more than 64 pages into a folio that is
split during the call.

It can also occur due to pages after a large folio being skipped if the
range starts in its second or later XArray slot and the call reschedules at
that folio.

Finally it can be caused by pages being counted twice due to folios in the
range being split or collapsed during the call.

filemap_cachestat() walks the page cache under the RCU read lock only,
and for each entry it counts the pages that lie in the range. It does
not track how far it has counted, so the counts are only right if the
walk returns each index once.

But a large folio can be split during the walk, or small folios can be
collapsed into a large one, and the next entries then cover indices that
were already counted.

Resolve the issue by tracking the current index as one past the last index
in the folio, and explicitly using xas_set() to set the current position in
the XArray.

Additionally, add a check for folio_last_index reaching or exceeding
last_index and break out of the loop if so.

Fixes: cf264e1329fb ("cachestat: implement cachestat syscall")
Cc: stable@vger.kernel.org
Signed-off-by: Chris Mason <mason@kernel.org>
[ ljs: updated to use curr_index, fixed comments, updated commit message ]
Reviewed-by: Gregory Price (Meta) <gourry@gourry.net>
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Reviewed-by: Johannes Weiner <hannes@cmpxchg.org>
Acked-by: Nhat Pham <nphamcs@gmail.com>
Reviewed-by: Pedro Falcato <pfalcato@suse.de>
Acked-by: Usama Arif <usama.arif@linux.dev>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/filemap.c | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/mm/filemap.c b/mm/filemap.c
index b04be64a5069..7c05e4fad081 100644
--- a/mm/filemap.c
+++ b/mm/filemap.c
@@ -4638,6 +4638,7 @@ static void filemap_cachestat(struct address_space *mapping,
 {
 	XA_STATE(xas, &mapping->i_pages, first_index);
 	struct folio *folio;
+	pgoff_t curr_index = first_index;
 
 	/* Flush stats (and potentially sleep) outside the RCU read section. */
 	mem_cgroup_flush_stats_ratelimited(NULL);
@@ -4667,13 +4668,19 @@ static void filemap_cachestat(struct address_space *mapping,
 		folio_first_index = round_down(xas.xa_index, 1 << order);
 		folio_last_index = folio_first_index + nr_pages - 1;
 
+		/* Skip already counted entries, e.g. after a split/collapse. */
+		if (folio_last_index < curr_index)
+			goto resched;
+
 		/* Folios might straddle the range boundaries, only count covered pages */
-		if (folio_first_index < first_index)
-			nr_pages -= first_index - folio_first_index;
+		if (folio_first_index < curr_index)
+			nr_pages -= curr_index - folio_first_index;
 
 		if (folio_last_index > last_index)
 			nr_pages -= folio_last_index - last_index;
 
+		curr_index = folio_last_index + 1;
+
 		if (xa_is_value(folio)) {
 			/* page is evicted */
 			void *shadow = (void *)folio;
@@ -4721,8 +4728,11 @@ static void filemap_cachestat(struct address_space *mapping,
 			cs->nr_writeback += nr_pages;
 
 resched:
+		if (folio_last_index >= last_index)
+			break;
+
 		if (need_resched()) {
-			xas_pause(&xas);
+			xas_set(&xas, curr_index);
 			cond_resched_rcu();
 		}
 	}

-- 
2.56.0

[PATCH v2 04/13] mm/migrate_device: allocate THP deferred split heads

[Lorenzo Stoakes (ARM)]

From: Josef Bacik <josef@toxicpanda.com>

A THP that device migration maps into an empty PMD can end up on the wrong
deferred split list when it is later queued for a deferred split.

This occurs if the memcg has no heads yet (resulting in a warning if
CONFIG_DEBUG_VM is set), where deferred_split_folio() will then queue the
folio on the parent's list.

Once a later THP fault in the same memcg allocates the heads,
__folio_unqueue_deferred_split() and the split path resolve folio_memcg()
to the memcg's own, now existing, list and take its lock.  They then
list_del_init() a folio that is linked on the parent's list, racing with
anything holding the parent's lock, and both lists' nr_items counts go
wrong.

Resolve the issue by calling folio_memcg_alloc_deferred() after the charge
in migrate_vma_insert_huge_pmd_page() and falling back like a failed charge
if that fails, as __do_huge_pmd_anonymous_page() does.

This was originally discovered by code inspection.

Fixes: fafaeceb89a5 ("mm: switch deferred split shrinker to list_lru")
Cc: stable@vger.kernel.org
Signed-off-by: Josef Bacik <josef@toxicpanda.com>
[ ljs: tweak commit message, remove unnecessary comment ]
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Acked-by: Balbir Singh <balbirs@nvidia.com>
Tested-by: Andrew Morton <akpm@linux-foundation.org>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 9e19fadfa65c..9e5dab8e1b77 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -833,6 +833,13 @@ static int migrate_vma_insert_huge_pmd_page(struct migrate_vma *migrate,
 		goto abort;
 	}
 
+	if (folio_memcg_alloc_deferred(folio)) {
+		count_vm_event(THP_FAULT_FALLBACK);
+		count_mthp_stat(HPAGE_PMD_ORDER, MTHP_STAT_ANON_FAULT_FALLBACK);
+		ret = -ENOMEM;
+		goto abort;
+	}
+
 	__folio_mark_uptodate(folio);
 
 	pgtable = pte_alloc_one(vma->vm_mm);

-- 
2.56.0

[PATCH v2 05/13] mm/migrate_device: don't collect a range twice

[Lorenzo Stoakes (ARM)]

From: Josef Bacik <josef@toxicpanda.com>

migrate_vma_collect_pmd() drops the page table lock to split a large folio
which it finds in the middle of a page table, and then starts over at the
"again" label. The entries for [start, addr) are in migrate->src[] and
migrate->dst[] by then.

While the page table lock is dropped the page table can go away, in which
case the restarted walk invokes migrate_vma_collect_huge_pmd(). However
this is passed start rather than addr, so [start, addr) is collected a
second time.

This can result in migrate->npages exceeding the number of pages the
caller asked for, and therefore in writes past the end of the
migrate->src[] and migrate->dst[] arrays.

Fix this by only handing over what has not yet been collected,
page-by-page, and flush the TLB for those entries which have been
previously replaced, similar to the other returns in the function.

Fixes: 022a12deda53 ("mm/migrate_device: handle partially mapped folios during collection")
Cc: stable@vger.kernel.org
Signed-off-by: Josef Bacik <josef@toxicpanda.com>
[ ljs: cleanup commit message ]
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Reviewed-by: Gregory Price (Meta) <gourry@gourry.net>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 9e5dab8e1b77..28ed706109da 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -257,13 +257,17 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 	pte_t *ptep;
 
 again:
+	/* [start, addr) is collected if we get here after splitting a folio */
 	if (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp)) {
-		int ret = migrate_vma_collect_huge_pmd(pmdp, start, end, walk, fault_folio);
+		int ret = migrate_vma_collect_huge_pmd(pmdp, addr, end, walk, fault_folio);
 
 		if (ret == -EAGAIN)
 			goto again;
-		if (ret == 0)
+		if (ret == 0) {
+			if (unmapped)
+				flush_tlb_range(walk->vma, start, end);
 			return 0;
+		}
 	}
 
 	ptep = pte_offset_map_lock(mm, pmdp, start, &ptl);

-- 
2.56.0

[PATCH v2 06/13] mm/migrate_device: don't let go of the fault page

[Lorenzo Stoakes (ARM)]

From: Josef Bacik <josef@toxicpanda.com>

When migrate_vma_split_folio() has split the folio of the fault page,
the lock and the reference which the fault handler has taken are at the
former head page.

It then moves them to the folio which now holds the fault page, by
taking a reference on and locking that folio, then unlocking and putting
the former head.

split_folio() has unlocked the new folio and has dropped its own
reference on it, so until folio_get() nothing but the page table entry
holds it.

Only the mmap lock is held for read, which does not prevent another thread
from zapping the entry with MADV_DONTNEED.

If this happens the folio is freed which results in folio_get() taking a
reference on a free device page, which the driver can hand out again at any
time.

In addition, folio_lock() can sleep here with the lock of the head page
held.

Utilise split_huge_page() which will leave the lock and reference at the
fault page correctly.

Fixes: 022a12deda53 ("mm/migrate_device: handle partially mapped folios during collection")
Cc: stable@vger.kernel.org
Signed-off-by: Josef Bacik <josef@toxicpanda.com>
[ ljs: tweaked commit message ]
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 39 ++++++++++-----------------------------
 1 file changed, 10 insertions(+), 29 deletions(-)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 28ed706109da..ae795fdf36f4 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -87,39 +87,20 @@ static int migrate_vma_split_folio(struct folio *folio,
 {
 	int ret;
 	struct folio *fault_folio = fault_page ? page_folio(fault_page) : NULL;
-	struct folio *new_fault_folio = NULL;
-
-	if (folio != fault_folio)
-		folio_lock(folio);
-
-	ret = split_folio(folio);
-	if (ret) {
-		if (folio != fault_folio) {
-			folio_unlock(folio);
-			folio_put(folio);
-		}
-		return ret;
-	}
-
-	new_fault_folio = fault_page ? page_folio(fault_page) : NULL;
 
 	/*
-	 * Ensure the lock is held on the correct
-	 * folio after the split
+	 * The lock and the reference of the fault handler stay with the
+	 * folio which holds the fault page after the split.
 	 */
-	if (!new_fault_folio) {
-		folio_unlock(folio);
-		folio_put(folio);
-	} else if (folio != new_fault_folio) {
-		if (new_fault_folio != fault_folio) {
-			folio_get(new_fault_folio);
-			folio_lock(new_fault_folio);
-		}
-		folio_unlock(folio);
-		folio_put(folio);
-	}
+	if (folio == fault_folio)
+		return split_huge_page(fault_page);
 
-	return 0;
+	folio_lock(folio);
+	ret = split_folio(folio);
+	folio_unlock(folio);
+	folio_put(folio);
+
+	return ret;
 }
 
 /** migrate_vma_collect_huge_pmd - collect THP pages without splitting the

-- 
2.56.0

[PATCH v2 07/13] mm/migrate_device: look up fault folio after a split

[Lorenzo Stoakes (ARM)]

From: Josef Bacik <josef@toxicpanda.com>

migrate_vma_collect_pmd() looks up the folio of migrate->fault_page once,
when it is entered, at which point the folio is locked.

The fault folio is assumed not to change after this, however if
migrate_vma_split_folio() splits a PTE-mapped large folio which contains
the fault page, the fault page ends up in a different folio.

Resolve the issue by looking up fault_folio again after each restart of the
walk.

Fixes: 022a12deda53 ("mm/migrate_device: handle partially mapped folios during collection")
Cc: stable@vger.kernel.org
Signed-off-by: Josef Bacik <josef@toxicpanda.com>
[ ljs: pared down commit message ]
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index ae795fdf36f4..45e571dbfb4c 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -233,11 +233,14 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 	struct mm_struct *mm = vma->vm_mm;
 	unsigned long addr = start, unmapped = 0;
 	spinlock_t *ptl;
-	struct folio *fault_folio = migrate->fault_page ?
-		page_folio(migrate->fault_page) : NULL;
+	struct folio *fault_folio;
 	pte_t *ptep;
 
 again:
+	/* The fault page is in another folio once its folio has been split */
+	fault_folio = migrate->fault_page ?
+		page_folio(migrate->fault_page) : NULL;
+
 	/* [start, addr) is collected if we get here after splitting a folio */
 	if (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp)) {
 		int ret = migrate_vma_collect_huge_pmd(pmdp, addr, end, walk, fault_folio);

-- 
2.56.0

[PATCH v2 08/13] mm/migrate_device: fix ksmd oops on a migrated KSM page

[Lorenzo Stoakes (ARM)]

From: Chris Mason <mason@kernel.org>

After a KSM page is migrated to device private memory, ksmd can oops,
and KSM then stops until reboot. A process that maps the page can also
loop in page faults when it next accesses the page.

The problem arises because KSM doesn't know not to try to look at device
private memory that does not exist in the direct map.

In addition, ksmd skips all device pages when it scans, so it drops its
rmap of a KSM page which was migrated to device private or device
coherent memory.

Disallow migrating a KSM page to either by explicitly testing for this
in __migrate_device_pages().

Fixes: a5430dda8a3a ("mm/migrate: support un-addressable ZONE_DEVICE page in migration")
Cc: stable@vger.kernel.org
Signed-off-by: Chris Mason <mason@kernel.org>
[ ljs: tweak commit message, comment, also refuse device coherent ]
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Acked-by: Balbir Singh <balbirs@nvidia.com>
Reviewed-by: Alistair Popple <apopple@nvidia.com>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 45e571dbfb4c..141d60d27e90 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -1207,6 +1207,11 @@ static void __migrate_device_pages(unsigned long *src_pfns,
 
 		if (folio_is_device_private(newfolio) ||
 		    folio_is_device_coherent(newfolio)) {
+			/* KSM does not scan or track device pages. */
+			if (folio_test_ksm(folio)) {
+				src_pfns[i] &= ~MIGRATE_PFN_MIGRATE;
+				goto next;
+			}
 			if (mapping) {
 				/*
 				 * For now only support anonymous memory migrating to

-- 
2.56.0

[PATCH v2 09/13] mm/migrate_device: fix use-after-free of the anon rmap

[Lorenzo Stoakes (ARM)]

A munmap() of anonymous memory can result in a use-after-free of the
anonymous rmap if a concurrent device memory migration occurs and the
migration runs without the mmap_lock held.

This can occur for a long-term pin of device-coherent memory, or a driver
evicting its device memory by PFN.

Consider:

    eviction                        munmap()
    ========                        ========
    migrate_device_range()
      migrate_device_unmap()
        try_to_migrate()
                                    zap the migration entry
                                    unlink_anon_vmas()
                                      last VMA: free the anon_vma
    migrate_device_finalize()
      remove_migration_ptes()
        /* locks and walks the freed anon_vma */

If munmap() runs before try_to_migrate() instead, the folio is no longer
mapped and gets no migration entry. The walk still runs, and follows the
stale folio->mapping to the freed anon_vma.

Resolve this by pinning the anon rmap of each folio which possesses one in
migrate_device_range(), migrate_device_pfns() and
migrate_device_coherent_folio() before its PTEs are replaced, and
unpinning it once the walk is done.

A folio which is no longer mapped cannot be pinned and has no migration
entries to remove, so skip it rather than migrating it.

Reported-by: Chris Mason <mason@kernel.org>
Fixes: e778406b40db ("mm/migrate_device.c: add migrate_device_range()")
Cc: stable@vger.kernel.org
Reviewed-by: Zi Yan <ziy@nvidia.com>
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 include/linux/migrate.h |  1 +
 mm/migrate_device.c     | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/include/linux/migrate.h b/include/linux/migrate.h
index 78424b3824c2..c6cce7925b9b 100644
--- a/include/linux/migrate.h
+++ b/include/linux/migrate.h
@@ -136,6 +136,7 @@ static inline int migrate_misplaced_folio(struct folio *folio, int node)
 #define MIGRATE_PFN_MIGRATE	(1UL << 1)
 #define MIGRATE_PFN_WRITE	(1UL << 3)
 #define MIGRATE_PFN_COMPOUND	(1UL << 4)
+#define MIGRATE_PFN_ANON_RMAP	(1UL << 5)	/* mm internal */
 #define MIGRATE_PFN_SHIFT	6
 
 static inline struct page *migrate_pfn_to_page(unsigned long mpfn)
diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 141d60d27e90..b4660eeb975a 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -552,6 +552,27 @@ static bool migrate_vma_check_page(struct page *page, struct page *fault_page)
 	return true;
 }
 
+/*
+ * Pin the anonymous rmap if present to allow migration entries to be safely
+ * removed. Returns true if the pin succeeded or no pin is required, and false
+ * if the folio is no longer mapped.
+ */
+static bool pin_anon_rmap(struct folio *folio, unsigned long *src_pfn)
+{
+	if (!folio_test_anon(folio) || folio_test_ksm(folio))
+		return true;
+	if (!folio_get_anon_vma(folio)) /* Pin the anon rmap. */
+		return false;
+	*src_pfn |= MIGRATE_PFN_ANON_RMAP;
+	return true;
+}
+
+static void unpin_anon_rmap(struct folio *folio, unsigned long src_pfn)
+{
+	if (src_pfn & MIGRATE_PFN_ANON_RMAP)
+		put_anon_vma(folio_anon_vma(folio)); /* Unpin the anon rmap. */
+}
+
 /*
  * Unmaps pages for migration. Returns number of source pfns marked as
  * migrating.
@@ -633,6 +654,7 @@ static unsigned long migrate_device_unmap(unsigned long *src_pfns,
 
 		folio = page_folio(page);
 		remove_migration_ptes(folio, folio, 0);
+		unpin_anon_rmap(folio, src_pfns[i]);
 
 		src_pfns[i] = 0;
 		if (fault_folio != folio)
@@ -1331,6 +1353,7 @@ static void __migrate_device_finalize(unsigned long *src_pfns,
 		if (!folio_is_zone_device(dst))
 			folio_add_lru(dst);
 		remove_migration_ptes(src, dst, 0);
+		unpin_anon_rmap(src, src_pfns[i]);
 		if (fault_folio != src)
 			folio_unlock(src);
 		folio_put(src);
@@ -1380,6 +1403,7 @@ EXPORT_SYMBOL(migrate_vma_finalize);
 
 static unsigned long migrate_device_pfn_lock(unsigned long pfn)
 {
+	unsigned long src_pfn = migrate_pfn(pfn) | MIGRATE_PFN_MIGRATE;
 	struct folio *folio;
 
 	folio = folio_get_nontail_page(pfn_to_page(pfn));
@@ -1391,7 +1415,13 @@ static unsigned long migrate_device_pfn_lock(unsigned long pfn)
 		return 0;
 	}
 
-	return migrate_pfn(pfn) | MIGRATE_PFN_MIGRATE;
+	if (!pin_anon_rmap(folio, &src_pfn)) {
+		folio_unlock(folio);
+		folio_put(folio);
+		return 0;
+	}
+
+	return src_pfn;
 }
 
 /**
@@ -1427,6 +1457,7 @@ int migrate_device_range(unsigned long *src_pfns, unsigned long start,
 		nr = folio_nr_pages(folio);
 		if (nr > npages - i) {
 			if (src_pfns[i] & MIGRATE_PFN_MIGRATE) {
+				unpin_anon_rmap(folio, src_pfns[i]);
 				folio_unlock(folio);
 				folio_put(folio);
 			}
@@ -1470,6 +1501,7 @@ int migrate_device_pfns(unsigned long *src_pfns, unsigned long npages)
 		nr = folio_nr_pages(folio);
 		if (nr > npages - i) {
 			if (src_pfns[i] & MIGRATE_PFN_MIGRATE) {
+				unpin_anon_rmap(folio, src_pfns[i]);
 				folio_unlock(folio);
 				folio_put(folio);
 			}
@@ -1505,6 +1537,11 @@ int migrate_device_coherent_folio(struct folio *folio)
 
 	folio_lock(folio);
 	src_pfn = migrate_pfn(folio_pfn(folio)) | MIGRATE_PFN_MIGRATE;
+	if (!pin_anon_rmap(folio, &src_pfn)) {
+		folio_unlock(folio);
+		folio_put(folio);
+		return -EBUSY;
+	}
 
 	/*
 	 * We don't have a VMA and don't need to walk the page tables to find

-- 
2.56.0

[PATCH v2 10/13] mm/migrate_device: refuse to migrate shadow stacks

[Lorenzo Stoakes (ARM)]

From: Chris Mason <mason@kernel.org>

The kernel blocks ordinary writes to a shadow stack. But pte_write() is
true for a shadow stack entry, and migrate_vma_setup() sets
MIGRATE_PFN_WRITE for each page where pte_write() is true.

Notably, the nouveau GPU driver takes that flag as permission to map its
copy of the page writable for the GPU, meaning that the next CPU access
migrates the page back into the shadow stack, and the page then holds what
the GPU wrote.

Be consistent and disallow this for migration.

Fixes: 54007f818206 ("mm: Introduce VM_SHADOW_STACK for shadow stack memory")
Cc: stable@vger.kernel.org
Signed-off-by: Chris Mason <mason@kernel.org>
[ ljs: removed unnecessary comment, refined commit msg ]
Reviewed-by: Alistair Popple <apopple@nvidia.com>
Reviewed-by: Balbir Singh <balbirs@nvidia.com>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/migrate_device.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index b4660eeb975a..e0bf6afe0768 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -759,6 +759,8 @@ int migrate_vma_setup(struct migrate_vma *args)
 	/* The migrate_vma page walk has no hugetlb support. */
 	if (!vma_is_mm_backed(vma) || vma_is_hugetlb(vma) || vma_is_dax(vma))
 		return -EINVAL;
+	if (vma_test_single_mask(vma, VMA_SHADOW_STACK))
+		return -EINVAL;
 	if (nr_pages <= 0)
 		return -EINVAL;
 	if (args->start < vma->vm_start || args->start >= vma->vm_end)

-- 
2.56.0

[PATCH v2 12/13] percpu: search the whole reserved chunk

[Lorenzo Stoakes (ARM)]

Loading a module with static percpu data can fail with -ENOMEM even
though the reserved chunk has a free area that fits the data.

pcpu_find_block_fit() does not scan the allocation bitmap, rather it
tests the request against hints, starting with the contig_hint of the
chunk, which describes its largest free area. If that area is too small
once aligned, the search gives up even if a smaller free area fits.

Consider:

    request:      256 bytes, aligned to 256 bytes
    contig_hint:  260 bytes, starts 4 bytes after a 256-byte boundary,
                  so 8 bytes remain after alignment
    also free:    256 bytes, starts on a 256-byte boundary

Ordinarily this is fine as the allocator moves on to the next chunk or
creates a new one, but the static percpu data of modules can only come
from the reserved chunk.

Fix this by searching the allocation bitmap of the reserved chunk
directly rather than using hints.

Reported-by: Chris Mason <mason@kernel.org>
Fixes: 13f966373f92 ("percpu: skip chunks if the alloc does not fit in the contig hint")
Fixes: b4c2116cfae6 ("percpu: update pcpu_find_block_fit to use an iterator")
Cc: stable@vger.kernel.org
Tested-by: Gregory Price (Meta) <gourry@gourry.net>
Reviewed-by: Gregory Price (Meta) <gourry@gourry.net>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/percpu.c | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mm/percpu.c b/mm/percpu.c
index dfd3ff0a1ac0..f87813e6ad73 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -1771,8 +1771,11 @@ void __percpu *pcpu_alloc_noprof(size_t size, size_t align, bool reserved,
 	if (reserved && pcpu_reserved_chunk) {
 		chunk = pcpu_reserved_chunk;
 
-		off = pcpu_find_block_fit(chunk, bits, bit_align, is_atomic);
-		if (off < 0) {
+		/* No other chunk to fall back on, so don't rely on hints. */
+		off = bitmap_find_next_zero_area(chunk->alloc_map,
+						 pcpu_chunk_map_bits(chunk), 0,
+						 bits, bit_align - 1);
+		if (off >= pcpu_chunk_map_bits(chunk)) {
 			err = "alloc from reserved chunk failed";
 			goto fail_unlock;
 		}

-- 
2.56.0

Re: [PATCH v2 07/13] mm/migrate_device: look up fault folio after a split

[Balbir Singh]

On Fri, Nov 13, 2026 at 09:45:25PM +0000, Lorenzo Stoakes (ARM) wrote:
> From: Josef Bacik <josef@toxicpanda.com>
> 
> migrate_vma_collect_pmd() looks up the folio of migrate->fault_page once,
> when it is entered, at which point the folio is locked.
> 
> The fault folio is assumed not to change after this, however if
> migrate_vma_split_folio() splits a PTE-mapped large folio which contains
> the fault page, the fault page ends up in a different folio.
> 
> Resolve the issue by looking up fault_folio again after each restart of the
> walk.
> 
> Fixes: 022a12deda53 ("mm/migrate_device: handle partially mapped folios during collection")
> Cc: stable@vger.kernel.org
> Signed-off-by: Josef Bacik <josef@toxicpanda.com>
> [ ljs: pared down commit message ]
> Tested-by: Gregory Price (Meta) <gourry@gourry.net>
> Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
> ---
>  mm/migrate_device.c | 7 +++++--
>  1 file changed, 5 insertions(+), 2 deletions(-)
> 
> diff --git a/mm/migrate_device.c b/mm/migrate_device.c
> index ae795fdf36f4..45e571dbfb4c 100644
> --- a/mm/migrate_device.c
> +++ b/mm/migrate_device.c
> @@ -233,11 +233,14 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
>  	struct mm_struct *mm = vma->vm_mm;
>  	unsigned long addr = start, unmapped = 0;
>  	spinlock_t *ptl;
> -	struct folio *fault_folio = migrate->fault_page ?
> -		page_folio(migrate->fault_page) : NULL;
> +	struct folio *fault_folio;
>  	pte_t *ptep;
>  
>  again:
> +	/* The fault page is in another folio once its folio has been split */
> +	fault_folio = migrate->fault_page ?
> +		page_folio(migrate->fault_page) : NULL;
> +
>  	/* [start, addr) is collected if we get here after splitting a folio */
>  	if (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp)) {
>  		int ret = migrate_vma_collect_huge_pmd(pmdp, addr, end, walk, fault_folio);
>

migrate_vma_split_folio(), already computes the new_fault_folio and gets
the right references and lock, but does not communicate it back

Reviewed-by: Balbir Singh <balbirs@nvidia.com>

Re: [PATCH v2 06/13] mm/migrate_device: don't let go of the fault page

[Balbir Singh]

On Fri, Nov 13, 2026 at 09:45:24PM +0000, Lorenzo Stoakes (ARM) wrote:
> From: Josef Bacik <josef@toxicpanda.com>
> 
> When migrate_vma_split_folio() has split the folio of the fault page,
> the lock and the reference which the fault handler has taken are at the
> former head page.
> 
> It then moves them to the folio which now holds the fault page, by
> taking a reference on and locking that folio, then unlocking and putting
> the former head.
> 
> split_folio() has unlocked the new folio and has dropped its own
> reference on it, so until folio_get() nothing but the page table entry
> holds it.
> 
> Only the mmap lock is held for read, which does not prevent another thread
> from zapping the entry with MADV_DONTNEED.
> 
> If this happens the folio is freed which results in folio_get() taking a
> reference on a free device page, which the driver can hand out again at any
> time.
> 
> In addition, folio_lock() can sleep here with the lock of the head page
> held.
> 
> Utilise split_huge_page() which will leave the lock and reference at the
> fault page correctly.
> 
> Fixes: 022a12deda53 ("mm/migrate_device: handle partially mapped folios during collection")
> Cc: stable@vger.kernel.org
> Signed-off-by: Josef Bacik <josef@toxicpanda.com>
> [ ljs: tweaked commit message ]
> Tested-by: Gregory Price (Meta) <gourry@gourry.net>
> Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
> ---
>  mm/migrate_device.c | 39 ++++++++++-----------------------------
>  1 file changed, 10 insertions(+), 29 deletions(-)
> 
> diff --git a/mm/migrate_device.c b/mm/migrate_device.c
> index 28ed706109da..ae795fdf36f4 100644
> --- a/mm/migrate_device.c
> +++ b/mm/migrate_device.c
> @@ -87,39 +87,20 @@ static int migrate_vma_split_folio(struct folio *folio,
>  {
>  	int ret;
>  	struct folio *fault_folio = fault_page ? page_folio(fault_page) : NULL;
> -	struct folio *new_fault_folio = NULL;
> -
> -	if (folio != fault_folio)
> -		folio_lock(folio);
> -
> -	ret = split_folio(folio);
> -	if (ret) {
> -		if (folio != fault_folio) {
> -			folio_unlock(folio);
> -			folio_put(folio);
> -		}
> -		return ret;
> -	}
> -
> -	new_fault_folio = fault_page ? page_folio(fault_page) : NULL;
>  
>  	/*
> -	 * Ensure the lock is held on the correct
> -	 * folio after the split
> +	 * The lock and the reference of the fault handler stay with the
> +	 * folio which holds the fault page after the split.
>  	 */
> -	if (!new_fault_folio) {
> -		folio_unlock(folio);
> -		folio_put(folio);
> -	} else if (folio != new_fault_folio) {
> -		if (new_fault_folio != fault_folio) {
> -			folio_get(new_fault_folio);
> -			folio_lock(new_fault_folio);
> -		}
> -		folio_unlock(folio);
> -		folio_put(folio);
> -	}
> +	if (folio == fault_folio)
> +		return split_huge_page(fault_page);
>  
> -	return 0;
> +	folio_lock(folio);
> +	ret = split_folio(folio);
> +	folio_unlock(folio);
> +	folio_put(folio);
> +
> +	return ret;
>  }
>  
>  /** migrate_vma_collect_huge_pmd - collect THP pages without splitting the
>

Very subtle race

Reviewed-by: Balbir Singh <balbirs@nvidia.com>