[PATCH bpf-next 11/16] bpf, arm64: Add support for arena percpu

[Xu Kuohai]

From: Xu Kuohai <xukuohai@xxxxxxxxxx>

Add arm64 JIT support for the arena percpu pseudo instruction.
Translate the address for each cpu by adding the offset
smp_processor_id() * chunk_sz to cpu 0's address.

Signed-off-by: Xu Kuohai <xukuohai@xxxxxxxxxx>
---
 arch/arm64/net/bpf_jit_comp.c | 51 ++++++++++++++++++++++++++++-------
 include/linux/bpf.h           |  2 +-
 kernel/bpf/arena.c            |  2 +-
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 1d5f4cf88d33..998806cc700e 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -500,6 +500,21 @@ static void emit_percpu_ptr(const u8 dst_reg, void __percpu *ptr,
 	emit(A64_ADD(1, dst_reg, dst_reg, tmp), ctx);
 }
 
+static void emit_current_cpu_id(u8 dst, struct jit_ctx *ctx)
+{
+	u8 tmp = bpf2a64[TMP_REG_1];
+	u8 tmp2 = bpf2a64[TMP_REG_2];
+	u32 cpu_offset = offsetof(struct thread_info, cpu);
+
+	emit(A64_MRS_SP_EL0(tmp), ctx);
+	if (is_lsi_offset(cpu_offset, 2)) {
+		emit(A64_LDR32I(dst, tmp, cpu_offset), ctx);
+	} else {
+		emit_a64_mov_i(1, tmp2, cpu_offset, ctx);
+		emit(A64_LDR32(dst, tmp, tmp2), ctx);
+	}
+}
+
 #define BTI_INSNS (IS_ENABLED(CONFIG_ARM64_BTI_KERNEL) ? 1 : 0)
 #define PAC_INSNS (IS_ENABLED(CONFIG_ARM64_PTR_AUTH_KERNEL) ? 1 : 0)
 
@@ -1742,19 +1757,10 @@ static int build_insn(const struct bpf_verifier_env *env, const struct bpf_insn
 		const u8 r0 = bpf2a64[BPF_REG_0];
 		bool func_addr_fixed;
 		u64 func_addr;
-		u32 cpu_offset;
 
 		/* Implement helper call to bpf_get_smp_processor_id() inline */
 		if (insn->src_reg == 0 && insn->imm == BPF_FUNC_get_smp_processor_id) {
-			cpu_offset = offsetof(struct thread_info, cpu);
-
-			emit(A64_MRS_SP_EL0(tmp), ctx);
-			if (is_lsi_offset(cpu_offset, 2)) {
-				emit(A64_LDR32I(r0, tmp, cpu_offset), ctx);
-			} else {
-				emit_a64_mov_i(1, tmp2, cpu_offset, ctx);
-				emit(A64_LDR32(r0, tmp, tmp2), ctx);
-			}
+			emit_current_cpu_id(r0, ctx);
 			break;
 		}
 
@@ -1829,6 +1835,26 @@ static int build_insn(const struct bpf_verifier_env *env, const struct bpf_insn
 		else
 			emit_a64_mov_i64(dst, imm64, ctx);
 
+#ifdef CONFIG_SMP
+		if (insn->src_reg == BPF_PSEUDO_ARENA_PERCPU) {
+			u32 chunk_sz = bpf_arena_percpu_chunk_sz(bpf_prog_arena(ctx->prog));
+
+			/* tmp = smp_processor_id() */
+			emit_current_cpu_id(tmp, ctx);
+			/*
+			 * tmp *= chunk_sz
+			 * arena_map_alloc() guarantees
+			 * chunk_sz * nr_cpu_ids <= vm_range <= 4 GiB and
+			 * cpu_number < nr_cpu_ids, so a 32-bit mul
+			 * cannot overflow here.
+			 */
+			emit_a64_mov_i(0, tmp2, chunk_sz, ctx);
+			emit(A64_MUL(0, tmp, tmp, tmp2), ctx);
+			/* dst += tmp */
+			emit(A64_ADD(1, dst, dst, tmp), ctx);
+		}
+#endif
+
 		return 1;
 	}
 
@@ -3446,6 +3472,11 @@ bool bpf_jit_supports_percpu_insn(void)
 	return true;
 }
 
+bool bpf_jit_supports_arena_percpu_insn(void)
+{
+	return true;
+}
+
 bool bpf_jit_bypass_spec_v4(void)
 {
 	/* In case of arm64, we rely on the firmware mitigation of Speculative
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 5ffefcd5f15e..4a7af8433c19 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -677,7 +677,7 @@ u64 bpf_arena_map_kern_vm_start(struct bpf_map *map);
 u32 bpf_arena_percpu_chunk_off(const struct bpf_map *map);
 u32 bpf_arena_percpu_chunk_sz(const struct bpf_map *map);
 bool bpf_arena_percpu_off_in_range(const struct bpf_map *map, u32 off);
-struct bpf_map *bpf_prog_arena(struct bpf_prog *prog);
+struct bpf_map *bpf_prog_arena(const struct bpf_prog *prog);
 int bpf_obj_name_cpy(char *dst, const char *src, unsigned int size);
 
 struct bpf_offload_dev;
diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 310573a2661b..4d2df31d1811 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -133,7 +133,7 @@ bool bpf_arena_percpu_off_in_range(const struct bpf_map *map, u32 off)
  * prog->aux->arena. Return that arena's underlying bpf_map, or NULL if
  * @prog does not reference an arena.
  */
-struct bpf_map *bpf_prog_arena(struct bpf_prog *prog)
+struct bpf_map *bpf_prog_arena(const struct bpf_prog *prog)
 {
 	struct bpf_arena *arena = prog->aux->arena;
 
-- 
2.47.3