[PATCH 21/24] arm64: Introduce mte_tag[_and_zero]_range()

["Pierre-Clément Tosi" <ptosi@google.com>]

Carve the KASAN-agnostic MTE logic out of the KASAN arch function, so
that it can be reused without needing to depend on a KASAN header.

Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte-kasan.h | 57 ++++----------------------------------
 arch/arm64/include/asm/mte.h       | 47 +++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 51 deletions(-)

diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index 66d72d0d3c43..81a8cf3b9a95 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -9,7 +9,6 @@
 #include <linux/const.h>
 #include <linux/kasan-tags.h>
 
-#include <asm/cputype.h>
 #include <asm/mte-def.h>
 
 #define __KASAN_TAG_MTE_MASK		GENMASK(MTE_TAG_SIZE - 1, 0)
@@ -32,6 +31,7 @@
 #include <linux/bitfield.h>
 #include <linux/types.h>
 
+#include <asm/mte.h>
 #include <asm/sysreg.h>
 
 #define KASAN_TAG_MTE_MASK		GENMASK_U8(MTE_TAG_SIZE - 1, 0)
@@ -81,58 +81,13 @@ static inline void *mte_set_random_tag(const void *addr)
 	return (void *)__tag_set(addr, kasan_tag_from_mte(mte_tag_of(addr)));
 }
 
-/*
- * Assign allocation tags for a region of memory based on the pointer tag.
- * Note: The address must be non-NULL and MTE_GRANULE_SIZE aligned and
- * size must be MTE_GRANULE_SIZE aligned.
- */
 static inline void mte_set_mem_tag_range(void *addr, size_t size, u8 tag,
-					 bool init)
-{
-	u64 curr, mask, dczid, dczid_bs, dczid_dzp, end1, end2, end3;
-
-	/* Read DC G(Z)VA block size from the system register. */
-	dczid = read_cpuid(DCZID_EL0);
-	dczid_bs = 4ul << (dczid & 0xf);
-	dczid_dzp = (dczid >> 4) & 1;
-
-	curr = (u64)mte_tagged(addr, kasan_tag_to_mte(tag));
-	mask = dczid_bs - 1;
-	/* STG/STZG up to the end of the first block. */
-	end1 = curr | mask;
-	end3 = curr + size;
-	/* DC GVA / GZVA in [end1, end2) */
-	end2 = end3 & ~mask;
-
-	/*
-	 * The following code uses STG on the first DC GVA block even if the
-	 * start address is aligned - it appears to be faster than an alignment
-	 * check + conditional branch. Also, if the range size is at least 2 DC
-	 * GVA blocks, the first two loops can use post-condition to save one
-	 * branch each.
-	 */
-#define SET_MEMTAG_RANGE(stg_post, dc_gva)		\
-	do {						\
-		if (!dczid_dzp && size >= 2 * dczid_bs) {\
-			do {				\
-				curr = stg_post(curr);	\
-			} while (curr < end1);		\
-							\
-			do {				\
-				dc_gva(curr);		\
-				curr += dczid_bs;	\
-			} while (curr < end2);		\
-		}					\
-							\
-		while (curr < end3)			\
-			curr = stg_post(curr);		\
-	} while (0)
-
-	if (init)
-		SET_MEMTAG_RANGE(__stzg_post, __dc_gzva);
+					 bool zero)
+{
+	if (zero)
+		mte_tag_and_zero_range(addr, size, kasan_tag_to_mte(tag));
 	else
-		SET_MEMTAG_RANGE(__stg_post, __dc_gva);
-#undef SET_MEMTAG_RANGE
+		mte_tag_range(addr, size, kasan_tag_to_mte(tag));
 }
 
 void mte_enable_kernel_sync(void);
diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 8826d64c8e33..9ba215b2b0fd 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -18,6 +18,7 @@
 
 #include <asm/barrier.h>
 #include <asm/cpufeature.h>
+#include <asm/cputype.h>
 #include <asm/pgtable-types.h>
 #include <asm/sysreg.h>
 
@@ -101,6 +102,52 @@ static inline bool task_uses_sync_mte_faults(struct task_struct *task)
 	return task->thread.sctlr_user & (1UL << SCTLR_EL1_TCF0_SHIFT);
 }
 
+/*
+ * Assign allocation tags for a region of memory (and optionally zero it).
+ * Note: The address must be non-NULL and MTE_GRANULE_SIZE aligned and
+ * size must be MTE_GRANULE_SIZE aligned.
+ */
+#define __mte_tag_range(addr, size, tag, stg_post, dc_gva)		\
+do {									\
+	u64 curr, mask, dczid, dczid_bs, dczid_dzp, end1, end2, end3;	\
+	/* Read DC G(Z)VA block size from the system register. */	\
+	dczid = read_cpuid(DCZID_EL0);					\
+	dczid_bs = 4ul << (dczid & 0xf);				\
+	dczid_dzp = (dczid >> 4) & 1;					\
+	curr = (u64)mte_tagged(addr, tag);				\
+	mask = dczid_bs - 1;						\
+	/* STG/STZG up to the end of the first block. */		\
+	end1 = curr | mask;						\
+	end3 = curr + size;						\
+	/* DC GVA / GZVA in [end1, end2) */				\
+	end2 = end3 & ~mask;						\
+	/* STG on first block is faster than alignment check + branch.*/\
+	if (!dczid_dzp && size >= 2 * dczid_bs) {			\
+		/* Post-condition saves one branch in each loop. */	\
+		do {							\
+			curr = stg_post(curr);				\
+		} while (curr < end1);					\
+									\
+		do {							\
+			dc_gva(curr);					\
+			curr += dczid_bs;				\
+		} while (curr < end2);					\
+	}								\
+	while (curr < end3)						\
+		curr = stg_post(curr);					\
+} while (0)
+
+static inline void mte_tag_range(void *addr, size_t size, u8 tag)
+{
+	__mte_tag_range(addr, size, tag, __stg_post, __dc_gva);
+}
+
+static inline void mte_tag_and_zero_range(void *addr, size_t size, u8 tag)
+{
+	__mte_tag_range(addr, size, tag, __stzg_post, __dc_gzva);
+}
+#undef __mte_tag_range
+
 #ifdef CONFIG_ARM64_MTE
 
 /* track which pages have valid allocation tags */

-- 
2.56.0.455.g6c5a503b38-goog