sync

linsyking · linsyking · commit 09dfb712ec26 · 2024-11-25T23:09:18.000-05:00
diff --git a/include/linux/bpf_ir.h b/include/linux/bpf_ir.h
@@ -1185,10 +1185,14 @@ void msan(struct bpf_ir_env *env, struct ir_function *fun, void *param);
 void bpf_ir_div_by_zero(struct bpf_ir_env *env, struct ir_function *fun,
 			void *param);
 
+void bpf_ir_optimize_code_compaction(struct bpf_ir_env *env,
+				     struct ir_function *fun, void *param);
+
 extern const struct builtin_pass_cfg bpf_ir_kern_add_counter_pass;
 extern const struct builtin_pass_cfg bpf_ir_kern_optimization_pass;
 extern const struct builtin_pass_cfg bpf_ir_kern_msan;
 extern const struct builtin_pass_cfg bpf_ir_kern_div_by_zero_pass;
+extern const struct builtin_pass_cfg bpf_ir_kern_compaction_pass;
 
 void translate_throw(struct bpf_ir_env *env, struct ir_function *fun,
 		     void *param);
diff --git a/kernel/bpf/ir/bpf_ir.c b/kernel/bpf/ir/bpf_ir.c
@@ -228,6 +228,8 @@ static const struct function_pass post_passes[] = {
 	DEF_FUNC_PASS(add_counter, "add_counter", false),
 	/* CG Preparation Passes */
 	DEF_NON_OVERRIDE_FUNC_PASS(translate_throw, "translate_throw"),
+	DEF_NON_OVERRIDE_FUNC_PASS(bpf_ir_optimize_code_compaction,
+				   "optimize_compaction"),
 	DEF_NON_OVERRIDE_FUNC_PASS(bpf_ir_optimize_ir, "optimize_ir"),
 	DEF_NON_OVERRIDE_FUNC_PASS(bpf_ir_cg_change_fun_arg, "change_fun_arg"),
 	DEF_NON_OVERRIDE_FUNC_PASS(bpf_ir_cg_change_call_pre_cg, "change_call"),
diff --git a/kernel/bpf/ir/code_compaction.c b/kernel/bpf/ir/code_compaction.c
@@ -15,8 +15,9 @@ void bpf_ir_optimize_code_compaction(struct bpf_ir_env *env,
 	{
 		struct ir_basic_block *bb = *pos;
 		list_for_each_entry(cur_insn, &bb->ir_insn_head, list_ptr) {
-			if (!prev_insn) {
+			if (prev_insn == NULL) {
 				prev_insn = cur_insn;
+				continue;
 			}
 			if (prev_insn->op == IR_INSN_LSH &&
 			    cur_insn->op == IR_INSN_RSH) {
@@ -26,35 +27,47 @@ void bpf_ir_optimize_code_compaction(struct bpf_ir_env *env,
 					    IR_VALUE_CONSTANT) {
 					continue;
 				}
+				if (prev_insn->values[1].data.constant_d !=
+					    32 ||
+				    cur_insn->values[1].data.constant_d != 32) {
+					continue;
+				}
 				if (prev_insn->values[0].type ==
 					    IR_VALUE_INSN &&
 				    cur_insn->values[0].type == IR_VALUE_INSN &&
-				    prev_insn->values[0].data.insn_d ==
+				    prev_insn ==
 					    cur_insn->values[0].data.insn_d) {
 					// Same!
 					bpf_ir_array_push(env, &opt_insns,
 							  &cur_insn);
-					bpf_ir_array_push(env, &opt_insns,
-							  &prev_insn);
 				}
 			}
 			prev_insn = cur_insn;
 		}
 	}
+	PRINT_LOG_DEBUG(env, "Found %d compaction opt\n", opt_insns.num_elem);
 	if (opt_insns.num_elem == 0) {
 		return;
 	}
 	struct ir_insn **pos2;
 	array_for(pos2, opt_insns)
 	{
 		struct ir_insn *insn = *pos2;
-		if (insn->op == IR_INSN_RSH) {
-			// Insert optimized code
-			struct ir_insn *ni = bpf_ir_create_assign_insn(
-				env, insn, insn->values[0], INSERT_BACK);
-			ni->vr_type = IR_VR_TYPE_32;
-		}
-		bpf_ir_erase_insn(env, insn);
+		struct ir_insn *lsh_insn = insn->values[0].data.insn_d;
+		DBGASSERT(insn->op == IR_INSN_RSH);
+		DBGASSERT(lsh_insn->op == IR_INSN_LSH);
+		// Insert optimized code
+		bpf_ir_change_value(env, insn, &insn->values[0],
+				    lsh_insn->values[0]);
+		insn->op = IR_INSN_ASSIGN;
+		insn->vr_type = IR_VR_TYPE_32;
+		insn->alu_op = IR_ALU_32;
+		insn->value_num = 1;
+		insn->raw_pos.valid = false;
+		bpf_ir_erase_insn(env, lsh_insn);
 	}
 	bpf_ir_array_free(&opt_insns);
-}
+}
+
+const struct builtin_pass_cfg bpf_ir_kern_compaction_pass =
+	DEF_BUILTIN_PASS_ENABLE_CFG("optimize_compaction", NULL, NULL);
