Re: [PATCH bpf-next 1/2] bpf, arm64: Add JIT support for timed may_goto

[<puranjay@kernel.org>]

Xu Kuohai <xukuohai@huaweicloud.com> writes:

> On 7/24/2025 8:54 PM, Puranjay Mohan wrote:
>> When verifier sees a timed may_goto instruction, it emits a call to
>> arch_bpf_timed_may_goto() with a stack offset in BPF_REG_AX (arm64 r9)
>> and expects a count value to be returned in the same register. The
>> verifier doesn't save or restore any registers before emitting this
>> call.
>> 
>> arch_bpf_timed_may_goto() should act as a trampoline to call
>> bpf_check_timed_may_goto() with AAPCS64 calling convention.
>> 
>> To support this custom calling convention, implement
>> arch_bpf_timed_may_goto() in assembly and make sure BPF caller saved
>> registers are saved and restored, call bpf_check_timed_may_goto with
>> arm64 calling convention where first argument and return value both are
>> in x0, then put the result back into BPF_REG_AX before returning.
>> 
>> Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
>> ---
>>   arch/arm64/net/Makefile             |  2 +-
>>   arch/arm64/net/bpf_jit_comp.c       | 13 ++++++++++-
>>   arch/arm64/net/bpf_timed_may_goto.S | 36 +++++++++++++++++++++++++++++
>>   3 files changed, 49 insertions(+), 2 deletions(-)
>>   create mode 100644 arch/arm64/net/bpf_timed_may_goto.S
>> 
>> diff --git a/arch/arm64/net/Makefile b/arch/arm64/net/Makefile
>> index 5c540efb7d9b9..3ae382bfca879 100644
>> --- a/arch/arm64/net/Makefile
>> +++ b/arch/arm64/net/Makefile
>> @@ -2,4 +2,4 @@
>>   #
>>   # ARM64 networking code
>>   #
>> -obj-$(CONFIG_BPF_JIT) += bpf_jit_comp.o
>> +obj-$(CONFIG_BPF_JIT) += bpf_jit_comp.o bpf_timed_may_goto.o
>> diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
>> index 89b1b8c248c62..6c954b36f57ea 100644
>> --- a/arch/arm64/net/bpf_jit_comp.c
>> +++ b/arch/arm64/net/bpf_jit_comp.c
>> @@ -1505,7 +1505,13 @@ static int build_insn(const struct bpf_insn *insn, struct jit_ctx *ctx,
>>   		if (ret < 0)
>>   			return ret;
>>   		emit_call(func_addr, ctx);
>> -		emit(A64_MOV(1, r0, A64_R(0)), ctx);
>> +		/*
>> +		 * Call to arch_bpf_timed_may_goto() is emitted by the
>> +		 * verifier and called with custom calling convention with
>> +		 * first argument and return value in BPF_REG_AX (x9).
>> +		 */
>> +		if (func_addr != (u64)arch_bpf_timed_may_goto)
>> +			emit(A64_MOV(1, r0, A64_R(0)), ctx);
>>   		break;
>>   	}
>>   	/* tail call */
>> @@ -2914,6 +2920,11 @@ bool bpf_jit_bypass_spec_v4(void)
>>   	return true;
>>   }
>>   
>> +bool bpf_jit_supports_timed_may_goto(void)
>> +{
>> +	return true;
>> +}
>> +
>>   bool bpf_jit_inlines_helper_call(s32 imm)
>>   {
>>   	switch (imm) {
>> diff --git a/arch/arm64/net/bpf_timed_may_goto.S b/arch/arm64/net/bpf_timed_may_goto.S
>> new file mode 100644
>> index 0000000000000..45f80e752345c
>> --- /dev/null
>> +++ b/arch/arm64/net/bpf_timed_may_goto.S
>> @@ -0,0 +1,36 @@
>> +/* SPDX-License-Identifier: GPL-2.0 */
>> +/* Copyright (c) 2025 Puranjay Mohan <puranjay@kernel.org> */
>> +
>> +#include <linux/linkage.h>
>> +
>> +SYM_FUNC_START(arch_bpf_timed_may_goto)
>> +	/* Allocate stack space and emit frame record */
>> +	stp     x29, x30, [sp, #-64]!
>> +	mov     x29, sp
>> +
>> +	/* Save BPF registers R0 - R5 (x7, x0-x4)*/
>> +	stp	x7, x0, [sp, #16]
>> +	stp	x1, x2, [sp, #32]
>> +	stp	x3, x4, [sp, #48]
>> +
>> +	/*
>> +	 * Stack depth was passed in BPF_REG_AX (x9), add it to the BPF_FP
>> +	 * (x25) to get the pointer to count and timestamp and pass it as the
>> +	 * first argument in x0.
>> +	 */
>> +	add	x0, x9, x25
>
> Whether BPF_REG_FP (x25) is set up by the arm64 jit depends on whether
> the jit detects any bpf instruction using it. Before generating the
> call to arch_bpf_timed_may_goto, the verifier generates a load
> instruction using FP, i.e. AX = *(u64 *)(FP - stack_off_cnt),
> so FP is always set up in this case.
>
> It seems a bit subtle. Maybe we should add a comment here?

Yes, a comment would be useful. I will add it in the next version.

>> +	bl	bpf_check_timed_may_goto
>> +	/* BPF_REG_AX(x9) will be stored into count, so move return value to it. */
>> +	mov	x9, x0
>> +
>> +
>
> Nit: one extra blank line
>

Thanks,
Puranjay