Re: [PATCH bpf-next 3/7] bpf, arm64: JIT __arena kfunc argument rebasing

On 8/11/2026 3:09 AM, Puranjay Mohan wrote:
> Implement arena argument rebasing for kfunc calls on arm64. x28 already
> holds kern_vm_start whenever the prog has an arena, and the newly added
> extended-register add zero-extends the 32-bit arena offset in place, so
> an unconditional argument costs a single instruction emitted right
> before the call:
> 
>    add xN, x28, wN, uxtw
> 
> A nullable argument first truncates into wN so that a zero offset leaves
> xN holding a real NULL, then tests it and jumps over the add:
> 
>    mov wN, wN
>    cbz wN, 1f
>    add xN, x28, wN, uxtw
> 1:
> 
> The rebase is native code generated after constant blinding has run on
> the BPF instruction stream, so blinding never sees it and needs no
> special handling. The emitted count depends only on the kfunc model, so
> it is identical across JIT passes.
> 
> bpf_jit_supports_arena_args() is not flipped yet; that happens when the
> struct_ops trampoline side is in place as well.
> 
> Signed-off-by: Puranjay Mohan <[email protected]>
> ---
>   arch/arm64/net/bpf_jit.h      | 11 +++++++++
>   arch/arm64/net/bpf_jit_comp.c | 42 +++++++++++++++++++++++++++++++++++
>   2 files changed, 53 insertions(+)
> 
> diff --git a/arch/arm64/net/bpf_jit.h b/arch/arm64/net/bpf_jit.h
> index d13de4222cfb7..b2fe6e6dcf447 100644
> --- a/arch/arm64/net/bpf_jit.h
> +++ b/arch/arm64/net/bpf_jit.h
> @@ -243,6 +243,17 @@
>   /* Rn - Rm; set condition flags */
>   #define A64_CMP(sf, Rn, Rm) A64_SUBS(sf, A64_ZR, Rn, Rm)
>   
> +/* Add/subtract (extended register) */
> +#define A64_ADDSUB_EREG(sf, Rd, Rn, Rm, ext, shift, type) \
> +	aarch64_insn_gen_add_sub_extended_reg(Rd, Rn, Rm, \
> +		AARCH64_INSN_EXTEND_##ext, shift, A64_VARIANT(sf), \
> +		AARCH64_INSN_ADSB_##type)
> +/* Rd = Rn + (EXT(Rm) << shift) */
> +#define A64_ADD_EXT(sf, Rd, Rn, Rm, ext, shift) \
> +	A64_ADDSUB_EREG(sf, Rd, Rn, Rm, ext, shift, ADD)
> +/* Rd = Rn + (u32)Rm */
> +#define A64_ADD_UXTW(Rd, Rn, Rm) A64_ADD_EXT(1, Rd, Rn, Rm, UXTW, 0)
> +
>   /* Data-processing (1 source) */
>   #define A64_DATA1(sf, Rd, Rn, type) aarch64_insn_gen_data1(Rd, Rn, \
>   	A64_VARIANT(sf), AARCH64_INSN_DATA1_##type)
> diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
> index 4af5a98b84e19..7aad17a51f006 100644
> --- a/arch/arm64/net/bpf_jit_comp.c
> +++ b/arch/arm64/net/bpf_jit_comp.c
> @@ -1249,6 +1249,43 @@ static void emit_stack_arg_store_imm(s32 imm, s16 bpf_off, const u8 tmp, struct
>   	}
>   }
>   
> +/*
> + * Rebase the __arena args of a kfunc call to arena kernel addresses,
> + * xN = kern_vm_start + (u32)xN, with the arena base register holding
> + * kern_vm_start. A nullable arg preserves NULL by skipping the add, tested
> + * on the truncated value as arena NULL is offset 0.
> + */
> +static int emit_kfunc_arena_args(struct jit_ctx *ctx, const struct bpf_insn *insn)
> +{
> +	const u8 arena_vm_base = bpf2a64[ARENA_VM_START];
> +	const struct btf_func_model *fm;
> +	int i;
> +
> +	fm = bpf_jit_find_kfunc_model(ctx->prog, insn);
> +	if (!fm)
> +		return -EINVAL;
> +
> +	for (i = 0; i < min_t(int, fm->nr_args, MAX_BPF_FUNC_REG_ARGS); i++) {
> +		const u8 reg = bpf2a64[BPF_REG_1 + i];
> +		u8 flags = fm->arg_flags[i];
> +
> +		if (!(flags & BTF_FMODEL_ARENA_ARG))
> +			continue;
> +		if (WARN_ON_ONCE(!ctx->arena_vm_start))
> +			return -EINVAL;
> +
> +		if (flags & BTF_FMODEL_NULLABLE_ARG) {
> +			/* 32-bit mov clears the upper 32 bits */
> +			emit(A64_MOV(0, reg, reg), ctx);
> +			/* skip the add so that NULL stays NULL */
> +			emit(A64_CBZ(0, reg, 2), ctx);
> +		}
> +		emit(A64_ADD_UXTW(reg, arena_vm_base, reg), ctx);
> +	}
> +
> +	return 0;
> +}
> +
>   /* JITs an eBPF instruction.
>    * Returns:
>    * 0  - successfully JITed an 8-byte eBPF instruction.
> @@ -1671,6 +1708,11 @@ static int build_insn(const struct bpf_verifier_env *env, const struct bpf_insn
>   					    &func_addr, &func_addr_fixed);
>   		if (ret < 0)
>   			return ret;
> +		if (insn->src_reg == BPF_PSEUDO_KFUNC_CALL) {
> +			ret = emit_kfunc_arena_args(ctx, insn);
> +			if (ret < 0)
> +				return ret;
> +		}
>   		emit_call(func_addr, ctx);
>   		/*
>   		 * Call to arch_bpf_timed_may_goto() is emitted by the

Reviewed-by: Xu Kuohai <[email protected]>