fix[venom]: track multi-output phi origins by variable

tests/unit/compiler/venom/test_phi_elimination.py

@@ -1,6 +1,9 @@
 import pytest
 
 from tests.venom_utils import PrePostChecker
+from vyper.venom.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRLabel
+from vyper.venom.parser import parse_venom
 from vyper.venom.passes import PhiEliminationPass
 
 _check_pre_post = PrePostChecker([PhiEliminationPass], default_hevm=False)
@@ -319,3 +322,81 @@ def test_phi_elim_two_phi_merges():
     """
 
     _check_pre_post(pre, post, hevm=True)
+
+
+def test_phi_elim_multi_output_invoke_distinct_outputs_kept():
+    # a phi over two different outputs of the same invoke instruction
+    # must not be collapsed: the origins trace to the same producing
+    # instruction, but to different variables (GH 5039)
+    pre = """
+    function main {
+    main:
+        %cond = source
+        %r1, %r2 = invoke @f
+        jnz %cond, @then, @else
+    then:
+        jmp @join
+    else:
+        jmp @join
+    join:
+        %z = phi @then, %r1, @else, %r2
+        sink %z
+    }
+
+    function f {
+    f:
+        %retpc = param
+        ret %retpc
+    }
+    """
+
+    ctx = parse_venom(pre)
+    fn = ctx.get_function(IRLabel("main"))
+    ac = IRAnalysesCache(fn)
+    PhiEliminationPass(ac, fn).run_pass()
+
+    insts = [inst for bb in fn.get_basic_blocks() for inst in bb.instructions]
+    phis = [inst for inst in insts if inst.opcode == "phi"]
+    assert len(phis) == 1, "phi over distinct invoke outputs must be kept"
+    assert phis[0].output.name == "%z"
+
+
+def test_phi_elim_multi_output_invoke_same_output_eliminated():
+    # Same producer instruction and same output variable: this phi is trivial
+    # even though the root instruction has multiple outputs.
+    source = """
+    function main {
+    main:
+        %cond = source
+        %r1, %r2 = invoke @f
+        %r2_copy = %r2
+        jnz %cond, @then, @else
+    then:
+        jmp @join
+    else:
+        jmp @join
+    join:
+        %z = phi @then, %r2, @else, %r2_copy
+        sink %z
+    }
+
+    function f {
+    f:
+        %retpc = param
+        ret %retpc
+    }
+    """
+
+    ctx = parse_venom(source)
+    fn = ctx.get_function(IRLabel("main"))
+    ac = IRAnalysesCache(fn)
+    PhiEliminationPass(ac, fn).run_pass()
+
+    insts = [inst for bb in fn.get_basic_blocks() for inst in bb.instructions]
+    phis = [inst for inst in insts if inst.opcode == "phi"]
+    assert len(phis) == 0, "phi over the same invoke output should be eliminated"
+
+    z_assigns = [inst for inst in insts if [out.name for out in inst.get_outputs()] == ["%z"]]
+    assert len(z_assigns) == 1
+    assert z_assigns[0].opcode == "assign"
+    assert z_assigns[0].operands[0].name == "%r2"

vyper/venom/passes/phi_elimination.py

@@ -1,10 +1,12 @@
 from vyper.venom.analysis import DFGAnalysis, LivenessAnalysis
-from vyper.venom.basicblock import IRInstruction, IRVariable
+from vyper.venom.basicblock import IRInstruction, IROperand, IRVariable
 from vyper.venom.passes.base_pass import InstUpdater, IRPass
 
 
 class PhiEliminationPass(IRPass):
-    phi_to_origins: dict[IRInstruction, set[IRInstruction]]
+    # phi -> set of root operands. The operand identifies both the
+    # producing instruction and the produced output slot.
+    phi_to_origins: dict[IRInstruction, set[IROperand]]
 
     def run_pass(self):
         self.dfg = self.analyses_cache.request_analysis(DFGAnalysis)
@@ -27,18 +29,10 @@ def _process_phi(self, inst: IRInstruction):
 
         # len > 1: multiple origins, phi is doing real work, keep it.
         if len(srcs) == 1:
-            src = next(iter(srcs))
-            if src == inst:
+            (var,) = srcs
+            if var == inst.output:
                 return
-            # the origin may be a multi-output instruction (invoke with a
-            # hidden adopted-FMP output, bump); `src.output` asserts
-            # single-output, and collapsing the phi to outputs[0] would be
-            # wrong when the phi'd value is a later output. Latent for
-            # multi-output invokes; common once FmpLoweringPass threads
-            # bump/invoke outputs through loop phis.
-            if len(src.get_outputs()) != 1:
-                return
-            self.updater.mk_assign(inst, src.output)
+            self.updater.mk_assign(inst, var)
 
     def _calculate_phi_origins(self):
         self.phi_to_origins = dict()
@@ -52,13 +46,14 @@ def _calculate_phi_origins(self):
     def _get_phi_origins(self, inst: IRInstruction):
         assert inst.opcode == "phi"  # sanity
         visited: set[IRInstruction] = set()
-        self.phi_to_origins[inst] = self._get_phi_origins_r(inst, visited)
+        self.phi_to_origins[inst] = self._get_phi_origins_r(inst, inst.output, visited)
 
-    # traverse chains of phis and stores to get the "root" instructions
-    # for phis.
+    # traverse chains of phis and stores to get the "root" definitions
+    # for phis. `var` is the ssa variable through which `inst` was
+    # reached; it identifies which output of `inst` is the origin.
     def _get_phi_origins_r(
-        self, inst: IRInstruction, visited: set[IRInstruction]
-    ) -> set[IRInstruction]:
+        self, inst: IRInstruction, var: IROperand, visited: set[IRInstruction]
+    ) -> set[IROperand]:
         if inst.opcode == "phi":
             if inst in self.phi_to_origins:
                 return self.phi_to_origins[inst]
@@ -72,12 +67,12 @@ def _get_phi_origins_r(
 
             visited.add(inst)
 
-            res: set[IRInstruction] = set()
+            res: set[IROperand] = set()
 
-            for _, var in inst.phi_operands:
-                next_inst = self.dfg.get_producing_instruction(var)
-                assert next_inst is not None, (inst, var)
-                res |= self._get_phi_origins_r(next_inst, visited)
+            for _, op_var in inst.phi_operands:
+                next_inst = self.dfg.get_producing_instruction(op_var)
+                assert next_inst is not None, (inst, op_var)
+                res |= self._get_phi_origins_r(next_inst, op_var, visited)
 
             if len(res) > 1:
                 # multi-origin phi: treat the phi itself as a "barrier" origin.
@@ -92,15 +87,17 @@ def _get_phi_origins_r(
                 #   %c = phi %a, %b  ; barrier (two origins)
                 #   %d = %c
                 #   %f = phi %d, %c  ; both paths lead to %c, so %f = %c
-                return set([inst])
+                return set([inst.output])
             return res
 
         if inst.opcode == "assign" and isinstance(inst.operands[0], IRVariable):
             # traverse assignment chain
-            var = inst.operands[0]
-            next_inst = self.dfg.get_producing_instruction(var)
+            src_var = inst.operands[0]
+            next_inst = self.dfg.get_producing_instruction(src_var)
             assert next_inst is not None
-            return self._get_phi_origins_r(next_inst, visited)
+            return self._get_phi_origins_r(next_inst, src_var, visited)
 
-        # root of the phi/assignment chain
-        return set([inst])
+        # root of the phi/assignment chain. note that for multi-output
+        # instructions (e.g. invoke), different outputs are distinct
+        # origins because their output operands are distinct.
+        return set([var])