feat[venom]: stack2mem pass implementation

tests/functional/codegen/features/test_clampers.py

@@ -5,7 +5,6 @@
 from eth_utils import keccak
 
 from tests.utils import ZERO_ADDRESS, decimal_to_int
-from vyper.exceptions import StackTooDeep
 from vyper.utils import int_bounds
 
 
@@ -502,7 +501,6 @@ def foo(b: DynArray[int128, 10]) -> DynArray[int128, 10]:
 
 
 @pytest.mark.parametrize("value", [0, 1, -1, 2**127 - 1, -(2**127)])
-@pytest.mark.venom_xfail(raises=StackTooDeep, reason="stack scheduler regression")
 def test_multidimension_dynarray_clamper_passing(get_contract, value):
     code = """
 @external

tests/unit/compiler/venom/test_mem_allocator.py

@@ -0,0 +1,116 @@
+import pytest
+
+from vyper.venom.mem_allocator import MemoryAllocator
+
+MEM_BLOCK_ADDRESS = 0x1000
+
+
+@pytest.fixture
+def allocator():
+    return MemoryAllocator(1024, MEM_BLOCK_ADDRESS)
+
+
+def test_initial_state(allocator):
+    assert allocator.get_free_memory() == 1024
+    assert allocator.get_allocated_memory() == 0
+
+
+def test_single_allocation(allocator):
+    addr = allocator.allocate(256)
+    assert addr == MEM_BLOCK_ADDRESS
+    assert allocator.get_free_memory() == 768
+    assert allocator.get_allocated_memory() == 256
+
+
+def test_multiple_allocations(allocator):
+    addr1 = allocator.allocate(256)
+    addr2 = allocator.allocate(128)
+    addr3 = allocator.allocate(64)
+
+    assert addr1 == MEM_BLOCK_ADDRESS
+    assert addr2 == MEM_BLOCK_ADDRESS + 256
+    assert addr3 == MEM_BLOCK_ADDRESS + 384
+    assert allocator.get_free_memory() == 576
+    assert allocator.get_allocated_memory() == 448
+
+
+def test_deallocation(allocator):
+    addr1 = allocator.allocate(256)
+    addr2 = allocator.allocate(128)
+
+    assert allocator.deallocate(addr1) is True
+    assert allocator.get_free_memory() == 896
+    assert allocator.get_allocated_memory() == 128
+
+    assert allocator.deallocate(addr2) is True
+    assert allocator.get_free_memory() == 1024
+    assert allocator.get_allocated_memory() == 0
+
+
+def test_allocation_after_deallocation(allocator):
+    addr1 = allocator.allocate(256)
+    allocator.deallocate(addr1)
+    addr2 = allocator.allocate(128)
+
+    assert addr2 == MEM_BLOCK_ADDRESS
+    assert allocator.get_free_memory() == 896
+    assert allocator.get_allocated_memory() == 128
+
+
+def test_out_of_memory(allocator):
+    allocator.allocate(1000)
+    with pytest.raises(MemoryError):
+        allocator.allocate(100)
+
+
+def test_invalid_deallocation(allocator):
+    assert allocator.deallocate(0x2000) is False
+
+
+def test_fragmentation_and_merging(allocator):
+    addr1 = allocator.allocate(256)
+    addr2 = allocator.allocate(256)
+    addr3 = allocator.allocate(256)
+
+    assert allocator.get_free_memory() == 256
+    assert allocator.get_allocated_memory() == 768
+
+    allocator.deallocate(addr1)
+    assert allocator.get_free_memory() == 512
+    assert allocator.get_allocated_memory() == 512
+
+    allocator.deallocate(addr3)
+    assert allocator.get_free_memory() == 768
+    assert allocator.get_allocated_memory() == 256
+
+    addr4 = allocator.allocate(512)
+    assert addr4 == MEM_BLOCK_ADDRESS + 512
+    assert allocator.get_free_memory() == 256
+    assert allocator.get_allocated_memory() == 768
+
+    allocator.deallocate(addr2)
+    assert allocator.get_free_memory() == 512
+    assert allocator.get_allocated_memory() == 512
+
+    allocator.deallocate(addr4)
+    assert allocator.get_free_memory() == 1024  # All blocks merged
+    assert allocator.get_allocated_memory() == 0
+
+    # Test if we can now allocate the entire memory
+    addr5 = allocator.allocate(1024)
+    assert addr5 == MEM_BLOCK_ADDRESS
+    assert allocator.get_free_memory() == 0
+    assert allocator.get_allocated_memory() == 1024
+
+
+def test_exact_fit_allocation(allocator):
+    addr1 = allocator.allocate(1024)
+    assert addr1 == MEM_BLOCK_ADDRESS
+    assert allocator.get_free_memory() == 0
+    assert allocator.get_allocated_memory() == 1024
+
+    allocator.deallocate(addr1)
+    addr2 = allocator.allocate(1024)
+    assert addr2 == MEM_BLOCK_ADDRESS
+    assert allocator.get_free_memory() == 0
+    assert allocator.get_allocated_memory() == 1024

vyper/exceptions.py

@@ -400,10 +400,6 @@ class CodegenPanic(VyperInternalException):
     """Invalid code generated during codegen phase"""
 
 
-class StackTooDeep(CodegenPanic):
-    """Stack too deep"""  # (should not happen)
-
-
 class UnexpectedNodeType(VyperInternalException):
     """Unexpected AST node type."""
 
@@ -424,6 +420,15 @@ class InvalidABIType(VyperInternalException):
     """An internal routine constructed an invalid ABI type"""
 
 
+class UnreachableStackException(VyperException):
+
+    """An unreachable stack operation was encountered."""
+
+    def __init__(self, message, op):
+        self.op = op
+        super().__init__(message)
+
+
 @contextlib.contextmanager
 def tag_exceptions(node, fallback_exception_type=CompilerPanic, note=None):
     try:

vyper/venom/__init__.py

@@ -6,19 +6,22 @@
 from vyper.codegen.ir_node import IRnode
 from vyper.compiler.settings import OptimizationLevel
 from vyper.venom.analysis.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRVariable
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
 from vyper.venom.ir_node_to_venom import ir_node_to_venom
 from vyper.venom.passes import (
     SCCP,
     AlgebraicOptimizationPass,
+    AllocaElimination,
     BranchOptimizationPass,
     DFTPass,
     FloatAllocas,
     MakeSSA,
     Mem2Var,
     RemoveUnusedVariablesPass,
     SimplifyCFGPass,
+    Stack2Mem,
     StoreElimination,
     StoreExpansionPass,
 )
@@ -51,6 +54,7 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     FloatAllocas(ac, fn).run_pass()
 
     SimplifyCFGPass(ac, fn).run_pass()
+    AllocaElimination(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
     Mem2Var(ac, fn).run_pass()
     MakeSSA(ac, fn).run_pass()
@@ -70,6 +74,7 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
 
     StoreExpansionPass(ac, fn).run_pass()
     DFTPass(ac, fn).run_pass()
+    Stack2Mem(ac, fn).run_pass()
 
 
 def generate_ir(ir: IRnode, optimize: OptimizationLevel) -> IRContext:

vyper/venom/context.py

@@ -2,6 +2,7 @@
 
 from vyper.venom.basicblock import IRInstruction, IRLabel, IROperand
 from vyper.venom.function import IRFunction
+from vyper.venom.mem_allocator import MemoryAllocator
 
 
 class IRContext:
@@ -10,13 +11,17 @@ class IRContext:
     immutables_len: Optional[int]
     data_segment: list[IRInstruction]
     last_label: int
+    mem_allocator: MemoryAllocator
 
     def __init__(self) -> None:
         self.functions = {}
         self.ctor_mem_size = None
         self.immutables_len = None
         self.data_segment = []
         self.last_label = 0
+        self.mem_allocator = MemoryAllocator(
+            4096, 0x100000
+        )  # TODO: Should get this from the original IR
 
     def add_function(self, fn: IRFunction) -> None:
         fn.ctx = self

vyper/venom/function.py

@@ -1,7 +1,9 @@
 from typing import Iterator, Optional
 
 from vyper.codegen.ir_node import IRnode
-from vyper.venom.basicblock import IRBasicBlock, IRLabel, IRVariable
+from vyper.utils import OrderedSet
+from vyper.venom.basicblock import CFG_ALTERING_INSTRUCTIONS, IRBasicBlock, IRLabel, IRVariable
+from vyper.venom.mem_allocator import MemoryAllocator
 
 
 class IRFunction:
@@ -15,6 +17,7 @@
     last_label: int
     last_variable: int
     _basic_block_dict: dict[str, IRBasicBlock]
+    _mem_allocator: MemoryAllocator
 
     # Used during code generation
     _ast_source_stack: list[IRnode]
@@ -31,6 +34,8 @@
         self._ast_source_stack = []
         self._error_msg_stack = []
 
+        self._mem_allocator = MemoryAllocator(0xFFFFFFFFFFFFFFFF, 32)
+
         self.append_basic_block(IRBasicBlock(name, self))
 
     @property

vyper/venom/mem_allocator.py

@@ -0,0 +1,61 @@
+from typing import List
+
+
+class MemoryBlock:
+    size: int
+    address: int
+    is_free: bool
+
+    def __init__(self, size: int, address: int):
+        self.size = size
+        self.address = address
+        self.is_free = True
+
+
+class MemoryAllocator:
+    total_size: int
+    start_address: int
+    blocks: List[MemoryBlock]
+
+    def __init__(self, total_size: int, start_address: int):
+        self.total_size = total_size
+        self.start_address = start_address
+        self.blocks = [MemoryBlock(total_size, 0)]
+
+    def allocate(self, size: int) -> int:
+        # print(f"Allocating {size} bytes with free memory {self.get_free_memory()}")
+        for block in self.blocks:
+            if block.is_free and block.size >= size:
+                if block.size > size:
+                    new_block = MemoryBlock(block.size - size, block.address + size)
+                    self.blocks.insert(self.blocks.index(block) + 1, new_block)
+                    block.size = size
+                block.is_free = False
+                return self.start_address + block.address
+        raise MemoryError(
+            f"Memory allocation failed for size {size} with free memory {self.get_free_memory()}"
+        )
+
+    def deallocate(self, address: int) -> bool:
+        relative_address = address - self.start_address
+        for block in self.blocks:
+            if block.address == relative_address:
+                block.is_free = True
+                self._merge_adjacent_free_blocks()
+                return True
+        return False  # invalid address
+
+    def _merge_adjacent_free_blocks(self) -> None:
+        i = 0
+        while i < len(self.blocks) - 1:
+            if self.blocks[i].is_free and self.blocks[i + 1].is_free:
+                self.blocks[i].size += self.blocks[i + 1].size
+                self.blocks.pop(i + 1)
+            else:
+                i += 1
+
+    def get_free_memory(self) -> int:
+        return sum(block.size for block in self.blocks if block.is_free)
+
+    def get_allocated_memory(self) -> int:
+        return sum(block.size for block in self.blocks if not block.is_free)

vyper/venom/passes/__init__.py

@@ -1,4 +1,5 @@
 from .algebraic_optimization import AlgebraicOptimizationPass
+from .alloca_elimination import AllocaElimination
 from .branch_optimization import BranchOptimizationPass
 from .dft import DFTPass
 from .float_allocas import FloatAllocas
@@ -8,5 +9,6 @@
 from .remove_unused_variables import RemoveUnusedVariablesPass
 from .sccp import SCCP
 from .simplify_cfg import SimplifyCFGPass
+from .stack2mem import Stack2Mem
 from .store_elimination import StoreElimination
 from .store_expansion import StoreExpansionPass

vyper/venom/passes/alloca_elimination.py

@@ -0,0 +1,21 @@
+from vyper.venom.basicblock import IRInstruction, IRLiteral
+from vyper.venom.passes.base_pass import IRPass
+
+
+class AllocaElimination(IRPass):
+    """
+    This pass eliminates alloca instructions by allocating memory for them
+    """
+
+    def run_pass(self):
+        for bb in self.function.get_basic_blocks():
+            for inst in bb.instructions:
+                if inst.opcode == "alloca":
+                    self._process_alloca(inst)
+
+    def _process_alloca(self, inst: IRInstruction):
+        offset, _size = inst.operands
+        address = inst.parent.parent._mem_allocator.allocate(_size.value)
+        inst.opcode = "store"
+        inst.operands = [IRLiteral(address)]
+        # print(f"Allocated address {address} for alloca {_size.value}")