[PATCH bpf-next v3 02/13] bpf: Add helpers to describe the R0:R2 return register pair

[Yonghong Song <yonghong.song@linux.dev>]

LLVM 23 added support for returning a value in two registers for an
__int128, or a struct/union whose size is greater than 8 but not more than
16 bytes: such a value comes back in the R0:R2 register pair, with R2
holding the upper half. See LLVM patches [1] and [2].

Later patches teach the JIT, precision backtracking, live register analysis
and the verifier itself about that convention. All of them need to answer
the same question: does this subprogram return its value in a register
pair? Add the shared helpers up front so that those patches can be ordered
independently of each other:

 - subprog_ret_type() resolves a subprogram's BTF return type. It is
   factored out of subprog_returns_void(). The verifier_bug_if(!func) and
   !func_proto checks it replaces are redundant, since
   check_btf_func_early() already rejects a func_info record whose type_id
   is not a BTF_KIND_FUNC pointing at a BTF_KIND_FUNC_PROTO. A check on
   prog->aux->{btf,func_info} is added instead: unlike
   subprog_returns_void(), which is only used for global subprograms, later
   callers ask about static subprograms too, and those may belong to a
   program loaded without BTF.

 - ret_regs_cnt() maps the size of a return value to the number of
   registers holding it.

 - bpf_ret_reg_pair() answers the question above. Its users query it at
   every subprogram call and at every subprogram exit, that is once per
   verifier state rather than once per subprogram, so the answer is
   precomputed into bpf_subprog_info->ret_reg_pair by
   bpf_compute_subprog_ret_regs() and the helper itself is a flag test.
   It lives in bpf_verifier.h because kernel/bpf/backtrack.c and
   kernel/bpf/liveness.c need it as well.

bpf_compute_subprog_ret_regs() runs in bpf_check() right before
bpf_compute_live_registers(), which is the first of those users: by then
BTF func_info has been validated and the subprogram list is final.

No functional change, bpf_ret_reg_pair() has no callers yet.

  [1] https://github.com/llvm/llvm-project/pull/190894
  [2] https://github.com/llvm/llvm-project/pull/206876

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 include/linux/bpf_verifier.h | 12 ++++++++
 kernel/bpf/verifier.c        | 55 ++++++++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 93f7c2075eea..79680f3b4c74 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -811,6 +811,8 @@ struct bpf_subprog_info {
 	bool is_async_cb: 1;
 	bool is_exception_cb: 1;
 	bool args_cached: 1;
+	/* true if the return value is passed in the R0:R2 register pair */
+	bool ret_reg_pair: 1;
 	/* true if bpf_fastcall stack region is used by functions that can't be inlined */
 	bool keep_fastcall_stack: 1;
 	bool changes_pkt_data: 1;
@@ -1044,6 +1046,16 @@ static inline struct bpf_subprog_info *subprog_info(struct bpf_verifier_env *env
 	return &env->subprog_info[subprog];
 }
 
+/*
+ * True if @subprog returns its value in the R0:R2 register pair. Cached by
+ * bpf_compute_subprog_ret_regs(), since this is queried on hot paths: at
+ * every subprogram call and at every subprogram exit.
+ */
+static inline bool bpf_ret_reg_pair(struct bpf_verifier_env *env, int subprog)
+{
+	return subprog_info(env, subprog)->ret_reg_pair;
+}
+
 struct bpf_call_summary {
 	u8 num_params;
 	bool is_void;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 40150390dd50..58a177d26c46 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -382,27 +382,57 @@ bool bpf_subprog_is_global(const struct bpf_verifier_env *env, int subprog)
 	return aux && aux[subprog].linkage == BTF_FUNC_GLOBAL;
 }
 
-static bool subprog_returns_void(struct bpf_verifier_env *env, int subprog)
+/* Return type of a subprogram, NULL if it cannot be resolved */
+static const struct btf_type *subprog_ret_type(struct bpf_verifier_env *env, int subprog)
 {
-	const struct btf_type *type, *func, *func_proto;
+	const struct btf_type *func, *func_proto;
 	const struct btf *btf = env->prog->aux->btf;
 	u32 btf_id;
 
+	if (!btf || !env->prog->aux->func_info)
+		return NULL;
+
 	btf_id = env->prog->aux->func_info[subprog].type_id;
 
+	/* Both already validated by prepare_btf_func() at prog load. */
 	func = btf_type_by_id(btf, btf_id);
-	if (verifier_bug_if(!func, env, "btf_id %u not found", btf_id))
-		return false;
-
 	func_proto = btf_type_by_id(btf, func->type);
-	if (!func_proto)
-		return false;
 
-	type = btf_type_skip_modifiers(btf, func_proto->type, NULL);
-	if (!type)
-		return false;
+	return btf_type_skip_modifiers(btf, func_proto->type, NULL);
+}
+
+static bool subprog_returns_void(struct bpf_verifier_env *env, int subprog)
+{
+	const struct btf_type *type = subprog_ret_type(env, subprog);
 
-	return btf_type_is_void(type);
+	return type && btf_type_is_void(type);
+}
+
+/*
+ * Number of registers holding a function return value: a value of up to 8
+ * bytes is returned in R0, a value of more than 8 bytes and no more than 16
+ * bytes (an __int128 or a struct/union of such size) is returned in the R0:R2
+ * register pair, with R2 holding the upper half.
+ */
+static u32 ret_regs_cnt(u32 size)
+{
+	return size > 8 && size <= 16 ? 2 : 1;
+}
+
+/*
+ * Resolve the return convention of every subprogram once, so that
+ * bpf_ret_reg_pair() is a plain flag test on the hot paths that use it.
+ */
+static void bpf_compute_subprog_ret_regs(struct bpf_verifier_env *env)
+{
+	const struct btf_type *type;
+	int subprog;
+
+	for (subprog = 0; subprog < env->subprog_cnt; subprog++) {
+		type = subprog_ret_type(env, subprog);
+		if (type && (btf_type_is_struct(type) || btf_type_is_scalar(type)))
+			subprog_info(env, subprog)->ret_reg_pair = ret_regs_cnt(type->size) > 1;
+	}
 }
 
 static const char *subprog_name(const struct bpf_verifier_env *env, int subprog)
@@ -20288,6 +20318,9 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr,
 	if (ret < 0)
 		goto skip_full_check;
 
+	/* must precede the first bpf_ret_reg_pair() user below */
+	bpf_compute_subprog_ret_regs(env);
+
 	ret = bpf_compute_live_registers(env);
 	if (ret < 0)
 		goto skip_full_check;
-- 
2.53.0-Meta