feat[venom]: optimize branching

tests/unit/compiler/venom/test_branch_optimizer.py

@@ -0,0 +1,38 @@
+from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRBasicBlock, IRLabel
+from vyper.venom.context import IRContext
+from vyper.venom.passes.branch_optimization import BranchOptimizationPass
+from vyper.venom.passes.make_ssa import MakeSSA
+
+
+def test_simple_jump_case():
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+
+    bb = fn.get_basic_block()
+
+    br1 = IRBasicBlock(IRLabel("then"), fn)
+    fn.append_basic_block(br1)
+    br2 = IRBasicBlock(IRLabel("else"), fn)
+    fn.append_basic_block(br2)
+
+    p1 = bb.append_instruction("param")
+    op1 = bb.append_instruction("store", p1)
+    op2 = bb.append_instruction("store", 64)
+    op3 = bb.append_instruction("add", op1, op2)
+    jnz_input = bb.append_instruction("iszero", op3)
+    bb.append_instruction("jnz", jnz_input, br1.label, br2.label)
+
+    br1.append_instruction("add", op3, 10)
+    br1.append_instruction("stop")
+    br2.append_instruction("add", op3, p1)
+    br2.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+    MakeSSA(ac, fn).run_pass()
+    BranchOptimizationPass(ac, fn).run_pass()
+
+    assert bb.instructions[-1].opcode == "jnz"
+    assert bb.instructions[-1].operands[0] == op3
+    assert bb.instructions[-1].operands[1] == br2.label
+    assert bb.instructions[-1].operands[2] == br1.label

vyper/venom/__init__.py

@@ -9,6 +9,7 @@
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
 from vyper.venom.ir_node_to_venom import ir_node_to_venom
+from vyper.venom.passes.branch_optimization import BranchOptimizationPass
 from vyper.venom.passes.dft import DFTPass
 from vyper.venom.passes.make_ssa import MakeSSA
 from vyper.venom.passes.mem2var import Mem2Var
@@ -49,6 +50,7 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     SCCP(ac, fn).run_pass()
     StoreElimination(ac, fn).run_pass()
     SimplifyCFGPass(ac, fn).run_pass()
+    BranchOptimizationPass(ac, fn).run_pass()
     RemoveUnusedVariablesPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()
 

vyper/venom/analysis/dfg.py

@@ -22,6 +22,10 @@ def get_uses(self, op: IRVariable) -> list[IRInstruction]:
     def get_producing_instruction(self, op: IRVariable) -> Optional[IRInstruction]:
         return self._dfg_outputs.get(op)
 
+    def add_use(self, op: IRVariable, inst: IRInstruction):
+        uses = self._dfg_inputs.setdefault(op, [])
+        uses.append(inst)
+
     def remove_use(self, op: IRVariable, inst: IRInstruction):
         uses = self._dfg_inputs.get(op, [])
         uses.remove(inst)

vyper/venom/passes/branch_optimization.py

@@ -0,0 +1,29 @@
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.passes.base_pass import IRPass
+
+
+class BranchOptimizationPass(IRPass):
+    """
+    This pass optimizes branches inverting jnz instructions where appropriate
+    """
+
+    def _optimize_branches(self) -> None:
+        fn = self.function
+        for bb in fn.get_basic_blocks():
+            term_inst = bb.instructions[-1]
+            if term_inst.opcode != "jnz":
+                continue
+
+            prev_inst = self.dfg.get_producing_instruction(term_inst.operands[0])
+            if prev_inst.opcode == "iszero":
+                new_cond = prev_inst.operands[0]
+                term_inst.operands = [new_cond, term_inst.operands[2], term_inst.operands[1]]
+
+                # Since the DFG update is simple we do in place to avoid invalidating the DFG
+                self.dfg.remove_use(prev_inst.output, term_inst)
+                self.dfg.add_use(new_cond, term_inst)
+
+    def run_pass(self):
+        self.dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+
+        self._optimize_branches()