feat[venom]: store expansion pass

vyper/venom/__init__.py

@@ -12,12 +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.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()
@@ -53,6 +55,8 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     SimplifyCFGPass(ac, fn).run_pass()
     AlgebraicOptimizationPass(ac, fn).run_pass()
     BranchOptimizationPass(ac, fn).run_pass()
+    ExtractLiteralsPass(ac, fn).run_pass()
+    StoreExpansionPass(ac, fn).run_pass()
     RemoveUnusedVariablesPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()
 

vyper/venom/passes/extract_literals.py

@@ -0,0 +1,38 @@
+from vyper.utils import OrderedSet
+from vyper.venom.analysis.dfg import DFGAnalysis
+from vyper.venom.analysis.liveness import LivenessAnalysis
+from vyper.venom.basicblock import IRInstruction, IRLiteral
+from vyper.venom.passes.base_pass import IRPass
+
+
+class ExtractLiteralsPass(IRPass):
+    """
+    This pass extracts literals so that they can be reordered by the DFT pass
+    """
+
+    def run_pass(self):
+        for bb in self.function.get_basic_blocks():
+            self._process_bb(bb)
+
+        self.analyses_cache.invalidate_analysis(DFGAnalysis)
+        self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+
+    def _process_bb(self, bb):
+        i = 0
+        while i < len(bb.instructions):
+            inst = bb.instructions[i]
+            if inst.opcode == "store":
+                i += 1
+                continue
+
+            for j, op in enumerate(inst.operands):
+                # first operand to log is magic
+                if inst.opcode == "log" and j == 0:
+                    continue
+
+                if isinstance(op, IRLiteral):
+                    var = self.function.get_next_variable()
+                    to_insert = IRInstruction("store", [op], var)
+                    bb.insert_instruction(to_insert, index=i)
+                    inst.operands[j] = var
+            i += 1

vyper/venom/passes/store_expansion.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 IRInstruction, IRVariable
+from vyper.venom.passes.base_pass import IRPass
+
+
+class StoreExpansionPass(IRPass):
+    """
+    This pass expands variables to their uses though `store` instructions,
+    reducing pressure on the stack scheduler
+    """
+
+    def run_pass(self):
+        self.analyses_cache.request_analysis(CFGAnalysis)
+        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+
+        for bb in self.function.get_basic_blocks():
+            for idx, inst in enumerate(bb.instructions):
+                if inst.opcode != "store":
+                    continue
+
+                self._process_inst(dfg, inst, idx)
+
+        self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+        self.analyses_cache.invalidate_analysis(DFGAnalysis)
+
+    def _process_inst(self, dfg, inst, idx):
+        """
+        Process store instruction. If the variable is only used by a load instruction,
+        forward the variable to the load instruction.
+        """
+        var = inst.output
+        uses = dfg.get_uses(var)
+
+        insertion_idx = idx + 1
+
+        for use_inst in uses:
+            if use_inst.parent != inst.parent:
+                continue
+
+            for i, operand in enumerate(use_inst.operands):
+                if operand == var:
+                    new_var = self.function.get_next_variable()
+                    new_inst = IRInstruction("store", [var], new_var)
+                    inst.parent.insert_instruction(new_inst, insertion_idx)
+                    insertion_idx += 1
+                    use_inst.operands[i] = new_var