diff --git a/tests/unit/compiler/venom/test_equivalent_variables.py b/tests/unit/compiler/venom/test_equivalent_variables.py new file mode 100644 index 0000000000..67b3065b6d --- /dev/null +++ b/tests/unit/compiler/venom/test_equivalent_variables.py @@ -0,0 +1,88 @@ +from collections import defaultdict + +from tests.venom_utils import parse_from_basic_block +from vyper.venom.analysis import IRAnalysesCache, VarEquivalenceAnalysis +from vyper.venom.basicblock import IRLiteral + + +def _check_expected(code, expected): + ctx = parse_from_basic_block(code) + fn = next(iter(ctx.functions.values())) + ac = IRAnalysesCache(fn) + eq = ac.request_analysis(VarEquivalenceAnalysis) + + tmp = defaultdict(list) + for var, bag in eq._bags.items(): + if not isinstance(var, IRLiteral): + tmp[bag].append(var) + + ret = [] + for varset in tmp.values(): + ret.append(tuple(var.value for var in varset)) + + assert tuple(ret) == expected + + +def test_simple_equivalent_vars(): + code = """ + main: + %1 = 5 + %2 = %1 + """ + expected = (("%1", "%2"),) + _check_expected(code, expected) + + +def test_equivalent_vars2(): + code = """ + main: + # graph with multiple edges from root: %1 => %2 and %1 => %3 + %1 = 5 + %2 = %1 + %3 = %1 + """ + expected = (("%1", "%2", "%3"),) + _check_expected(code, expected) + + +def test_equivalent_vars3(): + code = """ + main: + # even weirder graph + %1 = 5 + %2 = %1 + %3 = %2 + %4 = %2 + %5 = %1 + %6 = 7 ; disjoint + """ + expected = (("%1", "%2", "%3", "%4", "%5"), ("%6",)) + _check_expected(code, expected) + + +def test_equivalent_vars4(): + code = """ + main: + # even weirder graph + %1 = 5 + %2 = %1 + %3 = 5 ; not disjoint, equality on 5 + %4 = %3 + """ + expected = (("%1", "%2", "%3", "%4"),) + _check_expected(code, expected) + + +def test_equivalent_vars5(): + """ + Test with non-literal roots + """ + code = """ + main: + %1 = param + %2 = %1 + %3 = param ; disjoint + %4 = %3 + """ + expected = (("%1", "%2"), ("%3", "%4")) + _check_expected(code, expected) diff --git a/vyper/venom/analysis/equivalent_vars.py b/vyper/venom/analysis/equivalent_vars.py index 895895651a..880b16abf1 100644 --- a/vyper/venom/analysis/equivalent_vars.py +++ b/vyper/venom/analysis/equivalent_vars.py @@ -1,40 +1,79 @@ -from vyper.venom.analysis import DFGAnalysis, IRAnalysis -from vyper.venom.basicblock import IRVariable +from typing import Optional + +from vyper.venom.analysis import IRAnalysis +from vyper.venom.basicblock import IRInstruction, IRLiteral, IROperand, 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. + Generate equivalence sets of variables. This is used in passes so that + they can "peer through" store chains """ def analyze(self): - dfg = self.analyses_cache.request_analysis(DFGAnalysis) + # map from variables to "equivalence set" of variables, denoted + # by "bag" (an int). + self._bags: dict[IROperand, int] = {} + + # dict from bags to literal values + self._literals: dict[int, IRLiteral] = {} + + # the root of the store chain + self._root_instructions: dict[int, IRInstruction] = {} - equivalence_set: dict[IRVariable, int] = {} + bag = 0 + for bb in self.function.get_basic_blocks(): + for inst in bb.instructions: + if inst.output is None: + continue + if inst.opcode != "store": + self._handle_nonstore(inst, bag) + else: + self._handle_store(inst, bag) + bag += 1 - for bag, (var, inst) in enumerate(dfg._dfg_outputs.items()): - if inst.opcode != "store": - continue + def _handle_nonstore(self, inst: IRInstruction, bag: int): + assert inst.output is not None # help mypy + if bag in self._bags: + bag = self._bags[inst.output] + else: + self._bags[inst.output] = bag + self._root_instructions[bag] = inst - source = inst.operands[0] + def _handle_store(self, inst: IRInstruction, bag: int): + var = inst.output + 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 + assert var is not None # help mypy + assert isinstance(source, (IRVariable, IRLiteral)) + assert var not in self._bags # invariant - self._equivalence_set = equivalence_set + if source in self._bags: + bag = self._bags[source] + self._bags[var] = bag + else: + self._bags[source] = bag + self._bags[var] = bag - def equivalent(self, var1, var2): - if var1 not in self._equivalence_set: + if isinstance(source, IRLiteral): + self._literals[bag] = source + + def equivalent(self, var1: IROperand, var2: IROperand): + if var1 == var2: + return True + if var1 not in self._bags: return False - if var2 not in self._equivalence_set: + if var2 not in self._bags: return False - return self._equivalence_set[var1] == self._equivalence_set[var2] + return self._bags[var1] == self._bags[var2] + + def get_literal(self, var: IROperand) -> Optional[IRLiteral]: + if isinstance(var, IRLiteral): + return var + if (bag := self._bags.get(var)) is None: + return None + return self._literals.get(bag) + + def get_root_instruction(self, var: IROperand): + bag = self._bags[var] + return self._root_instructions[bag] diff --git a/vyper/venom/passes/load_elimination.py b/vyper/venom/passes/load_elimination.py index 6701b588fe..271188e45c 100644 --- a/vyper/venom/passes/load_elimination.py +++ b/vyper/venom/passes/load_elimination.py @@ -22,7 +22,7 @@ def run_pass(self): self.analyses_cache.invalidate_analysis(DFGAnalysis) def equivalent(self, op1, op2): - return op1 == op2 or self.equivalence.equivalent(op1, op2) + return 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;