fix[venom]: fix duplicate allocas (#4321)

this commit fixes a bug in the ir_node_to_venom translator. previously, `ir_node_to_venom` tried to detect unique allocas based on heuristics. this commit removes the heuristics and fixes the issue in the frontend by passing through a unique ID for each variable in the metadata. this ID is also passed into the `alloca` and `palloca` instructions to aid with debugging. note that this results in improved code, presumably due to more allocas being able to be reified. this commit makes a minor change to the `sqrt()`, builtin, which is to use `z_var.as_ir_node()` instead of `z_var.pos`, since `.as_ir_node()` correctly tags with the alloca metadata. to be maximally conservative, we could branch, only using `z_var.as_ir_node()` if we are using the venom pipeline, but the change should be correct for the legacy pipeline as well anyways. --------- Co-authored-by: Harry Kalogirou <[email protected]>
vyperlang · Nov 26, 2024 · cda634d · cda634d
1 parent f249c93
commit cda634d
Show file tree

Hide file tree

Showing 9 changed files with 89 additions and 37 deletions.
diff --git a/vyper/builtins/functions.py b/vyper/builtins/functions.py
@@ -2167,10 +2167,9 @@ def build_IR(self, expr, args, kwargs, context):
                 variables_2=variables_2,
                 memory_allocator=context.memory_allocator,
             )
+            z_ir = new_ctx.vars["z"].as_ir_node()
             ret = IRnode.from_list(
-                ["seq", placeholder_copy, sqrt_ir, new_ctx.vars["z"].pos],  # load x variable
-                typ=DecimalT(),
-                location=MEMORY,
+                ["seq", placeholder_copy, sqrt_ir, z_ir], typ=DecimalT(), location=MEMORY
             )
             return b1.resolve(ret)
 

diff --git a/vyper/codegen/context.py b/vyper/codegen/context.py
@@ -15,13 +15,26 @@ class Constancy(enum.Enum):
     Constant = 1
 
 
+_alloca_id = 0
+
+
+def _generate_alloca_id():
+    # note: this gets reset between compiler runs by codegen.core.reset_names
+    global _alloca_id
+
+    _alloca_id += 1
+    return _alloca_id
+
+
 @dataclass(frozen=True)
 class Alloca:
     name: str
     offset: int
     typ: VyperType
     size: int
 
+    _id: int
+
     def __post_init__(self):
         assert self.typ.memory_bytes_required == self.size
 
@@ -233,7 +246,9 @@ def _new_variable(
                 pos = f"$palloca_{ofst}_{size}"
             else:
                 pos = f"$alloca_{ofst}_{size}"
-            alloca = Alloca(name=name, offset=ofst, typ=typ, size=size)
+
+            alloca_id = _generate_alloca_id()
+            alloca = Alloca(name=name, offset=ofst, typ=typ, size=size, _id=alloca_id)
 
         var = VariableRecord(
             name=name,

diff --git a/vyper/codegen/core.py b/vyper/codegen/core.py
@@ -1,3 +1,4 @@
+import vyper.codegen.context as ctx
 from vyper.codegen.ir_node import Encoding, IRnode
 from vyper.compiler.settings import _opt_codesize, _opt_gas, _opt_none
 from vyper.evm.address_space import (
@@ -855,6 +856,9 @@ def reset_names():
     global _label
     _label = 0
 
+    # could be refactored
+    ctx._alloca_id = 0
+
 
 # returns True if t is ABI encoded and is a type that needs any kind of
 # validation

diff --git a/vyper/venom/README.md b/vyper/venom/README.md
@@ -209,15 +209,16 @@ Assembly can be inspected with `-f asm`, whereas an opcode view of the final byt
   - Effectively translates to `JUMP`, and marks the call site as a valid return destination (for callee to jump back to) by `JUMPDEST`.
 - `alloca`
   - ```
-    out = alloca size, offset
+    out = alloca size, offset, id
     ```
   - Allocates memory of a given `size` at a given `offset` in memory.
+  - The `id` argument is there to help debugging translation into venom
   - The output is the offset value itself.
   - Because the SSA form does not allow changing values of registers, handling mutable variables can be tricky. The `alloca` instruction is meant to simplify that.
   
 - `palloca`
   - ```
-    out = palloca size, offset
+    out = palloca size, offset, id
     ```
   - Like the `alloca` instruction but only used for parameters of internal functions which are passed by memory.
 - `iload`

diff --git a/vyper/venom/__init__.py b/vyper/venom/__init__.py
@@ -14,6 +14,7 @@
     AlgebraicOptimizationPass,
     BranchOptimizationPass,
     DFTPass,
+    FloatAllocas,
     MakeSSA,
     Mem2Var,
     RemoveUnusedVariablesPass,
@@ -47,6 +48,8 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
 
     ac = IRAnalysesCache(fn)
 
+    FloatAllocas(ac, fn).run_pass()
+
     SimplifyCFGPass(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
     Mem2Var(ac, fn).run_pass()

diff --git a/vyper/venom/ir_node_to_venom.py b/vyper/venom/ir_node_to_venom.py
@@ -107,18 +107,16 @@
 NOOP_INSTRUCTIONS = frozenset(["pass", "cleanup_repeat", "var_list", "unique_symbol"])
 
 SymbolTable = dict[str, Optional[IROperand]]
-_global_symbols: SymbolTable = None  # type: ignore
+_alloca_table: SymbolTable = None  # type: ignore
 MAIN_ENTRY_LABEL_NAME = "__main_entry"
-_external_functions: dict[int, SymbolTable] = None  # type: ignore
 
 
 # convert IRnode directly to venom
 def ir_node_to_venom(ir: IRnode) -> IRContext:
     _ = ir.unique_symbols  # run unique symbols check
 
-    global _global_symbols, _external_functions
-    _global_symbols = {}
-    _external_functions = {}
+    global _alloca_table
+    _alloca_table = {}
 
     ctx = IRContext()
     fn = ctx.create_function(MAIN_ENTRY_LABEL_NAME)
@@ -233,7 +231,7 @@ def pop_source(*args, **kwargs):
 def _convert_ir_bb(fn, ir, symbols):
     assert isinstance(ir, IRnode), ir
     # TODO: refactor these to not be globals
-    global _break_target, _continue_target, _global_symbols, _external_functions
+    global _break_target, _continue_target, _alloca_table
 
     # keep a map from external functions to all possible entry points
 
@@ -269,8 +267,8 @@ def _convert_ir_bb(fn, ir, symbols):
             if is_internal or len(re.findall(r"external.*__init__\(.*_deploy", current_func)) > 0:
                 # Internal definition
                 var_list = ir.args[0].args[1]
+                assert var_list.value == "var_list"
                 does_return_data = IRnode.from_list(["return_buffer"]) in var_list.args
-                _global_symbols = {}
                 symbols = {}
                 new_fn = _handle_internal_func(fn, ir, does_return_data, symbols)
                 for ir_node in ir.args[1:]:
@@ -298,8 +296,6 @@ def _convert_ir_bb(fn, ir, symbols):
         cont_ret = _convert_ir_bb(fn, cond, symbols)
         cond_block = fn.get_basic_block()
 
-        saved_global_symbols = _global_symbols.copy()
-
         then_block = IRBasicBlock(ctx.get_next_label("then"), fn)
         else_block = IRBasicBlock(ctx.get_next_label("else"), fn)
 
@@ -314,7 +310,6 @@ def _convert_ir_bb(fn, ir, symbols):
 
         # convert "else"
         cond_symbols = symbols.copy()
-        _global_symbols = saved_global_symbols.copy()
         fn.append_basic_block(else_block)
         else_ret_val = None
         if len(ir.args) == 3:
@@ -343,8 +338,6 @@ def _convert_ir_bb(fn, ir, symbols):
         if not then_block_finish.is_terminated:
             then_block_finish.append_instruction("jmp", exit_bb.label)
 
-        _global_symbols = saved_global_symbols
-
         return if_ret
 
     elif ir.value == "with":
@@ -385,13 +378,6 @@ def _convert_ir_bb(fn, ir, symbols):
                 data = _convert_ir_bb(fn, c, symbols)
                 ctx.append_data("db", [data])  # type: ignore
     elif ir.value == "label":
-        function_id_pattern = r"external (\d+)"
-        function_name = ir.args[0].value
-        m = re.match(function_id_pattern, function_name)
-        if m is not None:
-            function_id = m.group(1)
-            _global_symbols = _external_functions.setdefault(function_id, {})
-
         label = IRLabel(ir.args[0].value, True)
         bb = fn.get_basic_block()
         if not bb.is_terminated:
@@ -463,13 +449,11 @@ def _convert_ir_bb(fn, ir, symbols):
     elif ir.value == "repeat":
 
         def emit_body_blocks():
-            global _break_target, _continue_target, _global_symbols
+            global _break_target, _continue_target
             old_targets = _break_target, _continue_target
             _break_target, _continue_target = exit_block, incr_block
-            saved_global_symbols = _global_symbols.copy()
             _convert_ir_bb(fn, body, symbols.copy())
             _break_target, _continue_target = old_targets
-            _global_symbols = saved_global_symbols
 
         sym = ir.args[0]
         start, end, _ = _convert_ir_bb_list(fn, ir.args[1:4], symbols)
@@ -540,16 +524,25 @@ def emit_body_blocks():
     elif isinstance(ir.value, str) and ir.value.upper() in get_opcodes():
         _convert_ir_opcode(fn, ir, symbols)
     elif isinstance(ir.value, str):
-        if ir.value.startswith("$alloca") and ir.value not in _global_symbols:
+        if ir.value.startswith("$alloca"):
             alloca = ir.passthrough_metadata["alloca"]
-            ptr = fn.get_basic_block().append_instruction("alloca", alloca.offset, alloca.size)
-            _global_symbols[ir.value] = ptr
-        elif ir.value.startswith("$palloca") and ir.value not in _global_symbols:
+            if alloca._id not in _alloca_table:
+                ptr = fn.get_basic_block().append_instruction(
+                    "alloca", alloca.offset, alloca.size, alloca._id
+                )
+                _alloca_table[alloca._id] = ptr
+            return _alloca_table[alloca._id]
+
+        elif ir.value.startswith("$palloca"):
             alloca = ir.passthrough_metadata["alloca"]
-            ptr = fn.get_basic_block().append_instruction("palloca", alloca.offset, alloca.size)
-            _global_symbols[ir.value] = ptr
-
-        return _global_symbols.get(ir.value) or symbols.get(ir.value)
+            if alloca._id not in _alloca_table:
+                ptr = fn.get_basic_block().append_instruction(
+                    "palloca", alloca.offset, alloca.size, alloca._id
+                )
+                _alloca_table[alloca._id] = ptr
+            return _alloca_table[alloca._id]
+
+        return symbols.get(ir.value)
     elif ir.is_literal:
         return IRLiteral(ir.value)
     else:

diff --git a/vyper/venom/passes/__init__.py b/vyper/venom/passes/__init__.py
@@ -1,6 +1,7 @@
 from .algebraic_optimization import AlgebraicOptimizationPass
 from .branch_optimization import BranchOptimizationPass
 from .dft import DFTPass
+from .float_allocas import FloatAllocas
 from .make_ssa import MakeSSA
 from .mem2var import Mem2Var
 from .normalization import NormalizationPass

diff --git a/vyper/venom/passes/float_allocas.py b/vyper/venom/passes/float_allocas.py
@@ -0,0 +1,36 @@
+from vyper.venom.passes.base_pass import IRPass
+
+
+class FloatAllocas(IRPass):
+    """
+    This pass moves allocas to the entry basic block of a function
+    We could probably move them to the immediate dominator of the basic
+    block defining the alloca instead of the entry (which dominates all
+    basic blocks), but this is done for expedience.
+    Without this step, sccp fails, possibly because dominators are not
+    guaranteed to be traversed first.
+    """
+
+    def run_pass(self):
+        entry_bb = self.function.entry
+        assert entry_bb.is_terminated
+        tmp = entry_bb.instructions.pop()
+
+        for bb in self.function.get_basic_blocks():
+            if bb is entry_bb:
+                continue
+
+            # Extract alloca instructions
+            non_alloca_instructions = []
+            for inst in bb.instructions:
+                if inst.opcode in ("alloca", "palloca"):
+                    # note: order of allocas impacts bytecode.
+                    # TODO: investigate.
+                    entry_bb.insert_instruction(inst)
+                else:
+                    non_alloca_instructions.append(inst)
+
+            # Replace original instructions with filtered list
+            bb.instructions = non_alloca_instructions
+
+        entry_bb.instructions.append(tmp)
diff --git a/vyper/venom/passes/sccp/sccp.py b/vyper/venom/passes/sccp/sccp.py
@@ -252,7 +252,7 @@ def finalize(ret):
             if eval_result is LatticeEnum.BOTTOM:
                 return finalize(LatticeEnum.BOTTOM)
 
-            assert isinstance(eval_result, IROperand)
+            assert isinstance(eval_result, IROperand), (inst.parent.label, op, inst)
             ops.append(eval_result)
 
         # If we haven't found BOTTOM yet, evaluate the operation