s390/bpf: Change seen_reg to a mask

Using a mask instead of an array saves a small amount of memory and
allows marking multiple registers as seen with a simple "or". Another
positive side-effect is that it speeds up verification with jitterbug.

Signed-off-by: Ilya Leoshkevich <iii@linux.ibm.com>
Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>
Link: https://lore.kernel.org/bpf/20240703005047.40915-2-iii@linux.ibm.com

diff --git a/arch/s390/net/bpf_jit_comp.c b/arch/s390/net/bpf_jit_comp.c
index ddfc0e99872ef7354e2d7cf798adb641254d3234..945f2ee6511b8102d07593fd2ca404226e23d7ad 100644 (file)
--- a/arch/s390/net/bpf_jit_comp.c
+++ b/arch/s390/net/bpf_jit_comp.c
@@ -35,7 +35,7 @@
 
 struct bpf_jit {
        u32 seen;               /* Flags to remember seen eBPF instructions */
-       u32 seen_reg[16];       /* Array to remember which registers are used */
+       u16 seen_regs;          /* Mask to remember which registers are used */
        u32 *addrs;             /* Array with relative instruction addresses */
        u8 *prg_buf;            /* Start of program */
        int size;               /* Size of program and literal pool */
@@ -120,8 +120,8 @@ static inline void reg_set_seen(struct bpf_jit *jit, u32 b1)
 {
        u32 r1 = reg2hex[b1];
 
-       if (r1 >= 6 && r1 <= 15 && !jit->seen_reg[r1])
-               jit->seen_reg[r1] = 1;
+       if (r1 >= 6 && r1 <= 15)
+               jit->seen_regs |= (1 << r1);
 }
 
 #define REG_SET_SEEN(b1)                                       \
@@ -129,8 +129,6 @@ static inline void reg_set_seen(struct bpf_jit *jit, u32 b1)
        reg_set_seen(jit, b1);                                  \
 })
 
-#define REG_SEEN(b1) jit->seen_reg[reg2hex[(b1)]]
-
 /*
  * EMIT macros for code generation
  */
@@ -438,12 +436,12 @@ static void restore_regs(struct bpf_jit *jit, u32 rs, u32 re, u32 stack_depth)
 /*
  * Return first seen register (from start)
  */
-static int get_start(struct bpf_jit *jit, int start)
+static int get_start(u16 seen_regs, int start)
 {
        int i;
 
        for (i = start; i <= 15; i++) {
-               if (jit->seen_reg[i])
+               if (seen_regs & (1 << i))
                        return i;
        }
        return 0;
@@ -452,15 +450,15 @@ static int get_start(struct bpf_jit *jit, int start)
 /*
  * Return last seen register (from start) (gap >= 2)
  */
-static int get_end(struct bpf_jit *jit, int start)
+static int get_end(u16 seen_regs, int start)
 {
        int i;
 
        for (i = start; i < 15; i++) {
-               if (!jit->seen_reg[i] && !jit->seen_reg[i + 1])
+               if (!(seen_regs & (3 << i)))
                        return i - 1;
        }
-       return jit->seen_reg[15] ? 15 : 14;
+       return (seen_regs & (1 << 15)) ? 15 : 14;
 }
 
 #define REGS_SAVE      1
@@ -469,8 +467,10 @@ static int get_end(struct bpf_jit *jit, 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)
+static void save_restore_regs(struct bpf_jit *jit, int op, u32 stack_depth,
+                             u16 extra_regs)
 {
+       u16 seen_regs = jit->seen_regs | extra_regs;
        const int last = 15, save_restore_size = 6;
        int re = 6, rs;
 
@@ -484,10 +484,10 @@ static void save_restore_regs(struct bpf_jit *jit, int op, u32 stack_depth)
        }
 
        do {
-               rs = get_start(jit, re);
+               rs = get_start(seen_regs, re);
                if (!rs)
                        break;
-               re = get_end(jit, rs + 1);
+               re = get_end(seen_regs, rs + 1);
                if (op == REGS_SAVE)
                        save_regs(jit, rs, re);
                else
@@ -573,7 +573,7 @@ static void bpf_jit_prologue(struct bpf_jit *jit, struct bpf_prog *fp,
        /* Tail calls have to skip above initialization */
        jit->tail_call_start = jit->prg;
        /* Save registers */
-       save_restore_regs(jit, REGS_SAVE, stack_depth);
+       save_restore_regs(jit, REGS_SAVE, stack_depth, 0);
        /* Setup literal pool */
        if (is_first_pass(jit) || (jit->seen & SEEN_LITERAL)) {
                if (!is_first_pass(jit) &&
@@ -649,7 +649,7 @@ static void bpf_jit_epilogue(struct bpf_jit *jit, u32 stack_depth)
        /* Load exit code: lgr %r2,%b0 */
        EMIT4(0xb9040000, REG_2, BPF_REG_0);
        /* Restore registers */
-       save_restore_regs(jit, REGS_RESTORE, stack_depth);
+       save_restore_regs(jit, REGS_RESTORE, stack_depth, 0);
        if (nospec_uses_trampoline()) {
                jit->r14_thunk_ip = jit->prg;
                /* Generate __s390_indirect_jump_r14 thunk */
@@ -1847,7 +1847,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);
+               save_restore_regs(jit, REGS_RESTORE, stack_depth, 0);
 
                /*
                 * goto *(prog->bpf_func + tail_call_start);