Use DC GVA / DC GZVA to speed up KASan memory tagging in HW tags mode.
The first cacheline is always tagged using STG/STZG even if the address is
cacheline-aligned, as benchmarks show it is faster than a conditional
branch.
Signed-off-by: Evgenii Stepanov <[email protected]>
Co-developed-by: Peter Collingbourne <[email protected]>
Signed-off-by: Peter Collingbourne <[email protected]>
---
Changelog since v1:
- Added Co-developed-by.
Changelog since v2:
- Added Signed-off-by.
Changelog since v3:
- Move the implementation back to C with a bit of inline asm.
arch/arm64/include/asm/mte-kasan.h | 98 +++++++++++++++++++++---------
1 file changed, 70 insertions(+), 28 deletions(-)
diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
index ddd4d17cf9a0..34e23886f346 100644
--- a/arch/arm64/include/asm/mte-kasan.h
+++ b/arch/arm64/include/asm/mte-kasan.h
@@ -48,43 +48,85 @@ static inline u8 mte_get_random_tag(void)
return mte_get_ptr_tag(addr);
}
+static inline u64 __stg_post(u64 p)
+{
+ asm volatile(__MTE_PREAMBLE "stg %0, [%0], #16"
+ : "+r"(p)
+ :
+ : "memory");
+ return p;
+}
+
+static inline u64 __stzg_post(u64 p)
+{
+ asm volatile(__MTE_PREAMBLE "stzg %0, [%0], #16"
+ : "+r"(p)
+ :
+ : "memory");
+ return p;
+}
+
+static inline void __dc_gva(u64 p)
+{
+ asm volatile(__MTE_PREAMBLE "dc gva, %0" : : "r"(p) : "memory");
+}
+
+static inline void __dc_gzva(u64 p)
+{
+ asm volatile(__MTE_PREAMBLE "dc gzva, %0" : : "r"(p) : "memory");
+}
+
/*
* 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 non-zero and MTE_GRANULE_SIZE aligned.
+ * size must be MTE_GRANULE_SIZE aligned.
*/
-static inline void mte_set_mem_tag_range(void *addr, size_t size,
- u8 tag, bool init)
+static inline void mte_set_mem_tag_range(void *addr, size_t size, u8 tag,
+ bool init)
{
- u64 curr, end;
+ u64 curr, DCZID, mask, line_size, end1, end2, end3;
- if (!size)
- return;
+ /* Read DC G(Z)VA store size from the register. */
+ __asm__ __volatile__(__MTE_PREAMBLE "mrs %0, dczid_el0"
+ : "=r"(DCZID)::);
+ line_size = 4ul << (DCZID & 0xf);
curr = (u64)__tag_set(addr, tag);
- end = curr + size;
-
- /*
- * 'asm volatile' is required to prevent the compiler to move
- * the statement outside of the loop.
+ mask = line_size - 1;
+ /* STG/STZG up to the end of the first cache line. */
+ end1 = curr | mask;
+ end3 = curr + size;
+ /* DC GVA / GZVA in [end1, end2) */
+ end2 = end3 & ~mask;
+
+ /* The following code uses STG on the first cache line even if the start
+ * address is cache line aligned - it appears to be faster than an
+ * alignment check + conditional branch. Also, if the size is at least 2
+ * cache lines, the first two loops can use post-condition to save one
+ * branch each.
*/
- if (init) {
- do {
- asm volatile(__MTE_PREAMBLE "stzg %0, [%0]"
- :
- : "r" (curr)
- : "memory");
- curr += MTE_GRANULE_SIZE;
- } while (curr != end);
- } else {
- do {
- asm volatile(__MTE_PREAMBLE "stg %0, [%0]"
- :
- : "r" (curr)
- : "memory");
- curr += MTE_GRANULE_SIZE;
- } while (curr != end);
- }
+#define SET_MEMTAG_RANGE(stg_post, dc_gva) \
+ do { \
+ if (size >= 2 * line_size) { \
+ do { \
+ curr = stg_post(curr); \
+ } while (curr < end1); \
+ \
+ do { \
+ dc_gva(curr); \
+ curr += line_size; \
+ } while (curr < end2); \
+ } \
+ \
+ while (curr < end3) \
+ curr = stg_post(curr); \
+ } while (0)
+
+ if (init)
+ SET_MEMTAG_RANGE(__stzg_post, __dc_gzva);
+ else
+ SET_MEMTAG_RANGE(__stg_post, __dc_gva);
+#undef SET_MEMTAG_RANGE
}
void mte_enable_kernel_sync(void);
--
2.31.1.751.gd2f1c929bd-goog
On Wed, May 19, 2021 at 07:03:05PM -0700, Evgenii Stepanov wrote:
> Use DC GVA / DC GZVA to speed up KASan memory tagging in HW tags mode.
>
> The first cacheline is always tagged using STG/STZG even if the address is
> cacheline-aligned, as benchmarks show it is faster than a conditional
> branch.
>
> Signed-off-by: Evgenii Stepanov <[email protected]>
> Co-developed-by: Peter Collingbourne <[email protected]>
> Signed-off-by: Peter Collingbourne <[email protected]>
Some nitpicks below but it looks fine otherwise.
> diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
> index ddd4d17cf9a0..34e23886f346 100644
> --- a/arch/arm64/include/asm/mte-kasan.h
> +++ b/arch/arm64/include/asm/mte-kasan.h
> @@ -48,43 +48,85 @@ static inline u8 mte_get_random_tag(void)
> return mte_get_ptr_tag(addr);
> }
>
> +static inline u64 __stg_post(u64 p)
> +{
> + asm volatile(__MTE_PREAMBLE "stg %0, [%0], #16"
> + : "+r"(p)
> + :
> + : "memory");
> + return p;
> +}
> +
> +static inline u64 __stzg_post(u64 p)
> +{
> + asm volatile(__MTE_PREAMBLE "stzg %0, [%0], #16"
> + : "+r"(p)
> + :
> + : "memory");
> + return p;
> +}
> +
> +static inline void __dc_gva(u64 p)
> +{
> + asm volatile(__MTE_PREAMBLE "dc gva, %0" : : "r"(p) : "memory");
> +}
> +
> +static inline void __dc_gzva(u64 p)
> +{
> + asm volatile(__MTE_PREAMBLE "dc gzva, %0" : : "r"(p) : "memory");
> +}
> +
> /*
> * 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 non-zero and MTE_GRANULE_SIZE aligned.
> + * size must be MTE_GRANULE_SIZE aligned.
> */
> -static inline void mte_set_mem_tag_range(void *addr, size_t size,
> - u8 tag, bool init)
> +static inline void mte_set_mem_tag_range(void *addr, size_t size, u8 tag,
> + bool init)
> {
> - u64 curr, end;
> + u64 curr, DCZID, mask, line_size, end1, end2, end3;
Nitpick 1: please use lowercase for variables even if they match some
register.
>
> - if (!size)
> - return;
> + /* Read DC G(Z)VA store size from the register. */
> + __asm__ __volatile__(__MTE_PREAMBLE "mrs %0, dczid_el0"
> + : "=r"(DCZID)::);
> + line_size = 4ul << (DCZID & 0xf);
No need for __MTE_PREAMBLE here, this register has been available since
8.0. Even better, just use read_cpuid(DCZID_EL0) directly rather than
asm.
I'd also call this variable block_size (or dczid_bs etc.), it's not
necessarily a cache line size (we have CTR_EL0 for that), though most
implementations probably do just that. There are a few instances below
where the comments refer to cache lines.
> curr = (u64)__tag_set(addr, tag);
> - end = curr + size;
> -
> - /*
> - * 'asm volatile' is required to prevent the compiler to move
> - * the statement outside of the loop.
> + mask = line_size - 1;
> + /* STG/STZG up to the end of the first cache line. */
> + end1 = curr | mask;
> + end3 = curr + size;
> + /* DC GVA / GZVA in [end1, end2) */
> + end2 = end3 & ~mask;
> +
> + /* The following code uses STG on the first cache line even if the start
> + * address is cache line aligned - it appears to be faster than an
> + * alignment check + conditional branch. Also, if the size is at least 2
> + * cache lines, the first two loops can use post-condition to save one
> + * branch each.
> */
Nitpick 2: the multiline comments start with an empty /* (as per the
coding style doc).
> - if (init) {
> - do {
> - asm volatile(__MTE_PREAMBLE "stzg %0, [%0]"
> - :
> - : "r" (curr)
> - : "memory");
> - curr += MTE_GRANULE_SIZE;
> - } while (curr != end);
> - } else {
> - do {
> - asm volatile(__MTE_PREAMBLE "stg %0, [%0]"
> - :
> - : "r" (curr)
> - : "memory");
> - curr += MTE_GRANULE_SIZE;
> - } while (curr != end);
> - }
> +#define SET_MEMTAG_RANGE(stg_post, dc_gva) \
> + do { \
> + if (size >= 2 * line_size) { \
> + do { \
> + curr = stg_post(curr); \
> + } while (curr < end1); \
> + \
> + do { \
> + dc_gva(curr); \
> + curr += line_size; \
> + } while (curr < end2); \
> + } \
> + \
> + while (curr < end3) \
> + curr = stg_post(curr); \
> + } while (0)
> +
> + if (init)
> + SET_MEMTAG_RANGE(__stzg_post, __dc_gzva);
> + else
> + SET_MEMTAG_RANGE(__stg_post, __dc_gva);
> +#undef SET_MEMTAG_RANGE
> }
>
> void mte_enable_kernel_sync(void);
> --
> 2.31.1.751.gd2f1c929bd-goog
With the above fixed, feel free to add:
Reviewed-by: Catalin Marinas <[email protected]>
(aiming at 5.14)
--
Catalin