bpf, arm64: Optimize cast_user code generation

cast_user converts an arena offset into a user address by combining the
low 32 bits of the pointer with the upper 32 bits of user_vm_start, while
keeping a NULL pointer NULL. The current sequence always emits six
instructions: it materializes user_vm_start >> 32 into a register, shifts
it into place, and ORs in the offset.

The upper half of user_vm_start is a constant, so it can be written
directly onto the offset with MOVK. Move the 32-bit offset into dst
(which also zeroes the upper 32 bits), then MOVK the non-zero halfwords
of the upper address, branching over the MOVKs when the offset is zero so
NULL is preserved.

This emits at most four instructions, and only one when the upper half of
user_vm_start is zero. The generated code is equivalent.

Before:
  ; bpf_addr_space_cast(page1, 1, 0);
  7c:   mov     w10, w8
  80:   mov     w8, #1
  84:   lsl     x8, x8, #32
  88:   cbz     x10, 0xffff800087b80c20
  8c:   orr     x10, x8, x10
  90:   mov     x8, x10

After:
  ; bpf_addr_space_cast(page1, 1, 0);
  7c:   mov     w8, w8
  80:   cbz     w8, 0xffff800087b80c28
  84:   movk    x8, #1, lsl #32

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Acked-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260721105921.1070501-1-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

1 files changed, 19 insertions, 6 deletions

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index b0075ece4a6e..0f61662b900b 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -1284,12 +1284,25 @@ static int build_insn(const struct bpf_verifier_env *env, const struct bpf_insn
 	case BPF_ALU | BPF_MOV | BPF_X:
 	case BPF_ALU64 | BPF_MOV | BPF_X:
 		if (insn_is_cast_user(insn)) {
-			emit(A64_MOV(0, tmp, src), ctx); // 32-bit mov clears the upper 32 bits
-			emit_a64_mov_i(0, dst, ctx->user_vm_start >> 32, ctx);
-			emit(A64_LSL(1, dst, dst, 32), ctx);
-			emit(A64_CBZ(1, tmp, 2), ctx);
-			emit(A64_ORR(1, tmp, dst, tmp), ctx);
-			emit(A64_MOV(1, dst, tmp), ctx);
+			u32 upper = ctx->user_vm_start >> 32;
+			u16 upper_low = upper & 0xffff;
+			u16 upper_high = upper >> 16;
+			int nr_movk = !!upper_low + !!upper_high;
+
+			/*
+			 * Build the user address: the low 32 bits are the arena
+			 * offset, the upper 32 bits come from user_vm_start. A
+			 * zero offset must stay NULL, so branch over the MOVKs
+			 * when it is zero.
+			 */
+			emit(A64_MOV(0, dst, src), ctx); /* 32-bit mov clears the upper 32 bits */
+			if (nr_movk) {
+				emit(A64_CBZ(0, dst, nr_movk + 1), ctx);
+				if (upper_low)
+					emit(A64_MOVK(1, dst, upper_low, 32), ctx);
+				if (upper_high)
+					emit(A64_MOVK(1, dst, upper_high, 48), ctx);
+			}
 			break;
 		} else if (insn_is_mov_percpu_addr(insn)) {
 			if (dst != src)