feat[venom]: add store elimination pass (#4021)

Venom uses the `store` instruction to load a variable or literal to
another variable. This commit adds a new pass `StoreElimination` that
forwards the rhs for `store`s to their uses, eliminating the `store`
instruction in the process.

---------

Co-authored-by: Charles Cooper <[email protected]>

vyper/venom/__init__.py

@@ -16,6 +16,7 @@
 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.venom_to_assembly import VenomCompiler
 
 DEFAULT_OPT_LEVEL = OptimizationLevel.default()
@@ -45,8 +46,9 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     SimplifyCFGPass(ac, fn).run_pass()
     Mem2Var(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
+    StoreElimination(ac, fn).run_pass()
+    MakeSSA(ac, fn).run_pass()
     SCCP(ac, fn).run_pass()
-
     SimplifyCFGPass(ac, fn).run_pass()
     RemoveUnusedVariablesPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()

vyper/venom/passes/sccp/sccp.py

@@ -168,7 +168,9 @@ def _visit_phi(self, inst: IRInstruction):
                 continue
             in_vars.append(self._lookup_from_lattice(var))
         value = reduce(_meet, in_vars, LatticeEnum.TOP)  # type: ignore
-        assert inst.output in self.lattice, "Got undefined var for phi"
+
+        if inst.output not in self.lattice:
+            return
 
         if value != self._lookup_from_lattice(inst.output):
             self._set_lattice(inst.output, value)
@@ -327,26 +329,6 @@ def _replace_constants(self, inst: IRInstruction):
                 if isinstance(lat, IRLiteral):
                     inst.operands[i] = lat
 
-    def _propagate_variables(self):
-        """
-        Copy elimination. #NOTE: Not working yet, but it's also not needed atm.
-        """
-        for bb in self.dom.dfs_walk:
-            for inst in bb.instructions:
-                if inst.opcode == "store":
-                    uses = self._get_uses(inst.output)
-                    remove_inst = True
-                    for usage_inst in uses:
-                        if usage_inst.opcode == "phi":
-                            remove_inst = False
-                            continue
-                        for i, op in enumerate(usage_inst.operands):
-                            if op == inst.output:
-                                usage_inst.operands[i] = inst.operands[0]
-                    if remove_inst:
-                        inst.opcode = "nop"
-                        inst.operands = []
-
 
 def _meet(x: LatticeItem, y: LatticeItem) -> LatticeItem:
     if x == LatticeEnum.TOP:

vyper/venom/passes/store_elimination.py

@@ -0,0 +1,49 @@
+from vyper.venom.analysis.cfg import CFGAnalysis
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.analysis.dominators import DominatorTreeAnalysis
+from vyper.venom.analysis.liveness import LivenessAnalysis
+from vyper.venom.basicblock import IRVariable
+from vyper.venom.passes.base_pass import IRPass
+
+
+class StoreElimination(IRPass):
+    """
+    This pass forwards variables to their uses though `store` instructions,
+    and removes the `store` instruction.
+    """
+
+    def run_pass(self):
+        self.analyses_cache.request_analysis(CFGAnalysis)
+        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+
+        for var, inst in dfg.outputs.items():
+            if inst.opcode != "store":
+                continue
+            if not isinstance(inst.operands[0], IRVariable):
+                continue
+            if inst.operands[0].name in ["%ret_ofst", "%ret_size"]:
+                continue
+            if inst.output.name in ["%ret_ofst", "%ret_size"]:
+                continue
+            self._process_store(dfg, inst, var, inst.operands[0])
+
+        self.analyses_cache.invalidate_analysis(DominatorTreeAnalysis)
+        self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+        self.analyses_cache.invalidate_analysis(DFGAnalysis)
+
+    def _process_store(self, dfg, inst, var, new_var):
+        """
+        Process store instruction. If the variable is only used by a load instruction,
+        forward the variable to the load instruction.
+        """
+        uses = dfg.get_uses(var)
+
+        if any([inst.opcode == "phi" for inst in uses]):
+            return
+
+        for use_inst in uses:
+            for i, operand in enumerate(use_inst.operands):
+                if operand == var:
+                    use_inst.operands[i] = new_var
+
+        inst.parent.remove_instruction(inst)