feat[venom]: improve inter-basic block stack ordering

vyper/venom/analysis/__init__.py

@@ -7,4 +7,5 @@
 from .mem_alias import MemoryAliasAnalysis
 from .mem_ssa import MemSSA
 from .reachable import ReachableAnalysis
+from .stack_order import StackOrderAnalysis
 from .var_definition import VarDefinition

vyper/venom/analysis/analysis.py

@@ -75,7 +75,7 @@ def invalidate_analysis(self, analysis_cls: Type[IRAnalysis]):
         if analysis is not None:
             analysis.invalidate()
 
-    def force_analysis(self, analysis_cls: Type[IRAnalysis], *args, **kwargs):
+    def force_analysis(self, analysis_cls: Type[T], *args, **kwargs) -> T:
         """
         Force a specific analysis to be run on the IR even if it has already been run,
         and is cached.

vyper/venom/analysis/stack_order.py

@@ -0,0 +1,162 @@
+from vyper.venom.analysis import CFGAnalysis, LivenessAnalysis
+from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IROperand, IRVariable
+from vyper.venom.function import IRFunction
+
+# needed [top, ... , bottom]
+Needed = list[IRVariable]
+
+Stack = list[IROperand]
+
+
+def _swap(stack: Stack, op: IROperand):
+    top = len(stack) - 1
+    index = None
+    for i, item in reversed(list(enumerate(stack))):
+        if item == op:
+            index = i
+
+    assert index is not None
+
+    stack[index], stack[top] = stack[top], stack[index]
+
+
+def _swap_to(stack: Stack, depth: int):
+    top = len(stack) - 1
+    index = top - depth
+
+    stack[index], stack[top] = stack[top], stack[index]
+
+
+def _max_same_prefix(stack_a: Needed, stack_b: Needed):
+    res: Needed = []
+    for a, b in zip(stack_a, stack_b):
+        if a != b:
+            break
+        res.append(a)
+    return res
+
+
+class StackOrderAnalysis:
+    function: IRFunction
+    liveness: LivenessAnalysis
+    cfg: CFGAnalysis
+    _from_to: dict[tuple[IRBasicBlock, IRBasicBlock], Needed]
+
+    def __init__(self, ac: IRAnalysesCache):
+        self._from_to = dict()
+        self.ac = ac
+        self.liveness = ac.request_analysis(LivenessAnalysis)
+        self.cfg = ac.request_analysis(CFGAnalysis)
+
+    def analyze_bb(self, bb: IRBasicBlock) -> Needed:
+        self.needed: Needed = []
+        self.stack: Stack = []
+
+        for inst in bb.instructions:
+            if inst.opcode == "assign":
+                self._handle_assign(inst)
+            elif inst.opcode == "phi":
+                self._handle_inst(inst)
+            elif inst.is_bb_terminator:
+                self._handle_terminator(inst)
+            else:
+                self._handle_inst(inst)
+
+            if len(inst.operands) > 0:
+                if not inst.is_bb_terminator:
+                    assert self.stack[-len(inst.operands) :] == inst.operands, (
+                        inst,
+                        self.stack,
+                        inst.operands,
+                    )
+                self.stack = self.stack[: -len(inst.operands)]
+            if inst.output is not None:
+                self.stack.append(inst.output)
+
+        for pred in self.cfg.cfg_in(bb):
+            self._from_to[(pred, bb)] = self.needed.copy()
+
+        return self.needed
+
+    def get_stack(self, bb: IRBasicBlock) -> Needed:
+        succs = self.cfg.cfg_out(bb)
+        for succ in succs:
+            self.analyze_bb(succ)
+        orders = [self._from_to.get((bb, succ), []) for succ in succs]
+        return self._merge(orders)
+
+    def from_to(self, origin: IRBasicBlock, successor: IRBasicBlock) -> Needed:
+        target = self._from_to.get((origin, successor), []).copy()
+
+        for var in self.liveness.input_vars_from(origin, successor):
+            if var not in target:
+                target.append(var)
+
+        return target
+
+    def _handle_assign(self, inst: IRInstruction):
+        assert inst.opcode == "assign"
+        assert inst.output is not None
+
+        index = inst.parent.instructions.index(inst)
+        next_inst = inst.parent.instructions[index + 1]
+        next_live = self.liveness.live_vars_at(next_inst)
+
+        src = inst.operands[0]
+
+        if not isinstance(src, IRVariable):
+            self.stack.append(src)
+        elif src in next_live:
+            self.stack.append(src)
+            assert src in self.stack
+            self._add_needed(src)
+        else:
+            if src not in self.stack:
+                self.stack.append(src)
+                self._add_needed(src)
+            else:
+                _swap(self.stack, src)
+
+    def _add_needed(self, op: IRVariable):
+        if op not in self.needed:
+            self.needed.append(op)
+
+    def _reorder(self, target_stack: Stack):
+        count = len(target_stack)
+
+        for index, op in enumerate(target_stack):
+            depth = count - index - 1
+            _swap(self.stack, op)
+            _swap_to(self.stack, depth)
+
+        if len(target_stack) != 0:
+            assert target_stack == self.stack[-len(target_stack) :], (target_stack, self.stack)
+
+    def _handle_inst(self, inst: IRInstruction):
+        ops = inst.operands
+        for op in ops:
+            if isinstance(op, IRVariable) and op not in self.stack:
+                self._add_needed(op)
+            if op not in self.stack:
+                self.stack.append(op)
+        self._reorder(ops)
+
+    def _merge(self, orders: list[Needed]) -> Needed:
+        if len(orders) == 0:
+            return []
+        res = orders[0]
+        for order in orders:
+            res = _max_same_prefix(res, order)
+        return res
+
+    def _handle_terminator(self, inst: IRInstruction):
+        bb = inst.parent
+        orders = [self._from_to.get((bb, succ), []) for succ in self.cfg.cfg_out(bb)]
+        ops = (op for op in inst.operands if isinstance(op, IRVariable))
+        for op in ops:
+            if op not in self.stack:
+                self._add_needed(op)
+        for op in self._merge(orders):
+            if op not in self.stack:
+                self._add_needed(op)

vyper/venom/passes/dft.py

@@ -1,9 +1,10 @@
-from collections import defaultdict
+from collections import defaultdict, deque
 
 import vyper.venom.effects as effects
 from vyper.utils import OrderedSet
-from vyper.venom.analysis import DFGAnalysis, LivenessAnalysis
-from vyper.venom.basicblock import IRBasicBlock, IRInstruction
+from vyper.venom.analysis import CFGAnalysis, DFGAnalysis, LivenessAnalysis
+from vyper.venom.analysis.stack_order import StackOrderAnalysis
+from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRVariable
 from vyper.venom.function import IRFunction
 from vyper.venom.passes.base_pass import IRPass
 
@@ -17,15 +18,34 @@ class DFTPass(IRPass):
     # "effect dependency analysis"
     eda: dict[IRInstruction, OrderedSet[IRInstruction]]
 
+    stack_order: StackOrderAnalysis
+    cfg: CFGAnalysis
+
     def run_pass(self) -> None:
         self.data_offspring = {}
         self.visited_instructions: OrderedSet[IRInstruction] = OrderedSet()
 
         self.dfg = self.analyses_cache.force_analysis(DFGAnalysis)
+        self.cfg = self.analyses_cache.request_analysis(CFGAnalysis)
+        self.stack_order = StackOrderAnalysis(self.analyses_cache)
+
+        worklist = deque(self.cfg.dfs_post_walk)
 
-        for bb in self.function.get_basic_blocks():
+        last_order: dict[IRBasicBlock, list[IRVariable]] = dict()
+
+        while len(worklist) > 0:
+            bb = worklist.popleft()
+            self.stack_order.analyze_bb(bb)
+            order = self.stack_order.get_stack(bb)
+            if bb in last_order and last_order[bb] == order:
+                break
+            last_order[bb] = order
+            self.order = list(reversed(order))
             self._process_basic_block(bb)
 
+            for pred in self.cfg.cfg_in(bb):
+                worklist.append(pred)
+
         self.analyses_cache.invalidate_analysis(LivenessAnalysis)
 
     def _process_basic_block(self, bb: IRBasicBlock) -> None:
@@ -48,7 +68,6 @@ def _process_basic_block(self, bb: IRBasicBlock) -> None:
         self.visited_instructions = OrderedSet()
         for inst in entry_instructions_list:
             self._process_instruction_r(self.instructions, inst)
-
         bb.instructions = self.instructions
         assert bb.is_terminated, f"Basic block should be terminated {bb}"
 
@@ -63,12 +82,23 @@ def _process_instruction_r(self, instructions: list[IRInstruction], inst: IRInst
         children = list(self.dda[inst] | self.eda[inst])
 
         def cost(x: IRInstruction) -> int | float:
-            if x in self.eda[inst] or inst.flippable:
+            # intuition:
+            #   effect-only dependencies which have data dependencies
+            #   effect-only dependencies which have no data dependencies
+            #   indirect data dependencies (offspring of operands)
+            #   direct data dependencies (order of operands)
+
+            if (x not in self.dda[inst] and x in self.eda[inst]) or inst.flippable:
                 ret = -1 * int(len(self.data_offspring[x]) > 0)
+            elif x.output in inst.operands:
+                assert x in self.dda[inst]  # sanity check
+                assert x.output is not None  # help mypy
+                ret = inst.operands.index(x.output) + len(self.order)
             else:
                 assert x in self.dda[inst]  # sanity check
+                assert x.output in self.order
                 assert x.output is not None  # help mypy
-                ret = inst.operands.index(x.output)
+                ret = self.order.index(x.output)
             return ret
 
         # heuristic: sort by size of child dependency graph
@@ -97,6 +127,11 @@ def _calculate_dependency_graphs(self, bb: IRBasicBlock) -> None:
         all_read_effects: dict[effects.Effects, list[IRInstruction]] = defaultdict(list)
 
         for inst in non_phis:
+            if inst.is_bb_terminator:
+                for var in self.order:
+                    dep = self.dfg.get_producing_instruction(var)
+                    if dep is not None and dep.parent == bb:
+                        self.dda[inst].add(dep)
             for op in inst.operands:
                 dep = self.dfg.get_producing_instruction(op)
                 if dep is not None and dep.parent == bb:
@@ -111,7 +146,7 @@ def _calculate_dependency_graphs(self, bb: IRBasicBlock) -> None:
                     for read_inst in all_read_effects[write_effect]:
                         self.eda[inst].add(read_inst)
                 # prevent reordering write-after-write for the same effect
-                if write_effect in last_write_effects:
+                if (write_effect & ~effects.Effects.MSIZE) in last_write_effects:
                     self.eda[inst].add(last_write_effects[write_effect])
                 last_write_effects[write_effect] = inst
                 # clear previous read effects after a write

vyper/venom/passes/function_inliner.py

@@ -206,6 +206,7 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction):
 
         call_site_bb.instructions = call_site_bb.instructions[:call_idx]
         call_site_bb.append_instruction("jmp", func_copy.entry.label)
+
         self._fix_phi(call_site_bb, call_site_return)
 
     def _fix_phi(self, orig: IRBasicBlock, new: IRBasicBlock) -> None:

vyper/venom/passes/single_use_expansion.py

@@ -35,14 +35,15 @@ def _process_bb(self, bb):
                 i += 1
                 continue
 
-            for j, op in enumerate(inst.operands):
+            ops = inst.operands.copy()
+
+            for j, op in enumerate(ops):
                 # first operand to log is magic
                 if inst.opcode == "log" and j == 0:
                     continue
 
                 if isinstance(op, IRVariable):
                     uses = self.dfg.get_uses(op)
-                    # it's already only used once
                     if len(uses) == 1 and len([x for x in inst.operands if x == op]) == 1:
                         continue
 
@@ -54,7 +55,8 @@ def _process_bb(self, bb):
                 var = self.function.get_next_variable()
                 to_insert = IRInstruction("assign", [op], var)
                 bb.insert_instruction(to_insert, index=i)
-                inst.operands[j] = var
+                if len(inst.operands) > j:
+                    inst.operands[j] = var
                 i += 1
 
             i += 1

vyper/venom/venom_to_assembly.py

@@ -206,7 +206,9 @@ def _stack_reorder(
         if len(stack_ops) == 0:
             return 0
 
-        assert len(stack_ops) == len(set(stack_ops))  # precondition
+        assert len(stack_ops) == len(
+            set(stack_ops)
+        ), f"duplicated stack {stack_ops}"  # precondition
 
         cost = 0
         for i, op in enumerate(stack_ops):
@@ -271,7 +273,7 @@ def _emit_input_operands(
                 self.dup_op(assembly, stack, op)
 
             # guaranteed by store expansion
-            assert op not in seen, (op, seen)
+            assert op not in seen, (inst, op, seen)
             seen.add(op)
 
     def _prepare_stack_for_function(self, asm, fn: IRFunction, stack: StackModel):
@@ -462,8 +464,6 @@ def _generate_evm_for_instruction(
             assert len(self.cfg.cfg_in(next_bb)) > 1
 
             target_stack = self.liveness.input_vars_from(inst.parent, next_bb)
-            # NOTE: in general the stack can contain multiple copies of
-            # the same variable, however, before a jump that is not possible
             self._stack_reorder(assembly, stack, list(target_stack))
 
         if inst.is_commutative:
@@ -604,6 +604,11 @@ def _optimistic_swap(self, assembly, inst, next_liveness, stack):
         if DEBUG_SHOW_COST:
             stack0 = stack.copy()
 
+        next_index = inst.parent.instructions.index(inst)
+        next_inst = inst.parent.instructions[next_index + 1]
+
+        if next_inst.is_bb_terminator:
+            return
         # if there are no live vars at the next point, nothing to schedule
         if len(next_liveness) == 0:
             return