feat[venom]: implementation of the new calling convention

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/self_call.py

@@ -1,3 +1,6 @@
+import copy
+import dataclasses
+
 from vyper.codegen.core import _freshname, eval_once_check, make_setter
 from vyper.codegen.ir_node import IRnode
 from vyper.evm.address_space import MEMORY
@@ -66,7 +69,27 @@ def ir_for_self_call(stmt_expr, context):
 
     # 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/ir_node_to_venom.py

@@ -1,7 +1,9 @@
 import functools
 import re
+from collections import defaultdict
 from typing import Optional
 
+from vyper.codegen.context import Alloca
 from vyper.codegen.ir_node import IRnode
 from vyper.evm.opcodes import get_opcodes
 from vyper.venom.basicblock import (
@@ -13,7 +15,7 @@
     IRVariable,
 )
 from vyper.venom.context import IRContext
-from vyper.venom.function import IRFunction
+from vyper.venom.function import IRFunction, IRParameter
 
 # Instructions that are mapped to their inverse
 INVERSE_MAPPED_IR_INSTRUCTIONS = {"ne": "eq", "le": "gt", "sle": "sgt", "ge": "lt", "sge": "slt"}
@@ -63,7 +65,6 @@
         "gasprice",
         "gaslimit",
         "returndatasize",
-        "mload",
         "iload",
         "istore",
         "dload",
@@ -109,15 +110,18 @@
 
 SymbolTable = dict[str, IROperand]
 _alloca_table: dict[int, IROperand]
+# assumption: callsites (return pc labels) are globally unique.
+_callsites: dict[str, list[Alloca]]
 MAIN_ENTRY_LABEL_NAME = "__main_entry"
 
 
 # convert IRnode directly to venom
 def ir_node_to_venom(ir: IRnode) -> IRContext:
     _ = ir.unique_symbols  # run unique symbols check
 
-    global _alloca_table
+    global _alloca_table, _callsites
     _alloca_table = {}
+    _callsites = defaultdict(list)
 
     ctx = IRContext()
     fn = ctx.create_function(MAIN_ENTRY_LABEL_NAME)
@@ -160,66 +164,149 @@ def _append_return_args(fn: IRFunction, ofst: int = 0, size: int = 0):
 
 
 def _handle_self_call(fn: IRFunction, ir: IRnode, symbols: SymbolTable) -> Optional[IROperand]:
+    global _callsites
     setup_ir = ir.args[1]
     goto_ir = [ir for ir in ir.args if ir.value == "goto"][0]
-    target_label = goto_ir.args[0].value  # goto
-    ret_args: list[IROperand] = [IRLabel(target_label)]  # type: ignore
+    target_label = goto_ir.args[0].value  # target function
+    assert isinstance(target_label, str)  # help mypy
     func_t = ir.passthrough_metadata["func_t"]
     assert func_t is not None, "func_t not found in passthrough metadata"
 
+    returns_word = _returns_word(func_t)
+
+    stack_args: list[IROperand] = [IRLabel(target_label)]
+
     if setup_ir != goto_ir:
         _convert_ir_bb(fn, setup_ir, symbols)
 
     converted_args = _convert_ir_bb_list(fn, goto_ir.args[1:], symbols)
 
     callsite_op = converted_args[-1]
     assert isinstance(callsite_op, IRLabel), converted_args
+    callsite = callsite_op.value
 
     bb = fn.get_basic_block()
     return_buf = None
 
     if len(converted_args) > 1:
         return_buf = converted_args[0]
 
+    callsite_args = _callsites[callsite]
+
+    for alloca in callsite_args:
+        if not _is_word_type(alloca.typ):
+            continue
+        ptr = _alloca_table[alloca._id]
+        # to be optimized out by mem2var
+        stack_arg = bb.append_instruction("mload", ptr)
+        assert stack_arg is not None
+        stack_args.append(stack_arg)
+
+    # NOTE: order of return buf vis a vis stack args must be the same
+    # as the order of params!
     if return_buf is not None:
-        ret_args.append(return_buf)  # type: ignore
+        if not returns_word:
+            stack_args.append(return_buf)
 
-    bb.append_invoke_instruction(ret_args, returns=False)  # type: ignore
+    if returns_word:
+        ret_value = bb.append_invoke_instruction(stack_args, returns=True)  # type: ignore
+        assert ret_value is not None
+        assert isinstance(return_buf, IROperand)
+        # to be optimized out by mem2var
+        bb.append_instruction("mstore", ret_value, return_buf)
+        return return_buf
+
+    bb.append_invoke_instruction(stack_args, returns=False)  # type: ignore
 
     return return_buf
 
 
+_current_func_t = None
+_current_context = None
+
+
+def _is_word_type(typ):
+    return typ._is_prim_word
+    # return typ.memory_bytes_required == 32
+
+
+# func_t: ContractFunctionT
+def _returns_word(func_t) -> bool:
+    return_t = func_t.return_type
+    return return_t is not None and _is_word_type(return_t)
+
+
 def _handle_internal_func(
     # TODO: remove does_return_data, replace with `func_t.return_type is not None`
     fn: IRFunction,
     ir: IRnode,
     does_return_data: bool,
     symbols: SymbolTable,
 ) -> IRFunction:
-    global _alloca_table
+    global _alloca_table, _current_func_t, _current_context
+
+    _current_func_t = ir.passthrough_metadata["func_t"]
+    _current_context = ir.passthrough_metadata["context"]
+    func_t = _current_func_t
+    context = _current_context
 
     fn = fn.ctx.create_function(ir.args[0].args[0].value)
 
+    index = 0
+    if func_t.return_type is not None and not _returns_word(func_t):
+        index += 1
+    for arg in func_t.arguments:
+        var = context.lookup_var(arg.name)
+        if not _is_word_type(var.typ):
+            continue
+        venom_arg = IRParameter(
+            var.name, index, var.alloca.offset, var.alloca.size, None, None, None
+        )
+        fn.args.append(venom_arg)
+        index += 1
+
     bb = fn.get_basic_block()
 
     _saved_alloca_table = _alloca_table
     _alloca_table = {}
 
+    returns_word = _returns_word(func_t)
+
     # return buffer
     if does_return_data:
-        buf = bb.append_instruction("param")
-        bb.instructions[-1].annotation = "return_buffer"
+        if returns_word:
+            # this alloca should be stripped by mem2var. we can remove
+            # the hardcoded offset once we have proper memory allocator
+            # functionality in venom.
+            # (note frontend generates alloca IDs starting from 1)
+            buf = bb.append_instruction("alloca", IRLiteral(-1), IRLiteral(-1), IRLiteral(0))
+        else:
+            # NOTE: must match order of stack args!
+            buf = bb.append_instruction("param")
+            bb.instructions[-1].annotation = "return_buffer"
 
         assert buf is not None  # help mypy
         symbols["return_buffer"] = buf
 
+    for arg in fn.args:
+        ret = bb.append_instruction("param")
+        bb.instructions[-1].annotation = arg.name
+        assert ret is not None  # help mypy
+        symbols[arg.name] = ret
+        arg.func_var = ret
+
     # return address
     return_pc = bb.append_instruction("param")
     assert return_pc is not None  # help mypy
     symbols["return_pc"] = return_pc
 
     bb.instructions[-1].annotation = "return_pc"
 
+    for arg in fn.args:
+        var = IRVariable(arg.name)
+        bb.append_instruction("store", IRLiteral(arg.offset), ret=var)  # type: ignore
+        arg.addr_var = var
+
     _convert_ir_bb(fn, ir.args[0].args[2], symbols)
 
     _alloca_table = _saved_alloca_table
@@ -432,7 +519,12 @@ def _convert_ir_bb(fn, ir, symbols):
         if label.value == "return_pc":
             label = symbols.get("return_pc")
             # return label should be top of stack
-            bb.append_instruction("ret", label)
+            if _returns_word(_current_func_t):
+                buf = symbols["return_buffer"]
+                val = bb.append_instruction("mload", buf)
+                bb.append_instruction("ret", val, label)
+            else:
+                bb.append_instruction("ret", label)
         else:
             bb.append_instruction("jmp", label)
 
@@ -441,7 +533,28 @@ def _convert_ir_bb(fn, ir, symbols):
         # to fix upstream.
         val, ptr = _convert_ir_bb_list(fn, reversed(ir.args), symbols)
 
+        if isinstance(ptr, IRVariable):
+            # TODO: is this bad code?
+            param = fn.get_param_by_name(ptr)
+            if param is not None:
+                return fn.get_basic_block().append_instruction("store", val, ret=param.func_var)
+
+        if isinstance(ptr, IRLabel) and ptr.value.startswith("$palloca"):
+            symbol = symbols.get(ptr.annotation, None)
+            if symbol is not None:
+                return fn.get_basic_block().append_instruction("store", symbol)
+
         return fn.get_basic_block().append_instruction("mstore", val, ptr)
+    elif ir.value == "mload":
+        arg = ir.args[0]
+        ptr = _convert_ir_bb(fn, arg, symbols)
+
+        if isinstance(arg.value, str) and arg.value.startswith("$palloca"):
+            symbol = symbols.get(arg.annotation, None)
+            if symbol is not None:
+                return fn.get_basic_block().append_instruction("store", symbol)
+
+        return fn.get_basic_block().append_instruction("mload", ptr)
     elif ir.value == "ceil32":
         x = ir.args[0]
         expanded = IRnode.from_list(["and", ["add", x, 31], ["not", 31]])
@@ -551,13 +664,33 @@ def emit_body_blocks():
 
         elif ir.value.startswith("$palloca"):
             alloca = ir.passthrough_metadata["alloca"]
+            if fn.get_param_at_offset(alloca.offset) is not None:
+                return fn.get_param_at_offset(alloca.offset).addr_var
             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]
 
+        elif ir.value.startswith("$calloca"):
+            global _callsites
+            alloca = ir.passthrough_metadata["alloca"]
+            assert alloca._callsite is not None
+            if alloca._id not in _alloca_table:
+                bb = fn.get_basic_block()
+                if _is_word_type(alloca.typ):
+                    ptr = bb.append_instruction("alloca", alloca.offset, alloca.size, alloca._id)
+                else:
+                    ptr = IRLiteral(alloca.offset)
+
+                _alloca_table[alloca._id] = ptr
+            ret = _alloca_table[alloca._id]
+            # assumption: callocas appear in the same order as the
+            # order of arguments to the function.
+            _callsites[alloca._callsite].append(alloca)
+            return ret
+
         return symbols.get(ir.value)
     elif ir.is_literal:
         return IRLiteral(ir.value)

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)

vyper/venom/passes/function_inliner.py

@@ -123,6 +123,8 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
 
         func_copy = self._clone_function(func, prefix)
 
+        call_site_args = call_site.operands.copy()
+
         for bb in func_copy.get_basic_blocks():
             bb.parent = call_site_func
             call_site_func.append_basic_block(bb)
@@ -134,13 +136,18 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
                     # not valid venom code, but it will get removed in store elimination
                     # (or unused variable elimination)
                     inst.opcode = "store"
-                    val = call_site.operands[-param_idx - 1]
+                    val = call_site_args.pop()
                     inst.operands = [val]
                     param_idx += 1
                 elif inst.opcode == "palloca":
                     inst.opcode = "store"
                     inst.operands = [inst.operands[0]]
                 elif inst.opcode == "ret":
+                    if len(inst.operands) > 1:
+                        ret_value = inst.operands[0]
+                        bb.insert_instruction(
+                            IRInstruction("store", [ret_value], call_site.output), -1
+                        )
                     inst.opcode = "jmp"
                     inst.operands = [call_site_return.label]
                 elif inst.opcode == "revert":
@@ -152,6 +159,8 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
                 if not inst.annotation:
                     inst.annotation = f"from {func.name}"
 
+        assert len(call_site_args) == 0
+
         call_site_bb.instructions = call_site_bb.instructions[:call_idx]
         call_site_bb.append_instruction("jmp", func_copy.entry.label)
 

vyper/venom/passes/mem2var.py

@@ -61,6 +61,10 @@ def _process_alloca_var(self, dfg: DFGAnalysis, alloca_inst, var: IRVariable):
                 new_inst = IRInstruction("mstore", [var, inst.operands[1]])
                 bb.insert_instruction(new_inst, idx)
 
+            inst.annotation = f"mem2var alloca{alloca_id}"
+            if alloca_inst.annotation is not None:
+                inst.annotation += f" ({alloca_inst.annotation})"
+
     def _process_palloca_var(self, dfg: DFGAnalysis, palloca_inst: IRInstruction, var: IRVariable):
         """
         Process alloca allocated variable. If it is only used by mstore/mload