[PATCH bpf-next v2 1/2] bpf, riscv: Add BPF stack arguments support for RV64 JIT

[Feng Jiang <jiangfeng@kylinos.cn>]

Add bpf_jit_supports_stack_args() for the RV64 JIT so BPF subprograms
and kfuncs can receive more than 5 arguments via the stack
(BPF_REG_PARAMS / r11).

For BPF-to-BPF calls the caller writes outgoing arguments at the
bottom of its frame (SP-relative). The callee reads them with
FP-relative loads. Its FP is set to the caller SP in the prologue,
so the offsets match.

The RISC-V ABI puts arguments 6-8 in A5-A7 and arguments 9+ at
SP+0. Before each kfunc call, load arguments 6-8 from the outgoing
area into A5-A7 and copy any remaining arguments down so argument 9
lands at SP+0.

A5 (BPF_REG_0), A6 (TCC) and A7 are safe to clobber here: R0 is not
live before a call, TCC is backed up on the stack, and A7 is unused
by the JIT.

Limit the existing kfunc sign-extension loop to MAX_BPF_FUNC_REG_ARGS
iterations; otherwise idx >= 5 resolves to S1-S5 (BPF_R6-FP) and
corrupts callee-saved registers.

Signed-off-by: Feng Jiang <jiangfeng@kylinos.cn>
---
 arch/riscv/net/bpf_jit.h        |  1 +
 arch/riscv/net/bpf_jit_comp64.c | 79 +++++++++++++++++++++++++++++++++++++++--
 arch/riscv/net/bpf_jit_core.c   |  4 +++
 3 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/arch/riscv/net/bpf_jit.h b/arch/riscv/net/bpf_jit.h
index 419b9d795f2a..9eb4e149505d 100644
--- a/arch/riscv/net/bpf_jit.h
+++ b/arch/riscv/net/bpf_jit.h
@@ -82,6 +82,7 @@ struct rv_jit_context {
 	unsigned long flags;
 	int stack_size;
 	int tcc_offset;
+	u16 stack_arg_size;
 	u64 arena_vm_start;
 	u64 user_vm_start;
 };
diff --git a/arch/riscv/net/bpf_jit_comp64.c b/arch/riscv/net/bpf_jit_comp64.c
index 6b9972b07c1b..58cb3e5ff6b4 100644
--- a/arch/riscv/net/bpf_jit_comp64.c
+++ b/arch/riscv/net/bpf_jit_comp64.c
@@ -1815,18 +1815,43 @@ int bpf_jit_emit_insn(const struct bpf_insn *insn, struct rv_jit_context *ctx,
 
 		if (insn->src_reg == BPF_PSEUDO_KFUNC_CALL) {
 			const struct btf_func_model *fm;
-			int idx;
+			int idx, nargs;
 
 			fm = bpf_jit_find_kfunc_model(ctx->prog, insn);
 			if (!fm)
 				return -EINVAL;
 
-			for (idx = 0; idx < fm->nr_args; idx++) {
+			nargs = min_t(int, fm->nr_args, MAX_BPF_FUNC_REG_ARGS);
+			for (idx = 0; idx < nargs; idx++) {
 				u8 reg = bpf_to_rv_reg(BPF_REG_1 + idx, ctx);
 
 				if (fm->arg_size[idx] == sizeof(int))
 					emit_sextw(reg, reg, ctx);
 			}
+
+			/* BPF stack args -> RISC-V ABI: args 6-8 in A5-A7, 9+ at SP+0 */
+			if (fm->nr_args > MAX_BPF_FUNC_REG_ARGS) {
+				int n_stack = fm->nr_args - MAX_BPF_FUNC_REG_ARGS;
+				int n_reg = min_t(int, n_stack,
+						  RV_MAX_REG_ARGS - MAX_BPF_FUNC_REG_ARGS);
+
+				for (idx = 0; idx < n_reg; idx++) {
+					int sz = fm->arg_size[MAX_BPF_FUNC_REG_ARGS + idx];
+
+					emit_ld(RV_REG_A5 + idx, idx * 8, RV_REG_SP, ctx);
+					if (sz == sizeof(int))
+						emit_sextw(RV_REG_A5 + idx, RV_REG_A5 + idx, ctx);
+				}
+
+				for (idx = n_reg; idx < n_stack; idx++) {
+					int sz = fm->arg_size[MAX_BPF_FUNC_REG_ARGS + idx];
+
+					emit_ld(RV_REG_T1, idx * 8, RV_REG_SP, ctx);
+					if (sz == sizeof(int))
+						emit_sextw(RV_REG_T1, RV_REG_T1, ctx);
+					emit_sd(RV_REG_SP, (idx - n_reg) * 8, RV_REG_T1, ctx);
+				}
+			}
 		}
 
 		/* restore TCC to RV_REG_TCC before bpf2bpf call */
@@ -1891,6 +1916,21 @@ int bpf_jit_emit_insn(const struct bpf_insn *insn, struct rv_jit_context *ctx,
 	case BPF_LDX | BPF_MEM | BPF_H:
 	case BPF_LDX | BPF_MEM | BPF_W:
 	case BPF_LDX | BPF_MEM | BPF_DW:
+		if (insn->src_reg == BPF_REG_PARAMS) {
+			int idx = off / 8 - 1;
+
+			if (is_12b_int(idx * 8)) {
+				emit_ldx_insn(rd, idx * 8, RV_REG_FP, BPF_SIZE(code), false, ctx);
+			} else {
+				emit_imm(RV_REG_T1, idx * 8, ctx);
+				emit_add(RV_REG_T1, RV_REG_T1, RV_REG_FP, ctx);
+				emit_ldx_insn(rd, 0, RV_REG_T1, BPF_SIZE(code), false, ctx);
+			}
+			if (BPF_SIZE(code) != BPF_DW && insn_is_zext(&insn[1]))
+				return 1;
+			break;
+		}
+		fallthrough;
 	case BPF_LDX | BPF_PROBE_MEM | BPF_B:
 	case BPF_LDX | BPF_PROBE_MEM | BPF_H:
 	case BPF_LDX | BPF_PROBE_MEM | BPF_W:
@@ -1938,6 +1978,20 @@ int bpf_jit_emit_insn(const struct bpf_insn *insn, struct rv_jit_context *ctx,
 	case BPF_ST | BPF_MEM | BPF_H:
 	case BPF_ST | BPF_MEM | BPF_W:
 	case BPF_ST | BPF_MEM | BPF_DW:
+		if (insn->dst_reg == BPF_REG_PARAMS) {
+			int idx = -off / 8 - 1;
+
+			emit_imm(RV_REG_T1, imm, ctx);
+			if (is_12b_int(idx * 8)) {
+				emit_stx_insn(RV_REG_SP, idx * 8, RV_REG_T1, BPF_SIZE(code), ctx);
+			} else {
+				emit_imm(RV_REG_T2, idx * 8, ctx);
+				emit_add(RV_REG_T2, RV_REG_SP, RV_REG_T2, ctx);
+				emit_stx_insn(RV_REG_T2, 0, RV_REG_T1, BPF_SIZE(code), ctx);
+			}
+			break;
+		}
+		fallthrough;
 	/* ST | PROBE_MEM32: *(size *)(dst + RV_REG_ARENA + off) = imm */
 	case BPF_ST | BPF_PROBE_MEM32 | BPF_B:
 	case BPF_ST | BPF_PROBE_MEM32 | BPF_H:
@@ -1960,6 +2014,19 @@ int bpf_jit_emit_insn(const struct bpf_insn *insn, struct rv_jit_context *ctx,
 	case BPF_STX | BPF_MEM | BPF_H:
 	case BPF_STX | BPF_MEM | BPF_W:
 	case BPF_STX | BPF_MEM | BPF_DW:
+		if (insn->dst_reg == BPF_REG_PARAMS) {
+			int idx = -off / 8 - 1;
+
+			if (is_12b_int(idx * 8)) {
+				emit_stx_insn(RV_REG_SP, idx * 8, rs, BPF_SIZE(code), ctx);
+			} else {
+				emit_imm(RV_REG_T1, idx * 8, ctx);
+				emit_add(RV_REG_T1, RV_REG_SP, RV_REG_T1, ctx);
+				emit_stx_insn(RV_REG_T1, 0, rs, BPF_SIZE(code), ctx);
+			}
+			break;
+		}
+		fallthrough;
 	/* STX | PROBE_MEM32: *(size *)(dst + RV_REG_ARENA + off) = src */
 	case BPF_STX | BPF_PROBE_MEM32 | BPF_B:
 	case BPF_STX | BPF_PROBE_MEM32 | BPF_H:
@@ -2036,6 +2103,7 @@ void bpf_jit_build_prologue(struct rv_jit_context *ctx, bool is_subprog)
 
 	stack_adjust = round_up(stack_adjust, STACK_ALIGN);
 	stack_adjust += bpf_stack_adjust;
+	stack_adjust += ctx->stack_arg_size;
 
 	store_offset = stack_adjust - 8;
 
@@ -2093,7 +2161,7 @@ void bpf_jit_build_prologue(struct rv_jit_context *ctx, bool is_subprog)
 	emit_addi(RV_REG_FP, RV_REG_SP, stack_adjust, ctx);
 
 	if (bpf_stack_adjust)
-		emit_addi(RV_REG_S5, RV_REG_SP, bpf_stack_adjust, ctx);
+		emit_addi(RV_REG_S5, RV_REG_SP, ctx->stack_arg_size + bpf_stack_adjust, ctx);
 
 	ctx->stack_size = stack_adjust;
 
@@ -2171,3 +2239,8 @@ bool bpf_jit_supports_timed_may_goto(void)
 {
 	return true;
 }
+
+bool bpf_jit_supports_stack_args(void)
+{
+	return true;
+}
diff --git a/arch/riscv/net/bpf_jit_core.c b/arch/riscv/net/bpf_jit_core.c
index cbfcd287ea16..844a0f3e0fa9 100644
--- a/arch/riscv/net/bpf_jit_core.c
+++ b/arch/riscv/net/bpf_jit_core.c
@@ -72,6 +72,10 @@ struct bpf_prog *bpf_int_jit_compile(struct bpf_verifier_env *env, struct bpf_pr
 	ctx->arena_vm_start = bpf_arena_get_kern_vm_start(prog->aux->arena);
 	ctx->user_vm_start = bpf_arena_get_user_vm_start(prog->aux->arena);
 	ctx->prog = prog;
+
+	ctx->stack_arg_size = round_up(bpf_out_stack_arg_cnt(env, prog) *
+				       sizeof(u64), STACK_ALIGN);
+
 	ctx->offset = kvzalloc_objs(int, prog->len);
 	if (!ctx->offset)
 		goto out_offset;

-- 
2.53.0


_______________________________________________
linux-riscv mailing list
linux-riscv@lists.infradead.org
http://lists.infradead.org/mailman/listinfo/linux-riscv