Skip to content

Commit d8f4032

Browse files
feat[venom]: add small heuristic for cleaning input stack (#4251)
this commit adds a small heuristic for cleaning the input stack from `cfg_in`, which is to pop the shallowest items first. it also cleans up the code a little bit and adds stronger preconditions to the code. a further optimization would be to try to align the output stack as best as possible at the `jnz`, but that belongs in the DFT pass.
1 parent d079562 commit d8f4032

File tree

1 file changed

+22
-19
lines changed

1 file changed

+22
-19
lines changed

vyper/venom/venom_to_assembly.py

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,8 @@ def _generate_evm_for_basicblock_r(
295295
asm.append(f"_sym_{basicblock.label}")
296296
asm.append("JUMPDEST")
297297

298-
self.clean_stack_from_cfg_in(asm, basicblock, stack)
298+
if len(basicblock.cfg_in) == 1:
299+
self.clean_stack_from_cfg_in(asm, basicblock, stack)
299300

300301
all_insts = sorted(basicblock.instructions, key=lambda x: x.opcode != "param")
301302

@@ -321,26 +322,28 @@ def _generate_evm_for_basicblock_r(
321322
def clean_stack_from_cfg_in(
322323
self, asm: list, basicblock: IRBasicBlock, stack: StackModel
323324
) -> None:
324-
if len(basicblock.cfg_in) == 0:
325-
return
326-
327-
to_pop = OrderedSet[IRVariable]()
328-
for in_bb in basicblock.cfg_in:
329-
# inputs is the input variables we need from in_bb
330-
inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
331-
332-
# layout is the output stack layout for in_bb (which works
333-
# for all possible cfg_outs from the in_bb).
334-
layout = in_bb.out_vars
335-
336-
# pop all the stack items which in_bb produced which we don't need.
337-
to_pop |= layout.difference(inputs)
338-
325+
# the input block is a splitter block, like jnz or djmp
326+
assert len(basicblock.cfg_in) == 1
327+
in_bb = basicblock.cfg_in.first()
328+
assert len(in_bb.cfg_out) > 1
329+
330+
# inputs is the input variables we need from in_bb
331+
inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
332+
333+
# layout is the output stack layout for in_bb (which works
334+
# for all possible cfg_outs from the in_bb, in_bb is responsible
335+
# for making sure its output stack layout works no matter which
336+
# bb it jumps into).
337+
layout = in_bb.out_vars
338+
to_pop = list(layout.difference(inputs))
339+
340+
# small heuristic: pop from shallowest first.
341+
to_pop.sort(key=lambda var: -stack.get_depth(var))
342+
343+
# NOTE: we could get more fancy and try to optimize the swap
344+
# operations here, there is probably some more room for optimization.
339345
for var in to_pop:
340346
depth = stack.get_depth(var)
341-
# don't pop phantom phi inputs
342-
if depth is StackModel.NOT_IN_STACK:
343-
continue
344347

345348
if depth != 0:
346349
self.swap(asm, stack, depth)

0 commit comments

Comments
 (0)