[PATCH v2 4/4] arm64: mte: Add mte_clear_folio_tags()

[Fuad Tabba]

mte_sync_tags(), sanitise_mte_tags() and kvm_vm_ioctl_mte_copy_tags()
each open-code the loop that clears the tags of every page in a hugetlb
folio. Add mte_clear_folio_tags() and use it in all three.
kvm_vm_ioctl_mte_copy_tags() then tests for hugetlb first, as it does
where it sets the tagged flag below.

No functional change intended.

Suggested-by: Lorenzo Stoakes (ARM) <ljs at kernel.org>
Signed-off-by: Fuad Tabba <fuad.tabba at linux.dev>
---
 arch/arm64/include/asm/mte.h |  4 ++++
 arch/arm64/kernel/mte.c      | 15 ++++++++++-----
 arch/arm64/kvm/guest.c       | 10 +++++-----
 arch/arm64/kvm/mmu.c         |  3 +--
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 7f7b97e099968..5a2d2814c0416 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -96,6 +96,7 @@ static inline bool try_page_mte_tagging(struct page *page)
 }
 
 void mte_zero_clear_page_tags(void *addr);
+void mte_clear_folio_tags(struct folio *folio);
 void mte_sync_tags(pte_t pte, unsigned int nr_pages);
 void mte_copy_page_tags(void *kto, const void *kfrom);
 void mte_thread_init_user(void);
@@ -128,6 +129,9 @@ static inline bool try_page_mte_tagging(struct page *page)
 static inline void mte_zero_clear_page_tags(void *addr)
 {
 }
+static inline void mte_clear_folio_tags(struct folio *folio)
+{
+}
 static inline void mte_sync_tags(pte_t pte, unsigned int nr_pages)
 {
 }
diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index 1a9aad6ef22a0..fb2598f8c0ccd 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -36,6 +36,14 @@ DEFINE_STATIC_KEY_FALSE(mte_async_or_asymm_mode);
 EXPORT_SYMBOL_GPL(mte_async_or_asymm_mode);
 #endif
 
+void mte_clear_folio_tags(struct folio *folio)
+{
+	unsigned long i;
+
+	for (i = 0; i < folio_nr_pages(folio); i++)
+		mte_clear_page_tags(page_address(folio_page(folio, i)));
+}
+
 void mte_sync_tags(pte_t pte, unsigned int nr_pages)
 {
 	struct page *page = pte_page(pte);
@@ -43,12 +51,9 @@ void mte_sync_tags(pte_t pte, unsigned int nr_pages)
 	unsigned long i;
 
 	if (folio_test_hugetlb(folio)) {
-		unsigned long nr = folio_nr_pages(folio);
-
-		/* Hugetlb MTE flags are set for head page only */
+		/* A hugetlb folio is marked tagged as a whole, so clear every page */
 		if (folio_try_hugetlb_mte_tagging(folio)) {
-			for (i = 0; i < nr; i++, page++)
-				mte_clear_page_tags(page_address(page));
+			mte_clear_folio_tags(folio);
 			folio_set_hugetlb_mte_tagged(folio);
 		}
 
diff --git a/arch/arm64/kvm/guest.c b/arch/arm64/kvm/guest.c
index cda655fbff657..c75466384f8ac 100644
--- a/arch/arm64/kvm/guest.c
+++ b/arch/arm64/kvm/guest.c
@@ -1012,7 +1012,6 @@ int kvm_vm_ioctl_mte_copy_tags(struct kvm *kvm,
 		void *maddr;
 		unsigned long num_tags;
 		struct folio *folio;
-		unsigned long i;
 
 		if (!page) {
 			ret = -EFAULT;
@@ -1046,11 +1045,12 @@ int kvm_vm_ioctl_mte_copy_tags(struct kvm *kvm,
 			 * tags. A hugetlb folio is marked tagged as a whole,
 			 * so its first tagging clears every page.
 			 */
-			if (!folio_test_hugetlb(folio))
+			if (folio_test_hugetlb(folio)) {
+				if (folio_try_hugetlb_mte_tagging(folio))
+					mte_clear_folio_tags(folio);
+			} else {
 				try_page_mte_tagging(page);
-			else if (folio_try_hugetlb_mte_tagging(folio))
-				for (i = 0; i < folio_nr_pages(folio); i++)
-					mte_clear_page_tags(page_address(folio_page(folio, i)));
+			}
 			num_tags = mte_copy_tags_from_user(maddr, tags,
 							MTE_GRANULES_PER_PAGE);
 
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 6e10aa9913d71..f8139c0a54deb 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -1514,8 +1514,7 @@ static void sanitise_mte_tags(struct kvm *kvm, kvm_pfn_t pfn,
 	if (folio_test_hugetlb(folio)) {
 		/* A hugetlb folio is marked tagged as a whole, so clear every page */
 		if (folio_try_hugetlb_mte_tagging(folio)) {
-			for (i = 0; i < folio_nr_pages(folio); i++)
-				mte_clear_page_tags(page_address(folio_page(folio, i)));
+			mte_clear_folio_tags(folio);
 			folio_set_hugetlb_mte_tagged(folio);
 		}
 		return;
-- 
2.47.3