The tail call sequence moved the offsets of max_entries, ptrs and
bpf_func, and the tail call limit, into a temporary first, although
they all fit the immediate forms. It also zero extended and shifted the
index in separate instructions before the load. Use the immediate forms
and a register offset load with UXTW #3 instead.

That saves six instructions per tail call. A64_LDR64() loses its last
user and is removed.

Before/after for the checks and the lookup of a tail call:

  5280088a  mov w10, #0x44
  b86a682a  ldr w10, [x1, x10]
  2a0203e2  mov w2, w2
  6b0a005f  cmp w2, w10
  54000282  b.hs ...
  5280042a  mov w10, #0x21
  f940034c  ldr x12, [x26]
  eb0a019f  cmp x12, x10
  54000202  b.hs ...
  9100058c  add x12, x12, #0x1
  5280440a  mov w10, #0x220
  8b0a002a  add x10, x1, x10
  d37df04b  lsl x11, x2, #3
  f86b694b  ldr x11, [x10, x11]
  b400014b  cbz x11, ...

  b940442a  ldr w10, [x1, #68]
  6b0a005f  cmp w2, w10
  54000202  b.hs ...
  f940034c  ldr x12, [x26]
  f100859f  cmp x12, #0x21
  540001a2  b.hs ...
  9100058c  add x12, x12, #0x1
  9108802a  add x10, x1, #0x220
  f862594b  ldr x11, [x10, w2, uxtw #3]
  b400012b  cbz x11, ...

Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>
---
 arch/arm64/net/bpf_jit.h      |  3 ++-
 arch/arm64/net/bpf_jit_comp.c | 27 +++++++++++++++------------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/arch/arm64/net/bpf_jit.h b/arch/arm64/net/bpf_jit.h
index 882404b884e1..16c690c6cae9 100644
--- a/arch/arm64/net/bpf_jit.h
+++ b/arch/arm64/net/bpf_jit.h
@@ -63,7 +63,8 @@
 #define A64_LS_REG(Rt, Rn, Rm, size, type) \
 	A64_LS_REG_EXT(Rt, Rn, Rm, UXTX, 0, size, type)
 #define A64_LDR32(Wt, Xn, Xm) A64_LS_REG(Wt, Xn, Xm, 32, LOAD)
-#define A64_LDR64(Xt, Xn, Xm) A64_LS_REG(Xt, Xn, Xm, 64, LOAD)
+/* Xt = *(Xn + ((u32)Wm << 3)) */
+#define A64_LDR64_UXTW3(Xt, Xn, Wm) A64_LS_REG_EXT(Xt, Xn, Wm, UXTW, 3, 64, LOAD)
 
 /* Load/store register (immediate offset) */
 #define A64_LS_IMM(Rt, Rn, imm, size, type) \
diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 62d2f70cc412..a79c401c5262 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -909,13 +909,20 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
 	__le32 *branch2 = NULL;
 	__le32 *branch3 = NULL;
 
-	/* if (index >= array->map.max_entries)
+	/* The struct offsets fit the immediate forms, the count fits cmp */
+	BUILD_BUG_ON(offsetof(struct bpf_array, map.max_entries) % 4 ||
+		     offsetof(struct bpf_array, map.max_entries) > 0xfff * 4);
+	BUILD_BUG_ON(offsetof(struct bpf_array, ptrs) > 0xfff);
+	BUILD_BUG_ON(offsetof(struct bpf_prog, bpf_func) % 8 ||
+		     offsetof(struct bpf_prog, bpf_func) > 0xfff * 8);
+	BUILD_BUG_ON(MAX_TAIL_CALL_CNT > 0xfff);
+
+	/*
+	 * if ((u32)index >= array->map.max_entries)
 	 *     goto out;
 	 */
 	off = offsetof(struct bpf_array, map.max_entries);
-	emit_a64_mov_i64(tmp, off, ctx);
-	emit(A64_LDR32(tmp, r2, tmp), ctx);
-	emit(A64_MOV(0, r3, r3), ctx);
+	emit(A64_LDR32I(tmp, r2, off), ctx);
 	emit(A64_CMP(0, r3, tmp), ctx);
 	branch1 = ctx->image + ctx->idx;
 	emit(A64_NOP, ctx);
@@ -924,9 +931,8 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
 	 * if ((*tail_call_cnt_ptr) >= MAX_TAIL_CALL_CNT)
 	 *     goto out;
 	 */
-	emit_a64_mov_i64(tmp, MAX_TAIL_CALL_CNT, ctx);
 	emit(A64_LDR64I(tcc, ptr, 0), ctx);
-	emit(A64_CMP(1, tcc, tmp), ctx);
+	emit(A64_CMP_I(1, tcc, MAX_TAIL_CALL_CNT), ctx);
 	branch2 = ctx->image + ctx->idx;
 	emit(A64_NOP, ctx);
 
@@ -938,10 +944,8 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
 	 *     goto out;
 	 */
 	off = offsetof(struct bpf_array, ptrs);
-	emit_a64_mov_i64(tmp, off, ctx);
-	emit(A64_ADD(1, tmp, r2, tmp), ctx);
-	emit(A64_LSL(1, prg, r3, 3), ctx);
-	emit(A64_LDR64(prg, tmp, prg), ctx);
+	emit(A64_ADD_I(1, tmp, r2, off), ctx);
+	emit(A64_LDR64_UXTW3(prg, tmp, r3), ctx);
 	branch3 = ctx->image + ctx->idx;
 	emit(A64_NOP, ctx);
 
@@ -959,8 +963,7 @@ static int emit_bpf_tail_call(struct jit_ctx *ctx)
 
 	/* goto *(prog->bpf_func + prologue_offset); */
 	off = offsetof(struct bpf_prog, bpf_func);
-	emit_a64_mov_i64(tmp, off, ctx);
-	emit(A64_LDR64(tmp, prg, tmp), ctx);
+	emit(A64_LDR64I(tmp, prg, off), ctx);
 	emit(A64_ADD_I(1, tmp, tmp, sizeof(u32) * PROLOGUE_OFFSET), ctx);
 	emit(A64_BR(tmp), ctx);
 
-- 
2.43.0