s390/bpf: Centralize frame offset calculations

The calculation of the distance from %r15 to the caller-allocated
portion of the stack frame is copy-pasted into multiple places in the
JIT code.

Move it to bpf_jit_prog() and save the result into bpf_jit::frame_off,
so that the other parts of the JIT can use it.

Signed-off-by: Ilya Leoshkevich <iii@linux.ibm.com>
Link: https://lore.kernel.org/r/20250624121501.50536-2-iii@linux.ibm.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/arch/s390/net/bpf_jit_comp.c b/arch/s390/net/bpf_jit_comp.c
index c7f8313ba449716a8f18eafdeb6c77ed3b23f52e..76ad52a2444416fee76263783411a686c7c9ae2e 100644 (file)
--- a/arch/s390/net/bpf_jit_comp.c
+++ b/arch/s390/net/bpf_jit_comp.c
@@ -54,6 +54,7 @@ struct bpf_jit {
        int prologue_plt;       /* Start of prologue hotpatch PLT */
        int kern_arena;         /* Pool offset of kernel arena address */
        u64 user_arena;         /* User arena address */
+       u32 frame_off;          /* Offset of frame from %r15 */
 };
 
 #define SEEN_MEM       BIT(0)          /* use mem[] for temporary storage */
@@ -443,12 +444,9 @@ static void save_regs(struct bpf_jit *jit, u32 rs, u32 re)
 /*
  * Restore registers from "rs" (register start) to "re" (register end) on stack
  */
-static void restore_regs(struct bpf_jit *jit, u32 rs, u32 re, u32 stack_depth)
+static void restore_regs(struct bpf_jit *jit, u32 rs, u32 re)
 {
-       u32 off = STK_OFF_R6 + (rs - 6) * 8;
-
-       if (jit->seen & SEEN_STACK)
-               off += STK_OFF + stack_depth;
+       u32 off = jit->frame_off + STK_OFF_R6 + (rs - 6) * 8;
 
        if (rs == re)
                /* lg %rs,off(%r15) */
@@ -492,8 +490,7 @@ static int get_end(u16 seen_regs, int start)
  * Save and restore clobbered registers (6-15) on stack.
  * We save/restore registers in chunks with gap >= 2 registers.
  */
-static void save_restore_regs(struct bpf_jit *jit, int op, u32 stack_depth,
-                             u16 extra_regs)
+static void save_restore_regs(struct bpf_jit *jit, int op, u16 extra_regs)
 {
        u16 seen_regs = jit->seen_regs | extra_regs;
        const int last = 15, save_restore_size = 6;
@@ -516,7 +513,7 @@ static void save_restore_regs(struct bpf_jit *jit, int op, u32 stack_depth,
                if (op == REGS_SAVE)
                        save_regs(jit, rs, re);
                else
-                       restore_regs(jit, rs, re, stack_depth);
+                       restore_regs(jit, rs, re);
                re++;
        } while (re <= last);
 }
@@ -575,8 +572,7 @@ static void bpf_jit_plt(struct bpf_plt *plt, void *ret, void *target)
  * Save registers and create stack frame if necessary.
  * See stack frame layout description in "bpf_jit.h"!
  */
-static void bpf_jit_prologue(struct bpf_jit *jit, struct bpf_prog *fp,
-                            u32 stack_depth)
+static void bpf_jit_prologue(struct bpf_jit *jit, struct bpf_prog *fp)
 {
        /* No-op for hotpatching */
        /* brcl 0,prologue_plt */
@@ -609,7 +605,7 @@ static void bpf_jit_prologue(struct bpf_jit *jit, struct bpf_prog *fp,
                jit->seen_regs |= NVREGS;
        } else {
                /* Save registers */
-               save_restore_regs(jit, REGS_SAVE, stack_depth,
+               save_restore_regs(jit, REGS_SAVE,
                                  fp->aux->exception_boundary ? NVREGS : 0);
        }
        /* Setup literal pool */
@@ -631,8 +627,8 @@ static void bpf_jit_prologue(struct bpf_jit *jit, struct bpf_prog *fp,
                EMIT4(0xb9040000, REG_W1, REG_15);
                /* la %bfp,STK_160_UNUSED(%r15) (BPF frame pointer) */
                EMIT4_DISP(0x41000000, BPF_REG_FP, REG_15, STK_160_UNUSED);
-               /* aghi %r15,-STK_OFF */
-               EMIT4_IMM(0xa70b0000, REG_15, -(STK_OFF + stack_depth));
+               /* aghi %r15,-frame_off */
+               EMIT4_IMM(0xa70b0000, REG_15, -jit->frame_off);
                /* stg %w1,152(%r15) (backchain) */
                EMIT6_DISP_LH(0xe3000000, 0x0024, REG_W1, REG_0,
                              REG_15, 152);
@@ -669,13 +665,13 @@ static void call_r1(struct bpf_jit *jit)
 /*
  * Function epilogue
  */
-static void bpf_jit_epilogue(struct bpf_jit *jit, u32 stack_depth)
+static void bpf_jit_epilogue(struct bpf_jit *jit)
 {
        jit->exit_ip = jit->prg;
        /* Load exit code: lgr %r2,%b0 */
        EMIT4(0xb9040000, REG_2, BPF_REG_0);
        /* Restore registers */
-       save_restore_regs(jit, REGS_RESTORE, stack_depth, 0);
+       save_restore_regs(jit, REGS_RESTORE, 0);
        EMIT_JUMP_REG(14);
 
        jit->prg = ALIGN(jit->prg, 8);
@@ -857,7 +853,7 @@ static int sign_extend(struct bpf_jit *jit, int r, u8 size, u8 flags)
  * stack space for the large switch statement.
  */
 static noinline int bpf_jit_insn(struct bpf_jit *jit, struct bpf_prog *fp,
-                                int i, bool extra_pass, u32 stack_depth)
+                                int i, bool extra_pass)
 {
        struct bpf_insn *insn = &fp->insnsi[i];
        s32 branch_oc_off = insn->off;
@@ -1778,9 +1774,9 @@ static noinline int bpf_jit_insn(struct bpf_jit *jit, struct bpf_prog *fp,
                 * Note 2: We assume that the verifier does not let us call the
                 * main program, which clears the tail call counter on entry.
                 */
-               /* mvc STK_OFF_TCCNT(4,%r15),N(%r15) */
+               /* mvc STK_OFF_TCCNT(4,%r15),frame_off+STK_OFF_TCCNT(%r15) */
                _EMIT6(0xd203f000 | STK_OFF_TCCNT,
-                      0xf000 | (STK_OFF_TCCNT + STK_OFF + stack_depth));
+                      0xf000 | (jit->frame_off + STK_OFF_TCCNT));
 
                /* Sign-extend the kfunc arguments. */
                if (insn->src_reg == BPF_PSEUDO_KFUNC_CALL) {
@@ -1831,10 +1827,7 @@ static noinline int bpf_jit_insn(struct bpf_jit *jit, struct bpf_prog *fp,
                 *         goto out;
                 */
 
-               if (jit->seen & SEEN_STACK)
-                       off = STK_OFF_TCCNT + STK_OFF + stack_depth;
-               else
-                       off = STK_OFF_TCCNT;
+               off = jit->frame_off + STK_OFF_TCCNT;
                /* lhi %w0,1 */
                EMIT4_IMM(0xa7080000, REG_W0, 1);
                /* laal %w1,%w0,off(%r15) */
@@ -1864,7 +1857,7 @@ static noinline int bpf_jit_insn(struct bpf_jit *jit, struct bpf_prog *fp,
                /*
                 * Restore registers before calling function
                 */
-               save_restore_regs(jit, REGS_RESTORE, stack_depth, 0);
+               save_restore_regs(jit, REGS_RESTORE, 0);
 
                /*
                 * goto *(prog->bpf_func + tail_call_start);
@@ -2157,7 +2150,7 @@ static int bpf_set_addr(struct bpf_jit *jit, int i)
  * Compile eBPF program into s390x code
  */
 static int bpf_jit_prog(struct bpf_jit *jit, struct bpf_prog *fp,
-                       bool extra_pass, u32 stack_depth)
+                       bool extra_pass)
 {
        int i, insn_count, lit32_size, lit64_size;
        u64 kern_arena;
@@ -2166,24 +2159,28 @@ static int bpf_jit_prog(struct bpf_jit *jit, struct bpf_prog *fp,
        jit->lit64 = jit->lit64_start;
        jit->prg = 0;
        jit->excnt = 0;
+       if (is_first_pass(jit) || (jit->seen & SEEN_STACK))
+               jit->frame_off = STK_OFF + round_up(fp->aux->stack_depth, 8);
+       else
+               jit->frame_off = 0;
 
        kern_arena = bpf_arena_get_kern_vm_start(fp->aux->arena);
        if (kern_arena)
                jit->kern_arena = _EMIT_CONST_U64(kern_arena);
        jit->user_arena = bpf_arena_get_user_vm_start(fp->aux->arena);
 
-       bpf_jit_prologue(jit, fp, stack_depth);
+       bpf_jit_prologue(jit, fp);
        if (bpf_set_addr(jit, 0) < 0)
                return -1;
        for (i = 0; i < fp->len; i += insn_count) {
-               insn_count = bpf_jit_insn(jit, fp, i, extra_pass, stack_depth);
+               insn_count = bpf_jit_insn(jit, fp, i, extra_pass);
                if (insn_count < 0)
                        return -1;
                /* Next instruction address */
                if (bpf_set_addr(jit, i + insn_count) < 0)
                        return -1;
        }
-       bpf_jit_epilogue(jit, stack_depth);
+       bpf_jit_epilogue(jit);
 
        lit32_size = jit->lit32 - jit->lit32_start;
        lit64_size = jit->lit64 - jit->lit64_start;
@@ -2259,7 +2256,6 @@ static struct bpf_binary_header *bpf_jit_alloc(struct bpf_jit *jit,
  */
 struct bpf_prog *bpf_int_jit_compile(struct bpf_prog *fp)
 {
-       u32 stack_depth = round_up(fp->aux->stack_depth, 8);
        struct bpf_prog *tmp, *orig_fp = fp;
        struct bpf_binary_header *header;
        struct s390_jit_data *jit_data;
@@ -2312,7 +2308,7 @@ struct bpf_prog *bpf_int_jit_compile(struct bpf_prog *fp)
         *   - 3:   Calculate program size and addrs array
         */
        for (pass = 1; pass <= 3; pass++) {
-               if (bpf_jit_prog(&jit, fp, extra_pass, stack_depth)) {
+               if (bpf_jit_prog(&jit, fp, extra_pass)) {
                        fp = orig_fp;
                        goto free_addrs;
                }
@@ -2326,7 +2322,7 @@ struct bpf_prog *bpf_int_jit_compile(struct bpf_prog *fp)
                goto free_addrs;
        }
 skip_init_ctx:
-       if (bpf_jit_prog(&jit, fp, extra_pass, stack_depth)) {
+       if (bpf_jit_prog(&jit, fp, extra_pass)) {
                bpf_jit_binary_free(header);
                fp = orig_fp;
                goto free_addrs;