[PATCH v2 1/2] arm64: csum: Add fused copy and Internet checksum

[Demian Shulhan <demyansh@gmail.com>]

arm64 currently falls back to the generic csum_partial_copy_nocheck() in
lib/checksum.c, which performs memcpy() followed by a second pass in
csum_partial(). This double pass is hit on software GSO paths
(skb_segment() without SG or TX checksum offload, e.g. TCP over
WireGuard) and UDP/ICMP replies.

Implement a single-pass csum_partial_copy_nocheck() in
arch/arm64/lib/csum.c using four independent __uint128_t accumulators
over a 32-byte stride, followed by 16/8/4/2/1-byte tail handling. Each
word is read once into a GPR before storing and summing, emitting paired
ldp/stp and adds/cinc carry propagation with no stack frame.

Throughput on ARM Neoverse-V2 (Google Axion, 3.0 GHz, GCC 15.2 -O2,
warm L1 cache):

   len   memcpy + do_csum   csum_partial_copy_nocheck   speedup
  ----   ----------------   -------------------------   -------
    64         10245 MB/s                  15902 MB/s    +55.2%
   256         18851 MB/s                  26026 MB/s    +38.1%
   512         22465 MB/s                  28421 MB/s    +26.5%
  1024         24330 MB/s                  29900 MB/s    +22.9%
  1368         24192 MB/s                  30221 MB/s    +24.9%
  1448         24188 MB/s                  30279 MB/s    +25.2%
  1500         24165 MB/s                  30096 MB/s    +24.5%
  4096         25299 MB/s                  29946 MB/s    +18.4%

Suggested-by: David Laight <david.laight.linux@gmail.com>
Signed-off-by: Demian Shulhan <demyansh@gmail.com>
---
 arch/arm64/include/asm/checksum.h |   3 +
 arch/arm64/lib/csum.c             | 100 ++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)

diff --git a/arch/arm64/include/asm/checksum.h b/arch/arm64/include/asm/checksum.h
index dc52b733675d..9f766db01152 100644
--- a/arch/arm64/include/asm/checksum.h
+++ b/arch/arm64/include/asm/checksum.h
@@ -44,6 +44,9 @@ static inline __sum16 ip_fast_csum(const void *iph, unsigned int ihl)
 extern unsigned int do_csum(const unsigned char *buff, int len);
 #define do_csum do_csum
 
+#define _HAVE_ARCH_CSUM_AND_COPY
+__wsum csum_partial_copy_nocheck(const void *src, void *dst, int len);
+
 #include <asm-generic/checksum.h>
 
 #endif	/* __ASM_CHECKSUM_H */
diff --git a/arch/arm64/lib/csum.c b/arch/arm64/lib/csum.c
index 75db718e5b6a..a9617cde28c1 100644
--- a/arch/arm64/lib/csum.c
+++ b/arch/arm64/lib/csum.c
@@ -2,8 +2,10 @@
 // Copyright (C) 2019-2020 Arm Ltd.
 
 #include <linux/compiler.h>
+#include <linux/export.h>
 #include <linux/kasan-checks.h>
 #include <linux/kernel.h>
+#include <linux/unaligned.h>
 
 #include <net/checksum.h>
 
@@ -116,6 +118,104 @@ unsigned int __no_sanitize_address do_csum(const unsigned char *buff, int len)
 	return sum >> 16;
 }
 
+static __always_inline __wsum csum_fold128(__uint128_t sum)
+{
+	u64 s;
+
+	sum += (sum >> 64) | (sum << 64);
+	s = sum >> 64;
+	s += (s >> 32) | (s << 32);
+	return (__force __wsum)(u32)(s >> 32);
+}
+
+/*
+ * Single-pass copy and ones' complement sum. Each word is read once into a
+ * GPR before storing and summing, so the returned sum always matches @dst
+ * even if @src is concurrently modified.
+ */
+__wsum csum_partial_copy_nocheck(const void *src, void *dst, int len)
+{
+	const u8 *s = src;
+	u8 *d = dst;
+	__uint128_t sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
+	unsigned int n;
+
+	if (unlikely(len <= 0))
+		return 0;
+
+	n = (unsigned int)len >> 5;
+	if (n) {
+		do {
+			u64 v0, v1;
+
+			v0 = get_unaligned((const u64 *)s);
+			v1 = get_unaligned((const u64 *)(s + 8));
+			put_unaligned(v0, (u64 *)d);
+			put_unaligned(v1, (u64 *)(d + 8));
+			sum0 += v0;
+			sum1 += v1;
+
+			v0 = get_unaligned((const u64 *)(s + 16));
+			v1 = get_unaligned((const u64 *)(s + 24));
+			put_unaligned(v0, (u64 *)(d + 16));
+			put_unaligned(v1, (u64 *)(d + 24));
+			sum2 += v0;
+			sum3 += v1;
+
+			s += 32;
+			d += 32;
+		} while (--n);
+
+		sum0 += sum2;
+		sum1 += sum3;
+	}
+
+	if (len & 16) {
+		u64 v0 = get_unaligned((const u64 *)s);
+		u64 v1 = get_unaligned((const u64 *)(s + 8));
+
+		put_unaligned(v0, (u64 *)d);
+		put_unaligned(v1, (u64 *)(d + 8));
+		sum0 += v0;
+		sum1 += v1;
+		s += 16;
+		d += 16;
+	}
+	if (len & 8) {
+		u64 v = get_unaligned((const u64 *)s);
+
+		put_unaligned(v, (u64 *)d);
+		sum0 += v;
+		s += 8;
+		d += 8;
+	}
+	if (len & 4) {
+		u32 v = get_unaligned((const u32 *)s);
+
+		put_unaligned(v, (u32 *)d);
+		sum1 += v;
+		s += 4;
+		d += 4;
+	}
+	if (len & 2) {
+		u16 v = get_unaligned((const u16 *)s);
+
+		put_unaligned(v, (u16 *)d);
+		sum0 += v;
+		s += 2;
+		d += 2;
+	}
+	if (len & 1) {
+		u8 v = *s;
+
+		*d = v;
+		sum1 += v;
+	}
+
+	return csum_fold128(sum0 + sum1);
+}
+EXPORT_SYMBOL(csum_partial_copy_nocheck);
+
 __sum16 csum_ipv6_magic(const struct in6_addr *saddr,
 			const struct in6_addr *daddr,
 			__u32 len, __u8 proto, __wsum csum)
-- 
2.53.0