feat[venom]: add load elimination

tests/unit/compiler/venom/test_duplicate_operands.py

@@ -1,6 +1,8 @@
 from vyper.compiler.settings import OptimizationLevel
 from vyper.venom import generate_assembly_experimental
+from vyper.venom.analysis.analysis import IRAnalysesCache
 from vyper.venom.context import IRContext
+from vyper.venom.passes.store_expansion import StoreExpansionPass
 
 
 def test_duplicate_operands():
@@ -13,7 +15,7 @@ def test_duplicate_operands():
     %3 = mul %1, %2
     stop
 
-    Should compile to: [PUSH1, 10, DUP1, DUP1, DUP1, ADD, MUL, POP, STOP]
+    Should compile to: [PUSH1, 10, DUP1, DUP2, ADD, MUL, POP, STOP]
     """
     ctx = IRContext()
     fn = ctx.create_function("test")
@@ -23,5 +25,9 @@ def test_duplicate_operands():
     bb.append_instruction("mul", sum_, op)
     bb.append_instruction("stop")
 
-    asm = generate_assembly_experimental(ctx, optimize=OptimizationLevel.GAS)
-    assert asm == ["PUSH1", 10, "DUP1", "DUP1", "ADD", "MUL", "POP", "STOP"]
+    ac = IRAnalysesCache(fn)
+    StoreExpansionPass(ac, fn).run_pass()
+
+    optimize = OptimizationLevel.GAS
+    asm = generate_assembly_experimental(ctx, optimize=optimize)
+    assert asm == ["PUSH1", 10, "DUP1", "DUP2", "ADD", "MUL", "POP", "STOP"]

tests/unit/compiler/venom/test_load_elimination.py

@@ -0,0 +1,134 @@
+from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRLiteral, IRVariable
+from vyper.venom.context import IRContext
+from vyper.venom.passes.load_elimination import LoadElimination
+
+
+def test_simple_load_elimination():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+
+    bb = fn.get_basic_block()
+
+    ptr = IRLiteral(11)
+    bb.append_instruction("mload", ptr)
+    bb.append_instruction("mload", ptr)
+    bb.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+    LoadElimination(ac, fn).run_pass()
+
+    assert len([inst for inst in bb.instructions if inst.opcode == "mload"]) == 1
+
+    inst0, inst1, inst2 = bb.instructions
+
+    assert inst0.opcode == "mload"
+    assert inst1.opcode == "store"
+    assert inst1.operands[0] == inst0.output
+    assert inst2.opcode == "stop"
+
+
+def test_equivalent_var_elimination():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+
+    bb = fn.get_basic_block()
+
+    ptr1 = bb.append_instruction("store", IRLiteral(11))
+    ptr2 = bb.append_instruction("store", ptr1)
+    bb.append_instruction("mload", ptr1)
+    bb.append_instruction("mload", ptr2)
+    bb.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+    LoadElimination(ac, fn).run_pass()
+
+    assert len([inst for inst in bb.instructions if inst.opcode == "mload"]) == 1
+
+    inst0, inst1, inst2, inst3, inst4 = bb.instructions
+
+    assert inst0.opcode == "store"
+    assert inst1.opcode == "store"
+    assert inst2.opcode == "mload"
+    assert inst2.operands[0] == inst0.output
+    assert inst3.opcode == "store"
+    assert inst3.operands[0] == inst2.output
+    assert inst4.opcode == "stop"
+
+
+def test_elimination_barrier():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+
+    bb = fn.get_basic_block()
+
+    ptr = IRLiteral(11)
+    bb.append_instruction("mload", ptr)
+
+    arbitrary = IRVariable("%100")
+    # fence, writes to memory
+    bb.append_instruction("staticcall", arbitrary, arbitrary, arbitrary, arbitrary)
+
+    bb.append_instruction("mload", ptr)
+    bb.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+
+    instructions = bb.instructions.copy()
+    LoadElimination(ac, fn).run_pass()
+
+    assert instructions == bb.instructions  # no change
+
+
+def test_store_load_elimination():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+
+    bb = fn.get_basic_block()
+
+    val = IRLiteral(55)
+    ptr1 = bb.append_instruction("store", IRLiteral(11))
+    ptr2 = bb.append_instruction("store", ptr1)
+    bb.append_instruction("mstore", val, ptr1)
+    bb.append_instruction("mload", ptr2)
+    bb.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+    LoadElimination(ac, fn).run_pass()
+
+    assert len([inst for inst in bb.instructions if inst.opcode == "mload"]) == 0
+
+    inst0, inst1, inst2, inst3, inst4 = bb.instructions
+
+    assert inst0.opcode == "store"
+    assert inst1.opcode == "store"
+    assert inst2.opcode == "mstore"
+    assert inst3.opcode == "store"
+    assert inst3.operands[0] == inst2.operands[0]
+    assert inst4.opcode == "stop"
+
+
+def test_store_load_barrier():
+    ctx = IRContext()
+    fn = ctx.create_function("test")
+
+    bb = fn.get_basic_block()
+
+    val = IRLiteral(55)
+    ptr1 = bb.append_instruction("store", IRLiteral(11))
+    ptr2 = bb.append_instruction("store", ptr1)
+    bb.append_instruction("mstore", val, ptr1)
+
+    arbitrary = IRVariable("%100")
+    # fence, writes to memory
+    bb.append_instruction("staticcall", arbitrary, arbitrary, arbitrary, arbitrary)
+
+    bb.append_instruction("mload", ptr2)
+    bb.append_instruction("stop")
+
+    ac = IRAnalysesCache(fn)
+
+    instructions = bb.instructions.copy()
+    LoadElimination(ac, fn).run_pass()
+
+    assert instructions == bb.instructions

tests/unit/compiler/venom/test_stack_cleanup.py

@@ -9,7 +9,8 @@ def test_cleanup_stack():
     bb = fn.get_basic_block()
     ret_val = bb.append_instruction("param")
     op = bb.append_instruction("store", 10)
-    bb.append_instruction("add", op, op)
+    op2 = bb.append_instruction("store", op)
+    bb.append_instruction("add", op, op2)
     bb.append_instruction("ret", ret_val)
 
     asm = generate_assembly_experimental(ctx, optimize=OptimizationLevel.GAS)

vyper/venom/__init__.py

@@ -12,13 +12,14 @@
 from vyper.venom.passes.algebraic_optimization import AlgebraicOptimizationPass
 from vyper.venom.passes.branch_optimization import BranchOptimizationPass
 from vyper.venom.passes.dft import DFTPass
-from vyper.venom.passes.extract_literals import ExtractLiteralsPass
+from vyper.venom.passes.load_elimination import LoadElimination
 from vyper.venom.passes.make_ssa import MakeSSA
 from vyper.venom.passes.mem2var import Mem2Var
 from vyper.venom.passes.remove_unused_variables import RemoveUnusedVariablesPass
 from vyper.venom.passes.sccp import SCCP
 from vyper.venom.passes.simplify_cfg import SimplifyCFGPass
 from vyper.venom.passes.store_elimination import StoreElimination
+from vyper.venom.passes.store_expansion import StoreExpansionPass
 from vyper.venom.venom_to_assembly import VenomCompiler
 
 DEFAULT_OPT_LEVEL = OptimizationLevel.default()
@@ -52,10 +53,12 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     SCCP(ac, fn).run_pass()
     StoreElimination(ac, fn).run_pass()
     SimplifyCFGPass(ac, fn).run_pass()
+    LoadElimination(ac, fn).run_pass()
     AlgebraicOptimizationPass(ac, fn).run_pass()
     BranchOptimizationPass(ac, fn).run_pass()
-    ExtractLiteralsPass(ac, fn).run_pass()
     RemoveUnusedVariablesPass(ac, fn).run_pass()
+
+    StoreExpansionPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()
 
 

vyper/venom/analysis/equivalent_vars.py

@@ -0,0 +1,41 @@
+from vyper.venom.analysis.analysis import IRAnalysis
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.basicblock import IRVariable
+
+
+class VarEquivalenceAnalysis(IRAnalysis):
+    """
+    Generate equivalence sets of variables. This is used to avoid swapping
+    variables which are the same during venom_to_assembly. Theoretically,
+    the DFTPass should order variable declarations optimally, but, it is
+    not aware of the "pickaxe" heuristic in venom_to_assembly, so they can
+    interfere.
+    """
+
+    def analyze(self):
+        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+
+        equivalence_set: dict[IRVariable, int] = {}
+
+        for bag, (var, inst) in enumerate(dfg._dfg_outputs.items()):
+            if inst.opcode != "store":
+                continue
+
+            source = inst.operands[0]
+
+            assert var not in equivalence_set  # invariant
+            if source in equivalence_set:
+                equivalence_set[var] = equivalence_set[source]
+                continue
+            else:
+                equivalence_set[var] = bag
+                equivalence_set[source] = bag
+
+        self._equivalence_set = equivalence_set
+
+    def equivalent(self, var1, var2):
+        if var1 not in self._equivalence_set:
+            return False
+        if var2 not in self._equivalence_set:
+            return False
+        return self._equivalence_set[var1] == self._equivalence_set[var2]

vyper/venom/basicblock.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Iterator, Optional, Union
 
+import vyper.venom.effects as effects
 from vyper.codegen.ir_node import IRnode
 from vyper.utils import OrderedSet
 
@@ -240,6 +241,12 @@ def is_volatile(self) -> bool:
     def is_bb_terminator(self) -> bool:
         return self.opcode in BB_TERMINATORS
 
+    def get_read_effects(self):
+        return effects.reads.get(self.opcode, ())
+
+    def get_write_effects(self):
+        return effects.writes.get(self.opcode, ())
+
     def get_label_operands(self) -> Iterator[IRLabel]:
         """
         Get all labels in instruction.

vyper/venom/effects.py

@@ -0,0 +1,48 @@
+_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",
+}
+
+
+def _mktuple(x):
+    if not isinstance(x, tuple):
+        x = (x,)
+    return x
+
+
+writes = {k: _mktuple(v) for k, v in _writes.items()}
+reads = {k: _mktuple(v) for k, v in _reads.items()}

vyper/venom/passes/load_elimination.py

@@ -0,0 +1,74 @@
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.analysis.equivalent_vars import VarEquivalenceAnalysis
+from vyper.venom.analysis.liveness import LivenessAnalysis
+from vyper.venom.passes.base_pass import IRPass
+
+
+class LoadElimination(IRPass):
+    """
+    Eliminate sloads, mloads and tloads
+    """
+    # should this be renamed to EffectsElimination?
+
+    def run_pass(self):
+        self.equivalence = self.analyses_cache.request_analysis(VarEquivalenceAnalysis)
+
+        for bb in self.function.get_basic_blocks():
+            self._process_bb(bb)
+
+        self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+        self.analyses_cache.invalidate_analysis(DFGAnalysis)
+
+    def equivalent(self, op1, op2):
+        return op1 == op2 or self.equivalence.equivalent(op1, op2)
+
+    def _process_bb(self, bb):
+        transient = ()
+        storage = ()
+        memory = ()
+
+        for inst in bb.instructions:
+            if "memory" in inst.get_write_effects():
+                memory = ()
+            if "storage" in inst.get_write_effects():
+                storage = ()
+            if "transient" in inst.get_write_effects():
+                transient = ()
+
+            if inst.opcode == "mstore":
+                # mstore [val, ptr]
+                memory = (inst.operands[1], inst.operands[0])
+            if inst.opcode == "sstore":
+                storage = (inst.operands[1], inst.operands[0])
+            if inst.opcode == "tstore":
+                transient = (inst.operands[1], inst.operands[0])
+
+            if inst.opcode == "mload":
+                prev_memory = memory
+                memory = (inst.operands[0], inst.output)
+                if not prev_memory:
+                    continue
+                if not self.equivalent(inst.operands[0], prev_memory[0]):
+                    continue
+                inst.opcode = "store"
+                inst.operands = [prev_memory[1]]
+
+            if inst.opcode == "sload":
+                prev_storage = storage
+                storage = (inst.operands[0], inst.output)
+                if not prev_storage:
+                    continue
+                if not self.equivalent(inst.operands[0], prev_storage[0]):
+                    continue
+                inst.opcode = "store"
+                inst.operands = [prev_storage[1]]
+
+            if inst.opcode == "tload":
+                prev_transient = transient
+                transient = (inst.operands[0], inst.output)
+                if not prev_transient:
+                    continue
+                if not self.equivalent(inst.operands[0], prev_transient[0]):
+                    continue
+                inst.opcode = "store"
+                inst.operands = [prev_transient[1]]