feat[venom]: store expansion pass (#4068)

expand `extract_literals` pass (introduced in bb9129a) to
also extract variables and rename it to `store_expansion`, allowing for
the dft pass to reorder variable uses more effectively.

since this also gives us the guarantee that each variable is used
exactly once (besides by store instructions), this allows us to make
some simplifications in `venom_to_assembly.py`, since we no longer need
to account for the same variable occurring on the stack more than one
time (cf. for example 5d8280f).

this results in a stack scheduler improvement. for example:
- examples/tokens/ERC20.vy has a 20 byte codesize improvement
- examples/auctions/blind_auction.vy has a 145 byte codesize improvement
- examples/voting/ballot.vy has a 28 byte codesize improvement

across a range of contracts, the improvement seems to be between 1-2%,
but sometimes as high as 5%

since stack operands are now guaranteed to be unique, the old rule to
avoid swapping if two stack operands are the same no longer works. to
address this, this commit adds an equivalence analysis. this creates
equivalence sets of equivalent variables based on store chains, and
then generalizes the rule from "don't swap if two stack operands are
the same" to "don't swap if two stack operands are equivalent".

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_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)

tests/unit/compiler/venom/test_stack_reorder.py

@@ -1,5 +1,7 @@
 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_stack_reorder():
@@ -25,4 +27,7 @@ def test_stack_reorder():
 
     bb.append_instruction("ret", ret_val)
 
+    ac = IRAnalysesCache(fn)
+    StoreExpansionPass(ac, fn).run_pass()
+
     generate_assembly_experimental(ctx)

vyper/venom/__init__.py

@@ -12,13 +12,13 @@
 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()
@@ -54,8 +54,9 @@ 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()
     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/passes/extract_literals.py → vyper/venom/passes/store_expansion.py

@@ -1,12 +1,13 @@
 from vyper.venom.analysis.dfg import DFGAnalysis
 from vyper.venom.analysis.liveness import LivenessAnalysis
-from vyper.venom.basicblock import IRInstruction, IRLiteral
+from vyper.venom.basicblock import IRInstruction, IRLiteral, IRVariable
 from vyper.venom.passes.base_pass import IRPass
 
 
-class ExtractLiteralsPass(IRPass):
+class StoreExpansionPass(IRPass):
     """
-    This pass extracts literals so that they can be reordered by the DFT pass
+    This pass extracts literals and variables so that they can be
+    reordered by the DFT pass
     """
 
     def run_pass(self):
@@ -20,7 +21,7 @@ def _process_bb(self, bb):
         i = 0
         while i < len(bb.instructions):
             inst = bb.instructions[i]
-            if inst.opcode in ("store", "offset"):
+            if inst.opcode in ("store", "offset", "phi", "param"):
                 i += 1
                 continue
 
@@ -29,9 +30,11 @@ def _process_bb(self, bb):
                 if inst.opcode == "log" and j == 0:
                     continue
 
-                if isinstance(op, IRLiteral):
+                if isinstance(op, (IRVariable, 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
+
             i += 1

vyper/venom/stack_model.py

@@ -30,7 +30,7 @@ def push(self, op: IROperand) -> None:
     def pop(self, num: int = 1) -> None:
         del self._stack[len(self._stack) - num :]
 
-    def get_depth(self, op: IROperand, n: int = 1) -> int:
+    def get_depth(self, op: IROperand) -> int:
         """
         Returns the depth of the n-th matching operand in the stack map.
         If the operand is not in the stack map, returns NOT_IN_STACK.
@@ -39,10 +39,7 @@ def get_depth(self, op: IROperand, n: int = 1) -> int:
 
         for i, stack_op in enumerate(reversed(self._stack)):
             if stack_op.value == op.value:
-                if n <= 1:
-                    return -i
-                else:
-                    n -= 1
+                return -i
 
         return StackModel.NOT_IN_STACK  # type: ignore