feat[venom]: function inliner

vyper/codegen/context.py

@@ -35,6 +35,9 @@ class Alloca:
 
     _id: int
 
+    # special metadata for calloca. hint for venom to tie calloca to call site.
+    _callsite: Optional[str] = None
+
     def __post_init__(self):
         assert self.typ.memory_bytes_required == self.size
 

vyper/codegen/function_definitions/common.py

@@ -2,7 +2,7 @@
 from functools import cached_property
 from typing import Optional
 
-from vyper.codegen.context import Constancy, Context
+from vyper.codegen.context import Constancy, Context, VariableRecord
 from vyper.codegen.ir_node import IRnode
 from vyper.codegen.memory_allocator import MemoryAllocator
 from vyper.evm.opcodes import version_check
@@ -16,7 +16,7 @@
 class FrameInfo:
     frame_start: int
     frame_size: int
-    frame_vars: dict[str, tuple[int, VyperType]]
+    frame_vars: dict[str, VariableRecord]
 
     @property
     def mem_used(self):

vyper/codegen/self_call.py

@@ -1,5 +1,9 @@
+import copy
+import dataclasses
+
 from vyper.codegen.core import _freshname, eval_once_check, make_setter
 from vyper.codegen.ir_node import IRnode
+from vyper.codegen.memory_allocator import MemoryAllocator
 from vyper.evm.address_space import MEMORY
 from vyper.exceptions import StateAccessViolation
 from vyper.semantics.types.subscriptable import TupleT
@@ -66,7 +70,27 @@
 
     # note: dst_tuple_t != args_tuple_t
     dst_tuple_t = TupleT(tuple(func_t.argument_types))
-    args_dst = IRnode(func_t._ir_info.frame_info.frame_start, typ=dst_tuple_t, location=MEMORY)
+    if context.settings.experimental_codegen:
+        arg_items = ["multi"]
+        frame_info = func_t._ir_info.frame_info
+
+        for var in frame_info.frame_vars.values():
+            var = copy.copy(var)
+            alloca = var.alloca
+            assert alloca is not None
+            assert isinstance(var.pos, str)  # help mypy
+            if not var.pos.startswith("$palloca"):
+                continue
+            newname = var.pos.replace("$palloca", "$calloca")
+            var.pos = newname
+            alloca = dataclasses.replace(alloca, _callsite=return_label)
+            irnode = var.as_ir_node()
+            irnode.passthrough_metadata["alloca"] = alloca
+            arg_items.append(irnode)
+        args_dst = IRnode.from_list(arg_items, typ=dst_tuple_t)
+    else:
+        # legacy
+        args_dst = IRnode(func_t._ir_info.frame_info.frame_start, typ=dst_tuple_t, location=MEMORY)
 
     # if one of the arguments is a self call, the argument
     # buffer could get borked. to prevent against that,

vyper/venom/README.md

@@ -205,6 +205,11 @@ Assembly can be inspected with `-f asm`, whereas an opcode view of the final byt
     out = palloca size, offset, id
     ```
   - Like the `alloca` instruction but only used for parameters of internal functions which are passed by memory.
+- `calloca`
+  - ```
+    out = calloca size, offset, id, <callsite label>
+    ```
+  - Similar to the `calloca` instruction but only used for parameters of internal functions which are passed by memory. Used at the call-site of a call.
 - `iload`
   - ```
     out = iload offset

vyper/venom/__init__.py

@@ -14,6 +14,7 @@
     AlgebraicOptimizationPass,
     BranchOptimizationPass,
     DFTPass,
+    FunctionInlinerPass,
     FloatAllocas,
     MakeSSA,
     Mem2Var,
@@ -46,15 +47,22 @@ def generate_assembly_experimental(
 def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     # Run passes on Venom IR
     # TODO: Add support for optimization levels
+    ac = IRAnalysesCache(fn, optimize)
 
-    ac = IRAnalysesCache(fn)
+    FunctionInlinerPass(ac, fn).run_pass()
 
     FloatAllocas(ac, fn).run_pass()
 
     SimplifyCFGPass(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
+
     Mem2Var(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
+
+    # function inliner can insert bad variables, remove them before sccp
+
+    RemoveUnusedVariablesPass(ac, fn).run_pass()
+
     SCCP(ac, fn).run_pass()
     StoreElimination(ac, fn).run_pass()
     MemMergePass(ac, fn).run_pass()
@@ -77,11 +85,11 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
 def run_passes_on(ctx: IRContext, optimize: OptimizationLevel):
     for fn in ctx.functions.values():
         _run_passes(fn, optimize)
+    ctx.prune_unreachable_functions()
 
 
 def generate_ir(ir: IRnode, optimize: OptimizationLevel) -> IRContext:
     # Convert "old" IR to "new" IR
     ctx = ir_node_to_venom(ir)
     run_passes_on(ctx, optimize)
-
     return ctx

vyper/venom/analysis/analysis.py

@@ -1,5 +1,6 @@
 from typing import Type
 
+from vyper.compiler.settings import OptimizationLevel
 from vyper.venom.function import IRFunction
 
 
@@ -36,11 +37,16 @@ class IRAnalysesCache:
 
     function: IRFunction
     analyses_cache: dict[Type[IRAnalysis], IRAnalysis]
+    optimize: OptimizationLevel
 
-    def __init__(self, function: IRFunction):
+    def __init__(self, function: IRFunction, optimize=None):
         self.analyses_cache = {}
         self.function = function
 
+        if optimize is None:
+            optimize = OptimizationLevel.default()
+        self.optimize = optimize
+
     def request_analysis(self, analysis_cls: Type[IRAnalysis], *args, **kwargs):
         """
         Request a specific analysis to be run on the IR. The result is cached and

vyper/venom/basicblock.py

@@ -241,6 +241,13 @@ def __init__(
         self.ast_source = None
         self.error_msg = None
 
+    def copy(self):
+        cls = self.__class__
+        ret = cls.__new__(cls)
+        ret.__dict__ = self.__dict__.copy()
+        ret.operands = ret.operands.copy()
+        return ret
+
     @property
     def is_volatile(self) -> bool:
         return self.opcode in VOLATILE_INSTRUCTIONS

vyper/venom/context.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from vyper.utils import OrderedSet
 from vyper.venom.basicblock import IRInstruction, IRLabel, IROperand
 from vyper.venom.function import IRFunction
 
@@ -39,6 +40,23 @@ def get_next_label(self, suffix: str = "") -> IRLabel:
         self.last_label += 1
         return IRLabel(f"{self.last_label}{suffix}")
 
+    def prune_unreachable_functions(self):
+        entry = next(iter(self.functions.values()))
+        to_visit = OrderedSet([entry])
+        seen = OrderedSet()
+        while to_visit:
+            fn = to_visit.pop()
+            seen.add(fn)
+            for bb in fn.get_basic_blocks():
+                for inst in bb.instructions:
+                    if inst.opcode == "invoke":
+                        label = inst.operands[0]
+                        next_fn = self.get_function(label)
+                        if next_fn not in seen:
+                            to_visit.add(next_fn)
+
+        self.functions = {label: fn for label, fn in self.functions.items() if fn in seen}
+
     def chain_basic_blocks(self) -> None:
         """
         Chain basic blocks together. This is necessary for the IR to be valid, and is done after
@@ -47,6 +65,10 @@ def chain_basic_blocks(self) -> None:
         for fn in self.functions.values():
             fn.chain_basic_blocks()
 
+    def float_allocas(self) -> None:
+        for fn in self.functions.values():
+            fn.float_allocas()
+
     def append_data(self, opcode: str, args: list[IROperand]) -> None:
         """
         Append data

vyper/venom/function.py

@@ -179,6 +179,25 @@ def chain_basic_blocks(self) -> None:
             else:
                 bb.append_instruction("stop")
 
+    def float_allocas(self):
+        entry_bb = self.entry
+        assert entry_bb.is_terminated
+        tmp = entry_bb.instructions.pop()
+
+        for bb in self.get_basic_blocks():
+            if bb is entry_bb:
+                continue
+
+            # "fast" way to strip allocas from each basic block
+            def is_alloca(inst):
+                return inst.opcode in ("alloca", "palloca")
+
+            bb.instructions.sort(key=is_alloca)
+            while len(bb.instructions) > 0 and is_alloca(bb.instructions[-1]):
+                entry_bb.insert_instruction(bb.instructions.pop())
+
+        entry_bb.instructions.append(tmp)
+
     def copy(self):
         new = IRFunction(self.name)
         new._basic_block_dict = self._basic_block_dict.copy()

vyper/venom/ir_node_to_venom.py

@@ -125,6 +125,13 @@ def ir_node_to_venom(ir: IRnode) -> IRContext:
 
     ctx.chain_basic_blocks()
 
+    # float allocas to the front of the 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.
+    ctx.float_allocas()
+
     return ctx
 
 
@@ -177,9 +184,14 @@ def _handle_self_call(fn: IRFunction, ir: IRnode, symbols: SymbolTable) -> Optio
 def _handle_internal_func(
     fn: IRFunction, ir: IRnode, does_return_data: bool, symbols: SymbolTable
 ) -> IRFunction:
+    global _alloca_table
+
     fn = fn.ctx.create_function(ir.args[0].args[0].value)
     bb = fn.get_basic_block()
 
+    _saved_alloca_table = _alloca_table
+    _alloca_table = {}
+
     # return buffer
     if does_return_data:
         symbols["return_buffer"] = bb.append_instruction("param")
@@ -191,6 +203,7 @@ def _handle_internal_func(
 
     _convert_ir_bb(fn, ir.args[0].args[2], symbols)
 
+    _alloca_table = _saved_alloca_table
     return fn
 
 
@@ -539,6 +552,17 @@ def emit_body_blocks():
                 _alloca_table[alloca._id] = ptr
             return _alloca_table[alloca._id]
 
+        elif ir.value.startswith("$calloca"):
+            alloca = ir.passthrough_metadata["alloca"]
+            if alloca._id not in _alloca_table:
+                assert alloca._callsite is not None
+                callsite = IRLabel(alloca._callsite)
+                ptr = fn.get_basic_block().append_instruction(
+                    "calloca", alloca.offset, alloca.size, alloca._id, callsite
+                )
+                _alloca_table[alloca._id] = ptr
+            return _alloca_table[alloca._id]
+
         return symbols.get(ir.value)
     elif ir.is_literal:
         return IRLiteral(ir.value)

vyper/venom/passes/__init__.py

@@ -1,6 +1,7 @@
 from .algebraic_optimization import AlgebraicOptimizationPass
 from .branch_optimization import BranchOptimizationPass
 from .dft import DFTPass
+from .function_inliner import FunctionInlinerPass
 from .float_allocas import FloatAllocas
 from .make_ssa import MakeSSA
 from .mem2var import Mem2Var

vyper/venom/passes/float_allocas.py

@@ -23,7 +23,7 @@ def run_pass(self):
             # Extract alloca instructions
             non_alloca_instructions = []
             for inst in bb.instructions:
-                if inst.opcode in ("alloca", "palloca"):
+                if inst.opcode in ("alloca", "palloca", "calloca"):
                     # note: order of allocas impacts bytecode.
                     # TODO: investigate.
                     entry_bb.insert_instruction(inst)