[PATCHv2 bpf-next 3/3] bpf, x86: Add support for jit dry run

[Jiri Olsa <jolsa@kernel.org>]

Adding support to run jit code generation in dry_run mode that won't
store any code and only returns the jir code size.

The dry_run is enabled when __arch_prepare_bpf_trampoline is called
with rw_image argument as NULL.

It's used in arch_bpf_trampoline_size where it allows to skip the
image allocation, that gives speed up for tracing_multi attachment.

With current code:

  # ./test_progs -t tracing_multi_bench_attach -v
  ...
  serial_test_tracing_multi_bench_attach: found 55227 functions
  serial_test_tracing_multi_bench_attach: attached in   1.563s
  serial_test_tracing_multi_bench_attach: detached in   0.256s

With the fix:

  # ./test_progs -t tracing_multi_bench_attach -v
  ...
  serial_test_tracing_multi_bench_attach: found 55235 functions
  serial_test_tracing_multi_bench_attach: attached in   0.798s
  serial_test_tracing_multi_bench_attach: detached in   0.258s

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/net/bpf_jit_comp.c | 42 ++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 22 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index ee2ddeba3de1..d1af7dc5c5ce 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -25,6 +25,7 @@ static bool all_callee_regs_used[4] = {true, true, true, true};
 
 struct jit_emit_context {
 	u8 *prog;
+	bool dry_run;
 };
 
 static u8 *emit_code(u8 *ptr, u32 bytes, unsigned int len)
@@ -42,7 +43,10 @@ static u8 *emit_code(u8 *ptr, u32 bytes, unsigned int len)
 
 static void emit_code_jit(struct jit_emit_context *jit, u32 bytes, unsigned int len)
 {
-	jit->prog = emit_code(jit->prog, bytes, len);
+	if (jit->dry_run)
+		jit->prog += len;
+	else
+		jit->prog = emit_code(jit->prog, bytes, len);
 }
 
 #define EMIT(bytes, len) \
@@ -164,7 +168,8 @@ static int bpf_call_depth_emit_accounting(struct jit_emit_context *jit, void *fu
 	int size;
 
 	size = x86_call_depth_emit_accounting(insn_buff, func, ip);
-	memcpy(jit->prog, insn_buff, size);
+	if (!jit->dry_run)
+		memcpy(jit->prog, insn_buff, size);
 
 	jit->prog += size;
 	return size;
@@ -567,7 +572,8 @@ static int emit_patch(struct jit_emit_context *jit, void *func, void *ip, u8 opc
 	s64 offset;
 
 	offset = func - (ip + X86_PATCH_SIZE);
-	if (!is_simm32(offset)) {
+	/* We do not have meaningful ip value in the dry run, skip the check. */
+	if (!jit->dry_run && !is_simm32(offset)) {
 		pr_err("Target call %p is out of range\n", func);
 		return -ERANGE;
 	}
@@ -1647,6 +1653,7 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 	int err;
 
 	jit->prog = temp;
+	jit->dry_run = false;
 
 	stack_depth = bpf_prog->aux->stack_depth;
 	out_stack_arg_cnt = bpf_out_stack_arg_cnt(env, bpf_prog);
@@ -3190,8 +3197,10 @@ static int invoke_bpf_prog(const struct btf_func_model *m, struct jit_emit_conte
 		emit_stx(jit, BPF_DW, BPF_REG_FP, BPF_REG_0, -8);
 
 	/* replace 2 nops with JE insn, since jmp target is known */
-	jmp_insn[0] = X86_JE;
-	jmp_insn[1] = jit->prog - jmp_insn - 2;
+	if (!jit->dry_run) {
+		jmp_insn[0] = X86_JE;
+		jmp_insn[1] = jit->prog - jmp_insn - 2;
+	}
 
 	/* arg1: mov rdi, progs[i] */
 	emit_mov_imm64(jit, BPF_REG_1, (long) p >> 32, (u32) (long) p);
@@ -3472,6 +3481,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	}
 
 	jit->prog = rw_image;
+	jit->dry_run = !rw_image;
 
 	if (flags & BPF_TRAMP_F_INDIRECT) {
 		/*
@@ -3593,6 +3603,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		for (i = 0; i < fmod_ret->nr_nodes; i++) {
 			struct jit_emit_context branch_jit = {
 				.prog = branches[i],
+				.dry_run = jit->dry_run,
 			};
 
 			emit_cond_near_jump(&branch_jit, image + (jit->prog - (u8 *)rw_image),
@@ -3651,7 +3662,8 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	}
 	emit_return(jit, image + (jit->prog - (u8 *)rw_image));
 	/* Make sure the trampoline generation logic doesn't overflow */
-	if (WARN_ON_ONCE(jit->prog > (u8 *)rw_image_end - BPF_INSN_SAFETY)) {
+	if (!jit->dry_run &&
+	    WARN_ON_ONCE(jit->prog > (u8 *)rw_image_end - BPF_INSN_SAFETY)) {
 		ret = -EFAULT;
 		goto cleanup;
 	}
@@ -3710,23 +3722,9 @@ int arch_bpf_trampoline_size(const struct btf_func_model *m, u32 flags,
 			     struct bpf_tramp_nodes *tnodes, void *func_addr)
 {
 	struct bpf_tramp_image im;
-	void *image;
-	int ret;
 
-	/* Allocate a temporary buffer for __arch_prepare_bpf_trampoline().
-	 *
-	 * We cannot use kvmalloc here, because we need image to be in
-	 * module memory range.
-	 * Since it must be writable use bpf_jit_alloc_exec_rw().
-	 */
-	image = bpf_jit_alloc_exec_rw(PAGE_SIZE);
-	if (!image)
-		return -ENOMEM;
-
-	ret = __arch_prepare_bpf_trampoline(&im, image, image + PAGE_SIZE, image,
-					    m, flags, tnodes, func_addr);
-	bpf_jit_free_exec(image);
-	return ret;
+	return __arch_prepare_bpf_trampoline(&im, NULL, NULL, NULL, m, flags,
+					     tnodes, func_addr);
 }
 
 static int emit_bpf_dispatcher(struct jit_emit_context *jit, int a, int b, s64 *progs, u8 *image,
-- 
2.54.0