[PATCH v3 9/9] mm, sparc: batch arch_unmap_one()

[Dev Jain <dev.jain@arm.com>]

The anonymous large-folio unmap batching path can clear and replace
multiple PTEs at once, but architectures like sparc need arch_unmap_one()
to preserve per-page metadata before swap entries are installed.
sparc uses this hook to save ADI tags.

Extend arch_unmap_one() to operate on a pte range.

sparc uses the count to walk the PTE range and save ADI tags for entries
carrying _PAGE_MCD_4V. There is no reverse operation for a partially
saved range: saving ADI tags only records a copy for later restore, so
if a later entry fails and the core keeps the original PTEs, there is
no architectural state to unsave.

Signed-off-by: Dev Jain <dev.jain@arm.com>
---
 arch/sparc/include/asm/pgtable_64.h | 10 +++++++---
 arch/sparc/kernel/adi_64.c          | 21 +++++++++++++++++++++
 include/linux/pgtable.h             |  6 +++---
 mm/rmap.c                           | 11 +++--------
 4 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/arch/sparc/include/asm/pgtable_64.h b/arch/sparc/include/asm/pgtable_64.h
index 44d1333065a6a..f7dd1872884e4 100644
--- a/arch/sparc/include/asm/pgtable_64.h
+++ b/arch/sparc/include/asm/pgtable_64.h
@@ -1037,6 +1037,9 @@ void adi_restore_tags(struct mm_struct *mm, struct vm_area_struct *vma,
 int adi_save_tags(struct mm_struct *mm, struct vm_area_struct *vma,
 		  unsigned long addr, pte_t oldpte);
 
+int adi_save_tags_range(struct mm_struct *mm, struct vm_area_struct *vma,
+			unsigned long addr, pte_t oldpte, unsigned long nr);
+
 #define __HAVE_ARCH_DO_SWAP_PAGE
 static inline void arch_do_swap_page(struct mm_struct *mm,
 				     struct vm_area_struct *vma,
@@ -1057,10 +1060,11 @@ static inline void arch_do_swap_page(struct mm_struct *mm,
 #define __HAVE_ARCH_UNMAP_ONE
 static inline int arch_unmap_one(struct mm_struct *mm,
 				 struct vm_area_struct *vma,
-				 unsigned long addr, pte_t oldpte)
+				 unsigned long addr, pte_t oldpte,
+				 unsigned long nr)
 {
-	if (adi_state.enabled && (pte_val(oldpte) & _PAGE_MCD_4V))
-		return adi_save_tags(mm, vma, addr, oldpte);
+	if (adi_state.enabled)
+		return adi_save_tags_range(mm, vma, addr, oldpte, nr);
 	return 0;
 }
 
diff --git a/arch/sparc/kernel/adi_64.c b/arch/sparc/kernel/adi_64.c
index 18036a43cf568..ffa378accefe3 100644
--- a/arch/sparc/kernel/adi_64.c
+++ b/arch/sparc/kernel/adi_64.c
@@ -9,6 +9,7 @@
 #include <linux/init.h>
 #include <linux/slab.h>
 #include <linux/mm_types.h>
+#include <linux/pgtable.h>
 #include <asm/mdesc.h>
 #include <asm/adi_64.h>
 #include <asm/mmu_64.h>
@@ -394,3 +395,23 @@ int adi_save_tags(struct mm_struct *mm, struct vm_area_struct *vma,
 
 	return 0;
 }
+
+int adi_save_tags_range(struct mm_struct *mm, struct vm_area_struct *vma,
+			unsigned long addr, pte_t oldpte, unsigned long nr)
+{
+	unsigned long i;
+
+	for (i = 0; i < nr; i++, addr += PAGE_SIZE,
+	     oldpte = pte_next_pfn(oldpte)) {
+		int ret;
+
+		if (!(pte_val(oldpte) & _PAGE_MCD_4V))
+			continue;
+
+		ret = adi_save_tags(mm, vma, addr, oldpte);
+		if (ret)
+			return ret;
+	}
+
+	return 0;
+}
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index e3c8ab96941c5..d5c78dc938760 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -1431,12 +1431,12 @@ static inline void arch_do_swap_page_nr(struct mm_struct *mm,
  * restored when the page is swapped back in. SPARC M7 and newer
  * processors support an ADI (Application Data Integrity) tag for the
  * page as metadata for the page. arch_unmap_one() can save this
- * metadata on a swap-out of a page.
+ * metadata on swap-out of one or more pages.
  */
 static inline int arch_unmap_one(struct mm_struct *mm,
 				  struct vm_area_struct *vma,
-				  unsigned long addr,
-				  pte_t orig_pte)
+				  unsigned long addr, pte_t orig_pte,
+				  unsigned long nr)
 {
 	return 0;
 }
diff --git a/mm/rmap.c b/mm/rmap.c
index cdda9efd80bc3..774ddc75194b1 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1959,6 +1959,7 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 
 	if (flags & TTU_HWPOISON)
 		return 1;
+
 	if (!folio_test_large(folio))
 		return 1;
 
@@ -1969,11 +1970,6 @@ static inline unsigned int folio_unmap_pte_batch(struct folio *folio,
 	if (pte_unused(pte))
 		return 1;
 
-#ifdef __HAVE_ARCH_UNMAP_ONE
-	if (folio_test_anon(folio) && folio_test_swapbacked(folio))
-		return 1;
-#endif
-
 	/*
 	 * If unmap fails, we need to restore the ptes. To avoid accidentally
 	 * upgrading write permissions for ptes that were not originally
@@ -2168,8 +2164,7 @@ static bool __ttu_anon_swapbacked_folio(struct vm_area_struct *vma,
 	 * architectures where we could have PFN swap PTEs,
 	 * so we'll not check/care.
 	 */
-	if (arch_unmap_one(mm, vma, address, pteval) < 0) {
-		VM_WARN_ON(nr_pages != 1);
+	if (arch_unmap_one(mm, vma, address, pteval, nr_pages) < 0) {
 		folio_put_swap_pages(folio, page, nr_pages);
 		return false;
 	}
@@ -2746,7 +2741,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 * architectures where we could have PFN swap PTEs,
 			 * so we'll not check/care.
 			 */
-			if (arch_unmap_one(mm, vma, address, pteval) < 0) {
+			if (arch_unmap_one(mm, vma, address, pteval, 1) < 0) {
 				if (folio_test_hugetlb(folio))
 					set_huge_pte_at(mm, address, pvmw.pte,
 							pteval, hsz);
-- 
2.43.0