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

[PATCH v2 0/7] 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-only   + unmap         restore%   restore+unmap%
4K        2.96±0.15     3.02±0.14      2.99±0.17         +2.2%    +1.0%
64K       5.99±0.10     6.12±0.08      7.65±0.16         +2.2%   +27.7%
256K      6.48±0.09     6.67±0.08      9.23±0.14         +2.9%   +42.4%
1M        6.76±0.10     6.98±0.10      9.74±0.19         +3.1%   +44.0%
2M       10.50±0.20    10.44±0.16     10.67±0.21         -0.5%    +1.6%

In plain migration, the folio copy itself dominates 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 bottleneck and rmap-walk batching translates into much larger gains:

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

		    batch-copy                      dcbm+restore
folio    vanilla    -offload(dcbm)  dcbm+restore    +unmap          (rmap batch impact)
4K       3.05±0.13   2.76±0.11       2.77±0.12        2.74±0.09
64K      5.85±0.11  10.66±0.27      11.56±1.45       17.55±0.56
256K     6.26±0.08  13.64±2.73      13.68±0.16       34.76±0.57       (2.54x)
1M       6.58±0.08  14.25±0.07      15.27±0.77       44.28±0.27       (3.11x)
2M      10.61±0.24  65.26±0.96      65.25±1.65       65.66±0.87

Applies cleanly on mm-new+[2].

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

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
Changes in v2:

- Rebase onto mm-new, drop the borrowed folio_try_share_anon_rmap_ptes() patch
  and add it as dependency on Dev's series and adapt to recent migrate/rmap changes.
- Switch to recently added batched helpers - finish_folio_unmap, set_softleaf_ptes,
  and page_anon_exclusive_batch.
- Rename migration_softleaf_entry_to_pte() to migration_entry_to_pte(). (David)
- Split PTE construction from hugetlb handling into separate patch. (David)
  Mirror the split in mm/rmap.c patch as well.
- Rename make_migration_swp_pte() to make_migration_pte().
- Early exit the single-entry hugetlb page_vma_mapped_walk() paths - removes
  unnecessary nesting.
- Add the missing page_vma_mapped_walk_done() in remove_migration_pte_hugetlb(),
  which leaked the huge-PTE lock. (syzbot ci, Sashiko)
- Fix handling of hwpoisoned subpages in the batched unmap. Read PageHWPoison()
  once and skip batching when mapped page is poisoned, and use same result for
  hwpoison handling. (Sashiko)
- Remove the pte_unused(pteval) handling from try_to_migrate_hugetlb_one() (Sashiko).
- Link to v1: https://lore.kernel.org/r/20260712-migrate-rmap-batch-v1-0-872a734431d1@amd.com

---
Shivank Garg (7):
      mm: factor out generic PTE batch detection from swap_pte_batch()
      mm/migrate: factor out migration PTE construction
      mm/migrate: split remove_migration_pte_hugetlb() out of remove_migration_pte()
      mm/migrate: batch the restore-side migration rmap walk
      mm/rmap: factor out migration PTE construction
      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()

 mm/internal.h |  36 ++---
 mm/migrate.c  | 191 ++++++++++++++++++--------
 mm/rmap.c     | 419 ++++++++++++++++++++++++++++++++++++----------------------
 3 files changed, 416 insertions(+), 230 deletions(-)
---
base-commit: 480a31230b426efb005b6e71a14ef80f405f18b6
change-id: 20260701-migrate-rmap-batch-8ecc04de1905
prerequisite-message-id: 20260723070905.3422276-1-dev.jain@arm.com
prerequisite-patch-id: 5fdb03c3ea4a037f66fbd1c9b5d04e24e86406eb
prerequisite-patch-id: 033980575a1ef89d70e43a67c5530e2cf6cf97aa
prerequisite-patch-id: 6f7508334fc2bfc0df104ac50d36ed88072487e4
prerequisite-patch-id: 3f1aaae36e02d8b9343f7e719f732f10c1767572
prerequisite-patch-id: 98bc99c6bc274594946d022b3998f5f1dc70012f
prerequisite-patch-id: dd35ec00b82bf235f2e41e286d16eab6a3fd1e9e
prerequisite-patch-id: 552ee939dc6a7b3cce51373f4e75fa01a9d2b979
prerequisite-patch-id: 485cdbbe4e65b13494dc5104d4b8b7d0d23911a4

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

[PATCH v2 1/7] 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 a5d0488a54dc..1e1ab53b642e 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -560,6 +560,26 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
 	}
 }
 
+static inline int softleaf_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
+{
+	pte_t expected_pte = pte_next_softleaf_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_softleaf_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.
@@ -577,23 +597,9 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
  */
 static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {
-	pte_t expected_pte = pte_next_softleaf_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_softleaf_offset(expected_pte);
-		ptep++;
-	}
-
-	return ptep - start_ptep;
+	return softleaf_pte_batch(start_ptep, max_nr, pte);
 }
 #endif /* CONFIG_MMU */
 

-- 
2.43.0

[PATCH v2 2/7] mm/migrate: factor out migration PTE construction

[Shivank Garg]

remove_migration_pte() constructs a working-PTE from a migration entry
inline. Factor it into migration_entry_to_pte() so the logic can be
shared when the hugetlb restore path is separated.

No functional change intended.

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

diff --git a/mm/migrate.c b/mm/migrate.c
index 15b45832bcfa..a3362cc9ef66 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -345,6 +345,36 @@ struct rmap_walk_arg {
 	bool map_unused_to_zeropage;
 };
 
+static pte_t migration_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
  */
@@ -387,27 +417,8 @@ static bool remove_migration_pte(struct folio *folio,
 
 		folio_get(folio);
 		new = folio_page(folio, idx);
-		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;
+		pte = migration_entry_to_pte(folio, new, entry, old_pte, vma,
+					     &rmap_flags);
 
 		if (unlikely(is_device_private_page(new))) {
 			if (pte_write(pte))

-- 
2.43.0

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

[Shivank Garg]

remove_migration_pte() interleaves hugetlb handling with the regular
folio path. Move it into a dedicated callback selected by
remove_migration_ptes(), leaving the generic callback focused on regular
folios ahead of PTE batching.

With hugetlb folios routed separately, simplify the PMD mapping check to
a one-time warning.

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

diff --git a/mm/migrate.c b/mm/migrate.c
index a3362cc9ef66..ee1b8a55a2a4 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -375,6 +375,57 @@ static pte_t migration_entry_to_pte(struct folio *folio, struct page *new,
 	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);
+	struct page *new = folio_page(folio, 0);
+	rmap_t rmap_flags = RMAP_NONE;
+	pte_t old_pte, pte;
+	softleaf_t entry;
+
+	/* There is only a single mapping in a VMA. */
+	if (!page_vma_mapped_walk(&pvmw))
+		return true;
+
+	old_pte = huge_ptep_get(vma->vm_mm, pvmw.address, pvmw.pte);
+	entry = softleaf_from_pte(old_pte);
+	folio_get(folio);
+	pte = migration_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);
+	page_vma_mapped_walk_done(&pvmw);
+
+	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
  */
@@ -395,20 +446,14 @@ static bool remove_migration_pte(struct folio *folio,
 #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, folio);
 			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);
 		entry = softleaf_from_pte(old_pte);
-		if (folio_test_large(folio) && !folio_test_hugetlb(folio))
+		if (folio_test_large(folio))
 			idx = softleaf_to_pfn(entry) - pvmw.pfn;
 
 		if (rmap_walk_arg->map_unused_to_zeropage &&
@@ -434,30 +479,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();
 
@@ -484,7 +511,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 v2 4/7] 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 ee1b8a55a2a4..aad555453469 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_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)
@@ -433,11 +455,13 @@ 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)) {
 		rmap_t rmap_flags = RMAP_NONE;
 		unsigned long idx = 0;
+		unsigned int nr = 1;
 		softleaf_t entry;
 		struct page *new;
 		pte_t old_pte;
@@ -460,12 +484,19 @@ static bool remove_migration_pte(struct folio *folio,
 		    try_to_map_unused_to_zeropage(&pvmw, folio, old_pte, idx))
 			continue;
 
-		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);
 		new = folio_page(folio, idx);
 		pte = migration_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));
@@ -480,11 +511,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();
 
@@ -492,7 +523,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 v2 5/7] mm/rmap: factor out migration PTE construction

[Shivank Garg]

try_to_migrate_one() constructs the migration PTE that replaces a
present PTE or non-present swap PTE.

Factor it into make_migration_pte() which can be shared when the
hugetlb migration path is separated.

No functional change intended.

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

diff --git a/mm/rmap.c b/mm/rmap.c
index 0a4a0945dc65..8b89e32ae489 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2507,6 +2507,42 @@ void try_to_unmap(struct folio *folio, enum ttu_flags flags)
 		rmap_walk(folio, &rwc);
 }
 
+/* Build the migration PTE that replaces pteval for a page. */
+static pte_t make_migration_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;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument.
  *
@@ -2737,7 +2773,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;
 
 			/*
@@ -2779,32 +2814,8 @@ 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);
-			}
+			swp_pte = make_migration_pte(subpage, pteval,
+						     writable, anon_exclusive);
 			if (folio_test_hugetlb(folio))
 				set_huge_pte_at(mm, address, pvmw.pte, swp_pte,
 						hsz);

-- 
2.43.0

[PATCH v2 6/7] 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().

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 | 251 +++++++++++++++++++++++++++++++++++---------------------------
 1 file changed, 144 insertions(+), 107 deletions(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index 8b89e32ae489..35752a70f3a0 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2543,6 +2543,138 @@ static pte_t make_migration_pte(struct page *page, pte_t pteval,
 	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);
+
+	/* There is only a single mapping in a VMA. */
+	if (!page_vma_mapped_walk(&pvmw))
+		goto range_end;
+
+	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 {
+		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. A
+		 * hugetlb fault waits for migration to complete before retrying.
+		 */
+		swp_pte = make_migration_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);
+range_end:
+	mmu_notifier_invalidate_range_end(&range);
+
+	return ret;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument.
  *
@@ -2560,7 +2692,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,
@@ -2573,26 +2704,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)) {
@@ -2639,10 +2757,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 {
@@ -2653,66 +2768,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)) {
@@ -2751,14 +2813,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)) {
 			/*
@@ -2781,28 +2837,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);
@@ -2816,11 +2859,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			swp_pte = make_migration_pte(subpage, pteval,
 						     writable, anon_exclusive);
-			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);
+			set_pte_at(mm, address, pvmw.pte, swp_pte);
 			trace_set_migration_pte(address, pte_val(swp_pte),
 						folio_order(folio));
 			/*
@@ -2829,10 +2868,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);
@@ -2854,7 +2890,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 v2 7/7] 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 | 115 ++++++++++++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 86 insertions(+), 29 deletions(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index 35752a70f3a0..63b885c0b7ef 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2675,6 +2675,44 @@ 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;
+
+#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))
+		nr = page_anon_exclusive_batch(0, nr, subpage, anon_exclusive);
+
+	return nr;
+}
+
 /*
  * @arg: enum ttu_flags will be passed to this argument.
  *
@@ -2686,12 +2724,12 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 {
 	struct mm_struct *mm = vma->vm_mm;
 	DEFINE_FOLIO_VMA_WALK(pvmw, folio, vma, address, 0);
-	bool anon_exclusive, writable, ret = true;
+	bool anon_exclusive, hwpoison, writable, ret = true;
 	pte_t pteval;
 	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,
@@ -2744,11 +2782,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
 		}
@@ -2773,10 +2808,25 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 		subpage = folio_page(folio, pfn - folio_pfn(folio));
 		anon_exclusive = folio_test_anon(folio) &&
 				 PageAnonExclusive(subpage);
+		/*
+		 * memory_failure() can set PageHWPoison concurrently without holding
+		 * the folio lock. Snapshot the flag here to decide whether to batch
+		 * PTEs or install hwpoison entry.
+		 */
+		hwpoison = PageHWPoison(subpage);
 
+		nr_pages = 1;
 		if (likely(pte_present(pteval))) {
-			flush_cache_page(vma, address, pfn);
-			/* Nuke the page table entry. */
+			if (!hwpoison)
+				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
@@ -2786,11 +2836,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);
@@ -2809,7 +2857,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 		/* Update high watermark before we lower rss */
 		update_hiwater_rss(mm);
 
-		if (PageHWPoison(subpage)) {
+		if (hwpoison) {
+			VM_WARN_ON_ONCE(nr_pages != 1);
 			VM_WARN_ON_FOLIO(folio_is_device_private(folio), folio);
 
 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
@@ -2837,19 +2886,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;
 			}
 
 			/*
@@ -2859,19 +2904,31 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 */
 			swp_pte = make_migration_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. */
+			set_softleaf_ptes(mm, address, pvmw.pte, swp_pte, nr_pages);
 			/*
 			 * No need to invalidate here it will synchronize on
 			 * against the special swap migration pte.
 			 */
 		}
 
-		folio_remove_rmap_pte(folio, subpage, vma);
-		if (vma->vm_flags & VM_LOCKED)
-			mlock_drain_local();
-		folio_put(folio);
+		finish_folio_unmap(vma, folio, subpage, 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

Re: [PATCH v2 1/7] mm: factor out generic PTE batch detection from swap_pte_batch()

[David Hildenbrand (Arm)]

On 8/13/26 06:23, Shivank Garg wrote:
> 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 a5d0488a54dc..1e1ab53b642e 100644
> --- a/mm/internal.h
> +++ b/mm/internal.h
> @@ -560,6 +560,26 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
>  	}
>  }
>  
> +static inline int softleaf_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
> +{
> +	pte_t expected_pte = pte_next_softleaf_offset(pte);
> +	const pte_t *end_ptep = start_ptep + max_nr;
> +	pte_t *ptep = start_ptep + 1;
> +

Wouldn't we want some way to make sure that the softleaf entries actually carry
offsets or pfns of sorts?

> +	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_softleaf_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.
> @@ -577,23 +597,9 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
>   */
>  static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
>  {
> -	pte_t expected_pte = pte_next_softleaf_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)));

Any reason we cannot simply rename this thing and drop the VM_WARN_ON? I mean, a
caller can just do that.

-- 
Cheers,

David