LoongArch: BPF: Implement PROBE_MEM32 pseudo instructions

Add support for `{LDX,STX,ST} | PROBE_MEM32 | {B,H,W,DW}` instructions.
They are similar to PROBE_MEM instructions with the following differences:
 * PROBE_MEM32 supports store.
 * PROBE_MEM32 relies on the verifier to clear upper 32-bit of the
  src/dst register
 * PROBE_MEM32 adds 64-bit kern_vm_start address (which is stored in S6
  in the prologue). Due to bpf_arena constructions such S6 + reg +
  off16 access is guaranteed to be within arena virtual range, so no
  address check at run-time.
 * S6 is a free callee-saved register, so it is used to store arena_vm_start
 * PROBE_MEM32 allows ST and STX. If they fault the store is a nop. When
  LDX faults the destination register is zeroed.

To support these on LoongArch, we employ the t2/t3 registers to store the
intermediate results of reg_arena + src/dst reg and use the t2/t3 registers
as the new src/dst reg. This allows us to reuse most of the existing code.

Acked-by: Tiezhu Yang <yangtiezhu@loongson.cn>
Tested-by: Vincent Li <vincent.mc.li@gmail.com>
Signed-off-by: Hengqi Chen <hengqi.chen@gmail.com>
Signed-off-by: Huacai Chen <chenhuacai@loongson.cn>

Author: chenhengqi

arch/loongarch/net/bpf_jit.c

@@ -17,6 +17,7 @@
 #define LOONGARCH_BPF_FENTRY_NBYTES (LOONGARCH_LONG_JUMP_NINSNS * 4)
 
 #define REG_TCC		LOONGARCH_GPR_A6
+#define REG_ARENA	LOONGARCH_GPR_S6 /* For storing arena_vm_start */
 #define BPF_TAIL_CALL_CNT_PTR_STACK_OFF(stack) (round_up(stack, 16) - 80)
 
 static const int regmap[] = {
@@ -136,6 +137,9 @@ static void build_prologue(struct jit_ctx *ctx)
 	/* To store tcc and tcc_ptr */
 	stack_adjust += sizeof(long) * 2;
 
+	if (ctx->arena_vm_start)
+		stack_adjust += 8;
+
 	stack_adjust = round_up(stack_adjust, 16);
 	stack_adjust += bpf_stack_adjust;
 
@@ -178,6 +182,11 @@ static void build_prologue(struct jit_ctx *ctx)
 	store_offset -= sizeof(long);
 	emit_insn(ctx, std, LOONGARCH_GPR_S5, LOONGARCH_GPR_SP, store_offset);
 
+	if (ctx->arena_vm_start) {
+		store_offset -= sizeof(long);
+		emit_insn(ctx, std, REG_ARENA, LOONGARCH_GPR_SP, store_offset);
+	}
+
 	prepare_bpf_tail_call_cnt(ctx, &store_offset);
 
 	emit_insn(ctx, addid, LOONGARCH_GPR_FP, LOONGARCH_GPR_SP, stack_adjust);
@@ -186,6 +195,9 @@ static void build_prologue(struct jit_ctx *ctx)
 		emit_insn(ctx, addid, regmap[BPF_REG_FP], LOONGARCH_GPR_SP, bpf_stack_adjust);
 
 	ctx->stack_size = stack_adjust;
+
+	if (ctx->arena_vm_start)
+		move_imm(ctx, REG_ARENA, ctx->arena_vm_start, false);
 }
 
 static void __build_epilogue(struct jit_ctx *ctx, bool is_tail_call)
@@ -217,6 +229,11 @@ static void __build_epilogue(struct jit_ctx *ctx, bool is_tail_call)
 	load_offset -= sizeof(long);
 	emit_insn(ctx, ldd, LOONGARCH_GPR_S5, LOONGARCH_GPR_SP, load_offset);
 
+	if (ctx->arena_vm_start) {
+		load_offset -= sizeof(long);
+		emit_insn(ctx, ldd, REG_ARENA, LOONGARCH_GPR_SP, load_offset);
+	}
+
 	/*
 	 * When push into the stack, follow the order of tcc then tcc_ptr.
 	 * When pop from the stack, first pop tcc_ptr then followed by tcc.
@@ -442,14 +459,16 @@ static bool is_signed_bpf_cond(u8 cond)
 
 #define BPF_FIXUP_REG_MASK	GENMASK(31, 27)
 #define BPF_FIXUP_OFFSET_MASK	GENMASK(26, 0)
+#define REG_DONT_CLEAR_MARKER	0
 
 bool ex_handler_bpf(const struct exception_table_entry *ex,
 		    struct pt_regs *regs)
 {
 	int dst_reg = FIELD_GET(BPF_FIXUP_REG_MASK, ex->fixup);
 	off_t offset = FIELD_GET(BPF_FIXUP_OFFSET_MASK, ex->fixup);
 
-	regs->regs[dst_reg] = 0;
+	if (dst_reg != REG_DONT_CLEAR_MARKER)
+		regs->regs[dst_reg] = 0;
 	regs->csr_era = (unsigned long)&ex->fixup - offset;
 
 	return true;
@@ -468,7 +487,8 @@ static int add_exception_handler(const struct bpf_insn *insn,
 		return 0;
 
 	if (BPF_MODE(insn->code) != BPF_PROBE_MEM &&
-	    BPF_MODE(insn->code) != BPF_PROBE_MEMSX)
+	    BPF_MODE(insn->code) != BPF_PROBE_MEMSX &&
+	    BPF_MODE(insn->code) != BPF_PROBE_MEM32)
 		return 0;
 
 	if (WARN_ON_ONCE(ctx->num_exentries >= ctx->prog->aux->num_exentries))
@@ -528,8 +548,9 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx, bool ext
 	const u8 cond = BPF_OP(code);
 	const u8 t1 = LOONGARCH_GPR_T1;
 	const u8 t2 = LOONGARCH_GPR_T2;
-	const u8 src = regmap[insn->src_reg];
-	const u8 dst = regmap[insn->dst_reg];
+	const u8 t3 = LOONGARCH_GPR_T3;
+	u8 src = regmap[insn->src_reg];
+	u8 dst = regmap[insn->dst_reg];
 	const s16 off = insn->off;
 	const s32 imm = insn->imm;
 	const bool is32 = BPF_CLASS(insn->code) == BPF_ALU || BPF_CLASS(insn->code) == BPF_JMP32;
@@ -1035,8 +1056,19 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx, bool ext
 	case BPF_LDX | BPF_PROBE_MEMSX | BPF_B:
 	case BPF_LDX | BPF_PROBE_MEMSX | BPF_H:
 	case BPF_LDX | BPF_PROBE_MEMSX | BPF_W:
-		sign_extend = BPF_MODE(insn->code) == BPF_MEMSX ||
-			      BPF_MODE(insn->code) == BPF_PROBE_MEMSX;
+	/* LDX | PROBE_MEM32: dst = *(unsigned size *)(src + REG_ARENA + off) */
+	case BPF_LDX | BPF_PROBE_MEM32 | BPF_B:
+	case BPF_LDX | BPF_PROBE_MEM32 | BPF_H:
+	case BPF_LDX | BPF_PROBE_MEM32 | BPF_W:
+	case BPF_LDX | BPF_PROBE_MEM32 | BPF_DW:
+		sign_extend = BPF_MODE(code) == BPF_MEMSX ||
+			      BPF_MODE(code) == BPF_PROBE_MEMSX;
+
+		if (BPF_MODE(code) == BPF_PROBE_MEM32) {
+			emit_insn(ctx, addd, t2, src, REG_ARENA);
+			src = t2;
+		}
+
 		switch (BPF_SIZE(code)) {
 		case BPF_B:
 			if (is_signed_imm12(off)) {
@@ -1096,6 +1128,16 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx, bool ext
 	case BPF_ST | BPF_MEM | BPF_H:
 	case BPF_ST | BPF_MEM | BPF_W:
 	case BPF_ST | BPF_MEM | BPF_DW:
+	/* ST | PROBE_MEM32: *(size *)(dst + REG_ARENA + off) = imm */
+	case BPF_ST | BPF_PROBE_MEM32 | BPF_B:
+	case BPF_ST | BPF_PROBE_MEM32 | BPF_H:
+	case BPF_ST | BPF_PROBE_MEM32 | BPF_W:
+	case BPF_ST | BPF_PROBE_MEM32 | BPF_DW:
+		if (BPF_MODE(code) == BPF_PROBE_MEM32) {
+			emit_insn(ctx, addd, t3, dst, REG_ARENA);
+			dst = t3;
+		}
+
 		switch (BPF_SIZE(code)) {
 		case BPF_B:
 			move_imm(ctx, t1, imm, is32);
@@ -1138,13 +1180,27 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx, bool ext
 			}
 			break;
 		}
+
+		ret = add_exception_handler(insn, ctx, REG_DONT_CLEAR_MARKER);
+		if (ret)
+			return ret;
 		break;
 
 	/* *(size *)(dst + off) = src */
 	case BPF_STX | BPF_MEM | BPF_B:
 	case BPF_STX | BPF_MEM | BPF_H:
 	case BPF_STX | BPF_MEM | BPF_W:
 	case BPF_STX | BPF_MEM | BPF_DW:
+	/* STX | PROBE_MEM32: *(size *)(dst + REG_ARENA + off) = src */
+	case BPF_STX | BPF_PROBE_MEM32 | BPF_B:
+	case BPF_STX | BPF_PROBE_MEM32 | BPF_H:
+	case BPF_STX | BPF_PROBE_MEM32 | BPF_W:
+	case BPF_STX | BPF_PROBE_MEM32 | BPF_DW:
+		if (BPF_MODE(code) == BPF_PROBE_MEM32) {
+			emit_insn(ctx, addd, t2, dst, REG_ARENA);
+			dst = t2;
+		}
+
 		switch (BPF_SIZE(code)) {
 		case BPF_B:
 			if (is_signed_imm12(off)) {
@@ -1183,6 +1239,10 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx, bool ext
 			}
 			break;
 		}
+
+		ret = add_exception_handler(insn, ctx, REG_DONT_CLEAR_MARKER);
+		if (ret)
+			return ret;
 		break;
 
 	case BPF_STX | BPF_ATOMIC | BPF_W:
@@ -1894,6 +1954,7 @@ struct bpf_prog *bpf_int_jit_compile(struct bpf_prog *prog)
 
 	memset(&ctx, 0, sizeof(ctx));
 	ctx.prog = prog;
+	ctx.arena_vm_start = bpf_arena_get_kern_vm_start(prog->aux->arena);
 
 	ctx.offset = kvcalloc(prog->len + 1, sizeof(u32), GFP_KERNEL);
 	if (ctx.offset == NULL) {

arch/loongarch/net/bpf_jit.h

@@ -20,6 +20,7 @@ struct jit_ctx {
 	union loongarch_instruction *image;
 	union loongarch_instruction *ro_image;
 	u32 stack_size;
+	u64 arena_vm_start;
 };
 
 struct jit_data {