feat[venom]: add load elimination pass (#4265)

add primitive `sload`/`tload`/`mload` elimination. keeps latest
(m/s/t)load or (m/s/t)store in a one-element "lattice", and weakens the
(m/s/t)load to a `store` instruction if there is a hit on the same key.

---------

Co-authored-by: Harry Kalogirou <[email protected]>
Co-authored-by: HodanPlodky <[email protected]>

Author: 3 people

tests/unit/compiler/venom/test_load_elimination.py

@@ -0,0 +1,129 @@
+from tests.venom_utils import assert_ctx_eq, parse_from_basic_block
+from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.passes.load_elimination import LoadElimination
+
+
+def _check_pre_post(pre, post):
+    ctx = parse_from_basic_block(pre)
+
+    for fn in ctx.functions.values():
+        ac = IRAnalysesCache(fn)
+        LoadElimination(ac, fn).run_pass()
+
+    assert_ctx_eq(ctx, parse_from_basic_block(post))
+
+
+def _check_no_change(pre):
+    _check_pre_post(pre, pre)
+
+
+def test_simple_load_elimination():
+    pre = """
+    main:
+        %ptr = 11
+        %1 = mload %ptr
+
+        %2 = mload %ptr
+
+        stop
+    """
+    post = """
+    main:
+        %ptr = 11
+        %1 = mload %ptr
+
+        %2 = %1
+
+        stop
+    """
+    _check_pre_post(pre, post)
+
+
+def test_equivalent_var_elimination():
+    """
+    Test that the lattice can "peer through" equivalent vars
+    """
+    pre = """
+    main:
+        %1 = 11
+        %2 = %1
+        %3 = mload %1
+
+        %4 = mload %2
+
+        stop
+    """
+    post = """
+    main:
+        %1 = 11
+        %2 = %1
+        %3 = mload %1
+
+        %4 = %3  # %2 == %1
+
+        stop
+    """
+    _check_pre_post(pre, post)
+
+
+def test_elimination_barrier():
+    """
+    Check for barrier between load/load
+    """
+    pre = """
+    main:
+        %1 = 11
+        %2 = mload %1
+        %3 = %100
+        # fence - writes to memory
+        staticcall %3, %3, %3, %3
+        %4 = mload %1
+    """
+    _check_no_change(pre)
+
+
+def test_store_load_elimination():
+    """
+    Check that lattice stores the result of mstores (even through
+    equivalent variables)
+    """
+    pre = """
+    main:
+        %val = 55
+        %ptr1 = 11
+        %ptr2 = %ptr1
+        mstore %ptr1, %val
+
+        %3 = mload %ptr2
+
+        stop
+    """
+    post = """
+        main:
+        %val = 55
+        %ptr1 = 11
+        %ptr2 = %ptr1
+        mstore %ptr1, %val
+
+        %3 = %val
+
+        stop
+    """
+    _check_pre_post(pre, post)
+
+
+def test_store_load_barrier():
+    """
+    Check for barrier between store/load
+    """
+    pre = """
+    main:
+        %ptr = 11
+        %val = 55
+        mstore %ptr, %val
+        %3 = %100  ; arbitrary
+        # fence
+        staticcall %3, %3, %3, %3
+        %4 = mload %ptr
+    """
+    _check_no_change(pre)

vyper/venom/__init__.py

@@ -15,6 +15,7 @@
     BranchOptimizationPass,
     DFTPass,
     FloatAllocas,
+    LoadElimination,
     LowerDloadPass,
     MakeSSA,
     Mem2Var,
@@ -59,9 +60,12 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     Mem2Var(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
     SCCP(ac, fn).run_pass()
+
+    LoadElimination(ac, fn).run_pass()
     StoreElimination(ac, fn).run_pass()
     MemMergePass(ac, fn).run_pass()
     SimplifyCFGPass(ac, fn).run_pass()
+
     LowerDloadPass(ac, fn).run_pass()
     AlgebraicOptimizationPass(ac, fn).run_pass()
     # NOTE: MakeSSA is after algebraic optimization it currently produces

vyper/venom/passes/__init__.py

@@ -3,6 +3,7 @@
 from .dft import DFTPass
 from .float_allocas import FloatAllocas
 from .literals_codesize import ReduceLiteralsCodesize
+from .load_elimination import LoadElimination
 from .lower_dload import LowerDloadPass
 from .make_ssa import MakeSSA
 from .mem2var import Mem2Var

vyper/venom/passes/load_elimination.py

@@ -0,0 +1,50 @@
+from vyper.venom.analysis import DFGAnalysis, LivenessAnalysis, VarEquivalenceAnalysis
+from vyper.venom.effects import Effects
+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, Effects.MEMORY, "mload", "mstore")
+            self._process_bb(bb, Effects.TRANSIENT, "tload", "tstore")
+            self._process_bb(bb, Effects.STORAGE, "sload", "sstore")
+
+        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, eff, load_opcode, store_opcode):
+        # not really a lattice even though it is not really inter-basic block;
+        # we may generalize in the future
+        lattice = ()
+
+        for inst in bb.instructions:
+            if eff in inst.get_write_effects():
+                lattice = ()
+
+            if inst.opcode == store_opcode:
+                # mstore [val, ptr]
+                val, ptr = inst.operands
+                lattice = (ptr, val)
+
+            if inst.opcode == load_opcode:
+                prev_lattice = lattice
+                (ptr,) = inst.operands
+                lattice = (ptr, inst.output)
+                if not prev_lattice:
+                    continue
+                if not self.equivalent(ptr, prev_lattice[0]):
+                    continue
+                inst.opcode = "store"
+                inst.operands = [prev_lattice[1]]

vyper/venom/passes/store_elimination.py

@@ -9,6 +9,9 @@ class StoreElimination(IRPass):
     and removes the `store` instruction.
     """
 
+    # TODO: consider renaming `store` instruction, since it is confusing
+    # with LoadElimination
+
     def run_pass(self):
         self.analyses_cache.request_analysis(CFGAnalysis)
         dfg = self.analyses_cache.request_analysis(DFGAnalysis)