[PATCH 0/6] mm: batch rmap walks during large folio migration

[PATCH 0/6] mm: batch rmap walks during large folio migration

[Shivank Garg]

Migrating a PTE-mapped large folio currently walks the folio's rmap
one page at a time. It does so twice per migration:
1. try_to_migrate_one(): Replace present PTEs with migration entries
2. remove_migration_pte(): Restore migration entries to working PTEs.

For a folio with N contiguous PTEs, each walk repeats page_vma_mapped_walk()
and the associated operation per subpage of the folio. (256 iterations
for 1M folio).

This series enables both walks to process a batch of consecutive PTEs
that map the same large folio and do some refactoring and clean up on
the way.

Performance:

AMD EPYC ZEN 3, 2-socket system (1 node per socket).
move_pages() migrating 1 GB from node 0 -> node 1, throughput
in GB/s:
		         batch-	        batch-restore +
folio     vanilla        restore        unmap          restore%  restore+unmap%
4K        2.98±0.09      2.95±0.07      3.06±0.05        -1.0%    +2.6%
64K       6.16±0.09      6.24±0.08      7.80±0.08        +1.3%   +26.6%
256K      6.85±0.06      6.85±0.09      9.54±0.19        +0.0%   +39.3%
1M        7.06±0.09      7.17±0.09      9.89±0.06        +1.5%   +40.1%
2M       10.97±0.17     10.67±0.13     10.69±0.15        -2.7%    -2.6%

In plain migration, the folio copy itself dominate the cost, so these rmap walk
savings are only partly visible end-to-end.
When the folio copy is batched-copy and offloaded to DMA engine [1], the copy is
no longer the bottlneck and rmap-walk batching translates into much larger gains:

This series + [1], DMA Offload on PTDMA (DCBM), 16 channels:

		     batch-copy
folio   vanilla      -offload(dcbm)  dcbm+restore   dcbm+restore+unmap
4K      3.36 ± 0.10   4.16 ± 0.18     4.10 ± 0.14    4.12 ± 0.13
64K     6.16 ± 0.47   9.63 ± 0.23    10.00 ± 0.33   14.63 ± 0.14
256K    6.70 ± 0.09  12.98 ± 0.18    13.52 ± 0.15   32.30 ± 0.65    (2.38x)
1M      6.99 ± 0.02  14.32 ± 0.67    14.84 ± 0.17   38.97 ± 1.03    (2.62x)
2M     10.74 ± 0.29  64.25 ± 1.48    65.36 ± 1.00   64.97 ± 1.75

[1] https://lore.kernel.org/linux-mm/20260630-shivank-batch-migrate-offload-v6-0-da95d7e8b8a2@amd.com

Applies cleanly on mm-new 10-07-26 (61cccb8363f).

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
Dev Jain (1):
      mm/rmap: Add batched version of folio_try_share_anon_rmap_pte

Shivank Garg (5):
      mm: factor out generic PTE batch detection from swap_pte_batch()
      mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()
      mm/migrate: batch the restore-side migration rmap walk
      mm/rmap: split try_to_migrate_hugetlb_one() out of try_to_migrate_one()
      mm/rmap: batch the unmap of large folios in try_to_migrate_one()

 include/linux/rmap.h |  52 +++++--
 mm/internal.h        |  36 +++--
 mm/migrate.c         | 190 ++++++++++++++++-------
 mm/rmap.c            | 427 +++++++++++++++++++++++++++++++++------------------
 4 files changed, 464 insertions(+), 241 deletions(-)
---
base-commit: 61cccb8363fcc282d4ae0555b8739dd227f5ad0b
change-id: 20260701-migrate-rmap-batch-8ecc04de1905

Best regards,
-- 
Shivank Garg <shivankg@amd.com>

[PATCH 1/6] mm/rmap: Add batched version of folio_try_share_anon_rmap_pte

[Shivank Garg]

From: Dev Jain <dev.jain@arm.com>

To enable batched unmapping of anonymous folios, we need to handle the
sharing of exclusive pages. Hence, a batched version of
folio_try_share_anon_rmap_pte is required.

Currently, the sole purpose of nr_pages in __folio_try_share_anon_rmap is
to do some rmap sanity checks. Now, clear the PageAnonExclusive bit on a
batch of nr_pages. Refactor the function such that the clearing of the bit
can be done at one place without duplication.

Note that __folio_try_share_anon_rmap can receive nr_pages == HPAGE_PMD_NR
from the PMD path, but currently we only clear the bit on the head page.
Retain this behaviour by setting nr_pages = 1 in case the caller is
folio_try_share_anon_rmap_pmd.

While at it, convert nr_pages to unsigned long to future-proof from
overflow in case P4D-huge mappings etc get supported down the road.
I haven't made such a change in each function receiving nr_pages in
try_to_unmap_one - perhaps this can be done incrementally.

Borrrowed from [PATCH v4 00/12] Optimize anonymous large folio unmapping [1],
will drop and rebase once [1] is merged.

[1] https://lore.kernel.org/linux-mm/20260526063635.61721-10-dev.jain@arm.com/

Signed-off-by: Dev Jain <dev.jain@arm.com>
Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 include/linux/rmap.h | 52 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 35 insertions(+), 17 deletions(-)

diff --git a/include/linux/rmap.h b/include/linux/rmap.h
index 8dc0871e5f00..64929490a7cf 100644
--- a/include/linux/rmap.h
+++ b/include/linux/rmap.h
@@ -706,17 +706,18 @@ static inline int folio_try_dup_anon_rmap_pmd(struct folio *folio,
 }
 
 static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
-		struct page *page, int nr_pages, enum pgtable_level level)
+		struct page *page, unsigned long nr_pages, enum pgtable_level level)
 {
+	/* device private folios cannot get pinned via GUP. */
+	const bool pinnable = likely(!folio_is_device_private(folio));
+
 	VM_WARN_ON_FOLIO(!folio_test_anon(folio), folio);
 	VM_WARN_ON_FOLIO(!PageAnonExclusive(page), folio);
 	__folio_rmap_sanity_checks(folio, page, nr_pages, level);
 
-	/* device private folios cannot get pinned via GUP. */
-	if (unlikely(folio_is_device_private(folio))) {
-		ClearPageAnonExclusive(page);
-		return 0;
-	}
+	/* We only clear anon-exclusive from head page of PMD folio */
+	if (level == PGTABLE_LEVEL_PMD)
+		nr_pages = 1;
 
 	/*
 	 * We have to make sure that when we clear PageAnonExclusive, that
@@ -760,29 +761,38 @@ static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
 	 * so we use explicit ones here.
 	 */
 
-	/* Paired with the memory barrier in try_grab_folio(). */
-	if (IS_ENABLED(CONFIG_HAVE_GUP_FAST))
-		smp_mb();
+	if (pinnable) {
+		/* Paired with the memory barrier in try_grab_folio(). */
+		if (IS_ENABLED(CONFIG_HAVE_GUP_FAST))
+			smp_mb();
 
-	if (unlikely(folio_maybe_dma_pinned(folio)))
-		return -EBUSY;
-	ClearPageAnonExclusive(page);
+		if (unlikely(folio_maybe_dma_pinned(folio)))
+			return -EBUSY;
+	}
+
+	for (;;) {
+		ClearPageAnonExclusive(page);
+		if (--nr_pages == 0)
+			break;
+		page++;
+	}
 
 	/*
 	 * This is conceptually a smp_wmb() paired with the smp_rmb() in
 	 * gup_must_unshare().
 	 */
-	if (IS_ENABLED(CONFIG_HAVE_GUP_FAST))
+	if (pinnable && IS_ENABLED(CONFIG_HAVE_GUP_FAST))
 		smp_mb__after_atomic();
 	return 0;
 }
 
 /**
- * folio_try_share_anon_rmap_pte - try marking an exclusive anonymous page
- *				   mapped by a PTE possibly shared to prepare
+ * folio_try_share_anon_rmap_ptes - try marking exclusive anonymous pages
+ *				   mapped by PTEs possibly shared to prepare
  *				   for KSM or temporary unmapping
  * @folio:	The folio to share a mapping of
- * @page:	The mapped exclusive page
+ * @page:	The first mapped exclusive page of the batch in the folio
+ * @nr_pages:	The number of pages to share in the folio (batch size)
  *
  * The caller needs to hold the page table lock and has to have the page table
  * entries cleared/invalidated.
@@ -797,11 +807,19 @@ static __always_inline int __folio_try_share_anon_rmap(struct folio *folio,
  *
  * Returns 0 if marking the mapped page possibly shared succeeded. Returns
  * -EBUSY otherwise.
+ *
+ * The caller needs to hold the page table lock.
  */
+static inline int folio_try_share_anon_rmap_ptes(struct folio *folio,
+		struct page *page, unsigned long nr_pages)
+{
+	return __folio_try_share_anon_rmap(folio, page, nr_pages, PGTABLE_LEVEL_PTE);
+}
+
 static inline int folio_try_share_anon_rmap_pte(struct folio *folio,
 		struct page *page)
 {
-	return __folio_try_share_anon_rmap(folio, page, 1, PGTABLE_LEVEL_PTE);
+	return folio_try_share_anon_rmap_ptes(folio, page, 1);
 }
 
 /**

-- 
2.43.0

[PATCH 2/6] mm: factor out generic PTE batch detection from swap_pte_batch()

[Shivank Garg]

Factor out the generic PTE batch detection logic from swap_pte_batch()
into softleaf_pte_batch() helper, so that it can be reused by upcoming
patch that adds restore-side batching of migration entries.

swap_pte_batch() now keeps only its swap-specific concerns and
delegates the batch detection to softleaf_pte_batch().

No functional changes intended.

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/internal.h | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/mm/internal.h b/mm/internal.h
index 874be94cf257..a90721f8c687 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -505,6 +505,26 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
 	return pte_move_swp_offset(pte, 1);
 }
 
+static inline int softleaf_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
+{
+	pte_t expected_pte = pte_next_swp_offset(pte);
+	const pte_t *end_ptep = start_ptep + max_nr;
+	pte_t *ptep = start_ptep + 1;
+
+	VM_WARN_ON(max_nr < 1);
+
+	while (ptep < end_ptep) {
+		pte = ptep_get(ptep);
+
+		if (!pte_same(pte, expected_pte))
+			break;
+		expected_pte = pte_next_swp_offset(expected_pte);
+		ptep++;
+	}
+
+	return ptep - start_ptep;
+}
+
 /**
  * swap_pte_batch - detect a PTE batch for a set of contiguous swap entries
  * @start_ptep: Page table pointer for the first entry.
@@ -522,23 +542,9 @@ static inline pte_t pte_next_swp_offset(pte_t pte)
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {
-	pte_t expected_pte = pte_next_swp_offset(pte);
-	const pte_t *end_ptep = start_ptep + max_nr;
-	pte_t *ptep = start_ptep + 1;
-
-	VM_WARN_ON(max_nr < 1);
 	VM_WARN_ON(!softleaf_is_swap(softleaf_from_pte(pte)));
 
-	while (ptep < end_ptep) {
-		pte = ptep_get(ptep);
-
-		if (!pte_same(pte, expected_pte))
-			break;
-		expected_pte = pte_next_swp_offset(expected_pte);
-		ptep++;
-	}
-
-	return ptep - start_ptep;
+	return softleaf_pte_batch(start_ptep, max_nr, pte);
 }
 #endif /* CONFIG_MMU */
 

-- 
2.43.0

[PATCH 3/6] mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()

[Shivank Garg]

remove_migration_pte() interleaves the regular folio path with the
hugetlb path via folio_test_hugetlb() and CONFIG_HUGETLB_PAGE
special cases. Simplify this by seprating the hugetlb handling into
remove_migration_pte_hugetlb().

Factor out working-PTE construction from migration entry into
migration_softleaf_entry_to_pte(), used by both paths.

While here, convert the VM_BUG_ON_FOLIO() to VM_WARN_ON_ONCE_FOLIO().

No functional change intended.

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/migrate.c | 153 ++++++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 98 insertions(+), 55 deletions(-)

diff --git a/mm/migrate.c b/mm/migrate.c
index 222c8c15f782..8a098402b8c9 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -345,6 +345,89 @@ struct rmap_walk_arg {
 	bool map_unused_to_zeropage;
 };
 
+static pte_t migration_softleaf_entry_to_pte(struct folio *folio, struct page *new,
+		softleaf_t entry, pte_t old_pte, struct vm_area_struct *vma,
+		rmap_t *rmap_flags)
+{
+	pte_t pte = mk_pte(new, READ_ONCE(vma->vm_page_prot));
+
+	if (!softleaf_is_migration_young(entry))
+		pte = pte_mkold(pte);
+	if (folio_test_dirty(folio) && softleaf_is_migration_dirty(entry))
+		pte = pte_mkdirty(pte);
+	if (pte_swp_soft_dirty(old_pte))
+		pte = pte_mksoft_dirty(pte);
+	else
+		pte = pte_clear_soft_dirty(pte);
+
+	if (softleaf_is_migration_write(entry))
+		pte = pte_mkwrite(pte, vma);
+	else if (pte_swp_uffd(old_pte))
+		pte = pte_mkuffd(pte);
+
+	/* See do_swap_page(): restore PAGE_NONE for RWP */
+	if (pte_swp_uffd(old_pte) && userfaultfd_rwp(vma))
+		pte = pte_modify(pte, PAGE_NONE);
+
+	if (folio_test_anon(folio) && !softleaf_is_migration_read(entry))
+		*rmap_flags |= RMAP_EXCLUSIVE;
+
+	return pte;
+}
+
+/*
+ * Restore a potential migration pte to a working pte entry for hugetlb folios.
+ */
+#ifdef CONFIG_HUGETLB_PAGE
+static bool remove_migration_pte_hugetlb(struct folio *folio,
+		struct vm_area_struct *vma, unsigned long addr, void *arg)
+{
+	struct rmap_walk_arg *rmap_walk_arg = arg;
+	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
+	struct hstate *h = hstate_vma(vma);
+	unsigned int shift = huge_page_shift(h);
+	unsigned long psize = huge_page_size(h);
+
+	if (page_vma_mapped_walk(&pvmw)) {
+		rmap_t rmap_flags = RMAP_NONE;
+		struct page *new = folio_page(folio, 0);
+		pte_t old_pte, pte;
+		softleaf_t entry;
+
+		old_pte = huge_ptep_get(vma->vm_mm, pvmw.address, pvmw.pte);
+		entry = softleaf_from_pte(old_pte);
+
+		folio_get(folio);
+		pte = migration_softleaf_entry_to_pte(folio, new, entry, old_pte,
+						      vma, &rmap_flags);
+		pte = arch_make_huge_pte(pte, shift, vma->vm_flags);
+		if (folio_test_anon(folio))
+			hugetlb_add_anon_rmap(folio, vma, pvmw.address,
+					      rmap_flags);
+		else
+			hugetlb_add_file_rmap(folio);
+		set_huge_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte, psize);
+
+		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
+			mlock_drain_local();
+
+		trace_remove_migration_pte(pvmw.address, pte_val(pte),
+					   compound_order(new));
+
+		/* No need to invalidate - it was non-present before */
+		update_mmu_cache(vma, pvmw.address, pvmw.pte);
+	}
+
+	return true;
+}
+#else
+static bool remove_migration_pte_hugetlb(struct folio *folio,
+		struct vm_area_struct *vma, unsigned long addr, void *arg)
+{
+	return false;
+}
+#endif /* CONFIG_HUGETLB_PAGE */
+
 /*
  * Restore a potential migration pte to a working pte entry
  */
@@ -363,52 +446,28 @@ static bool remove_migration_pte(struct folio *folio,
 		unsigned long idx = 0;
 
 		/* pgoff is invalid for ksm pages, but they are never large */
-		if (folio_test_large(folio) && !folio_test_hugetlb(folio))
+		if (folio_test_large(folio))
 			idx = linear_page_index(vma, pvmw.address) - pvmw.pgoff;
 		new = folio_page(folio, idx);
 
 #ifdef CONFIG_ARCH_HAS_PMD_SOFTLEAVES
 		/* PMD-mapped THP migration entry */
 		if (!pvmw.pte) {
-			VM_BUG_ON_FOLIO(folio_test_hugetlb(folio) ||
-					!folio_test_pmd_mappable(folio), folio);
+			VM_WARN_ON_ONCE_FOLIO(!folio_test_pmd_mappable(folio), folio);
 			remove_migration_pmd(&pvmw, new);
 			continue;
 		}
 #endif
-		if (folio_test_hugetlb(folio))
-			old_pte = huge_ptep_get(vma->vm_mm, pvmw.address,
-						pvmw.pte);
-		else
-			old_pte = ptep_get(pvmw.pte);
+		old_pte = ptep_get(pvmw.pte);
 		if (rmap_walk_arg->map_unused_to_zeropage &&
 		    try_to_map_unused_to_zeropage(&pvmw, folio, old_pte, idx))
 			continue;
 
-		folio_get(folio);
-		pte = mk_pte(new, READ_ONCE(vma->vm_page_prot));
-
 		entry = softleaf_from_pte(old_pte);
-		if (!softleaf_is_migration_young(entry))
-			pte = pte_mkold(pte);
-		if (folio_test_dirty(folio) && softleaf_is_migration_dirty(entry))
-			pte = pte_mkdirty(pte);
-		if (pte_swp_soft_dirty(old_pte))
-			pte = pte_mksoft_dirty(pte);
-		else
-			pte = pte_clear_soft_dirty(pte);
 
-		if (softleaf_is_migration_write(entry))
-			pte = pte_mkwrite(pte, vma);
-		else if (pte_swp_uffd(old_pte))
-			pte = pte_mkuffd(pte);
-
-		/* See do_swap_page(): restore PAGE_NONE for RWP */
-		if (pte_swp_uffd(old_pte) && userfaultfd_rwp(vma))
-			pte = pte_modify(pte, PAGE_NONE);
-
-		if (folio_test_anon(folio) && !softleaf_is_migration_read(entry))
-			rmap_flags |= RMAP_EXCLUSIVE;
+		folio_get(folio);
+		pte = migration_softleaf_entry_to_pte(folio, new, entry, old_pte,
+						      vma, &rmap_flags);
 
 		if (unlikely(is_device_private_page(new))) {
 			if (pte_write(pte))
@@ -424,30 +483,12 @@ static bool remove_migration_pte(struct folio *folio,
 				pte = pte_swp_mkuffd(pte);
 		}
 
-#ifdef CONFIG_HUGETLB_PAGE
-		if (folio_test_hugetlb(folio)) {
-			struct hstate *h = hstate_vma(vma);
-			unsigned int shift = huge_page_shift(h);
-			unsigned long psize = huge_page_size(h);
-
-			pte = arch_make_huge_pte(pte, shift, vma->vm_flags);
-			if (folio_test_anon(folio))
-				hugetlb_add_anon_rmap(folio, vma, pvmw.address,
-						      rmap_flags);
-			else
-				hugetlb_add_file_rmap(folio);
-			set_huge_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte,
-					psize);
-		} else
-#endif
-		{
-			if (folio_test_anon(folio))
-				folio_add_anon_rmap_pte(folio, new, vma,
-							pvmw.address, rmap_flags);
-			else
-				folio_add_file_rmap_pte(folio, new, vma);
-			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
-		}
+		if (folio_test_anon(folio))
+			folio_add_anon_rmap_pte(folio, new, vma,
+						pvmw.address, rmap_flags);
+		else
+			folio_add_file_rmap_pte(folio, new, vma);
+		set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
 		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
 			mlock_drain_local();
 
@@ -474,7 +515,9 @@ void remove_migration_ptes(struct folio *src, struct folio *dst,
 	};
 
 	struct rmap_walk_control rwc = {
-		.rmap_one = remove_migration_pte,
+		.rmap_one = folio_test_hugetlb(src) ?
+				remove_migration_pte_hugetlb :
+				remove_migration_pte,
 		.arg = &rmap_walk_arg,
 	};
 

-- 
2.43.0

[PATCH 4/6] mm/migrate: batch the restore-side migration rmap walk

[Shivank Garg]

remove_migration_pte() restores migration entries one PTE at a time. For
a PTE-mapped large folio, it repeats calls to building the PTE, add rmap,
set_pte_at() and page_vma_mapped_walk() per base page (256 times for
a 1M folio).

Add migration_pte_batch() to detect batch of contiguous migration entries
that map consecutive subpages, and have identical PTE bits.

While at it, use folio_is_device_private() instead of
is_device_private_page().

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/migrate.c | 49 ++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 42 insertions(+), 7 deletions(-)

diff --git a/mm/migrate.c b/mm/migrate.c
index 8a098402b8c9..6d05d3f41095 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -345,6 +345,28 @@ struct rmap_walk_arg {
 	bool map_unused_to_zeropage;
 };
 
+/*
+ * Detect a batch of contiguous migration entries: consecutive (non-present) PTEs
+ * containing migration entries with consecutive offsets and matching pte bits.
+ */
+static unsigned int migration_pte_batch(struct page_vma_mapped_walk *pvmw,
+		struct folio *folio, pte_t first_pte, unsigned long start_idx)
+{
+	struct vm_area_struct *vma = pvmw->vma;
+	unsigned long end_addr = pmd_addr_end(pvmw->address, vma->vm_end);
+	unsigned int folio_nr = folio_nr_pages(folio);
+	unsigned int max_nr;
+
+	VM_WARN_ON(!softleaf_is_migration(softleaf_from_pte(first_pte)));
+
+	/* Bound by VMA / PMD end and by the remaining subpages of the folio. */
+	max_nr = min((end_addr - pvmw->address) >> PAGE_SHIFT, folio_nr - start_idx);
+	if (max_nr <= 1)
+		return 1;
+
+	return softleaf_pte_batch(pvmw->pte, max_nr, first_pte);
+}
+
 static pte_t migration_softleaf_entry_to_pte(struct folio *folio, struct page *new,
 		softleaf_t entry, pte_t old_pte, struct vm_area_struct *vma,
 		rmap_t *rmap_flags)
@@ -435,6 +457,7 @@ static bool remove_migration_pte(struct folio *folio,
 		struct vm_area_struct *vma, unsigned long addr, void *arg)
 {
 	struct rmap_walk_arg *rmap_walk_arg = arg;
+	bool is_devpriv = folio_is_device_private(folio);
 	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
 
 	while (page_vma_mapped_walk(&pvmw)) {
@@ -444,6 +467,7 @@ static bool remove_migration_pte(struct folio *folio,
 		softleaf_t entry;
 		struct page *new;
 		unsigned long idx = 0;
+		unsigned int nr = 1;
 
 		/* pgoff is invalid for ksm pages, but they are never large */
 		if (folio_test_large(folio))
@@ -465,11 +489,18 @@ static bool remove_migration_pte(struct folio *folio,
 
 		entry = softleaf_from_pte(old_pte);
 
-		folio_get(folio);
+		/*
+		 * Try to restore nr>1 contiguous PTEs in one shot. Falls back
+		 * to original per-PTE path (nr=1) if batching is not possible.
+		 */
+		if (!rmap_walk_arg->map_unused_to_zeropage && likely(!is_devpriv))
+			nr = migration_pte_batch(&pvmw, folio, old_pte, idx);
+
+		folio_ref_add(folio, nr);
 		pte = migration_softleaf_entry_to_pte(folio, new, entry, old_pte,
 						      vma, &rmap_flags);
 
-		if (unlikely(is_device_private_page(new))) {
+		if (unlikely(is_devpriv)) {
 			if (pte_write(pte))
 				entry = make_writable_device_private_entry(
 							page_to_pfn(new));
@@ -484,11 +515,11 @@ static bool remove_migration_pte(struct folio *folio,
 		}
 
 		if (folio_test_anon(folio))
-			folio_add_anon_rmap_pte(folio, new, vma,
-						pvmw.address, rmap_flags);
+			folio_add_anon_rmap_ptes(folio, new, nr, vma,
+						 pvmw.address, rmap_flags);
 		else
-			folio_add_file_rmap_pte(folio, new, vma);
-		set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
+			folio_add_file_rmap_ptes(folio, new, nr, vma);
+		set_ptes(vma->vm_mm, pvmw.address, pvmw.pte, pte, nr);
 		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
 			mlock_drain_local();
 
@@ -496,7 +527,11 @@ static bool remove_migration_pte(struct folio *folio,
 					   compound_order(new));
 
 		/* No need to invalidate - it was non-present before */
-		update_mmu_cache(vma, pvmw.address, pvmw.pte);
+		update_mmu_cache_range(NULL, vma, pvmw.address, pvmw.pte, nr);
+
+		/* Skip the batched PTEs */
+		pvmw.pte += nr - 1;
+		pvmw.address += (nr - 1) * PAGE_SIZE;
 	}
 
 	return true;

-- 
2.43.0

[PATCH 5/6] mm/rmap: split try_to_migrate_hugetlb_one() out of try_to_migrate_one()

[Shivank Garg]

try_to_migrate_one() interleaves hugetlb special handling through the
regular PTE/PMD migration path with folio_test_hugetlb() checks
scattered across the walk. Simplify this by moving hugetlb handling
into its own rmap_one callback, try_to_migrate_hugetlb_one().

Factor the migration swap-pte construction shared by both hugetlb
and regular folios path into make_migration_swp_pte().

Few minor changes:
- Since hugetlb uses separate counters, update_hiwater_rss() is not
needed for hugetlb path. Also, use huge_pte_dirty instead of pte_dirty
- No PVMW_SYNC flag is needed for hugetlb
- hugetlb folio has single mapping per VMA, so no loop is needed.
- Convert VM_BUG_ON to VM_WARN_ON_ONCE and not use irrelevant WARN check.

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/rmap.c | 329 +++++++++++++++++++++++++++++++++++++-------------------------
 1 file changed, 195 insertions(+), 134 deletions(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index e854679553b9..ef68adba89c4 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2405,6 +2405,187 @@ void try_to_unmap(struct folio *folio, enum ttu_flags flags)
 		rmap_walk(folio, &rwc);
 }
 
+/* Build the migration swap PTE that replaces pteval for a page. */
+static pte_t make_migration_swp_pte(struct page *page, pte_t pteval,
+		bool writable, bool anon_exclusive)
+{
+	swp_entry_t entry;
+	pte_t swp_pte;
+
+	if (writable)
+		entry = make_writable_migration_entry(page_to_pfn(page));
+	else if (anon_exclusive)
+		entry = make_readable_exclusive_migration_entry(
+						page_to_pfn(page));
+	else
+		entry = make_readable_migration_entry(page_to_pfn(page));
+
+	if (likely(pte_present(pteval))) {
+		if (pte_young(pteval))
+			entry = make_migration_entry_young(entry);
+		if (pte_dirty(pteval))
+			entry = make_migration_entry_dirty(entry);
+		swp_pte = swp_entry_to_pte(entry);
+		if (pte_soft_dirty(pteval))
+			swp_pte = pte_swp_mksoft_dirty(swp_pte);
+		if (pte_uffd(pteval))
+			swp_pte = pte_swp_mkuffd(swp_pte);
+	} else {
+		swp_pte = swp_entry_to_pte(entry);
+		if (pte_swp_soft_dirty(pteval))
+			swp_pte = pte_swp_mksoft_dirty(swp_pte);
+		if (pte_swp_uffd(pteval))
+			swp_pte = pte_swp_mkuffd(swp_pte);
+	}
+
+	return swp_pte;
+}
+
+/*
+ * Replace the hugetlb page table mapping for folio with migration entries.
+ *
+ * @arg: enum ttu_flags will be passed to this argument.
+ */
+static bool try_to_migrate_hugetlb_one(struct folio *folio,
+		struct vm_area_struct *vma, unsigned long address, void *arg)
+{
+	const unsigned long hsz = huge_page_size(hstate_vma(vma));
+	DEFINE_FOLIO_VMA_WALK(pvmw, folio, vma, address, 0);
+	enum ttu_flags flags = (enum ttu_flags)(long)arg;
+	bool anon_exclusive, writable, ret = true;
+	bool anon = folio_test_anon(folio);
+	struct mm_struct *mm = vma->vm_mm;
+	struct mmu_notifier_range range;
+	unsigned long pfn;
+	struct page *page;
+	pte_t pteval;
+
+	range.end = vma_address_end(&pvmw);
+	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, vma->vm_mm,
+				address, range.end);
+	/* If sharing is possible, start and end will be adjusted accordingly. */
+	adjust_range_if_pmd_sharing_possible(vma, &range.start, &range.end);
+	mmu_notifier_invalidate_range_start(&range);
+
+	if (page_vma_mapped_walk(&pvmw)) {
+		address = pvmw.address;
+		pteval = huge_ptep_get(mm, address, pvmw.pte);
+		if (likely(pte_present(pteval)))
+			pfn = pte_pfn(pteval);
+		else
+			pfn = softleaf_to_pfn(softleaf_from_pte(pteval));
+
+		page = folio_page(folio, pfn - folio_pfn(folio));
+		anon_exclusive = anon && PageAnonExclusive(page);
+
+		/*
+		 * huge_pmd_unshare may unmap an entire PMD page. There is no
+		 * way of knowing exactly which PMDs may be cached for this mm,
+		 * so we must flush them all. start/end were already adjusted
+		 * above to cover this range.
+		 */
+		flush_cache_range(vma, range.start, range.end);
+
+		/*
+		 * To call huge_pmd_unshare, i_mmap_rwsem must be held in write
+		 * mode. Caller needs to explicitly do this outside rmap
+		 * routines. We also must hold hugetlb vma_lock in write mode.
+		 * Lock order dictates acquiring vma_lock BEFORE i_mmap_rwsem.
+		 * We can only try lock here and fail if unsuccessful.
+		 */
+		if (!anon) {
+			struct mmu_gather tlb;
+
+			VM_WARN_ON_ONCE(!(flags & TTU_RMAP_LOCKED));
+			if (!hugetlb_vma_trylock_write(vma)) {
+				ret = false;
+				goto walk_done;
+			}
+
+			tlb_gather_mmu_vma(&tlb, vma);
+			if (huge_pmd_unshare(&tlb, vma, address, pvmw.pte)) {
+				hugetlb_vma_unlock_write(vma);
+				huge_pmd_unshare_flush(&tlb, vma);
+				tlb_finish_mmu(&tlb);
+				/*
+				 * The PMD table was unmapped,
+				 * consequently unmapping the folio.
+				 */
+				goto walk_done;
+			}
+			hugetlb_vma_unlock_write(vma);
+			tlb_finish_mmu(&tlb);
+		}
+
+		/* Nuke the hugetlb page table entry */
+		pteval = huge_ptep_clear_flush(vma, address, pvmw.pte);
+		if (huge_pte_dirty(pteval))
+			folio_mark_dirty(folio);
+		writable = pte_write(pteval);
+
+		VM_WARN_ON_FOLIO(writable && anon && !anon_exclusive, folio);
+
+		if (PageHWPoison(page)) {
+			pteval = swp_entry_to_pte(make_hwpoison_entry(page));
+			hugetlb_count_sub(folio_nr_pages(folio), mm);
+			set_huge_pte_at(mm, address, pvmw.pte, pteval, hsz);
+		} else if (likely(pte_present(pteval)) && pte_unused(pteval) &&
+			   !userfaultfd_armed(vma)) {
+			/*
+			 * The guest indicated that the page content is of no
+			 * interest anymore. Simply discard the pte, vmscan
+			 * will take care of the rest.
+			 */
+			dec_mm_counter(mm, mm_counter(folio));
+		} else {
+			pte_t swp_pte;
+
+			/*
+			 * arch_unmap_one() is expected to be a NOP on
+			 * architectures where we could have PFN swap PTEs,
+			 * so we'll not check/care.
+			 */
+			if (arch_unmap_one(mm, vma, address, pteval) < 0) {
+				set_huge_pte_at(mm, address, pvmw.pte,
+						pteval, hsz);
+				ret = false;
+				goto walk_done;
+			}
+
+			/* See folio_try_share_anon_rmap_pte(): clear PTE first. */
+			if (anon_exclusive &&
+			    hugetlb_try_share_anon_rmap(folio)) {
+				set_huge_pte_at(mm, address, pvmw.pte,
+						pteval, hsz);
+				ret = false;
+				goto walk_done;
+			}
+
+			/*
+			 * Store the pfn of the page in a special migration
+			 * pte. do_swap_page() will wait until the migration
+			 * pte is removed and then restart fault handling.
+			 */
+			swp_pte = make_migration_swp_pte(page, pteval,
+							 writable, anon_exclusive);
+			set_huge_pte_at(mm, address, pvmw.pte, swp_pte, hsz);
+			trace_set_migration_pte(address, pte_val(swp_pte),
+						folio_order(folio));
+		}
+
+		hugetlb_remove_rmap(folio);
+		if (vma->vm_flags & VM_LOCKED)
+			mlock_drain_local();
+		folio_put(folio);
+walk_done:
+		page_vma_mapped_walk_done(&pvmw);
+	}
+
+	mmu_notifier_invalidate_range_end(&range);
+
+	return ret;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument.
  *
@@ -2422,7 +2603,6 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 	struct mmu_notifier_range range;
 	enum ttu_flags flags = (enum ttu_flags)(long)arg;
 	unsigned long pfn;
-	unsigned long hsz = 0;
 
 	/*
 	 * When racing against e.g. zap_pte_range() on another cpu,
@@ -2435,26 +2615,13 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 
 	/*
 	 * For THP, we have to assume the worse case ie pmd for invalidation.
-	 * For hugetlb, it could be much worse if we need to do pud
-	 * invalidation in the case of pmd sharing.
 	 *
 	 * Note that the page can not be free in this function as call of
-	 * try_to_unmap() must hold a reference on the page.
+	 * try_to_migrate() must hold a reference on the page.
 	 */
 	range.end = vma_address_end(&pvmw);
 	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, vma->vm_mm,
 				address, range.end);
-	if (folio_test_hugetlb(folio)) {
-		/*
-		 * If sharing is possible, start and end will be adjusted
-		 * accordingly.
-		 */
-		adjust_range_if_pmd_sharing_possible(vma, &range.start,
-						     &range.end);
-
-		/* We need the huge page size for set_huge_pte_at() */
-		hsz = huge_page_size(hstate_vma(vma));
-	}
 	mmu_notifier_invalidate_range_start(&range);
 
 	while (page_vma_mapped_walk(&pvmw)) {
@@ -2501,10 +2668,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 		VM_BUG_ON_FOLIO(!pvmw.pte, folio);
 
 		address = pvmw.address;
-		if (folio_test_hugetlb(folio))
-			pteval = huge_ptep_get(mm, address, pvmw.pte);
-		else
-			pteval = ptep_get(pvmw.pte);
+		pteval = ptep_get(pvmw.pte);
 		if (likely(pte_present(pteval))) {
 			pfn = pte_pfn(pteval);
 		} else {
@@ -2515,66 +2679,13 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			const softleaf_t entry = softleaf_from_pte(pteval);
 
 			pfn = softleaf_to_pfn(entry);
-			VM_WARN_ON_FOLIO(folio_test_hugetlb(folio), folio);
 		}
 
 		subpage = folio_page(folio, pfn - folio_pfn(folio));
 		anon_exclusive = folio_test_anon(folio) &&
 				 PageAnonExclusive(subpage);
 
-		if (folio_test_hugetlb(folio)) {
-			bool anon = folio_test_anon(folio);
-
-			/*
-			 * huge_pmd_unshare may unmap an entire PMD page.
-			 * There is no way of knowing exactly which PMDs may
-			 * be cached for this mm, so we must flush them all.
-			 * start/end were already adjusted above to cover this
-			 * range.
-			 */
-			flush_cache_range(vma, range.start, range.end);
-
-			/*
-			 * To call huge_pmd_unshare, i_mmap_rwsem must be
-			 * held in write mode.  Caller needs to explicitly
-			 * do this outside rmap routines.
-			 *
-			 * We also must hold hugetlb vma_lock in write mode.
-			 * Lock order dictates acquiring vma_lock BEFORE
-			 * i_mmap_rwsem.  We can only try lock here and
-			 * fail if unsuccessful.
-			 */
-			if (!anon) {
-				struct mmu_gather tlb;
-
-				VM_BUG_ON(!(flags & TTU_RMAP_LOCKED));
-				if (!hugetlb_vma_trylock_write(vma)) {
-					page_vma_mapped_walk_done(&pvmw);
-					ret = false;
-					break;
-				}
-
-				tlb_gather_mmu_vma(&tlb, vma);
-				if (huge_pmd_unshare(&tlb, vma, address, pvmw.pte)) {
-					hugetlb_vma_unlock_write(vma);
-					huge_pmd_unshare_flush(&tlb, vma);
-					tlb_finish_mmu(&tlb);
-					/*
-					 * The PMD table was unmapped,
-					 * consequently unmapping the folio.
-					 */
-					page_vma_mapped_walk_done(&pvmw);
-					break;
-				}
-				hugetlb_vma_unlock_write(vma);
-				tlb_finish_mmu(&tlb);
-			}
-			/* Nuke the hugetlb page table entry */
-			pteval = huge_ptep_clear_flush(vma, address, pvmw.pte);
-			if (pte_dirty(pteval))
-				folio_mark_dirty(folio);
-			writable = pte_write(pteval);
-		} else if (likely(pte_present(pteval))) {
+		if (likely(pte_present(pteval))) {
 			flush_cache_page(vma, address, pfn);
 			/* Nuke the page table entry. */
 			if (should_defer_flush(mm, flags)) {
@@ -2613,14 +2724,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			VM_WARN_ON_FOLIO(folio_is_device_private(folio), folio);
 
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
-			if (folio_test_hugetlb(folio)) {
-				hugetlb_count_sub(folio_nr_pages(folio), mm);
-				set_huge_pte_at(mm, address, pvmw.pte, pteval,
-						hsz);
-			} else {
-				dec_mm_counter(mm, mm_counter(folio));
-				set_pte_at(mm, address, pvmw.pte, pteval);
-			}
+			dec_mm_counter(mm, mm_counter(folio));
+			set_pte_at(mm, address, pvmw.pte, pteval);
 		} else if (likely(pte_present(pteval)) && pte_unused(pteval) &&
 			   !userfaultfd_armed(vma)) {
 			/*
@@ -2635,7 +2740,6 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			dec_mm_counter(mm, mm_counter(folio));
 		} else {
-			swp_entry_t entry;
 			pte_t swp_pte;
 
 			/*
@@ -2644,28 +2748,15 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 * so we'll not check/care.
 			 */
 			if (arch_unmap_one(mm, vma, address, pteval) < 0) {
-				if (folio_test_hugetlb(folio))
-					set_huge_pte_at(mm, address, pvmw.pte,
-							pteval, hsz);
-				else
-					set_pte_at(mm, address, pvmw.pte, pteval);
+				set_pte_at(mm, address, pvmw.pte, pteval);
 				ret = false;
 				page_vma_mapped_walk_done(&pvmw);
 				break;
 			}
 
 			/* See folio_try_share_anon_rmap_pte(): clear PTE first. */
-			if (folio_test_hugetlb(folio)) {
-				if (anon_exclusive &&
-				    hugetlb_try_share_anon_rmap(folio)) {
-					set_huge_pte_at(mm, address, pvmw.pte,
-							pteval, hsz);
-					ret = false;
-					page_vma_mapped_walk_done(&pvmw);
-					break;
-				}
-			} else if (anon_exclusive &&
-				   folio_try_share_anon_rmap_pte(folio, subpage)) {
+			if (anon_exclusive &&
+			    folio_try_share_anon_rmap_pte(folio, subpage)) {
 				set_pte_at(mm, address, pvmw.pte, pteval);
 				ret = false;
 				page_vma_mapped_walk_done(&pvmw);
@@ -2677,37 +2768,9 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 * pte. do_swap_page() will wait until the migration
 			 * pte is removed and then restart fault handling.
 			 */
-			if (writable)
-				entry = make_writable_migration_entry(
-							page_to_pfn(subpage));
-			else if (anon_exclusive)
-				entry = make_readable_exclusive_migration_entry(
-							page_to_pfn(subpage));
-			else
-				entry = make_readable_migration_entry(
-							page_to_pfn(subpage));
-			if (likely(pte_present(pteval))) {
-				if (pte_young(pteval))
-					entry = make_migration_entry_young(entry);
-				if (pte_dirty(pteval))
-					entry = make_migration_entry_dirty(entry);
-				swp_pte = swp_entry_to_pte(entry);
-				if (pte_soft_dirty(pteval))
-					swp_pte = pte_swp_mksoft_dirty(swp_pte);
-				if (pte_uffd(pteval))
-					swp_pte = pte_swp_mkuffd(swp_pte);
-			} else {
-				swp_pte = swp_entry_to_pte(entry);
-				if (pte_swp_soft_dirty(pteval))
-					swp_pte = pte_swp_mksoft_dirty(swp_pte);
-				if (pte_swp_uffd(pteval))
-					swp_pte = pte_swp_mkuffd(swp_pte);
-			}
-			if (folio_test_hugetlb(folio))
-				set_huge_pte_at(mm, address, pvmw.pte, swp_pte,
-						hsz);
-			else
-				set_pte_at(mm, address, pvmw.pte, swp_pte);
+			swp_pte = make_migration_swp_pte(subpage, pteval,
+							 writable, anon_exclusive);
+			set_pte_at(mm, address, pvmw.pte, swp_pte);
 			trace_set_migration_pte(address, pte_val(swp_pte),
 						folio_order(folio));
 			/*
@@ -2716,10 +2779,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 		}
 
-		if (unlikely(folio_test_hugetlb(folio)))
-			hugetlb_remove_rmap(folio);
-		else
-			folio_remove_rmap_pte(folio, subpage, vma);
+		folio_remove_rmap_pte(folio, subpage, vma);
 		if (vma->vm_flags & VM_LOCKED)
 			mlock_drain_local();
 		folio_put(folio);
@@ -2741,7 +2801,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 void try_to_migrate(struct folio *folio, enum ttu_flags flags)
 {
 	struct rmap_walk_control rwc = {
-		.rmap_one = try_to_migrate_one,
+		.rmap_one = folio_test_hugetlb(folio) ?
+				try_to_migrate_hugetlb_one : try_to_migrate_one,
 		.arg = (void *)flags,
 		.done = folio_not_mapped,
 		.anon_lock = folio_lock_anon_vma_read,

-- 
2.43.0

[PATCH 6/6] mm/rmap: batch the unmap of large folios in try_to_migrate_one()

[Shivank Garg]

try_to_migrate_one() converts present PTEs to migration entries one at a
time. For a PTE-mapped large folio, this repeat calls to ptep clear+flush,
the migration entry build and set, folio_remove_rmap_pte() and folio_put(),
each re-entering page_vma_mapped_walk() once per base page (256 times for
1M folio).

Mirror try_to_unmap_one() to introduce folio_migrate_pte_batch() to detect
eligible batch for PTEs mapping conseuctive subpages of a large folios,
and convert the whole batch in one shot using the batched helpers.

A side-effect of this change is trace_set_migration_pte() will record
one event per batched run instead of earlier behavior of one per base page.

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/rmap.c | 110 ++++++++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 85 insertions(+), 25 deletions(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index ef68adba89c4..cfac85f6f68a 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2586,6 +2586,48 @@ static bool try_to_migrate_hugetlb_one(struct folio *folio,
 	return ret;
 }
 
+static inline unsigned int folio_migrate_pte_batch(struct folio *folio,
+		struct page_vma_mapped_walk *pvmw, pte_t pte,
+		struct page *subpage, bool anon_exclusive)
+{
+	unsigned long end_addr, addr = pvmw->address;
+	struct vm_area_struct *vma = pvmw->vma;
+	unsigned int max_nr, nr, i;
+
+#ifdef __HAVE_ARCH_UNMAP_ONE
+	/* Cannot batch unmap if arch_unmap_one() is defined. */
+	return 1;
+#endif
+
+	if (!folio_test_large(folio))
+		return 1;
+	if (folio_is_zone_device(folio) || folio_test_has_hwpoisoned(folio))
+		return 1;
+	if (pte_unused(pte))
+		return 1;
+
+	/* We may only batch within a single VMA and a single page table. */
+	end_addr = pmd_addr_end(addr, vma->vm_end);
+	max_nr = (end_addr - addr) >> PAGE_SHIFT;
+	/*
+	 * If unmap fails, we need to restore the ptes. To avoid accidentally
+	 * upgrading write permissions for ptes that were not originally writable,
+	 * and to avoid losing the soft-dirty bit, use the appropriate FPB flags.
+	 */
+	nr = folio_pte_batch_flags(folio, vma, pvmw->pte, &pte, max_nr,
+				   FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY);
+
+	/* Limit possible batch count to a uniform PageAnonExclusive value */
+	if (folio_test_anon(folio)) {
+		for (i = 1; i < nr; i++)
+			if (PageAnonExclusive(subpage + i) != anon_exclusive)
+				break;
+		nr = i;
+	}
+
+	return nr;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument.
  *
@@ -2602,7 +2644,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 	struct page *subpage;
 	struct mmu_notifier_range range;
 	enum ttu_flags flags = (enum ttu_flags)(long)arg;
-	unsigned long pfn;
+	unsigned long pfn, end_addr, nr_pages;
 
 	/*
 	 * When racing against e.g. zap_pte_range() on another cpu,
@@ -2655,11 +2697,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			VM_BUG_ON_FOLIO(folio_test_hugetlb(folio) ||
 					!folio_test_pmd_mappable(folio), folio);
 
-			if (set_pmd_migration_entry(&pvmw, subpage)) {
-				ret = false;
-				page_vma_mapped_walk_done(&pvmw);
-				break;
-			}
+			if (set_pmd_migration_entry(&pvmw, subpage))
+				goto walk_abort;
 			continue;
 #endif
 		}
@@ -2685,9 +2724,16 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 		anon_exclusive = folio_test_anon(folio) &&
 				 PageAnonExclusive(subpage);
 
+		nr_pages = 1;
 		if (likely(pte_present(pteval))) {
-			flush_cache_page(vma, address, pfn);
-			/* Nuke the page table entry. */
+			nr_pages = folio_migrate_pte_batch(folio, &pvmw, pteval,
+							   subpage, anon_exclusive);
+
+			end_addr = address + nr_pages * PAGE_SIZE;
+			flush_cache_range(vma, address, end_addr);
+
+			/* Nuke the page table entries. */
+			pteval = get_and_clear_ptes(mm, address, pvmw.pte, nr_pages);
 			if (should_defer_flush(mm, flags)) {
 				/*
 				 * We clear the PTE but do not flush so potentially
@@ -2697,11 +2743,9 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 				 * transition on a cached TLB entry is written through
 				 * and traps if the PTE is unmapped.
 				 */
-				pteval = ptep_get_and_clear(mm, address, pvmw.pte);
-
-				set_tlb_ubc_flush_pending(mm, pteval, address, address + PAGE_SIZE);
+				set_tlb_ubc_flush_pending(mm, pteval, address, end_addr);
 			} else {
-				pteval = ptep_clear_flush(vma, address, pvmw.pte);
+				flush_tlb_range(vma, address, end_addr);
 			}
 			if (pte_dirty(pteval))
 				folio_mark_dirty(folio);
@@ -2741,6 +2785,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			dec_mm_counter(mm, mm_counter(folio));
 		} else {
 			pte_t swp_pte;
+			unsigned int i;
 
 			/*
 			 * arch_unmap_one() is expected to be a NOP on
@@ -2748,19 +2793,15 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 * so we'll not check/care.
 			 */
 			if (arch_unmap_one(mm, vma, address, pteval) < 0) {
-				set_pte_at(mm, address, pvmw.pte, pteval);
-				ret = false;
-				page_vma_mapped_walk_done(&pvmw);
-				break;
+				set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
+				goto walk_abort;
 			}
 
-			/* See folio_try_share_anon_rmap_pte(): clear PTE first. */
+			/* See folio_try_share_anon_rmap_ptes(): clear PTE first. */
 			if (anon_exclusive &&
-			    folio_try_share_anon_rmap_pte(folio, subpage)) {
-				set_pte_at(mm, address, pvmw.pte, pteval);
-				ret = false;
-				page_vma_mapped_walk_done(&pvmw);
-				break;
+			    folio_try_share_anon_rmap_ptes(folio, subpage, nr_pages)) {
+				set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
+				goto walk_abort;
 			}
 
 			/*
@@ -2770,19 +2811,38 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			swp_pte = make_migration_swp_pte(subpage, pteval,
 							 writable, anon_exclusive);
-			set_pte_at(mm, address, pvmw.pte, swp_pte);
 			trace_set_migration_pte(address, pte_val(swp_pte),
 						folio_order(folio));
+
+			/* Set nr_pages migration entries, advancing the PFN. */
+			for (i = 0; i < nr_pages; i++) {
+				set_pte_at(mm, address + (unsigned long)i * PAGE_SIZE,
+					   pvmw.pte + i, swp_pte);
+				swp_pte = pte_next_swp_offset(swp_pte);
+			}
 			/*
 			 * No need to invalidate here it will synchronize on
 			 * against the special swap migration pte.
 			 */
 		}
 
-		folio_remove_rmap_pte(folio, subpage, vma);
+		folio_remove_rmap_ptes(folio, subpage, nr_pages, vma);
 		if (vma->vm_flags & VM_LOCKED)
 			mlock_drain_local();
-		folio_put(folio);
+		folio_put_refs(folio, nr_pages);
+
+		/*
+		 * If we batched the entire folio, there is nothing left to
+		 * walk; stop right here.
+		 */
+		if (nr_pages == folio_nr_pages(folio))
+			goto walk_done;
+		continue;
+walk_abort:
+		ret = false;
+walk_done:
+		page_vma_mapped_walk_done(&pvmw);
+		break;
 	}
 
 	mmu_notifier_invalidate_range_end(&range);

-- 
2.43.0

[syzbot ci] Re: mm: batch rmap walks during large folio migration

[syzbot ci]

syzbot ci has tested the following series

[v1] mm: batch rmap walks during large folio migration
https://lore.kernel.org/all/20260712-migrate-rmap-batch-v1-0-872a734431d1@amd.com
* [PATCH 1/6] mm/rmap: Add batched version of folio_try_share_anon_rmap_pte
* [PATCH 2/6] mm: factor out generic PTE batch detection from swap_pte_batch()
* [PATCH 3/6] mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()
* [PATCH 4/6] mm/migrate: batch the restore-side migration rmap walk
* [PATCH 5/6] mm/rmap: split try_to_migrate_hugetlb_one() out of try_to_migrate_one()
* [PATCH 6/6] mm/rmap: batch the unmap of large folios in try_to_migrate_one()

and found the following issues:
* BUG: scheduling while atomic in exit_to_user_mode_loop
* BUG: sleeping function called from invalid context in migrate_pages
* BUG: sleeping function called from invalid context in migrate_pages_batch
* BUG: sleeping function called from invalid context in move_pages
* BUG: sleeping function called from invalid context in store_status

Full report is available here:
https://ci.syzbot.org/series/f6ed0914-22d0-4ad7-be5e-be4ed8fb9961

***

BUG: scheduling while atomic in exit_to_user_mode_loop

tree:      mm-new
URL:       https://kernel.googlesource.com/pub/scm/linux/kernel/git/akpm/mm.git
base:      61cccb8363fcc282d4ae0555b8739dd227f5ad0b
arch:      amd64
compiler:  Debian clang version 22.1.6 (++20260514074242+fc4aad7b5db3-1~exp1~20260514074407.73), Debian LLD 22.1.6
config:    https://ci.syzbot.org/builds/a27db994-58f5-48f4-a7f5-528b63f5a7c6/config
syz repro: https://ci.syzbot.org/findings/72fb0597-196f-4679-88a2-7d2dafda8454/syz_repro

BUG: scheduling while atomic: syz.0.17/5797/0x00000002
1 lock held by syz.0.17/5797:
 #0: ffff88810aa64e58 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff88810aa64e58 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
Modules linked in:
Preemption disabled at:
[<0000000000000000>] 0x0
Kernel panic - not syncing: scheduling while atomic: panic_on_warn set ...
CPU: 0 UID: 0 PID: 5797 Comm: syz.0.17 Not tainted syzkaller #0 PREEMPT(full) 
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 vpanic+0x56c/0xa60 kernel/panic.c:651
 panic+0xc5/0xd0 kernel/panic.c:788
 check_panic_on_warn+0x89/0xb0 kernel/panic.c:525
 __schedule_bug+0xf6/0x150 kernel/sched/core.c:6061
 schedule_debug kernel/sched/core.c:6090 [inline]
 __schedule+0x189d/0x56c0 kernel/sched/core.c:7083
 __schedule_loop kernel/sched/core.c:7311 [inline]
 schedule+0x164/0x2b0 kernel/sched/core.c:7326
 __exit_to_user_mode_loop kernel/entry/common.c:55 [inline]
 exit_to_user_mode_loop+0x83/0x730 kernel/entry/common.c:101
 __exit_to_user_mode_prepare include/linux/irq-entry-common.h:207 [inline]
 syscall_exit_to_user_mode_prepare include/linux/irq-entry-common.h:230 [inline]
 syscall_exit_to_user_mode include/linux/entry-common.h:318 [inline]
 do_syscall_64+0x353/0x580 arch/x86/entry/syscall_64.c:100
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7f1b3439ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007ffc9098d528 EFLAGS: 00000246 ORIG_RAX: 00000000000000ed
RAX: 0000000000000000 RBX: 00007f1b34615fa0 RCX: 00007f1b3439ce59
RDX: 0000000000004000 RSI: 0000000000002000 RDI: 0000200000195000
RBP: 00007f1b34432e6f R08: 0000000000000008 R09: 0000000000000003
R10: 0000200000000400 R11: 0000000000000246 R12: 0000000000000000
R13: 00007f1b34615fac R14: 00007f1b34615fa0 R15: 00007f1b34615fa0
 </TASK>
Kernel Offset: disabled
Rebooting in 86400 seconds..


***

BUG: sleeping function called from invalid context in migrate_pages

tree:      mm-new
URL:       https://kernel.googlesource.com/pub/scm/linux/kernel/git/akpm/mm.git
base:      61cccb8363fcc282d4ae0555b8739dd227f5ad0b
arch:      amd64
compiler:  Debian clang version 22.1.6 (++20260514074242+fc4aad7b5db3-1~exp1~20260514074407.73), Debian LLD 22.1.6
config:    https://ci.syzbot.org/builds/a27db994-58f5-48f4-a7f5-528b63f5a7c6/config
syz repro: https://ci.syzbot.org/findings/9dd2d2c8-a699-4e7c-8f73-4400d3057853/syz_repro

BUG: sleeping function called from invalid context at mm/migrate.c:1741
in_atomic(): 1, irqs_disabled(): 0, non_block: 0, pid: 5809, name: syz.0.17
preempt_count: 1, expected: 0
RCU nest depth: 0, expected: 0
1 lock held by syz.0.17/5809:
 #0: ffff88817766d6d8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff88817766d6d8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
Preemption disabled at:
[<0000000000000000>] 0x0
CPU: 1 UID: 0 PID: 5809 Comm: syz.0.17 Not tainted syzkaller #0 PREEMPT(full) 
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 __might_resched+0x378/0x4d0 kernel/sched/core.c:9197
 migrate_hugetlbs mm/migrate.c:1741 [inline]
 migrate_pages+0x49b/0x27e0 mm/migrate.c:2196
 do_mbind mm/mempolicy.c:1617 [inline]
 kernel_mbind mm/mempolicy.c:1760 [inline]
 __do_sys_mbind mm/mempolicy.c:1834 [inline]
 __se_sys_mbind+0xe82/0x10e0 mm/mempolicy.c:1830
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7fc92ad9ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007fc92bc65028 EFLAGS: 00000246 ORIG_RAX: 00000000000000ed
RAX: ffffffffffffffda RBX: 00007fc92b015fa0 RCX: 00007fc92ad9ce59
RDX: 0000000000000001 RSI: 0000000000600000 RDI: 0000200000000000
RBP: 00007fc92ae32e6f R08: 0000000000000000 R09: 0000000000000003
R10: 0000000000000000 R11: 0000000000000246 R12: 0000000000000000
R13: 00007fc92b016038 R14: 00007fc92b015fa0 R15: 00007ffed35b2a68
 </TASK>

======================================================


***

BUG: sleeping function called from invalid context in migrate_pages_batch

tree:      mm-new
URL:       https://kernel.googlesource.com/pub/scm/linux/kernel/git/akpm/mm.git
base:      61cccb8363fcc282d4ae0555b8739dd227f5ad0b
arch:      amd64
compiler:  Debian clang version 22.1.6 (++20260514074242+fc4aad7b5db3-1~exp1~20260514074407.73), Debian LLD 22.1.6
config:    https://ci.syzbot.org/builds/a27db994-58f5-48f4-a7f5-528b63f5a7c6/config
syz repro: https://ci.syzbot.org/findings/1c235813-ad03-46d5-93c7-2702e3513739/syz_repro

BUG: sleeping function called from invalid context at mm/migrate.c:1923
in_atomic(): 1, irqs_disabled(): 0, non_block: 0, pid: 5820, name: syz.1.20
preempt_count: 1, expected: 0
RCU nest depth: 0, expected: 0
1 lock held by syz.1.20/5820:
 #0: ffff88810903b618 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff88810903b618 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
Preemption disabled at:
[<0000000000000000>] 0x0
CPU: 1 UID: 0 PID: 5820 Comm: syz.1.20 Not tainted syzkaller #0 PREEMPT(full) 
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 __might_resched+0x378/0x4d0 kernel/sched/core.c:9197
 migrate_pages_batch+0x5c5/0x5020 mm/migrate.c:1923
 migrate_pages_sync mm/migrate.c:2115 [inline]
 migrate_pages+0x1a8b/0x27e0 mm/migrate.c:2224
 migrate_to_node mm/mempolicy.c:1322 [inline]
 do_migrate_pages+0x656/0x8b0 mm/mempolicy.c:1421
 kernel_migrate_pages mm/mempolicy.c:1935 [inline]
 __do_sys_migrate_pages mm/mempolicy.c:1953 [inline]
 __se_sys_migrate_pages+0xaa2/0xc90 mm/mempolicy.c:1949
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7f1175d9ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007f1176c21028 EFLAGS: 00000246 ORIG_RAX: 0000000000000100
RAX: ffffffffffffffda RBX: 00007f1176015fa0 RCX: 00007f1175d9ce59
RDX: 0000200000000000 RSI: 0000000000000004 RDI: 0000000000000000
RBP: 00007f1175e32e6f R08: 0000000000000000 R09: 0000000000000000
R10: 0000200000000300 R11: 0000000000000246 R12: 0000000000000000
R13: 00007f1176016038 R14: 00007f1176015fa0 R15: 00007ffd17a9f8a8
 </TASK>

======================================================


***

BUG: sleeping function called from invalid context in move_pages

tree:      mm-new
URL:       https://kernel.googlesource.com/pub/scm/linux/kernel/git/akpm/mm.git
base:      61cccb8363fcc282d4ae0555b8739dd227f5ad0b
arch:      amd64
compiler:  Debian clang version 22.1.6 (++20260514074242+fc4aad7b5db3-1~exp1~20260514074407.73), Debian LLD 22.1.6
config:    https://ci.syzbot.org/builds/a27db994-58f5-48f4-a7f5-528b63f5a7c6/config
syz repro: https://ci.syzbot.org/findings/5cf317a6-38fc-4f7b-aa15-e8cca68ee0f8/syz_repro

BUG: sleeping function called from invalid context at mm/migrate.c:2320
in_atomic(): 1, irqs_disabled(): 0, non_block: 0, pid: 5872, name: syz.0.21
preempt_count: 1, expected: 0
RCU nest depth: 0, expected: 0
1 lock held by syz.0.21/5872:
 #0: ffff88810e2682b8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff88810e2682b8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
Preemption disabled at:
[<0000000000000000>] 0x0
CPU: 0 UID: 0 PID: 5872 Comm: syz.0.21 Not tainted syzkaller #0 PREEMPT(full) 
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 __might_resched+0x378/0x4d0 kernel/sched/core.c:9197
 __might_fault+0x76/0x130 mm/memory.c:7425
 store_status mm/migrate.c:2320 [inline]
 move_pages_and_store_status mm/migrate.c:2431 [inline]
 do_pages_move mm/migrate.c:2525 [inline]
 kernel_move_pages mm/migrate.c:2704 [inline]
 __do_sys_move_pages mm/migrate.c:2718 [inline]
 __se_sys_move_pages+0x1b5d/0x1c80 mm/migrate.c:2713
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7f364019ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007f3640f79028 EFLAGS: 00000246 ORIG_RAX: 0000000000000117
RAX: ffffffffffffffda RBX: 00007f3640415fa0 RCX: 00007f364019ce59
RDX: 00002000000003c0 RSI: 0000000000000002 RDI: 0000000000000000
RBP: 00007f3640232e6f R08: 0000000000000000 R09: 0000000000000000
R10: 0000200000000000 R11: 0000000000000246 R12: 0000000000000000
R13: 00007f3640416038 R14: 00007f3640415fa0 R15: 00007fff4d7049e8
 </TASK>

=============================
[ BUG: Invalid wait context ]
syzkaller #0 Tainted: G        W          
-----------------------------
syz.0.21/5872 is trying to lock:
ffff88816dd71078 (&mm->mmap_lock){++++}-{4:4}, at: __might_fault+0xaf/0x130 mm/memory.c:7427
other info that might help us debug this:
context-{5:5}
1 lock held by syz.0.21/5872:
 #0: ffff88810e2682b8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff88810e2682b8 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
stack backtrace:
CPU: 0 UID: 0 PID: 5872 Comm: syz.0.21 Tainted: G        W           syzkaller #0 PREEMPT(full) 
Tainted: [W]=WARN
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 print_lock_invalid_wait_context kernel/locking/lockdep.c:4830 [inline]
 check_wait_context kernel/locking/lockdep.c:4902 [inline]
 __lock_acquire+0xef0/0x2e50 kernel/locking/lockdep.c:5188
 lock_acquire+0x115/0x350 kernel/locking/lockdep.c:5890
 __might_fault+0xcb/0x130 mm/memory.c:7427
 store_status mm/migrate.c:2320 [inline]
 move_pages_and_store_status mm/migrate.c:2431 [inline]
 do_pages_move mm/migrate.c:2525 [inline]
 kernel_move_pages mm/migrate.c:2704 [inline]
 __do_sys_move_pages mm/migrate.c:2718 [inline]
 __se_sys_move_pages+0x1b5d/0x1c80 mm/migrate.c:2713
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7f364019ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007f3640f79028 EFLAGS: 00000246 ORIG_RAX: 0000000000000117
RAX: ffffffffffffffda RBX: 00007f3640415fa0 RCX: 00007f364019ce59
RDX: 00002000000003c0 RSI: 0000000000000002 RDI: 0000000000000000
RBP: 00007f3640232e6f R08: 0000000000000000 R09: 0000000000000000
R10: 0000200000000000 R11: 0000000000000246 R12: 0000000000000000
R13: 00007f3640416038 R14: 00007f3640415fa0 R15: 00007fff4d7049e8
 </TASK>
BUG: scheduling while atomic: syz.0.21/5872/0x00000002
INFO: lockdep is turned off.
Modules linked in:
Preemption disabled at:
[<0000000000000000>] 0x0


***

BUG: sleeping function called from invalid context in store_status

tree:      mm-new
URL:       https://kernel.googlesource.com/pub/scm/linux/kernel/git/akpm/mm.git
base:      61cccb8363fcc282d4ae0555b8739dd227f5ad0b
arch:      amd64
compiler:  Debian clang version 22.1.6 (++20260514074242+fc4aad7b5db3-1~exp1~20260514074407.73), Debian LLD 22.1.6
config:    https://ci.syzbot.org/builds/a27db994-58f5-48f4-a7f5-528b63f5a7c6/config
syz repro: https://ci.syzbot.org/findings/bd9708eb-bf3f-4bb4-acb1-82ff8e5bb3c2/syz_repro

BUG: sleeping function called from invalid context at mm/migrate.c:2320
in_atomic(): 1, irqs_disabled(): 0, non_block: 0, pid: 5813, name: syz.1.18
preempt_count: 1, expected: 0
RCU nest depth: 0, expected: 0
1 lock held by syz.1.18/5813:
 #0: ffff888176f9ac78 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff888176f9ac78 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
Preemption disabled at:
[<0000000000000000>] 0x0
CPU: 1 UID: 0 PID: 5813 Comm: syz.1.18 Not tainted syzkaller #0 PREEMPT(full) 
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 __might_resched+0x378/0x4d0 kernel/sched/core.c:9197
 __might_fault+0x76/0x130 mm/memory.c:7425
 store_status+0x40/0xb0 mm/migrate.c:2320
 move_pages_and_store_status mm/migrate.c:2431 [inline]
 do_pages_move mm/migrate.c:2513 [inline]
 kernel_move_pages mm/migrate.c:2704 [inline]
 __do_sys_move_pages mm/migrate.c:2718 [inline]
 __se_sys_move_pages+0x10e4/0x1c80 mm/migrate.c:2713
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7fb06159ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007fb062390028 EFLAGS: 00000246 ORIG_RAX: 0000000000000117
RAX: ffffffffffffffda RBX: 00007fb061815fa0 RCX: 00007fb06159ce59
RDX: 00002000000017c0 RSI: 0000000000000003 RDI: 0000000000000000
RBP: 00007fb061632e6f R08: 0000200000001840 R09: 0000000000000002
R10: 0000200000001800 R11: 0000000000000246 R12: 0000000000000000
R13: 00007fb061816038 R14: 00007fb061815fa0 R15: 00007ffcb4b262e8
 </TASK>

=============================
[ BUG: Invalid wait context ]
syzkaller #0 Tainted: G        W          
-----------------------------
syz.1.18/5813 is trying to lock:
ffff88817695d038 (&mm->mmap_lock){++++}-{4:4}, at: __might_fault+0xaf/0x130 mm/memory.c:7427
other info that might help us debug this:
context-{5:5}
1 lock held by syz.1.18/5813:
 #0: ffff888176f9ac78 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: spin_lock include/linux/spinlock.h:342 [inline]
 #0: ffff888176f9ac78 (ptlock_ptr(ptdesc)){+.+.}-{3:3}, at: huge_pte_lock+0x17b/0x1e0 include/linux/hugetlb.h:1318
stack backtrace:
CPU: 1 UID: 0 PID: 5813 Comm: syz.1.18 Tainted: G        W           syzkaller #0 PREEMPT(full) 
Tainted: [W]=WARN
Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.2-debian-1.16.2-1 04/01/2014
Call Trace:
 <TASK>
 dump_stack_lvl+0xe8/0x150 lib/dump_stack.c:120
 print_lock_invalid_wait_context kernel/locking/lockdep.c:4830 [inline]
 check_wait_context kernel/locking/lockdep.c:4902 [inline]
 __lock_acquire+0xef0/0x2e50 kernel/locking/lockdep.c:5188
 lock_acquire+0x115/0x350 kernel/locking/lockdep.c:5890
 __might_fault+0xcb/0x130 mm/memory.c:7427
 store_status+0x40/0xb0 mm/migrate.c:2320
 move_pages_and_store_status mm/migrate.c:2431 [inline]
 do_pages_move mm/migrate.c:2513 [inline]
 kernel_move_pages mm/migrate.c:2704 [inline]
 __do_sys_move_pages mm/migrate.c:2718 [inline]
 __se_sys_move_pages+0x10e4/0x1c80 mm/migrate.c:2713
 do_syscall_x64 arch/x86/entry/syscall_64.c:63 [inline]
 do_syscall_64+0x174/0x580 arch/x86/entry/syscall_64.c:94
 entry_SYSCALL_64_after_hwframe+0x77/0x7f
RIP: 0033:0x7fb06159ce59
Code: ff c3 66 2e 0f 1f 84 00 00 00 00 00 0f 1f 44 00 00 48 89 f8 48 89 f7 48 89 d6 48 89 ca 4d 89 c2 4d 89 c8 4c 8b 4c 24 08 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 c7 c1 e8 ff ff ff f7 d8 64 89 01 48
RSP: 002b:00007fb062390028 EFLAGS: 00000246 ORIG_RAX: 0000000000000117
RAX: ffffffffffffffda RBX: 00007fb061815fa0 RCX: 00007fb06159ce59
RDX: 00002000000017c0 RSI: 0000000000000003 RDI: 0000000000000000
RBP: 00007fb061632e6f R08: 0000200000001840 R09: 0000000000000002
R10: 0000200000001800 R11: 0000000000000246 R12: 0000000000000000
R13: 00007fb061816038 R14: 00007fb061815fa0 R15: 00007ffcb4b262e8
 </TASK>
BUG: scheduling while atomic: syz.1.18/5813/0x00000002
INFO: lockdep is turned off.
Modules linked in:
Preemption disabled at:
[<0000000000000000>] 0x0


***

If these findings have caused you to resend the series or submit a
separate fix, please add the following tag to your commit message:
  Tested-by: syzbot@syzkaller.appspotmail.com

---
This report is generated by a bot. It may contain errors.
syzbot ci engineers can be reached at syzkaller@googlegroups.com.

To test a patch for this bug, please reply with `#syz test`
(should be on a separate line).

The patch should be attached to the email.
Note: arguments like custom git repos and branches are not supported.

Re: [PATCH 0/6] mm: batch rmap walks during large folio migration

[Andrew Morton]

On Sun, 12 Jul 2026 07:38:46 +0000 Shivank Garg <shivankg@amd.com> wrote:

> Performance:
> 
> AMD EPYC ZEN 3, 2-socket system (1 node per socket).
> move_pages() migrating 1 GB from node 0 -> node 1, throughput
> in GB/s:
> 		         batch-	        batch-restore +
> folio     vanilla        restore        unmap          restore%  restore+unmap%
> 4K        2.98±0.09      2.95±0.07      3.06±0.05        -1.0%    +2.6%
> 64K       6.16±0.09      6.24±0.08      7.80±0.08        +1.3%   +26.6%
> 256K      6.85±0.06      6.85±0.09      9.54±0.19        +0.0%   +39.3%
> 1M        7.06±0.09      7.17±0.09      9.89±0.06        +1.5%   +40.1%
> 2M       10.97±0.17     10.67±0.13     10.69±0.15        -2.7%    -2.6%
> 
> In plain migration, the folio copy itself dominate the cost, so these rmap walk
> savings are only partly visible end-to-end.
> When the folio copy is batched-copy and offloaded to DMA engine [1], the copy is
> no longer the bottlneck and rmap-walk batching translates into much larger gains:
> 
> This series + [1], DMA Offload on PTDMA (DCBM), 16 channels:
> 
> 		     batch-copy
> folio   vanilla      -offload(dcbm)  dcbm+restore   dcbm+restore+unmap
> 4K      3.36 ± 0.10   4.16 ± 0.18     4.10 ± 0.14    4.12 ± 0.13
> 64K     6.16 ± 0.47   9.63 ± 0.23    10.00 ± 0.33   14.63 ± 0.14
> 256K    6.70 ± 0.09  12.98 ± 0.18    13.52 ± 0.15   32.30 ± 0.65    (2.38x)
> 1M      6.99 ± 0.02  14.32 ± 0.67    14.84 ± 0.17   38.97 ± 1.03    (2.62x)
> 2M     10.74 ± 0.29  64.25 ± 1.48    65.36 ± 1.00   64.97 ± 1.75

Well dang.

Thanks.  I'll await reviewer input on this series.

Sashiko said things:
	https://sashiko.dev/#/patchset/20260712-migrate-rmap-batch-v1-0-872a734431d1@amd.com

Re: [PATCH 0/6] mm: batch rmap walks during large folio migration

[Garg, Shivank]

On 7/14/2026 7:35 AM, Andrew Morton wrote:
> On Sun, 12 Jul 2026 07:38:46 +0000 Shivank Garg <shivankg@amd.com> wrote:
> 
>> Performance:
>>
>> AMD EPYC ZEN 3, 2-socket system (1 node per socket).
>> move_pages() migrating 1 GB from node 0 -> node 1, throughput
>> in GB/s:
>> 		         batch-	        batch-restore +
>> folio     vanilla        restore        unmap          restore%  restore+unmap%
>> 4K        2.98±0.09      2.95±0.07      3.06±0.05        -1.0%    +2.6%
>> 64K       6.16±0.09      6.24±0.08      7.80±0.08        +1.3%   +26.6%
>> 256K      6.85±0.06      6.85±0.09      9.54±0.19        +0.0%   +39.3%
>> 1M        7.06±0.09      7.17±0.09      9.89±0.06        +1.5%   +40.1%
>> 2M       10.97±0.17     10.67±0.13     10.69±0.15        -2.7%    -2.6%
>>
>> In plain migration, the folio copy itself dominate the cost, so these rmap walk
>> savings are only partly visible end-to-end.
>> When the folio copy is batched-copy and offloaded to DMA engine [1], the copy is
>> no longer the bottlneck and rmap-walk batching translates into much larger gains:
>>
>> This series + [1], DMA Offload on PTDMA (DCBM), 16 channels:
>>
>> 		     batch-copy
>> folio   vanilla      -offload(dcbm)  dcbm+restore   dcbm+restore+unmap
>> 4K      3.36 ± 0.10   4.16 ± 0.18     4.10 ± 0.14    4.12 ± 0.13
>> 64K     6.16 ± 0.47   9.63 ± 0.23    10.00 ± 0.33   14.63 ± 0.14
>> 256K    6.70 ± 0.09  12.98 ± 0.18    13.52 ± 0.15   32.30 ± 0.65    (2.38x)
>> 1M      6.99 ± 0.02  14.32 ± 0.67    14.84 ± 0.17   38.97 ± 1.03    (2.62x)
>> 2M     10.74 ± 0.29  64.25 ± 1.48    65.36 ± 1.00   64.97 ± 1.75
> 
> Well dang.
> 
> Thanks.  I'll await reviewer input on this series.
> 
> Sashiko said things:
> 	https://sashiko.dev/#/patchset/20260712-migrate-rmap-batch-v1-0-872a734431d1@amd.com

Thanks Andrew.
I'm working on this and the syzbot reported issue.

Best regards,
Shivank

Re: [PATCH 0/6] mm: batch rmap walks during large folio migration

[David Hildenbrand (Arm)]

On 7/12/26 09:38, Shivank Garg wrote:
> Migrating a PTE-mapped large folio currently walks the folio's rmap
> one page at a time. It does so twice per migration:
> 1. try_to_migrate_one(): Replace present PTEs with migration entries
> 2. remove_migration_pte(): Restore migration entries to working PTEs.
> 
> For a folio with N contiguous PTEs, each walk repeats page_vma_mapped_walk()
> and the associated operation per subpage of the folio. (256 iterations
> for 1M folio).
> 
> This series enables both walks to process a batch of consecutive PTEs
> that map the same large folio and do some refactoring and clean up on
> the way.
> 
> Performance:
> 
> AMD EPYC ZEN 3, 2-socket system (1 node per socket).
> move_pages() migrating 1 GB from node 0 -> node 1, throughput
> in GB/s:
> 		         batch-	        batch-restore +
> folio     vanilla        restore        unmap          restore%  restore+unmap%
> 4K        2.98±0.09      2.95±0.07      3.06±0.05        -1.0%    +2.6%
> 64K       6.16±0.09      6.24±0.08      7.80±0.08        +1.3%   +26.6%
> 256K      6.85±0.06      6.85±0.09      9.54±0.19        +0.0%   +39.3%
> 1M        7.06±0.09      7.17±0.09      9.89±0.06        +1.5%   +40.1%
> 2M       10.97±0.17     10.67±0.13     10.69±0.15        -2.7%    -2.6%
> 
> In plain migration, the folio copy itself dominate the cost, so these rmap walk
> savings are only partly visible end-to-end.
> When the folio copy is batched-copy and offloaded to DMA engine [1], the copy is
> no longer the bottlneck and rmap-walk batching translates into much larger gains:
> 
> This series + [1], DMA Offload on PTDMA (DCBM), 16 channels:
> 
> 		     batch-copy
> folio   vanilla      -offload(dcbm)  dcbm+restore   dcbm+restore+unmap
> 4K      3.36 ± 0.10   4.16 ± 0.18     4.10 ± 0.14    4.12 ± 0.13
> 64K     6.16 ± 0.47   9.63 ± 0.23    10.00 ± 0.33   14.63 ± 0.14
> 256K    6.70 ± 0.09  12.98 ± 0.18    13.52 ± 0.15   32.30 ± 0.65    (2.38x)
> 1M      6.99 ± 0.02  14.32 ± 0.67    14.84 ± 0.17   38.97 ± 1.03    (2.62x)
> 2M     10.74 ± 0.29  64.25 ± 1.48    65.36 ± 1.00   64.97 ± 1.75
> 
> [1] https://lore.kernel.org/linux-mm/20260630-shivank-batch-migrate-offload-v6-0-da95d7e8b8a2@amd.com
> 
> Applies cleanly on mm-new 10-07-26 (61cccb8363f).
> 
> Signed-off-by: Shivank Garg <shivankg@amd.com>
> ---
> Dev Jain (1):
>       mm/rmap: Add batched version of folio_try_share_anon_rmap_pte

This would ideally be based on Dev's series instead if including a patch from
his series.

-- 
Cheers,

David

Re: [PATCH 3/6] mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()

[David Hildenbrand (Arm)]

On 7/12/26 09:38, Shivank Garg wrote:
> remove_migration_pte() interleaves the regular folio path with the
> hugetlb path via folio_test_hugetlb() and CONFIG_HUGETLB_PAGE
> special cases. Simplify this by seprating the hugetlb handling into
> remove_migration_pte_hugetlb().
> 
> Factor out working-PTE construction from migration entry into
> migration_softleaf_entry_to_pte(), used by both paths.
> 
> While here, convert the VM_BUG_ON_FOLIO() to VM_WARN_ON_ONCE_FOLIO().
> 
> No functional change intended.
> 
> Signed-off-by: Shivank Garg <shivankg@amd.com>
> ---
>  mm/migrate.c | 153 ++++++++++++++++++++++++++++++++++++++---------------------
>  1 file changed, 98 insertions(+), 55 deletions(-)
> 
> diff --git a/mm/migrate.c b/mm/migrate.c
> index 222c8c15f782..8a098402b8c9 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -345,6 +345,89 @@ struct rmap_walk_arg {
>  	bool map_unused_to_zeropage;
>  };
>  
> +static pte_t migration_softleaf_entry_to_pte(struct folio *folio, struct page *new,

"migration_entry" should be sufficient, right? It implies "softleaf".

> +		softleaf_t entry, pte_t old_pte, struct vm_area_struct *vma,
> +		rmap_t *rmap_flags)
> +{
> +	pte_t pte = mk_pte(new, READ_ONCE(vma->vm_page_prot));
> +
> +	if (!softleaf_is_migration_young(entry))
> +		pte = pte_mkold(pte);
> +	if (folio_test_dirty(folio) && softleaf_is_migration_dirty(entry))
> +		pte = pte_mkdirty(pte);
> +	if (pte_swp_soft_dirty(old_pte))
> +		pte = pte_mksoft_dirty(pte);
> +	else
> +		pte = pte_clear_soft_dirty(pte);
> +
> +	if (softleaf_is_migration_write(entry))
> +		pte = pte_mkwrite(pte, vma);
> +	else if (pte_swp_uffd(old_pte))
> +		pte = pte_mkuffd(pte);
> +
> +	/* See do_swap_page(): restore PAGE_NONE for RWP */
> +	if (pte_swp_uffd(old_pte) && userfaultfd_rwp(vma))
> +		pte = pte_modify(pte, PAGE_NONE);
> +
> +	if (folio_test_anon(folio) && !softleaf_is_migration_read(entry))
> +		*rmap_flags |= RMAP_EXCLUSIVE;
> +
> +	return pte;
> +}

Can you move splitting that out into a separate patch? IIUC you can do that even
without messing with the hugetlb stuff.

-- 
Cheers,

David