[PATCH bpf-next v1 0/6] Eliminate spill overhead for bpf_throw()

[PATCH bpf-next v1 0/6] Eliminate spill overhead for bpf_throw()

[Kumar Kartikeya Dwivedi]

A program that may throw is an exception boundary: its prologue pushes
every callee-saved register the JIT could clobber, five on x86-64 and
ten on arm64, whether the program uses them or not, and the exception
callback pops them all from the boundary frame before it runs. Every
invocation of such a program pays the pushes and pops: 18 bytes of code
on x86-64 for a program that uses none of the registers, five stp/ldp
pairs and 80 bytes of stack on arm64.

This series moves that work to the throw path. Each JITed function
records where its prologue saved the caller's callee-saved registers,
and bpf_throw() enters through a native entry that captures the
callee-saved registers of the throw site. The unwinder walks from the
throw site outwards, starting from the captured values and taking in
each frame the values the frame saved of its caller's; past the main
frame they are the registers the program's caller expects back. The
exception callback receives them and loads them before it reuses the
main frame. The boundary prologue is then an ordinary prologue on x86-64
and arm64.

Architectures without a native entry, s390 and powerpc, keep their
boundary scheme unchanged. arm64 is compile-tested only.

Kumar Kartikeya Dwivedi (6):
  bpf, x86: Account for the r9 push in exception boundary program
  bpf: Record where JITed subprograms keep callee-saved registers
  bpf: Add bpf_throw_unwind() and a hook for bpf_throw() entry
  bpf, x86: Add a native bpf_throw() entry that captures R6 to R9
  bpf, arm64: Add a native bpf_throw() entry that captures R6 to R9
  bpf: Restore callee-saved registers from the frames being unwound

 arch/arm64/net/Makefile                       |   1 +
 arch/arm64/net/bpf_jit_comp.c                 | 151 ++++++++++--------
 arch/arm64/net/bpf_throw.S                    |  30 ++++
 arch/x86/net/Makefile                         |   1 +
 arch/x86/net/bpf_jit_comp.c                   | 134 +++++++++-------
 arch/x86/net/bpf_throw.S                      |  39 +++++
 include/linux/bpf.h                           |  17 +-
 include/linux/bpf_throw.h                     |  49 ++++++
 kernel/bpf/core.c                             |   9 ++
 kernel/bpf/fixups.c                           |   1 -
 kernel/bpf/helpers.c                          |  47 +++++-
 kernel/bpf/verifier.c                         |   5 +
 .../bpf/progs/verifier_private_stack.c        |  12 +-
 13 files changed, 356 insertions(+), 140 deletions(-)
 create mode 100644 arch/arm64/net/bpf_throw.S
 create mode 100644 arch/x86/net/bpf_throw.S
 create mode 100644 include/linux/bpf_throw.h


base-commit: b71326388a8a24674846c1d8d3b6dd831d9843a9
-- 
2.53.0

[PATCH bpf-next v1 1/6] bpf, x86: Account for the r9 push in exception boundary program

[Kumar Kartikeya Dwivedi]

A program with a private stack pushes r9, which holds the private frame
pointer, around every call. The exception callback of an exception
boundary program resets the stack pointer to the value the boundary frame
had at the call it was thrown from, plus stack_arg_sp_adjust, and pops
the callee-saved registers from there. With a private stack that value
is 8 bytes short, so the registers are restored from the wrong slots and
the program returns to the kernel with corrupted callee-saved registers.
Include the push in stack_arg_sp_adjust.

A tracing program with 64 bytes or more of stack runs on a private
stack, so any such program that calls bpf_throw() is affected. The two
private stack exception selftests run such a program; they pass only
because nothing observes the restored registers after the return into
the kernel. The case was found while reworking how the callback restores
the callee-saved registers.

Fixes: 7d1cd70d4b16 ("bpf, x86: Support private stack in jit")
Cc: Yonghong Song <yonghong.song@linux.dev>
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/x86/net/bpf_jit_comp.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 72a2a0a89e03..ea97e068efda 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -2625,7 +2625,8 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 	 */
 	outgoing_rsp = out_stack_arg_cnt > 1 ? (out_stack_arg_cnt - 1) * 8 : 0;
 	if (bpf_prog->aux->exception_boundary)
-		bpf_prog->aux->stack_arg_sp_adjust = outgoing_rsp;
+		/* With a private stack every call site also has r9 pushed. */
+		bpf_prog->aux->stack_arg_sp_adjust = outgoing_rsp + (priv_frame_ptr ? 8 : 0);
 	emit_sub_rsp(&prog, outgoing_rsp);
 
 	if (arena_vm_start)
-- 
2.53.0

[PATCH bpf-next v1 2/6] bpf: Record where JITed subprograms keep callee-saved registers

[Kumar Kartikeya Dwivedi]

When bpf_throw() unwinds to the exception callback of the main program,
the frames between the throw site and the main program are discarded
without their epilogues running, so the callee-saved registers they
clobbered have to be restored some other way. Doing that from the frames
themselves needs, for each JITed program, where its prologue saved the
caller's R6..R9: walking from the throw site outwards, a register is
either still in the physical register, when no deeper frame saved it, or
sits in the prologue spill area of the nearest deeper frame that did.

Add callee_reg_off[4] to struct bpf_prog_aux, holding the offset of each
saved register from the native frame pointer that arch_bpf_stack_walk()
reports for the frame, or 0 when the register is not saved, and fill it
in from the x86-64 and arm64 JITs. Nothing consumes it yet.

On x86-64 the frame below rbp holds, from high to low addresses, the
program stack of round_up(stack_depth, 8) bytes (none with a private
stack), the tail call count and its pointer when tail_call_reachable,
r12 when the program is an exception boundary or uses an arena, and then
rbx, r13, r14 and r15 in push order, each only when it is pushed. An
exception boundary pushes all four.

On arm64 the frame record sits at x29, the tail call count pair right
below it, and the callee-saved registers follow as 16 byte pairs in the
order find_used_callee_regs() lists them, with the first register of a
pair at the lower address. The i-th pushed register is thus at
-32 - 16 * (i / 2) + 8 * (i % 2). An exception boundary pushes x19..x28
in order, so R6..R9 take the first four slots.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/bpf_jit_comp.c | 37 +++++++++++++++++++++++++++++++++++
 arch/x86/net/bpf_jit_comp.c   | 33 +++++++++++++++++++++++++++++++
 include/linux/bpf.h           |  6 ++++++
 3 files changed, 76 insertions(+)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 1d5f4cf88d33..6a99f883a091 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -421,6 +421,42 @@ static void find_used_callee_regs(struct jit_ctx *ctx)
 	ctx->nr_used_callee_reg = i;
 }
 
+/*
+ * Offset from the frame pointer of the @idx-th register pushed by push_callee_regs(). The
+ * frame record sits at the frame pointer, the tail call count pair right below it, and
+ * then the pushed pairs, each 16 bytes with its first register at the lower address.
+ */
+static int callee_reg_slot_off(int idx)
+{
+	return -32 - 16 * (idx / 2) + 8 * (idx % 2);
+}
+
+/*
+ * Record where push_callee_regs() saved the caller's R6..R9, relative to the frame
+ * pointer that arch_bpf_stack_walk() reports for this frame, for the exception unwinder.
+ * An exception boundary pushes x19..x28 in order, so R6..R9 take the first four slots.
+ */
+static void record_callee_reg_off(struct jit_ctx *ctx)
+{
+	struct bpf_prog_aux *aux = ctx->prog->aux;
+	int i, j;
+
+	if (aux->exception_boundary) {
+		for (i = 0; i < 4; i++)
+			aux->callee_reg_off[i] = callee_reg_slot_off(i);
+		return;
+	}
+	for (i = 0; i < 4; i++) {
+		aux->callee_reg_off[i] = 0;
+		for (j = 0; j < ctx->nr_used_callee_reg; j++) {
+			if (ctx->used_callee_reg[j] != bpf2a64[BPF_REG_6 + i])
+				continue;
+			aux->callee_reg_off[i] = callee_reg_slot_off(j);
+			break;
+		}
+	}
+}
+
 /* Save callee-saved registers */
 static void push_callee_regs(struct jit_ctx *ctx)
 {
@@ -451,6 +487,7 @@ static void push_callee_regs(struct jit_ctx *ctx)
 			emit(A64_PUSH(reg1, A64_ZR, A64_SP), ctx);
 		}
 	}
+	record_callee_reg_off(ctx);
 }
 
 /* Restore callee-saved registers */
diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index ea97e068efda..2d7495438d18 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -427,6 +427,24 @@ static void pop_callee_regs(u8 **pprog, bool *callee_regs_used)
 	*pprog = prog;
 }
 
+/*
+ * Record where push_callee_regs() saved the caller's R6..R9 relative to rbp, for the
+ * exception unwinder. @off is the offset of rsp from rbp right before the first push.
+ */
+static void record_callee_reg_off(struct bpf_prog *prog, s32 off, bool *callee_regs_used)
+{
+	int i;
+
+	for (i = 0; i < 4; i++) {
+		if (!callee_regs_used[i]) {
+			prog->aux->callee_reg_off[i] = 0;
+			continue;
+		}
+		off -= 8;
+		prog->aux->callee_reg_off[i] = off;
+	}
+}
+
 /* add rsp, depth */
 static void emit_add_rsp(u8 **pprog, u32 depth)
 {
@@ -2530,6 +2548,7 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 	u32 stack_depth;
 	int callee_saved_size;
 	s32 outgoing_arg_base;
+	s32 callee_regs_off;
 	int err;
 
 	stack_depth = bpf_prog->aux->stack_depth;
@@ -2591,6 +2610,20 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 		push_callee_regs(&prog, callee_regs_used);
 	}
 
+	/*
+	 * Record where the pushes above saved the caller's R6..R9, relative to the
+	 * rbp that arch_bpf_stack_walk() reports for this frame. Below rbp lie, from
+	 * high to low addresses: round_up(stack_depth, 8) bytes of program stack
+	 * (none with a private stack), the tail call count and its pointer (16 bytes)
+	 * when tail_call_reachable, r12 when exception_boundary or arena_vm_start,
+	 * and then rbx, r13, r14 and r15 in push order, each only when pushed.
+	 */
+	callee_regs_off = -(s32)round_up(stack_depth, 8) - (tail_call_reachable ? 16 : 0);
+	if (bpf_prog->aux->exception_boundary || arena_vm_start)
+		callee_regs_off -= 8;
+	record_callee_reg_off(bpf_prog, callee_regs_off,
+			      bpf_prog->aux->exception_boundary ? all_callee_regs_used : callee_regs_used);
+
 	/* Compute callee-saved register area size. */
 	callee_saved_size = 0;
 	if (bpf_prog->aux->exception_boundary || arena_vm_start)
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 7750df973494..010f14bb6c77 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1939,6 +1939,12 @@ struct bpf_prog_aux {
 	u64 (*bpf_exception_cb)(u64 cookie, u64 sp, u64 bp, u64, u64);
 	u16 stack_arg_sp_adjust;
 	u16 freplace_link_cnt; /* counts freplace links extending this prog */
+	/*
+	 * Offsets from the native frame pointer of this JITed program at which
+	 * its prologue saved the caller's R6..R9, or 0 when that register is not
+	 * saved. Filled by the JIT; consumed by the exception unwinder.
+	 */
+	s16 callee_reg_off[4];
 #ifdef CONFIG_SECURITY
 	void *security;
 #endif
-- 
2.53.0

[PATCH bpf-next v1 3/6] bpf: Add bpf_throw_unwind() and a hook for bpf_throw() entry

[Kumar Kartikeya Dwivedi]

bpf_throw() walks the stack to the frame of the main program and calls
its exception callback, which resets the stack to that frame and returns
from the program. Restoring the callee-saved registers of the frames the
walk discards needs their values at the throw site, and by the time C
code runs, the compiler may have clobbered them.

Move the body of bpf_throw() into bpf_throw_unwind(), which takes the
callee-saved registers of the throw site, R6..R9, along with the cookie.
An architecture that provides a native entry for the kfunc captures them
before calling it. bpf_arch_throw_entry() returns that entry, and the
verifier retargets the call to bpf_throw() there when it specializes the
kfunc, so the JIT keeps finding the function model of the call through
the specialized descriptor. Without an entry, bpf_throw() calls the
unwinder with NULL, as before. Nothing consumes the registers yet.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/bpf.h       |  1 +
 include/linux/bpf_throw.h | 33 +++++++++++++++++++++++++++++++++
 kernel/bpf/core.c         |  9 +++++++++
 kernel/bpf/helpers.c      | 16 +++++++++++++---
 kernel/bpf/verifier.c     |  5 +++++
 5 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 include/linux/bpf_throw.h

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 010f14bb6c77..839c28d3743f 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -7,6 +7,7 @@
 #include <uapi/linux/bpf.h>
 #include <uapi/linux/filter.h>
 #include <linux/bpf_defs.h>
+#include <linux/bpf_throw.h>
 
 #include <crypto/sha2.h>
 #include <linux/workqueue.h>
diff --git a/include/linux/bpf_throw.h b/include/linux/bpf_throw.h
new file mode 100644
index 000000000000..663776715b33
--- /dev/null
+++ b/include/linux/bpf_throw.h
@@ -0,0 +1,33 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * The interface between the architecture entries of bpf_throw() and the
+ * generic unwinder.
+ */
+#ifndef _LINUX_BPF_THROW_H
+#define _LINUX_BPF_THROW_H
+
+/* The callee-saved registers the unwinder tracks: R6..R9. */
+#define BPF_THROW_NR_CALLEE_SAVED	4
+/* The size of struct bpf_throw_regs, for the entries. */
+#define BPF_THROW_REGS_SIZE		(8 * BPF_THROW_NR_CALLEE_SAVED)
+
+#ifndef __ASSEMBLER__
+#include <linux/build_bug.h>
+#include <linux/types.h>
+
+/*
+ * The callee-saved registers captured at a throw site, indexed like
+ * bpf_prog_aux::callee_reg_off.
+ */
+struct bpf_throw_regs {
+	u64 saved[BPF_THROW_NR_CALLEE_SAVED];
+};
+static_assert(sizeof(struct bpf_throw_regs) == BPF_THROW_REGS_SIZE);
+
+void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie);
+/* The native entry of bpf_throw(), on an architecture that has one. */
+void bpf_throw_entry(void);
+void *bpf_arch_throw_entry(void);
+#endif /* __ASSEMBLER__ */
+
+#endif /* _LINUX_BPF_THROW_H */
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index fc22d66d194a..0205a687e1b5 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3601,6 +3601,15 @@ void __weak arch_bpf_stack_walk(bool (*consume_fn)(void *cookie, u64 ip, u64 sp,
 {
 }
 
+/*
+ * The native entry of bpf_throw(), which captures the BPF register file of
+ * the throw site, or NULL when the architecture calls the kfunc directly.
+ */
+void * __weak bpf_arch_throw_entry(void)
+{
+	return NULL;
+}
+
 bool __weak bpf_jit_supports_timed_may_goto(void)
 {
 	return false;
diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 648abe1ba63a..7c78630e4dd3 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -3427,7 +3427,12 @@ static bool bpf_stack_walker(void *cookie, u64 ip, u64 sp, u64 bp)
 	return false;
 }
 
-__bpf_kfunc void bpf_throw(u64 cookie)
+/*
+ * Unwind the stack to the exception callback of the main program. @regs holds
+ * the callee-saved registers of the throw site, or is NULL when the
+ * architecture has no native entry and the kfunc was called directly.
+ */
+void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie)
 {
 	struct bpf_throw_ctx ctx = {};
 
@@ -3439,8 +3444,8 @@ __bpf_kfunc void bpf_throw(u64 cookie)
 	WARN_ON_ONCE(!ctx.cnt);
 	/*
 	 * Prevent KASAN false positives for CONFIG_KASAN_STACK by unpoisoning
-	 * deeper stack depths than ctx.sp as we do not return from bpf_throw,
-	 * which skips compiler generated instrumentation to do the same. Some
+	 * deeper stack depths than ctx.sp as we do not return from here, which
+	 * skips compiler generated instrumentation to do the same. Some
 	 * architectures cannot recover sp while unwinding, so fall back to bp.
 	 */
 	kasan_unpoison_task_stack_below((void *)(long)(ctx.sp ?: ctx.bp));
@@ -3448,6 +3453,11 @@ __bpf_kfunc void bpf_throw(u64 cookie)
 	WARN(1, "A call to BPF exception callback should never return\n");
 }
 
+__bpf_kfunc void bpf_throw(u64 cookie)
+{
+	bpf_throw_unwind(NULL, cookie);
+}
+
 __bpf_kfunc int bpf_wq_init(struct bpf_wq *wq, void *p__const_map, unsigned int flags)
 {
 	struct bpf_async_kern *async = (struct bpf_async_kern *)wq;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 01434e1d23b5..d980530e58bb 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -23162,6 +23162,11 @@ static int specialize_kfunc(struct bpf_verifier_env *env, struct bpf_kfunc_desc
 	} else if (func_id == special_kfunc_list[KF_bpf_arena_free_pages]) {
 		if (env->insn_aux_data[insn_idx].non_sleepable)
 			addr = (unsigned long)bpf_arena_free_pages_non_sleepable;
+	} else if (func_id == special_kfunc_list[KF_bpf_throw]) {
+		void *entry = bpf_arch_throw_entry();
+
+		if (entry)
+			addr = (unsigned long)entry;
 	}
 	desc->addr = addr;
 	return 0;
-- 
2.53.0

[PATCH bpf-next v1 4/6] bpf, x86: Add a native bpf_throw() entry that captures R6 to R9

[Kumar Kartikeya Dwivedi]

Add bpf_throw_entry, which the verifier calls in place of bpf_throw(). It
saves R6..R9 to the stack in the layout of struct bpf_throw_regs, passes
the cookie from R1 and calls bpf_throw_unwind(), which does not return,
so a ud2 follows the call.

The JIT calls the entry directly through __bpf_call_base + imm, so it
needs no ENDBR. objtool generates ORC for its frame, so the stack walk of
bpf_throw_unwind() passes through the entry to the BPF frames, where the
frame pointer fallback of the ORC unwinder takes over as before.

The entry calls bpf_throw_unwind(), which is built with the syscall, so
it is built with CONFIG_BPF_SYSCALL as well and the JIT returns it only
then. The JIT itself builds with CONFIG_BPF_JIT alone, for cBPF.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/x86/net/Makefile       |  1 +
 arch/x86/net/bpf_jit_comp.c |  8 ++++++++
 arch/x86/net/bpf_throw.S    | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+)
 create mode 100644 arch/x86/net/bpf_throw.S

diff --git a/arch/x86/net/Makefile b/arch/x86/net/Makefile
index dddbefc0f439..daef0aa8b28c 100644
--- a/arch/x86/net/Makefile
+++ b/arch/x86/net/Makefile
@@ -7,4 +7,5 @@ ifeq ($(CONFIG_X86_32),y)
         obj-$(CONFIG_BPF_JIT) += bpf_jit_comp32.o
 else
         obj-$(CONFIG_BPF_JIT) += bpf_jit_comp.o bpf_timed_may_goto.o
+        obj-$(CONFIG_BPF_SYSCALL) += bpf_throw.o
 endif
diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 2d7495438d18..1a865313d09e 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -5171,6 +5171,14 @@ void arch_bpf_stack_walk(bool (*consume_fn)(void *cookie, u64 ip, u64 sp, u64 bp
 #endif
 }
 
+#ifdef CONFIG_BPF_SYSCALL
+/* arch/x86/net/bpf_throw.S, which calls bpf_throw_unwind(); both need the syscall. */
+void *bpf_arch_throw_entry(void)
+{
+	return (void *)bpf_throw_entry;
+}
+#endif
+
 void bpf_arch_poke_desc_update(struct bpf_jit_poke_descriptor *poke,
 			       struct bpf_prog *new, struct bpf_prog *old)
 {
diff --git a/arch/x86/net/bpf_throw.S b/arch/x86/net/bpf_throw.S
new file mode 100644
index 000000000000..c70a4002cc8a
--- /dev/null
+++ b/arch/x86/net/bpf_throw.S
@@ -0,0 +1,37 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
+
+#include <linux/linkage.h>
+#include <asm/nospec-branch.h>
+
+	.code64
+	.section .text, "ax"
+
+/*
+ * Native entry of the bpf_throw() kfunc. The verifier retargets the call to
+ * bpf_throw() here, so the BPF register file of the throw site is still intact
+ * on entry. Capture its callee-saved registers on the stack as a struct
+ * bpf_throw_regs and hand it to bpf_throw_unwind(), which never returns.
+ */
+SYM_FUNC_START(bpf_throw_entry)
+	ANNOTATE_NOENDBR
+
+	/* Setup frame. */
+	pushq %rbp
+	movq %rsp, %rbp
+
+	/* Save R9..R6 so that [rsp + 8 * i] is saved[i]. */
+	pushq %r15	/* R9 */
+	pushq %r14	/* R8 */
+	pushq %r13	/* R7 */
+	pushq %rbx	/* R6 */
+
+	/* bpf_throw_unwind(regs, cookie); the cookie is R1 of the caller. */
+	movq %rdi, %rsi
+	movq %rsp, %rdi
+
+	/* Emit call depth accounting for call below. */
+	CALL_DEPTH_ACCOUNT
+	call bpf_throw_unwind
+	ud2
+SYM_FUNC_END(bpf_throw_entry)
-- 
2.53.0

[PATCH bpf-next v1 5/6] bpf, arm64: Add a native bpf_throw() entry that captures R6 to R9

[Kumar Kartikeya Dwivedi]

Give arm64 the native entry of bpf_throw() that x86-64 has: the verifier
retargets the call to bpf_throw() to bpf_throw_entry(), which pushes a
frame record whose return address is the one into the program and saves
R6-R9 of the throw site, x19-x22, below it as a struct bpf_throw_regs,
then calls bpf_throw_unwind() with it. Nothing consumes the registers
yet.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/Makefile       |  1 +
 arch/arm64/net/bpf_jit_comp.c |  8 ++++++++
 arch/arm64/net/bpf_throw.S    | 26 ++++++++++++++++++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 arch/arm64/net/bpf_throw.S

diff --git a/arch/arm64/net/Makefile b/arch/arm64/net/Makefile
index 3ae382bfca87..70d18f067d01 100644
--- a/arch/arm64/net/Makefile
+++ b/arch/arm64/net/Makefile
@@ -3,3 +3,4 @@
 # ARM64 networking code
 #
 obj-$(CONFIG_BPF_JIT) += bpf_jit_comp.o bpf_timed_may_goto.o
+obj-$(CONFIG_BPF_SYSCALL) += bpf_throw.o
diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 6a99f883a091..cc0ee9b3726b 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -10,6 +10,7 @@
 #include <linux/arm-smccc.h>
 #include <linux/bitfield.h>
 #include <linux/bpf.h>
+#include <linux/bpf_throw.h>
 #include <linux/cfi.h>
 #include <linux/filter.h>
 #include <linux/memory.h>
@@ -3499,6 +3500,13 @@ bool bpf_jit_supports_timed_may_goto(void)
 	return true;
 }
 
+#ifdef CONFIG_BPF_SYSCALL
+void *bpf_arch_throw_entry(void)
+{
+	return (void *)bpf_throw_entry;
+}
+#endif
+
 bool bpf_jit_inlines_helper_call(s32 imm)
 {
 	switch (imm) {
diff --git a/arch/arm64/net/bpf_throw.S b/arch/arm64/net/bpf_throw.S
new file mode 100644
index 000000000000..aac1c5e967e1
--- /dev/null
+++ b/arch/arm64/net/bpf_throw.S
@@ -0,0 +1,26 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
+
+#include <linux/linkage.h>
+#include <linux/bpf_throw.h>
+
+/*
+ * Native entry of the bpf_throw() kfunc. The verifier retargets the call to
+ * bpf_throw() here, so the BPF register file of the throw site is still intact
+ * on entry and x30 is the return address into the program. Push a frame record
+ * whose link register is that address and save R6-R9, x19-x22, below it as a
+ * struct bpf_throw_regs: saved[i] sits at [sp + 16 + 8 * i]. Then hand it to
+ * bpf_throw_unwind(), which never returns.
+ */
+SYM_FUNC_START(bpf_throw_entry)
+	stp	x29, x30, [sp, #-(16 + BPF_THROW_REGS_SIZE)]!
+	mov	x29, sp
+	stp	x19, x20, [sp, #16]
+	stp	x21, x22, [sp, #32]
+
+	/* bpf_throw_unwind(regs, cookie); the cookie is R1 of the caller. */
+	mov	x1, x0
+	add	x0, sp, #16
+	bl	bpf_throw_unwind
+	brk	#0x1
+SYM_FUNC_END(bpf_throw_entry)
-- 
2.53.0

Re: [PATCH bpf-next v1 5/6] bpf, arm64: Add a native bpf_throw() entry that captures R6 to R9

[bot+bpf-ci]

[-- Attachment #1: Type: text/plain, Size: 1755 bytes --]

> bpf, arm64: Add a native bpf_throw() entry that captures R6 to R9

TL;DR: Nit only: the open-coded "brk #0x1" after bl bpf_throw_unwind() in the
arm64 bpf_throw_entry could be ASM_BUG(), which would give a clearer report if
the call ever returned. Not a bug.

> diff --git a/arch/arm64/net/bpf_throw.S b/arch/arm64/net/bpf_throw.S
> new file mode 100644
> index 000000000000..aac1c5e967e1
> --- /dev/null
> +++ b/arch/arm64/net/bpf_throw.S

[ ... ]

> +SYM_FUNC_START(bpf_throw_entry)
> +	stp	x29, x30, [sp, #-(16 + BPF_THROW_REGS_SIZE)]!

[ ... ]

> +	/* bpf_throw_unwind(regs, cookie); the cookie is R1 of the caller. */
> +	mov	x1, x0
> +	add	x0, sp, #16
> +	bl	bpf_throw_unwind
> +	brk	#0x1
> +SYM_FUNC_END(bpf_throw_entry)

This isn't a bug, but would ASM_BUG() fit better here than the open-coded
brk #0x1 after the call to bpf_throw_unwind()?

Immediate 0x1 is not one of the BRK immediates listed in
arch/arm64/include/asm/brk-imm.h, so if bpf_throw_unwind() ever did return,
the result would be an anonymous "Oops - BRK" from do_el1_brk64().

Other arm64 assembly with the same "bl <noreturn>; trap" pattern uses
ASM_BUG() from <asm/asm-bug.h>:

arch/arm64/kernel/entry.S:
	bl	handle_bad_stack
	ASM_BUG()

arch/arm64/kernel/head.S:
	bl	start_kernel
	ASM_BUG()

ASM_BUG() emits brk BUG_BRK_IMM plus a __bug_table entry, so the report
would name bpf_throw.S and the line.

The follow-up commit 9a710c167b82 ("bpf: Restore callee-saved registers from
the frames being unwound") keeps the brk #0x1 as well.


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/39472833568

[PATCH bpf-next v1 6/6] bpf: Restore callee-saved registers from the frames being unwound

[Kumar Kartikeya Dwivedi]

A program that may throw is an exception boundary: its prologue pushes
every callee-saved register the JIT could clobber, five on x86-64 and
ten on arm64, whether the program uses them or not, and the exception
callback pops them all from the boundary frame before it runs. That is
what restores the registers a discarded frame clobbered, since the
epilogues of the frames between the throw site and the boundary never
run. It costs every invocation of such a program the pushes and pops,
18 bytes of code on x86-64 for a program that uses none of them, and
80 bytes of stack on arm64.

The unwinder can do that work instead, on the throw path only. Each
JITed function already records where its prologue saved the caller's
R6..R9; extend the record to every other callee-saved register, r12 on
x86-64, which holds the arena base, and x23..x28 on arm64: x25..x28 are
the BPF frame pointer, the tail call count pointer, the private stack
pointer and the arena base, and x23 and x24, which no JITed code
touches, may still be in use by the C code of the unwinder when it
enters the callback, whose epilogue never runs. The native entries
capture the extra registers along with R6..R9, before any C code runs.
Walking from the throw site outwards, the unwinder starts from the
values the site had and takes, in each frame, the values that frame
saved of its caller's; past the main frame, they are the registers the
program's caller expects back. The exception callback receives them as
its fourth argument and loads them before it reuses the main frame, in
place of popping the boundary's saves, and builds the rest of its frame
like any function.

The boundary prologue is then an ordinary prologue, and the x86-64
stack pointer adjustment the callback needed to find the boundary's
saves goes away with the pops.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/bpf_jit_comp.c                 | 126 ++++++++----------
 arch/arm64/net/bpf_throw.S                    |  10 +-
 arch/x86/net/bpf_jit_comp.c                   | 120 +++++++----------
 arch/x86/net/bpf_throw.S                      |   6 +-
 include/linux/bpf.h                           |  16 ++-
 include/linux/bpf_throw.h                     |  24 +++-
 kernel/bpf/fixups.c                           |   1 -
 kernel/bpf/helpers.c                          |  33 ++++-
 .../bpf/progs/verifier_private_stack.c        |  12 +-
 9 files changed, 174 insertions(+), 174 deletions(-)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index cc0ee9b3726b..1afcf10c49ba 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -433,24 +433,29 @@ static int callee_reg_slot_off(int idx)
 }
 
 /*
- * Record where push_callee_regs() saved the caller's R6..R9, relative to the frame
- * pointer that arch_bpf_stack_walk() reports for this frame, for the exception unwinder.
- * An exception boundary pushes x19..x28 in order, so R6..R9 take the first four slots.
+ * Record where the prologue saved the caller's callee-saved registers, relative to
+ * the frame pointer that arch_bpf_stack_walk() reports for this frame, for the
+ * exception unwinder. The registers are x19..x28 in the order of struct
+ * bpf_throw_regs, R6..R9 first: the tail call count pointer is saved with the frame
+ * record, the others where push_callee_regs() put them, and x23 and x24 nowhere,
+ * since JITed code leaves them alone.
  */
 static void record_callee_reg_off(struct jit_ctx *ctx)
 {
 	struct bpf_prog_aux *aux = ctx->prog->aux;
 	int i, j;
 
-	if (aux->exception_boundary) {
-		for (i = 0; i < 4; i++)
-			aux->callee_reg_off[i] = callee_reg_slot_off(i);
-		return;
-	}
-	for (i = 0; i < 4; i++) {
+	BUILD_BUG_ON(BPF_THROW_NR_CALLEE_SAVED != 10);
+	for (i = 0; i < BPF_THROW_NR_CALLEE_SAVED; i++) {
+		u8 reg = A64_R(19 + i);
+
 		aux->callee_reg_off[i] = 0;
+		if (reg == bpf2a64[TCCNT_PTR]) {
+			aux->callee_reg_off[i] = -8;
+			continue;
+		}
 		for (j = 0; j < ctx->nr_used_callee_reg; j++) {
-			if (ctx->used_callee_reg[j] != bpf2a64[BPF_REG_6 + i])
+			if (ctx->used_callee_reg[j] != reg)
 				continue;
 			aux->callee_reg_off[i] = callee_reg_slot_off(j);
 			break;
@@ -458,35 +463,25 @@ static void record_callee_reg_off(struct jit_ctx *ctx)
 	}
 }
 
-/* Save callee-saved registers */
+/*
+ * Save the callee-saved registers this function uses. An exception boundary
+ * saves nothing more: the unwinder recovers the registers of the frames it
+ * discards and hands them to the exception callback.
+ */
 static void push_callee_regs(struct jit_ctx *ctx)
 {
 	int reg1, reg2, i;
 
-	/*
-	 * Program acting as exception boundary should save all ARM64
-	 * Callee-saved registers as the exception callback needs to recover
-	 * all ARM64 Callee-saved registers in its epilogue.
-	 */
-	if (ctx->prog->aux->exception_boundary) {
-		emit(A64_PUSH(A64_R(19), A64_R(20), A64_SP), ctx);
-		emit(A64_PUSH(A64_R(21), A64_R(22), A64_SP), ctx);
-		emit(A64_PUSH(A64_R(23), A64_R(24), A64_SP), ctx);
-		emit(A64_PUSH(A64_R(25), A64_R(26), A64_SP), ctx);
-		emit(A64_PUSH(A64_R(27), A64_R(28), A64_SP), ctx);
-		ctx->fp_used = true;
-	} else {
-		find_used_callee_regs(ctx);
-		for (i = 0; i + 1 < ctx->nr_used_callee_reg; i += 2) {
-			reg1 = ctx->used_callee_reg[i];
-			reg2 = ctx->used_callee_reg[i + 1];
-			emit(A64_PUSH(reg1, reg2, A64_SP), ctx);
-		}
-		if (i < ctx->nr_used_callee_reg) {
-			reg1 = ctx->used_callee_reg[i];
-			/* keep SP 16-byte aligned */
-			emit(A64_PUSH(reg1, A64_ZR, A64_SP), ctx);
-		}
+	find_used_callee_regs(ctx);
+	for (i = 0; i + 1 < ctx->nr_used_callee_reg; i += 2) {
+		reg1 = ctx->used_callee_reg[i];
+		reg2 = ctx->used_callee_reg[i + 1];
+		emit(A64_PUSH(reg1, reg2, A64_SP), ctx);
+	}
+	if (i < ctx->nr_used_callee_reg) {
+		reg1 = ctx->used_callee_reg[i];
+		/* keep SP 16-byte aligned */
+		emit(A64_PUSH(reg1, A64_ZR, A64_SP), ctx);
 	}
 	record_callee_reg_off(ctx);
 }
@@ -494,34 +489,19 @@ static void push_callee_regs(struct jit_ctx *ctx)
 /* Restore callee-saved registers */
 static void pop_callee_regs(struct jit_ctx *ctx)
 {
-	struct bpf_prog_aux *aux = ctx->prog->aux;
 	int reg1, reg2, i;
 
-	/*
-	 * Program acting as exception boundary pushes R23 and R24 in addition
-	 * to BPF callee-saved registers. Exception callback uses the boundary
-	 * program's stack frame, so recover these extra registers in the above
-	 * two cases.
-	 */
-	if (aux->exception_boundary || aux->exception_cb) {
-		emit(A64_POP(A64_R(27), A64_R(28), A64_SP), ctx);
-		emit(A64_POP(A64_R(25), A64_R(26), A64_SP), ctx);
-		emit(A64_POP(A64_R(23), A64_R(24), A64_SP), ctx);
-		emit(A64_POP(A64_R(21), A64_R(22), A64_SP), ctx);
-		emit(A64_POP(A64_R(19), A64_R(20), A64_SP), ctx);
-	} else {
-		i = ctx->nr_used_callee_reg - 1;
-		if (ctx->nr_used_callee_reg % 2 != 0) {
-			reg1 = ctx->used_callee_reg[i];
-			emit(A64_POP(reg1, A64_ZR, A64_SP), ctx);
-			i--;
-		}
-		while (i > 0) {
-			reg1 = ctx->used_callee_reg[i - 1];
-			reg2 = ctx->used_callee_reg[i];
-			emit(A64_POP(reg1, reg2, A64_SP), ctx);
-			i -= 2;
-		}
+	i = ctx->nr_used_callee_reg - 1;
+	if (ctx->nr_used_callee_reg % 2 != 0) {
+		reg1 = ctx->used_callee_reg[i];
+		emit(A64_POP(reg1, A64_ZR, A64_SP), ctx);
+		i--;
+	}
+	while (i > 0) {
+		reg1 = ctx->used_callee_reg[i - 1];
+		reg2 = ctx->used_callee_reg[i];
+		emit(A64_POP(reg1, reg2, A64_SP), ctx);
+		i -= 2;
 	}
 }
 
@@ -622,24 +602,22 @@ static int build_prologue(struct jit_ctx *ctx, bool ebpf_from_cbpf)
 			/* BTI landing pad for the tail call, done with a BR */
 			emit_bti(A64_BTI_J, ctx);
 		}
-		push_callee_regs(ctx);
 	} else {
 		/*
-		 * Exception callback receives FP of Main Program as third
-		 * parameter
+		 * The exception callback receives the main program's frame
+		 * pointer as its third parameter and, as its fourth, the
+		 * callee-saved registers x19..x28 the unwinder recovered from
+		 * the frames it discarded: the values the main program's caller
+		 * expects back. Load them, then build the rest of the frame
+		 * below the main program's frame record like any function.
 		 */
+		for (int i = 0; i < BPF_THROW_NR_CALLEE_SAVED; i++)
+			emit(A64_LDR64I(A64_R(19 + i), A64_R(3), 8 * i), ctx);
 		emit(A64_MOV(1, A64_FP, A64_R(2)), ctx);
-		/*
-		 * Main Program already pushed the frame record and the
-		 * callee-saved registers. The exception callback will not push
-		 * anything and re-use the main program's stack.
-		 *
-		 * 12 registers are on the stack
-		 */
-		emit(A64_SUB_I(1, A64_SP, A64_FP, 96), ctx);
-		/* The callback may use its own BPF stack, set up fp for it. */
-		ctx->fp_used = true;
+		emit(A64_MOV(1, A64_SP, A64_FP), ctx);
+		prepare_bpf_tail_call_cnt(ctx);
 	}
+	push_callee_regs(ctx);
 
 	/* Stack must be multiples of 16B */
 	ctx->stack_size = round_up(prog->aux->stack_depth, 16);
diff --git a/arch/arm64/net/bpf_throw.S b/arch/arm64/net/bpf_throw.S
index aac1c5e967e1..7010ba2388cd 100644
--- a/arch/arm64/net/bpf_throw.S
+++ b/arch/arm64/net/bpf_throw.S
@@ -8,15 +8,19 @@
  * Native entry of the bpf_throw() kfunc. The verifier retargets the call to
  * bpf_throw() here, so the BPF register file of the throw site is still intact
  * on entry and x30 is the return address into the program. Push a frame record
- * whose link register is that address and save R6-R9, x19-x22, below it as a
- * struct bpf_throw_regs: saved[i] sits at [sp + 16 + 8 * i]. Then hand it to
- * bpf_throw_unwind(), which never returns.
+ * whose link register is that address and save R6-R9, x19-x22, and the extra
+ * callee-saved registers x23-x28 below it as a struct bpf_throw_regs: saved[i]
+ * sits at [sp + 16 + 8 * i]. Then hand it to bpf_throw_unwind(), which never
+ * returns.
  */
 SYM_FUNC_START(bpf_throw_entry)
 	stp	x29, x30, [sp, #-(16 + BPF_THROW_REGS_SIZE)]!
 	mov	x29, sp
 	stp	x19, x20, [sp, #16]
 	stp	x21, x22, [sp, #32]
+	stp	x23, x24, [sp, #48]
+	stp	x25, x26, [sp, #64]
+	stp	x27, x28, [sp, #80]
 
 	/* bpf_throw_unwind(regs, cookie); the cookie is R1 of the caller. */
 	mov	x1, x0
diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 1a865313d09e..be9bfb9c8336 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -34,7 +34,6 @@ void __asan_load8(void *p);
 void __asan_store8(void *p);
 #endif
 
-static bool all_callee_regs_used[4] = {true, true, true, true};
 
 static void *trampoline_size_image;
 
@@ -428,13 +427,19 @@ static void pop_callee_regs(u8 **pprog, bool *callee_regs_used)
 }
 
 /*
- * Record where push_callee_regs() saved the caller's R6..R9 relative to rbp, for the
- * exception unwinder. @off is the offset of rsp from rbp right before the first push.
+ * Record where the prologue saved the caller's callee-saved registers relative to rbp,
+ * for the exception unwinder: r12 when @r12 was pushed, then rbx, r13, r14 and r15
+ * as push_callee_regs() pushed them. @off is the offset of rsp from rbp right before
+ * the first push.
  */
-static void record_callee_reg_off(struct bpf_prog *prog, s32 off, bool *callee_regs_used)
+static void record_callee_reg_off(struct bpf_prog *prog, s32 off, bool *callee_regs_used,
+				  bool r12)
 {
 	int i;
 
+	prog->aux->callee_reg_off[4] = r12 ? off - 8 : 0;
+	if (r12)
+		off -= 8;
 	for (i = 0; i < 4; i++) {
 		if (!callee_regs_used[i]) {
 			prog->aux->callee_reg_off[i] = 0;
@@ -614,18 +619,21 @@ static void emit_prologue(u8 **pprog, u8 *ip, u32 stack_depth, bool ebpf_from_cb
 			/* Keep the same instruction layout. */
 			emit_nops(&prog, 3);     /* nop3 */
 	}
-	/* Exception callback receives FP as third parameter */
+	/*
+	 * The exception callback receives the main program's frame pointer as
+	 * its third parameter and, as its fourth, the callee-saved registers
+	 * the unwinder recovered from the frames it discarded, rbx, r13, r14,
+	 * r15 and r12 in that order: the values the main program's caller
+	 * expects back. Load them, then reuse the main program's frame.
+	 */
 	if (is_exception_cb) {
-		EMIT3(0x48, 0x89, 0xF4); /* mov rsp, rsi */
-		EMIT3(0x48, 0x89, 0xD5); /* mov rbp, rdx */
-		/* The main frame must have exception_boundary as true, so we
-		 * first restore those callee-saved regs from stack, before
-		 * reusing the stack frame.
-		 */
-		pop_callee_regs(&prog, all_callee_regs_used);
-		pop_r12(&prog);
-		/* Reset the stack frame. */
-		EMIT3(0x48, 0x89, 0xEC); /* mov rsp, rbp */
+		EMIT3(0x48, 0x8B, 0x19);       /* mov rbx, [rcx] */
+		EMIT4(0x4C, 0x8B, 0x69, 0x08); /* mov r13, [rcx + 8] */
+		EMIT4(0x4C, 0x8B, 0x71, 0x10); /* mov r14, [rcx + 16] */
+		EMIT4(0x4C, 0x8B, 0x79, 0x18); /* mov r15, [rcx + 24] */
+		EMIT4(0x4C, 0x8B, 0x61, 0x20); /* mov r12, [rcx + 32] */
+		EMIT3(0x48, 0x89, 0xD5);       /* mov rbp, rdx */
+		EMIT3(0x48, 0x89, 0xEC);       /* mov rsp, rbp */
 	} else {
 		EMIT1(0x55);             /* push rbp */
 		EMIT3(0x48, 0x89, 0xE5); /* mov rbp, rsp */
@@ -901,14 +909,9 @@ static void emit_bpf_tail_call_indirect(struct bpf_prog *bpf_prog,
 	/* Inc tail_call_cnt if the slot is populated. */
 	EMIT4(0x48, 0x83, 0x00, 0x01);            /* add qword ptr [rax], 1 */
 
-	if (bpf_prog->aux->exception_boundary) {
-		pop_callee_regs(&prog, all_callee_regs_used);
+	pop_callee_regs(&prog, callee_regs_used);
+	if (bpf_prog_uses_arena_insns(bpf_prog))
 		pop_r12(&prog);
-	} else {
-		pop_callee_regs(&prog, callee_regs_used);
-		if (bpf_prog_uses_arena_insns(bpf_prog))
-			pop_r12(&prog);
-	}
 
 	/* Pop tail_call_cnt_ptr. */
 	EMIT1(0x58);                              /* pop rax */
@@ -967,14 +970,9 @@ static void emit_bpf_tail_call_direct(struct bpf_prog *bpf_prog,
 	/* Inc tail_call_cnt if the slot is populated. */
 	EMIT4(0x48, 0x83, 0x00, 0x01);                /* add qword ptr [rax], 1 */
 
-	if (bpf_prog->aux->exception_boundary) {
-		pop_callee_regs(&prog, all_callee_regs_used);
+	pop_callee_regs(&prog, callee_regs_used);
+	if (bpf_prog_uses_arena_insns(bpf_prog))
 		pop_r12(&prog);
-	} else {
-		pop_callee_regs(&prog, callee_regs_used);
-		if (bpf_prog_uses_arena_insns(bpf_prog))
-			pop_r12(&prog);
-	}
 
 	/* Pop tail_call_cnt_ptr. */
 	EMIT1(0x58);                                  /* pop rax */
@@ -2594,49 +2592,31 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 
 	bpf_prog->aux->ksym.fp_start = prog - temp;
 
-	/* Exception callback will clobber callee regs for its own use, and
-	 * restore the original callee regs from main prog's stack frame.
+	/*
+	 * Save the callee-saved registers this function uses. An exception
+	 * boundary saves nothing more: the unwinder recovers the registers of
+	 * the frames it discards and hands them to the exception callback.
 	 */
-	if (bpf_prog->aux->exception_boundary) {
-		/* We also need to save r12, which is not mapped to any BPF
-		 * register, as we throw after entry into the kernel, which may
-		 * overwrite r12.
-		 */
+	if (arena_vm_start)
 		push_r12(&prog);
-		push_callee_regs(&prog, all_callee_regs_used);
-	} else {
-		if (arena_vm_start)
-			push_r12(&prog);
-		push_callee_regs(&prog, callee_regs_used);
-	}
+	push_callee_regs(&prog, callee_regs_used);
 
 	/*
-	 * Record where the pushes above saved the caller's R6..R9, relative to the
-	 * rbp that arch_bpf_stack_walk() reports for this frame. Below rbp lie, from
-	 * high to low addresses: round_up(stack_depth, 8) bytes of program stack
-	 * (none with a private stack), the tail call count and its pointer (16 bytes)
-	 * when tail_call_reachable, r12 when exception_boundary or arena_vm_start,
-	 * and then rbx, r13, r14 and r15 in push order, each only when pushed.
+	 * Record where the pushes above saved the caller's registers, relative to
+	 * the rbp that arch_bpf_stack_walk() reports for this frame. Below rbp lie,
+	 * from high to low addresses: round_up(stack_depth, 8) bytes of program
+	 * stack (none with a private stack), the tail call count and its pointer
+	 * (16 bytes) when tail_call_reachable, r12 when arena_vm_start, and then
+	 * rbx, r13, r14 and r15 in push order, each only when pushed.
 	 */
 	callee_regs_off = -(s32)round_up(stack_depth, 8) - (tail_call_reachable ? 16 : 0);
-	if (bpf_prog->aux->exception_boundary || arena_vm_start)
-		callee_regs_off -= 8;
-	record_callee_reg_off(bpf_prog, callee_regs_off,
-			      bpf_prog->aux->exception_boundary ? all_callee_regs_used : callee_regs_used);
+	record_callee_reg_off(bpf_prog, callee_regs_off, callee_regs_used, arena_vm_start);
 
 	/* Compute callee-saved register area size. */
-	callee_saved_size = 0;
-	if (bpf_prog->aux->exception_boundary || arena_vm_start)
-		callee_saved_size += 8; /* r12 */
-	if (bpf_prog->aux->exception_boundary) {
-		callee_saved_size += 4 * 8; /* rbx, r13, r14, r15 */
-	} else {
-		int j;
-
-		for (j = 0; j < 4; j++)
-			if (callee_regs_used[j])
-				callee_saved_size += 8;
-	}
+	callee_saved_size = arena_vm_start ? 8 : 0; /* r12 */
+	for (int j = 0; j < 4; j++)
+		if (callee_regs_used[j])
+			callee_saved_size += 8;
 	/*
 	 * Base offset from rbp for translating BPF outgoing args 7+
 	 * to native offsets. BPF uses negative offsets from r11
@@ -2657,9 +2637,6 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 	 * Arg 6 goes into r9 register, not on stack.
 	 */
 	outgoing_rsp = out_stack_arg_cnt > 1 ? (out_stack_arg_cnt - 1) * 8 : 0;
-	if (bpf_prog->aux->exception_boundary)
-		/* With a private stack every call site also has r9 pushed. */
-		bpf_prog->aux->stack_arg_sp_adjust = outgoing_rsp + (priv_frame_ptr ? 8 : 0);
 	emit_sub_rsp(&prog, outgoing_rsp);
 
 	if (arena_vm_start)
@@ -3822,14 +3799,9 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 			}
 			/* Deallocate outgoing args 7+ area. */
 			emit_add_rsp(&prog, outgoing_rsp);
-			if (bpf_prog->aux->exception_boundary) {
-				pop_callee_regs(&prog, all_callee_regs_used);
+			pop_callee_regs(&prog, callee_regs_used);
+			if (arena_vm_start)
 				pop_r12(&prog);
-			} else {
-				pop_callee_regs(&prog, callee_regs_used);
-				if (arena_vm_start)
-					pop_r12(&prog);
-			}
 			EMIT1(0xC9);         /* leave */
 			bpf_prog->aux->ksym.fp_end = prog - temp;
 
diff --git a/arch/x86/net/bpf_throw.S b/arch/x86/net/bpf_throw.S
index c70a4002cc8a..7d249e099b67 100644
--- a/arch/x86/net/bpf_throw.S
+++ b/arch/x86/net/bpf_throw.S
@@ -11,7 +11,8 @@
  * Native entry of the bpf_throw() kfunc. The verifier retargets the call to
  * bpf_throw() here, so the BPF register file of the throw site is still intact
  * on entry. Capture its callee-saved registers on the stack as a struct
- * bpf_throw_regs and hand it to bpf_throw_unwind(), which never returns.
+ * bpf_throw_regs, along with r12, the callee-saved register no BPF register
+ * maps to, and hand it to bpf_throw_unwind(), which never returns.
  */
 SYM_FUNC_START(bpf_throw_entry)
 	ANNOTATE_NOENDBR
@@ -20,7 +21,8 @@ SYM_FUNC_START(bpf_throw_entry)
 	pushq %rbp
 	movq %rsp, %rbp
 
-	/* Save R9..R6 so that [rsp + 8 * i] is saved[i]. */
+	/* Save r12, then R9..R6 so that [rsp + 8 * i] is saved[i]. */
+	pushq %r12	/* saved[4] */
 	pushq %r15	/* R9 */
 	pushq %r14	/* R8 */
 	pushq %r13	/* R7 */
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 839c28d3743f..b4a1edfb1df4 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1937,15 +1937,21 @@ struct bpf_prog_aux {
 	int cgroup_atype; /* enum cgroup_bpf_attach_type */
 	struct bpf_map *cgroup_storage[MAX_BPF_CGROUP_STORAGE_TYPE];
 	char name[BPF_OBJ_NAME_LEN];
-	u64 (*bpf_exception_cb)(u64 cookie, u64 sp, u64 bp, u64, u64);
-	u16 stack_arg_sp_adjust;
+	/*
+	 * The exception callback of the main program: @cookie is the value of
+	 * the throw, @sp and @bp the native stack and frame pointers of the main
+	 * program's frame, @saved the callee-saved registers the unwinder
+	 * restored, R6..R9 then the architecture's own, see struct bpf_throw_regs.
+	 */
+	u64 (*bpf_exception_cb)(u64 cookie, u64 sp, u64 bp, u64 saved, u64);
 	u16 freplace_link_cnt; /* counts freplace links extending this prog */
 	/*
 	 * Offsets from the native frame pointer of this JITed program at which
-	 * its prologue saved the caller's R6..R9, or 0 when that register is not
-	 * saved. Filled by the JIT; consumed by the exception unwinder.
+	 * its prologue saved the caller's callee-saved registers, R6..R9 then
+	 * the architecture's own, or 0 when that register is not saved. Filled
+	 * by the JIT; consumed by the exception unwinder.
 	 */
-	s16 callee_reg_off[4];
+	s16 callee_reg_off[BPF_THROW_NR_CALLEE_SAVED];
 #ifdef CONFIG_SECURITY
 	void *security;
 #endif
diff --git a/include/linux/bpf_throw.h b/include/linux/bpf_throw.h
index 663776715b33..e3daa2875415 100644
--- a/include/linux/bpf_throw.h
+++ b/include/linux/bpf_throw.h
@@ -6,8 +6,24 @@
 #ifndef _LINUX_BPF_THROW_H
 #define _LINUX_BPF_THROW_H
 
-/* The callee-saved registers the unwinder tracks: R6..R9. */
-#define BPF_THROW_NR_CALLEE_SAVED	4
+/*
+ * The callee-saved registers of the architecture that no BPF register maps
+ * to, which the unwinder restores along with R6..R9 when it discards frames:
+ * r12 on x86-64, which holds the arena base, and x23..x28 on arm64, of which
+ * x25..x28 are the BPF frame pointer, the tail call count pointer, the
+ * private stack pointer and the arena base, while x23 and x24 are left alone
+ * by JITed code but may be in use by the C code of the unwinder when it
+ * enters the exception callback, whose epilogue never runs.
+ */
+#if defined(CONFIG_X86_64)
+#define BPF_THROW_NR_EXTRA	1
+#elif defined(CONFIG_ARM64)
+#define BPF_THROW_NR_EXTRA	6
+#else
+#define BPF_THROW_NR_EXTRA	0
+#endif
+/* The callee-saved registers the unwinder tracks: R6..R9, then the extra ones. */
+#define BPF_THROW_NR_CALLEE_SAVED	(4 + BPF_THROW_NR_EXTRA)
 /* The size of struct bpf_throw_regs, for the entries. */
 #define BPF_THROW_REGS_SIZE		(8 * BPF_THROW_NR_CALLEE_SAVED)
 
@@ -16,8 +32,8 @@
 #include <linux/types.h>
 
 /*
- * The callee-saved registers captured at a throw site, indexed like
- * bpf_prog_aux::callee_reg_off.
+ * The callee-saved registers captured at a throw site, R6..R9 then the extra
+ * ones of the architecture, indexed like bpf_prog_aux::callee_reg_off.
  */
 struct bpf_throw_regs {
 	u64 saved[BPF_THROW_NR_CALLEE_SAVED];
diff --git a/kernel/bpf/fixups.c b/kernel/bpf/fixups.c
index 409b9104c723..8b9944cac47a 100644
--- a/kernel/bpf/fixups.c
+++ b/kernel/bpf/fixups.c
@@ -1584,7 +1584,6 @@ static int jit_subprogs(struct bpf_verifier_env *env)
 	prog->aux->real_func_cnt = env->subprog_cnt;
 	prog->aux->bpf_exception_cb = (void *)func[env->exception_callback_subprog]->bpf_func;
 	prog->aux->exception_boundary = func[0]->aux->exception_boundary;
-	prog->aux->stack_arg_sp_adjust = func[0]->aux->stack_arg_sp_adjust;
 	bpf_prog_jit_attempt_done(prog);
 	return 0;
 out_free:
diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 7c78630e4dd3..d898aaf33b97 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -3401,12 +3401,21 @@ struct bpf_throw_ctx {
 	u64 sp;
 	u64 bp;
 	int cnt;
+	const struct bpf_throw_regs *regs;	/* register file of the throw site, or NULL */
+	/*
+	 * The callee-saved registers of the frame being walked, R6-R9 then the
+	 * architecture's own: what the throw site had, then what each frame
+	 * recovered of its caller's. After the main frame they are the ones
+	 * its caller expects back.
+	 */
+	u64 saved[BPF_THROW_NR_CALLEE_SAVED];
 };
 
 static bool bpf_stack_walker(void *cookie, u64 ip, u64 sp, u64 bp)
 {
 	struct bpf_throw_ctx *ctx = cookie;
 	struct bpf_prog *prog;
+	int i;
 
 	/*
 	 * The RCU read lock is held to safely traverse the latch tree, but we
@@ -3418,6 +3427,20 @@ static bool bpf_stack_walker(void *cookie, u64 ip, u64 sp, u64 bp)
 	rcu_read_unlock();
 	if (!prog)
 		return !ctx->cnt;
+	if (ctx->regs) {
+		/* The symbol of a program JITed as a set of functions covers its first function. */
+		const struct bpf_prog_aux *fn = !prog->is_func && prog->aux->func ?
+						prog->aux->func[0]->aux : prog->aux;
+
+		/*
+		 * Recover the caller's callee-saved registers: those this frame's
+		 * prologue saved sit in its frame, the rest are still the values
+		 * the frame itself has.
+		 */
+		for (i = 0; i < BPF_THROW_NR_CALLEE_SAVED; i++)
+			if (fn->callee_reg_off[i])
+				ctx->saved[i] = *(u64 *)(long)(bp + fn->callee_reg_off[i]);
+	}
 	ctx->cnt++;
 	if (bpf_is_subprog(prog))
 		return true;
@@ -3430,12 +3453,16 @@ static bool bpf_stack_walker(void *cookie, u64 ip, u64 sp, u64 bp)
 /*
  * Unwind the stack to the exception callback of the main program. @regs holds
  * the callee-saved registers of the throw site, or is NULL when the
- * architecture has no native entry and the kfunc was called directly.
+ * architecture has no native entry and the kfunc was called directly; with it,
+ * the callee-saved registers of the frames unwound are restored for the
+ * callback, without it the main program's prologue has saved them all.
  */
 void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie)
 {
-	struct bpf_throw_ctx ctx = {};
+	struct bpf_throw_ctx ctx = { .regs = regs };
 
+	if (regs)
+		memcpy(ctx.saved, regs->saved, sizeof(ctx.saved));
 	arch_bpf_stack_walk(bpf_stack_walker, &ctx);
 	WARN_ON_ONCE(!ctx.aux);
 	if (ctx.aux)
@@ -3449,7 +3476,7 @@ void bpf_throw_unwind(struct bpf_throw_regs *regs, u64 cookie)
 	 * architectures cannot recover sp while unwinding, so fall back to bp.
 	 */
 	kasan_unpoison_task_stack_below((void *)(long)(ctx.sp ?: ctx.bp));
-	ctx.aux->bpf_exception_cb(cookie, ctx.sp + ctx.aux->stack_arg_sp_adjust, ctx.bp, 0, 0);
+	ctx.aux->bpf_exception_cb(cookie, ctx.sp, ctx.bp, (u64)(long)ctx.saved, 0);
 	WARN(1, "A call to BPF exception callback should never return\n");
 }
 
diff --git a/tools/testing/selftests/bpf/progs/verifier_private_stack.c b/tools/testing/selftests/bpf/progs/verifier_private_stack.c
index 3d9994a52464..76a4e1823c6c 100644
--- a/tools/testing/selftests/bpf/progs/verifier_private_stack.c
+++ b/tools/testing/selftests/bpf/progs/verifier_private_stack.c
@@ -212,11 +212,7 @@ __jited("	stp	x29, x30, [sp, #-0x10]!")
 __jited("	mov	x29, sp")
 __jited("	stp	xzr, x26, [sp, #-0x10]!")
 __jited("	mov	x26, sp")
-__jited("	stp	x19, x20, [sp, #-0x10]!")
-__jited("	stp	x21, x22, [sp, #-0x10]!")
-__jited("	stp	x23, x24, [sp, #-0x10]!")
-__jited("	stp	x25, x26, [sp, #-0x10]!")
-__jited("	stp	x27, x28, [sp, #-0x10]!")
+__jited("	stp	x25, x27, [sp, #-0x10]!")
 __jited("	mov	x27, {{.*}}")
 __jited("	movk	x27, {{.*}}, lsl #16")
 __jited("	movk	x27, {{.*}}")
@@ -228,7 +224,7 @@ __jited("	str	x0, [x27]")
 __jited("	mov	x0, #0x0")
 __jited("	bl	0x{{.*}}")
 __jited("	mov	x8, x0")
-__jited("	ldp	x27, x28, [sp], #0x10")
+__jited("	ldp	x25, x27, [sp], #0x10")
 int private_stack_exception_main_prog(void)
 {
 	asm volatile ("			\
@@ -256,7 +252,7 @@ __jited("...")
 __jited("	callq")
 __jited("	popq	%r9")
 __arch_arm64
-__jited("	stp	x27, x28, [sp, #-0x10]!")
+__jited("	stp	x25, x27, [sp, #-0x10]!")
 __jited("	mov	x27, {{.*}}")
 __jited("	movk	x27, {{.*}}, lsl #16")
 __jited("	movk	x27, {{.*}}")
@@ -267,7 +263,7 @@ __jited("	mov	x0, #0x2a")
 __jited("	str	x0, [x27]")
 __jited("	bl	0x{{.*}}")
 __jited("	mov	x8, x0")
-__jited("	ldp	x27, x28, [sp], #0x10")
+__jited("	ldp	x25, x27, [sp], #0x10")
 int private_stack_exception_sub_prog(void)
 {
 	asm volatile ("			\
-- 
2.53.0