selftests/bpf: libarena: Optimize and make public arena_memset

Clang currently provides no __builtin_{memset, memcpy, memcmp}
for its BPF backend. This is especially an issue for arena code
that is more likely to do these operations on buffers with user-provided
bounds. One example is the arena ASAN implementation that uses
memset to update the shadow bitmap. Arena ASAN actually already
has a naive implementation of this operation. Another user would
be a calloc() call that has to zero the memory it returns.

Introduce a more optimized version of the memset() operation for
arena memory and make it public to all libarena users. The operation
uses word-sized assignments to speed up the function for larger sizes.
We expose the function through common.h to allow for inlining from
the callers.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
Link: https://lore.kernel.org/bpf/20260817191616.11071-7-emil@etsalapatis.com
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

2 files changed, 56 insertions, 2 deletions

diff --git a/tools/testing/selftests/bpf/libarena/Makefile b/tools/testing/selftests/bpf/libarena/Makefile
index 6d62eda34920..91164d43bd61 100644
--- a/tools/testing/selftests/bpf/libarena/Makefile
+++ b/tools/testing/selftests/bpf/libarena/Makefile
@@ -60,6 +60,7 @@ override BPF_CFLAGS += -O2 -g
 override BPF_CFLAGS += -Wno-incompatible-pointer-types-discards-qualifiers
 # Required for suppressing harmless vmlinux.h-related warnings.
 override BPF_CFLAGS += -Wno-missing-declarations
+override BPF_CFLAGS += -fno-strict-aliasing
 override BPF_CFLAGS += $(INCLUDES)
 
 CFLAGS = -O2 -no-pie
diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/common.h b/tools/testing/selftests/bpf/libarena/include/libarena/common.h
index d32a51ff5e7f..a6220355a697 100644
--- a/tools/testing/selftests/bpf/libarena/include/libarena/common.h
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/common.h
@@ -71,10 +71,63 @@ void arena_free(void __arena *ptr);
  */
 static inline int arena_memset(s8 __arena *dst, s8 val, size_t size)
 {
+	size_t headalign;
+	size_t tailalign;
+	u8 uval = (u8)val;
+	size_t val64;
 	size_t i;
 
-	for (i = zero; i < size && can_loop; i++)
-		dst[i] = val;
+	/*
+	 * Calculate how many bytes to the next word-aligned one.
+	 * We get this by truncating the 2s complement of the
+	 * pointer to the last 3 bits. Intuitively, since
+	 *
+	 * The N LSBs of dst and -dst add to 1 << N, which
+	 * is why dst + (-dst) = 0x0ULL through overflow. So the
+	 * last N = 3 bits of the negative are the number of
+	 * bytes to align dst on the last 3 bits.
+	 *
+	 */
+	headalign = -(u64)dst & (sizeof(u64) - 1);
+	if (!headalign || size < headalign)
+		goto ptraligned;
+
+	for (i = zero; i < headalign && can_loop; i++)
+		dst[i] = uval;
+
+	dst += headalign;
+	size -= headalign;
+
+ptraligned:
+
+	/*
+	 * Make a word with all bytes equal to the byte we are setting.
+	 * Since 1 byte -> 2 hex digits.
+	 *
+	 * Shifting the value by a 0 bytes is equal to multiplication by 0x01
+	 * Shifting by 1 bytes is equal to multiplication by 0x01 << 8,
+	 * ...
+	 * Shifting by 7 bytes is equal to multiplication by 0x01 << 56.
+	 *
+	 * End operation to replicate the byte into all the bytes of a word
+	 * is (since a | b = a + b when a & b == 0):
+	 *
+	 * val + val * (1UL << 8) + val * (1UL << 16) + .. + val * (1UL << 56)
+	 * = val * (1UL << 56 + 1UL << 48 + ... + 1UL << 0)
+	 * = val * (0x01UL << 56 | 0x01UL << 48 + ... + 1UL << 0)
+	 * = val * 0x0101 0101 0101 0101
+	 */
+	val64 = (u8)val * 0x0101010101010101ULL;
+
+	/* Pointer is now aligned, use word-aligned assignments. */
+	for (i = zero; i < size / sizeof(u64) && can_loop; i++)
+		((u64 __arena *)dst)[i] = val64;
+
+	/* Go back to byte-aligned for the tail. */
+	tailalign = size % sizeof(u64);
+	dst += size - tailalign;
+	for (i = zero; i < tailalign && can_loop; i++)
+		dst[i] = uval;
 
 	return 0;
 }