feat[venom]: common subexpression elimination pass

tests/unit/compiler/venom/test_common_subexpression_elimination.py

@@ -0,0 +1,23 @@
+from vyper.venom.context import IRContext
+from vyper.venom.passes.common_subexpression_elimination import CSE
+from vyper.venom.passes.extract_literals import ExtractLiteralsPass
+from vyper.venom.analysis.analysis import IRAnalysesCache
+
+def test_common_subexpression_elimination():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+    bb = fn.get_basic_block()
+    op = bb.append_instruction("store", 10)
+    sum_1 = bb.append_instruction("add", op, 10)
+    bb.append_instruction("mul", sum_1, 10)
+    sum_2 = bb.append_instruction("add", op, 10)
+    bb.append_instruction("mul", sum_2, 10)
+    bb.append_instruction("stop")
+
+
+    ac = IRAnalysesCache(fn)
+    CSE(ac, fn).run_pass()
+    ExtractLiteralsPass(ac, fn).run_pass()
+
+    assert sum(1 for inst in bb.instructions if inst.opcode == "add") == 1, "wrong number of adds"
+    assert sum(1 for inst in bb.instructions if inst.opcode == "mul") == 1, "wrong number of muls"

vyper/venom/__init__.py

@@ -20,6 +20,7 @@
 from vyper.venom.passes.simplify_cfg import SimplifyCFGPass
 from vyper.venom.passes.store_elimination import StoreElimination
 from vyper.venom.venom_to_assembly import VenomCompiler
+from vyper.venom.passes.common_subexpression_elimination import CSE
 
 DEFAULT_OPT_LEVEL = OptimizationLevel.default()
 
@@ -54,6 +55,7 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     SimplifyCFGPass(ac, fn).run_pass()
     AlgebraicOptimizationPass(ac, fn).run_pass()
     BranchOptimizationPass(ac, fn).run_pass()
+    CSE(ac, fn).run_pass()
     ExtractLiteralsPass(ac, fn).run_pass()
     RemoveUnusedVariablesPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()

vyper/venom/analysis/available_expression.py

@@ -0,0 +1,196 @@
+from vyper.venom.analysis.analysis import IRAnalysis
+from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.analysis.cfg import CFGAnalysis
+from vyper.venom.context import IRFunction
+from vyper.venom.basicblock import IRBasicBlock
+from vyper.venom.basicblock import IRInstruction
+from vyper.venom.basicblock import IRLiteral
+from vyper.venom.basicblock import IRVariable
+from vyper.venom.basicblock import IROperand
+from vyper.venom.basicblock import BB_TERMINATORS
+from vyper.utils import OrderedSet
+from dataclasses import dataclass
+from collections import deque
+
+@dataclass
+class _Expression:
+    first_inst : IRInstruction
+    opcode: str
+    operands : list[IROperand]
+
+    def __eq__(self, other):
+        if not isinstance(other, _Expression):
+            return False
+        return self.first_inst == other.first_inst
+
+    def __hash__(self) -> int:
+        res : int = hash(self.opcode)
+        for op in self.operands:
+            res ^= hash(op)
+        return res
+
+    def __repr__(self) -> str:
+        if self.opcode == "store":
+            assert len(self.operands) == 1, "wrong store"
+            return repr(self.operands[0])
+        res = self.opcode + " [ "
+        for op in self.operands:
+            res += repr(op) + " "
+        res += "]"
+        return res
+
+    def contains_expr(self, expr : "_Expression") -> bool:
+        for op in self.operands:
+            if op == expr:
+                return True
+            if isinstance(op, _Expression) and op.contains_expr(expr):
+                return True
+        return False
+
+class _BBLattice:
+    data : dict[IRInstruction, OrderedSet[_Expression]]
+    out : OrderedSet[_Expression]
+    in_cache: OrderedSet[_Expression] | None
+
+    def __init__(self, bb : IRBasicBlock):
+        self.data = dict()
+        self.out = OrderedSet()
+        self.in_cache = None
+        for inst in bb.instructions:
+            self.data[inst] = OrderedSet()
+
+_UNINTRESTING_OPCODES = [
+    "store",
+    "param",
+    "offset",
+    "phi",
+    "nop",
+]
+
+_ALL = ("storage", "transient", "memory", "immutables", "balance", "returndata")
+
+writes = {
+    "sstore": ("storage"),
+    "tstore": ("transient"),
+    "mstore": ("memory"),
+    "istore": ("immutables"),
+    "call": _ALL,
+    "delegatecall": _ALL,
+    "staticcall": ("memory"),
+    "create": _ALL,
+    "create2": _ALL,
+    "invoke": _ALL,  # could be smarter, look up the effects of the invoked function
+    "dloadbytes": ("memory"),
+    "returndatacopy": ("memory"),
+    "calldatacopy": ("memory"),
+    "codecopy": ("memory"),
+    "extcodecopy": ("memory"),
+    "mcopy": ("memory"),
+}
+reads = {
+    "sload": ("storage"),
+    "tload": ("transient"),
+    "iload": ("immutables"),
+    "mload": ("memory"),
+    "mcopy": ("memory"),
+    "call": _ALL,
+    "delegatecall": _ALL,
+    "staticcall": _ALL,
+    "returndatasize": ("returndata"),
+    "returndatacopy": ("returndata"),
+    "balance": ("balance"),
+    "selfbalance": ("balance"),
+    "log": ("memory"),
+    "revert": ("memory"),
+    "return": ("memory"),
+    "sha3": ("memory"),
+}
+
+
+class _FunctionLattice:
+    data : dict[IRBasicBlock, _BBLattice]
+
+    def __init__(self, function : IRFunction):
+        self.data = dict()
+        for bb in function.get_basic_blocks():
+            self.data[bb] = _BBLattice(bb)
+
+class AvailableExpressionAnalysis(IRAnalysis):
+    expressions : OrderedSet[_Expression] = OrderedSet()
+    inst_to_expr: dict[IRInstruction, _Expression] = dict()
+    dfg: DFGAnalysis
+    lattice : _FunctionLattice
+
+    def __init__(self, analyses_cache: IRAnalysesCache, function: IRFunction):
+        super().__init__(analyses_cache, function)
+        self.analyses_cache.request_analysis(CFGAnalysis)
+        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+        assert isinstance(dfg, DFGAnalysis)
+        self.dfg = dfg
+
+        self.lattice = _FunctionLattice(function)
+
+    def analyze(self, *args, **kwargs):
+        worklist = deque()
+        worklist.append(self.function.entry)
+        while len(worklist) > 0:
+            bb : IRBasicBlock = worklist.popleft()
+            changed = self._handle_bb(bb)
+
+            if changed:
+                for out in bb.cfg_out:
+                    if out not in worklist:
+                        worklist.append(out)
+
+    def _handle_bb(self, bb : IRBasicBlock) -> bool:
+        available_expr : OrderedSet[_Expression] = OrderedSet()
+        if len(bb.cfg_in) > 0:
+            available_expr = OrderedSet.intersection(*(self.lattice.data[in_bb].out for in_bb in bb.cfg_in))
+
+        bb_lat = self.lattice.data[bb]
+        if bb_lat.in_cache is not None and available_expr == bb_lat.in_cache:
+            return False
+        bb_lat.in_cache = available_expr
+        change = False
+        for inst in bb.instructions:
+            if (inst.opcode in _UNINTRESTING_OPCODES 
+                or inst.opcode in BB_TERMINATORS):
+                continue
+            if available_expr != bb_lat.data[inst]:
+                bb_lat.data[inst] = available_expr.copy()
+                change |= True
+
+
+            inst_expr = self.get_expression(inst, available_expr)
+            write_effects = writes.get(inst_expr.opcode, ())
+            for expr in available_expr.copy():
+                if expr.contains_expr(inst_expr):
+                    available_expr.remove(expr)
+                read_effects = reads.get(expr.opcode, ())
+                if any(eff in write_effects for eff in read_effects):
+                    available_expr.remove(expr)
+
+            if "call" not in inst.opcode and inst.opcode != "invoke":
+                available_expr.add(inst_expr)
+
+        if available_expr != bb_lat.out:
+            bb_lat.out = available_expr.copy()
+            change |= True
+
+        return change
+
+    def get_expression(self, inst: IRInstruction, available_exprs : OrderedSet[_Expression] | None = None) -> _Expression:
+        if available_exprs is None:
+            available_exprs = self.lattice.data[inst.parent].data[inst]
+        operands: list[IROperand] = inst.operands.copy()
+        expr = _Expression(inst, inst.opcode, operands)
+        for e in available_exprs:
+            if e.opcode == expr.opcode and e.operands == expr.operands:
+                return e
+
+        return expr
+
+    def get_available(self, inst : IRInstruction) -> OrderedSet[_Expression]:
+        return self.lattice.data[inst.parent].data[inst]
+

vyper/venom/passes/common_subexpression_elimination.py

@@ -0,0 +1,65 @@
+from vyper.venom.passes.base_pass import IRPass
+from vyper.venom.analysis.available_expression import AvailableExpressionAnalysis
+from vyper.venom.analysis.liveness import LivenessAnalysis
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.basicblock import IRInstruction
+from vyper.venom.basicblock import IRVariable
+from vyper.venom.basicblock import IRBasicBlock
+from vyper.utils import OrderedSet
+
+class CSE(IRPass):
+    available_expression_analysis : AvailableExpressionAnalysis
+
+    def run_pass(self, *args, **kwargs):
+        available_expression_analysis = self.analyses_cache.request_analysis(AvailableExpressionAnalysis)
+        assert isinstance(available_expression_analysis, AvailableExpressionAnalysis)
+        self.available_expression_analysis = available_expression_analysis
+
+        while True:
+            replace_dict = self._find_replaceble()
+            if len(replace_dict) == 0:
+                return
+            self._replace(replace_dict)
+            self.analyses_cache.invalidate_analysis(DFGAnalysis)
+            self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+
+    # return instruction and to which instruction it could
+    # replaced by
+    def _find_replaceble(self) -> dict[IRInstruction, IRInstruction]:
+        res : dict[IRInstruction, IRInstruction] = dict()
+        for bb in self.function.get_basic_blocks():
+            for inst in bb.instructions:
+                inst_expr = self.available_expression_analysis.get_expression(inst)
+                avail = self.available_expression_analysis.get_available(inst) 
+                if inst_expr in avail:
+                    res[inst] = inst_expr.first_inst
+
+        return res
+
+    def _replace(self, replace_dict : dict[IRInstruction, IRInstruction]):
+        for (orig, to) in replace_dict.items():
+            while to in replace_dict.keys():
+                to = replace_dict[to]
+            self._replace_inst(orig, to)
+
+    def _replace_inst(self, orig_inst : IRInstruction, to_inst : IRInstruction):
+        visited: OrderedSet[IRBasicBlock] = OrderedSet()
+        if orig_inst.output is not None:
+            assert isinstance(orig_inst.output, IRVariable), f"not var {orig_inst}"
+            assert isinstance(to_inst.output, IRVariable), f"not var {to_inst}"
+            self._replace_inst_r(orig_inst.parent, orig_inst.output, to_inst.output, visited)
+        orig_inst.parent.remove_instruction(orig_inst)
+
+    def _replace_inst_r(self, bb : IRBasicBlock, orig : IRVariable, to : IRVariable, visited : OrderedSet[IRBasicBlock]):
+        if bb in visited:
+            return
+        visited.add(bb)
+
+        for inst in bb.instructions:
+            for i in range(len(inst.operands)):
+                op = inst.operands[i]
+                if op == orig:
+                    inst.operands[i] = to
+
+        for out in bb.cfg_out:
+            self._replace_inst_r(out, orig, to, visited)