feat[venom]: add function inliner (vyperlang#4478)

This commit adds a function inliner to the Venom backend, including:

- A basic heuristic to determine when inlining is beneficial
  (placeholder for future sophisticated analysis)
- new machinery for global (inter-function) passes over Venom IR to
  perform inlining
- A new function call graph analysis to support inlining
- Removal of the variable equivalence analysis (now provided directly by
  DFG analysis)

misc/refactor:
- add new `all2` utility function
- add dfs preorder traversal to CFG
- add ensure_well_formed util to IRBasicBlock
- machinery for IRParameter (may be useful in the future)
- add `copy()` machinery for venom data structures (may be useful in the
  future)

---------

Co-authored-by: Charles Cooper <[email protected]>

tests/functional/codegen/calling_convention/test_inlineable_functions.py

@@ -0,0 +1,66 @@
+"""
+Test functionality of internal functions which may be inlined
+"""
+# note for refactor: this may be able to be merged with
+# calling_convention/test_internal_call.py
+
+
+def test_call_in_call(get_contract):
+    code = """
+@internal
+def _foo(a: uint256,) -> uint256:
+    return 1 + a
+
+@internal
+def _foo2() -> uint256:
+    return 4
+
+@external
+def foo() -> uint256:
+    return self._foo(self._foo2())
+    """
+
+    c = get_contract(code)
+    assert c.foo() == 5
+
+
+def test_call_in_call_with_raise(get_contract, tx_failed):
+    code = """
+@internal
+def sum(a: uint256) -> uint256:
+    if a > 1:
+        return a + 1
+    raise
+
+@internal
+def middle(a: uint256) -> uint256:
+    return self.sum(a)
+
+@external
+def test(a: uint256) -> uint256:
+    return self.middle(a)
+    """
+
+    c = get_contract(code)
+
+    assert c.test(2) == 3
+
+    with tx_failed():
+        c.test(0)
+
+
+def test_inliner_with_unused_param(get_contract):
+    code = """
+data: public(uint256)
+
+@internal
+def _foo(start: uint256, length: uint256):
+    self.data = start
+
+@external
+def foo(x: uint256, y: uint256):
+    self._foo(x, y)
+"""
+
+    c = get_contract(code)
+    c.foo(1, 2)

tests/unit/compiler/asm/test_asm_optimizer.py

@@ -120,7 +120,9 @@ def foo():
     input_bundle = make_input_bundle({"library.vy": library})
     res = compile_code(code, input_bundle=input_bundle, output_formats=["asm"])
     asm = res["asm"]
-    assert "some_function()" in asm
+
+    if not experimental_codegen:
+        assert "some_function()" in asm  # Venom function inliner will remove this
 
     assert "unused1()" not in asm
     assert "unused2()" not in asm

tests/unit/compiler/test_source_map.py

@@ -37,17 +37,25 @@ def test_jump_map(optimize, experimental_codegen):
     pos_map = source_map["pc_pos_map"]
     jump_map = source_map["pc_jump_map"]
 
-    expected_jumps = 1
     if optimize == OptimizationLevel.NONE:
         # some jumps which don't get optimized out when optimizer is off
         # (slightly different behavior depending if venom pipeline is enabled):
-        if not experimental_codegen:
+        if experimental_codegen:
+            expected_jumps = 0
+            expected_internals = 0
+        else:
             expected_jumps = 3
+            expected_internals = 2
+    else:
+        if experimental_codegen:
+            expected_jumps = 0
+            expected_internals = 0
         else:
-            expected_jumps = 2
+            expected_jumps = 1
+            expected_internals = 2
 
     assert len([v for v in jump_map.values() if v == "o"]) == expected_jumps
-    assert len([v for v in jump_map.values() if v == "i"]) == 2
+    assert len([v for v in jump_map.values() if v == "i"]) == expected_internals
 
     code_lines = [i + "\n" for i in TEST_CODE.split("\n")]
     for pc in [k for k, v in jump_map.items() if v == "o"]:

tests/unit/compiler/venom/test_variable_equivalence.py

@@ -0,0 +1,32 @@
+import itertools
+
+from tests.venom_utils import parse_from_basic_block
+from vyper.venom.analysis import DFGAnalysis, IRAnalysesCache
+from vyper.venom.basicblock import IRVariable
+
+
+def test_variable_equivalence_dfg_order():
+    a_code = """
+    main:
+    %1 = 1
+    %2 = %1
+    %3 = %2
+    """
+    # technically invalid code, but variable equivalence should handle
+    # it either way
+    b_code = """
+    main:
+    %3 = %2
+    %2 = %1
+    %1 = 1
+    """
+    fn1 = parse_from_basic_block(a_code).entry_function
+    fn2 = parse_from_basic_block(b_code).entry_function
+
+    dfg1 = IRAnalysesCache(fn1).request_analysis(DFGAnalysis)
+    dfg2 = IRAnalysesCache(fn2).request_analysis(DFGAnalysis)
+
+    vars_ = map(IRVariable, ("%1", "%2", "%3"))
+    for var1, var2 in itertools.combinations(vars_, 2):
+        assert dfg1.are_equivalent(var1, var2)
+        assert dfg2.are_equivalent(var1, var2)

vyper/codegen/function_definitions/internal_function.py

@@ -78,6 +78,11 @@ def generate_ir_for_internal_function(
 
     ir_node = IRnode.from_list(["seq", body, cleanup_routine])
 
+    # add function signature to passthru metadata so that venom
+    # has more information to work with
+    ir_node.passthrough_metadata["func_t"] = func_t
+    ir_node.passthrough_metadata["context"] = context
+
     # tag gas estimate and frame info
     func_t._ir_info.gas_estimate = ir_node.gas
     tag_frame_info(func_t, context)

vyper/codegen/self_call.py

@@ -109,4 +109,6 @@ def ir_for_self_call(stmt_expr, context):
     )
     o.is_self_call = True
     o.invoked_function_ir = func_t._ir_info.func_ir
+    o.passthrough_metadata["func_t"] = func_t
+    o.passthrough_metadata["args_ir"] = args_ir
     return o

vyper/compiler/phases.py

@@ -256,8 +256,8 @@ def function_signatures(self) -> dict[str, ContractFunctionT]:
     @cached_property
     def venom_functions(self):
         deploy_ir, runtime_ir = self._ir_output
-        deploy_venom = generate_ir(deploy_ir, self.settings.optimize)
-        runtime_venom = generate_ir(runtime_ir, self.settings.optimize)
+        deploy_venom = generate_ir(deploy_ir, self.settings)
+        runtime_venom = generate_ir(runtime_ir, self.settings)
         return deploy_venom, runtime_venom
 
     @cached_property

vyper/utils.py

@@ -693,3 +693,17 @@ def safe_relpath(path):
         # on Windows, if path and curdir are on different drives, an exception
         # can be thrown
         return path
+
+
+def all2(iterator):
+    """
+    This function checks if all elements in the given `iterable` are truthy,
+    similar to Python's built-in `all()` function. However, `all2` differs
+    in the case where there are no elements in the iterable. `all()` returns
+    `True` for the empty iterable, but `all2()` returns False.
+    """
+    try:
+        s = next(iterator)
+    except StopIteration:
+        return False
+    return bool(s) and all(iterator)

vyper/venom/__init__.py

@@ -4,7 +4,7 @@
 from typing import Optional
 
 from vyper.codegen.ir_node import IRnode
-from vyper.compiler.settings import OptimizationLevel
+from vyper.compiler.settings import OptimizationLevel, Settings
 from vyper.venom.analysis.analysis import IRAnalysesCache
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
@@ -15,6 +15,7 @@
     BranchOptimizationPass,
     DFTPass,
     FloatAllocas,
+    FunctionInlinerPass,
     LoadElimination,
     LowerDloadPass,
     MakeSSA,
@@ -46,12 +47,10 @@ def generate_assembly_experimental(
     return compiler.generate_evm(optimize == OptimizationLevel.NONE)
 
 
-def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
+def _run_passes(fn: IRFunction, optimize: OptimizationLevel, ac: IRAnalysesCache) -> None:
     # Run passes on Venom IR
     # TODO: Add support for optimization levels
 
-    ac = IRAnalysesCache(fn)
-
     FloatAllocas(ac, fn).run_pass()
 
     SimplifyCFGPass(ac, fn).run_pass()
@@ -94,14 +93,31 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     DFTPass(ac, fn).run_pass()
 
 
-def run_passes_on(ctx: IRContext, optimize: OptimizationLevel):
+def _run_global_passes(ctx: IRContext, optimize: OptimizationLevel, ir_analyses: dict) -> None:
+    FunctionInlinerPass(ir_analyses, ctx, optimize).run_pass()
+
+
+def run_passes_on(ctx: IRContext, optimize: OptimizationLevel) -> None:
+    ir_analyses = {}
     for fn in ctx.functions.values():
-        _run_passes(fn, optimize)
+        ir_analyses[fn] = IRAnalysesCache(fn)
 
+    _run_global_passes(ctx, optimize, ir_analyses)
+
+    ir_analyses = {}
+    for fn in ctx.functions.values():
+        ir_analyses[fn] = IRAnalysesCache(fn)
 
-def generate_ir(ir: IRnode, optimize: OptimizationLevel) -> IRContext:
+    for fn in ctx.functions.values():
+        _run_passes(fn, optimize, ir_analyses[fn])
+
+
+def generate_ir(ir: IRnode, settings: Settings) -> IRContext:
     # Convert "old" IR to "new" IR
     ctx = ir_node_to_venom(ir)
+
+    optimize = settings.optimize
+    assert optimize is not None  # help mypy
     run_passes_on(ctx, optimize)
 
     return ctx

vyper/venom/analysis/__init__.py

@@ -2,5 +2,5 @@
 from .cfg import CFGAnalysis
 from .dfg import DFGAnalysis
 from .dominators import DominatorTreeAnalysis
-from .equivalent_vars import VarEquivalenceAnalysis
+from .fcg import FCGAnalysis
 from .liveness import LivenessAnalysis

vyper/venom/analysis/analysis.py

@@ -72,4 +72,5 @@ def force_analysis(self, analysis_cls: Type[IRAnalysis], *args, **kwargs):
         assert issubclass(analysis_cls, IRAnalysis), f"{analysis_cls} is not an IRAnalysis"
         if analysis_cls in self.analyses_cache:
             self.invalidate_analysis(analysis_cls)
+
         return self.request_analysis(analysis_cls, *args, **kwargs)

vyper/venom/analysis/cfg.py

@@ -34,29 +34,48 @@ def analyze(self) -> None:
                     bb.add_cfg_out(next_bb)
                     next_bb.add_cfg_in(bb)
 
-        self._compute_dfs_r(self.function.entry)
+        self._compute_dfs_post_r(self.function.entry)
 
-    def _compute_dfs_r(self, bb):
+    def _compute_dfs_post_r(self, bb):
         if bb.is_reachable:
             return
         bb.is_reachable = True
 
         for out_bb in bb.cfg_out:
-            self._compute_dfs_r(out_bb)
+            self._compute_dfs_post_r(out_bb)
 
         self._dfs.add(bb)
 
     @property
-    def dfs_walk(self) -> Iterator[IRBasicBlock]:
+    def dfs_pre_walk(self) -> Iterator[IRBasicBlock]:
+        visited: OrderedSet[IRBasicBlock] = OrderedSet()
+
+        def _visit_dfs_pre_r(bb: IRBasicBlock):
+            if bb in visited:
+                return
+            visited.add(bb)
+
+            yield bb
+
+            for out_bb in bb.cfg_out:
+                yield from _visit_dfs_pre_r(out_bb)
+
+        yield from _visit_dfs_pre_r(self.function.entry)
+
+    @property
+    def dfs_post_walk(self) -> Iterator[IRBasicBlock]:
         return iter(self._dfs)
 
     def invalidate(self):
         from vyper.venom.analysis import DFGAnalysis, DominatorTreeAnalysis, LivenessAnalysis
 
+        fn = self.function
+        for bb in fn.get_basic_blocks():
+            bb.cfg_in = OrderedSet()
+            bb.cfg_out = OrderedSet()
+            bb.out_vars = OrderedSet()
+
         self.analyses_cache.invalidate_analysis(DominatorTreeAnalysis)
         self.analyses_cache.invalidate_analysis(LivenessAnalysis)
-
-        self._dfs = None
-
-        # be conservative - assume cfg invalidation invalidates dfg
         self.analyses_cache.invalidate_analysis(DFGAnalysis)
+        self._dfs = None

vyper/venom/analysis/dfg.py

@@ -3,7 +3,7 @@
 from vyper.utils import OrderedSet
 from vyper.venom.analysis.analysis import IRAnalysesCache, IRAnalysis
 from vyper.venom.analysis.liveness import LivenessAnalysis
-from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRVariable
+from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IROperand, IRVariable
 from vyper.venom.function import IRFunction
 
 
@@ -41,6 +41,23 @@ def remove_use(self, op: IRVariable, inst: IRInstruction):
         uses: OrderedSet = self._dfg_inputs.get(op, OrderedSet())
         uses.remove(inst)
 
+    def are_equivalent(self, var1: IROperand, var2: IROperand) -> bool:
+        if var1 == var2:
+            return True
+
+        if isinstance(var1, IRVariable) and isinstance(var2, IRVariable):
+            var1 = self._traverse_store_chain(var1)
+            var2 = self._traverse_store_chain(var2)
+
+        return var1 == var2
+
+    def _traverse_store_chain(self, var: IRVariable) -> IRVariable:
+        while True:
+            inst = self.get_producing_instruction(var)
+            if inst is None or inst.opcode != "store":
+                return var
+            var = inst.operands[0]  # type: ignore
+
     @property
     def outputs(self) -> dict[IRVariable, IRInstruction]:
         return self._dfg_outputs