[PATCH v2] lib/crc: arm64: Add NEON PMULL+EOR3 crc64_nvme implementation

[Shreesh Adiga <16567adigashreesh@gmail.com>]

Add a 4x-folding implementation for CPUs with FEAT_SHA3 which benefits
newer ARMv9 CPUs like Cortex-X925.
Usage of EOR3 from SHA3 extension improves performance by 20% compared
to the regular PMULL without EOR3 on Cortex-X925 measure in userspace
application via libbenchmark for 512 byte input size.
Measured about 3x speedup for 512 byte input compared with the existing
NEON single fold implementation on Cortex-X925 in userspace.

Signed-off-by: Shreesh Adiga <16567adigashreesh@gmail.com>
---
Changes in v2:
  Improved the 4x to 1x reduction by using x^255 and x^319 x2 reduction.
  Added static keys for the CPU feature checks instead of regular
  branches.

 lib/crc/arm64/crc64.h | 24 +++++++++++--
 lib/crc/crc64-neon.c  | 83 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/lib/crc/arm64/crc64.h b/lib/crc/arm64/crc64.h
index c7a69e1f3d8f..e5be0ef21140 100644
--- a/lib/crc/arm64/crc64.h
+++ b/lib/crc/arm64/crc64.h
@@ -9,20 +9,38 @@
 #include <linux/sizes.h>
 
 u64 crc64_nvme_neon(u64 crc, const u8 *p, size_t len);
+u64 crc64_nvme_neon_pmull_eor3(u64 crc, const u8 *p, size_t len);
+
+static __ro_after_init DEFINE_STATIC_KEY_FALSE(have_eor3);
+static __ro_after_init DEFINE_STATIC_KEY_FALSE(have_pmull);
 
 #define crc64_be_arch crc64_be_generic
 
 static inline u64 crc64_nvme_arch(u64 crc, const u8 *p, size_t len)
 {
-	if (len >= 128 && cpu_have_named_feature(PMULL) &&
+	if (len >= 128 && static_branch_likely(&have_pmull) &&
 	    likely(may_use_simd())) {
 		size_t chunk = len & ~15;
 
-		scoped_ksimd()
-			crc = crc64_nvme_neon(crc, p, chunk);
+		scoped_ksimd() {
+			if (static_branch_likely(&have_eor3))
+				crc = crc64_nvme_neon_pmull_eor3(crc, p, chunk);
+			else
+				crc = crc64_nvme_neon(crc, p, chunk);
+		}
 
 		p += chunk;
 		len &= 15;
 	}
 	return crc64_nvme_generic(crc, p, len);
 }
+
+#define crc64_mod_init_arch crc64_mod_init_arch
+static void crc64_mod_init_arch(void)
+{
+	if (cpu_have_named_feature(PMULL)) {
+		static_branch_enable(&have_pmull);
+		if (cpu_have_named_feature(SHA3))
+			static_branch_enable(&have_eor3);
+	}
+}
diff --git a/lib/crc/crc64-neon.c b/lib/crc/crc64-neon.c
index 4753fb94a4be..767e55edc0e0 100644
--- a/lib/crc/crc64-neon.c
+++ b/lib/crc/crc64-neon.c
@@ -10,6 +10,17 @@
 
 u64 crc64_nvme_neon(u64 crc, const u8 *p, size_t len);
 
+#ifdef CONFIG_ARM64
+u64 crc64_nvme_neon_pmull_eor3(u64 crc, const u8 *p, size_t len);
+
+/* x^575 mod G, x^511 mod G */
+static const u64 fold_consts_4x_val[2] = { 0x0c32cdb31e18a84aULL,
+					   0x62242240ace5045aULL };
+/* x^319 mod G, x^255 mod G */
+static const u64 fold_consts_2x_val[2] = { 0xb0bc2e589204f500ULL,
+					   0xe1e0bb9d45d7a44cULL };
+#endif
+
 /* x^191 mod G, x^127 mod G */
 static const u64 fold_consts_val[2] = { 0xeadc41fd2ba3d420ULL,
 					0x21e9761e252621acULL };
@@ -45,3 +56,75 @@ u64 crc64_nvme_neon(u64 crc, const u8 *p, size_t len)
 
 	return vgetq_lane_u64(v0, 1);
 }
+
+#ifdef CONFIG_ARM64
+#define SHA3_TARGET __attribute__((target("arch=armv8.2-a+crypto+sha3")))
+
+SHA3_TARGET
+static inline uint64x2_t fold_single_round(uint64x2_t fold,
+					   uint64x2_t temp,
+					   uint64x2_t consts)
+{
+	uint64x2_t temp1 = pmull64(consts, fold);
+	uint64x2_t temp2 = pmull64_high(consts, fold);
+
+	return veor3q_u64(temp, temp1, temp2);
+}
+
+SHA3_TARGET
+u64 crc64_nvme_neon_pmull_eor3(u64 crc, const u8 *p, size_t len)
+{
+	/* Assumes len >= 64. */
+	uint64x2_t fold_consts = vld1q_u64(fold_consts_val);
+	uint64x2_t fold_consts_4x = vld1q_u64(fold_consts_4x_val);
+	uint64x2_t fold_consts_2x = vld1q_u64(fold_consts_2x_val);
+	uint64x2_t v0 = { crc, 0 };
+	uint64x2_t zero = { };
+	uint64x2_t v1, v2, v3;
+
+	v0 ^= vreinterpretq_u64_u8(vld1q_u8(p));
+	v1  = vreinterpretq_u64_u8(vld1q_u8(p + 16));
+	v2  = vreinterpretq_u64_u8(vld1q_u8(p + 32));
+	v3  = vreinterpretq_u64_u8(vld1q_u8(p + 48));
+
+	p += 64;
+	len -= 64;
+	while (len >= 64) {
+		uint64x2_t t0 = vreinterpretq_u64_u8(vld1q_u8(p +  0));
+		uint64x2_t t1 = vreinterpretq_u64_u8(vld1q_u8(p + 16));
+		uint64x2_t t2 = vreinterpretq_u64_u8(vld1q_u8(p + 32));
+		uint64x2_t t3 = vreinterpretq_u64_u8(vld1q_u8(p + 48));
+
+		v0 = fold_single_round(v0, t0, fold_consts_4x);
+		v1 = fold_single_round(v1, t1, fold_consts_4x);
+		v2 = fold_single_round(v2, t2, fold_consts_4x);
+		v3 = fold_single_round(v3, t3, fold_consts_4x);
+		p += 64;
+		len -= 64;
+	}
+
+	v0 = fold_single_round(v0, v2, fold_consts_2x);
+	v1 = fold_single_round(v1, v3, fold_consts_2x);
+	v0 = fold_single_round(v0, v1, fold_consts);
+
+	while (len >= 16) {
+		uint64x2_t t0 = vreinterpretq_u64_u8(vld1q_u8(p));
+
+		v0 = fold_single_round(v0, t0, fold_consts);
+		p += 16;
+		len -= 16;
+	}
+
+	/* Multiply the 128-bit value by x^64 and reduce it back to 128 bits. */
+	v0 = vextq_u64(v0, zero, 1) ^ pmull64_hi_lo(fold_consts, v0);
+
+	/* Final Barrett reduction */
+	uint64x2_t bconsts = vld1q_u64(bconsts_val);
+	uint64x2_t final = pmull64(bconsts, v0);
+
+	v0 = veor3q_u64(v0, vextq_u64(zero, final, 1),
+			pmull64_hi_lo(bconsts, final));
+
+	return vgetq_lane_u64(v0, 1);
+}
+#endif
-- 
2.56.1