[RFC bpf-next 01/11] bpf: use branch predictions in opt_hard_wire_dead_code_branches()

[Eduard Zingerman <eddyz87@gmail.com>]

Consider dead code elimination problem for program like below:

    main:
      1: r1 = 42
      2: call <subprogram>;
      3: exit

    subprogram:
      4: r0 = 1
      5: if r1 != 42 goto +1
      6: r0 = 2
      7: exit;

Here verifier would visit every instruction and thus
bpf_insn_aux_data->seen flag would be set for both true (7)
and falltrhough (6) branches of conditional (5).
Hence opt_hard_wire_dead_code_branches() will not replace
conditional (5) with unconditional jump.

To cover such cases:
- add two fields in struct bpf_insn_aux_data:
  - true_branch_taken;
  - false_branch_taken;
- adjust check_cond_jmp_op() to set the fields according to jump
  predictions;
- handle these flags in the opt_hard_wire_dead_code_branches():
  - true_branch_taken && !false_branch_taken
    always jump, replace instruction with 'goto off';
  - !true_branch_taken && false_branch_taken
    always falltrhough, replace with 'goto +0';
  - true_branch_taken && false_branch_taken
    jump and falltrhough are possible, don't change the instruction;
  - !true_branch_taken && !false_branch_taken
    neither jump, nor falltrhough are possible, if condition itself
    must be a dead code (should be removed by opt_remove_dead_code).

Signed-off-by: Eduard Zingerman <eddyz87@gmail.com>
---
 include/linux/bpf_verifier.h |  4 +++-
 kernel/bpf/verifier.c        | 16 ++++++++++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 4513372c5bc8..ed4eacfd4db7 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -570,7 +570,9 @@ struct bpf_insn_aux_data {
 	struct btf_struct_meta *kptr_struct_meta;
 	u64 map_key_state; /* constant (32 bit) key tracking for maps */
 	int ctx_field_size; /* the ctx field size for load insn, maybe 0 */
-	u32 seen; /* this insn was processed by the verifier at env->pass_cnt */
+	bool seen; /* this insn was processed by the verifier at env->pass_cnt */
+	bool true_branch_taken; /* for cond jumps, set if verifier ever took true branch */
+	bool false_branch_taken; /* for cond jumps, set if verifier ever took false branch */
 	bool sanitize_stack_spill; /* subject to Spectre v4 sanitation */
 	bool zext_dst; /* this insn zero extends dst reg */
 	bool needs_zext; /* alu op needs to clear upper bits */
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 7958d6ff6b73..3bae0bbc1da9 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -13265,7 +13265,7 @@ static void sanitize_mark_insn_seen(struct bpf_verifier_env *env)
 	 * rewrite/sanitize them.
 	 */
 	if (!vstate->speculative)
-		env->insn_aux_data[env->insn_idx].seen = env->pass_cnt;
+		env->insn_aux_data[env->insn_idx].seen = env->pass_cnt > 0;
 }
 
 static int sanitize_err(struct bpf_verifier_env *env,
@@ -15484,6 +15484,7 @@ static int check_cond_jmp_op(struct bpf_verifier_env *env,
 {
 	struct bpf_verifier_state *this_branch = env->cur_state;
 	struct bpf_verifier_state *other_branch;
+	struct bpf_insn_aux_data *aux = &env->insn_aux_data[*insn_idx];
 	struct bpf_reg_state *regs = this_branch->frame[this_branch->curframe]->regs;
 	struct bpf_reg_state *dst_reg, *other_branch_regs, *src_reg = NULL;
 	struct bpf_reg_state *eq_branch_regs;
@@ -15510,6 +15511,8 @@ static int check_cond_jmp_op(struct bpf_verifier_env *env,
 				insn->off, insn->imm);
 			return -EINVAL;
 		}
+		aux->true_branch_taken = true;
+		aux->false_branch_taken = true;
 		prev_st = find_prev_entry(env, cur_st->parent, idx);
 
 		/* branch out 'fallthrough' insn as a new state to explore */
@@ -15579,6 +15582,7 @@ static int check_cond_jmp_op(struct bpf_verifier_env *env,
 		 * the fall-through branch for simulation under speculative
 		 * execution.
 		 */
+		aux->true_branch_taken = true;
 		if (!env->bypass_spec_v1 &&
 		    !sanitize_speculative_path(env, insn, *insn_idx + 1,
 					       *insn_idx))
@@ -15592,6 +15596,7 @@ static int check_cond_jmp_op(struct bpf_verifier_env *env,
 		 * program will go. If needed, push the goto branch for
 		 * simulation under speculative execution.
 		 */
+		aux->false_branch_taken = true;
 		if (!env->bypass_spec_v1 &&
 		    !sanitize_speculative_path(env, insn,
 					       *insn_idx + insn->off + 1,
@@ -15602,6 +15607,9 @@ static int check_cond_jmp_op(struct bpf_verifier_env *env,
 		return 0;
 	}
 
+	aux->true_branch_taken = true;
+	aux->false_branch_taken = true;
+
 	/* Push scalar registers sharing same ID to jump history,
 	 * do this before creating 'other_branch', so that both
 	 * 'this_branch' and 'other_branch' share this history
@@ -19288,7 +19296,7 @@ static void adjust_insn_aux_data(struct bpf_verifier_env *env,
 {
 	struct bpf_insn_aux_data *old_data = env->insn_aux_data;
 	struct bpf_insn *insn = new_prog->insnsi;
-	u32 old_seen = old_data[off].seen;
+	bool old_seen = old_data[off].seen;
 	u32 prog_len;
 	int i;
 
@@ -19608,9 +19616,9 @@ static void opt_hard_wire_dead_code_branches(struct bpf_verifier_env *env)
 		if (!insn_is_cond_jump(insn->code))
 			continue;
 
-		if (!aux_data[i + 1].seen)
+		if (aux_data[i].true_branch_taken && !aux_data[i].false_branch_taken)
 			ja.off = insn->off;
-		else if (!aux_data[i + 1 + insn->off].seen)
+		else if (!aux_data[i].true_branch_taken && aux_data[i].false_branch_taken)
 			ja.off = 0;
 		else
 			continue;
-- 
2.47.0