feat[venom]: add codesize optimization pass

tests/unit/compiler/venom/test_literals_codesize.py

@@ -0,0 +1,77 @@
+import pytest
+
+from vyper.utils import evm_not
+from vyper.venom.analysis import IRAnalysesCache
+from vyper.venom.basicblock import IRLiteral
+from vyper.venom.context import IRContext
+from vyper.venom.passes import ReduceLiteralsCodesize
+
+
+@pytest.mark.parametrize("orig_value", [0xFFFF << 240, 2**256 - 1])
+def test_ff_inversion(orig_value):
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+    bb = fn.get_basic_block()
+
+    bb.append_instruction("store", IRLiteral(orig_value))
+    bb.append_instruction("stop")
+    ac = IRAnalysesCache(fn)
+    ReduceLiteralsCodesize(ac, fn).run_pass()
+
+    assert bb.instructions[0].opcode == "not"
+    assert evm_not(bb.instructions[0].operands[0].value) == orig_value
+
+
+should_not_invert = [1, 0xFE << 248 | (2**248 - 1)]  # 0xfeff...ff
+
+
+@pytest.mark.parametrize("orig_value", should_not_invert)
+def test_no_inversion(orig_value):
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+    bb = fn.get_basic_block()
+
+    bb.append_instruction("store", IRLiteral(orig_value))
+    bb.append_instruction("stop")
+    ac = IRAnalysesCache(fn)
+    ReduceLiteralsCodesize(ac, fn).run_pass()
+
+    assert bb.instructions[0].opcode == "store"
+    assert bb.instructions[0].operands[0].value == orig_value
+
+
+should_shl = [0x01_000000]  # saves 3 bytes
+
+
+@pytest.mark.parametrize("orig_value", should_shl)
+def test_shl(orig_value):
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+    bb = fn.get_basic_block()
+
+    bb.append_instruction("store", IRLiteral(orig_value))
+    bb.append_instruction("stop")
+    ac = IRAnalysesCache(fn)
+    ReduceLiteralsCodesize(ac, fn).run_pass()
+
+    assert bb.instructions[0].opcode == "shl"
+    op0, op1 = bb.instructions[0].operands
+    assert op0.value << op1.value == orig_value
+
+
+should_not_shl = [0x01_00]  # only saves 2 bytes
+
+
+@pytest.mark.parametrize("orig_value", should_not_shl)
+def test_no_shl(orig_value):
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+    bb = fn.get_basic_block()
+
+    bb.append_instruction("store", IRLiteral(orig_value))
+    bb.append_instruction("stop")
+    ac = IRAnalysesCache(fn)
+    ReduceLiteralsCodesize(ac, fn).run_pass()
+
+    assert bb.instructions[0].opcode == "store"
+    assert bb.instructions[0].operands[0].value == orig_value

vyper/utils.py

@@ -391,6 +391,11 @@ def evm_twos_complement(x: int) -> int:
     return ((2**256 - 1) ^ x) + 1
 
 
+def evm_not(val: int) -> int:
+    assert 0 <= val <= SizeLimits.MAX_UINT256, "Value out of bounds"
+    return SizeLimits.MAX_UINT256 ^ val
+
+
 # EVM div semantics as a python function
 def evm_div(x, y):
     if y == 0:

vyper/venom/__init__.py

@@ -16,6 +16,7 @@
     DFTPass,
     MakeSSA,
     Mem2Var,
+    ReduceLiteralsCodesize,
     RemoveUnusedVariablesPass,
     SimplifyCFGPass,
     StoreElimination,
@@ -66,6 +67,10 @@ def _run_passes(fn: IRFunction, optimize: OptimizationLevel) -> None:
     RemoveUnusedVariablesPass(ac, fn).run_pass()
 
     StoreExpansionPass(ac, fn).run_pass()
+
+    if optimize == OptimizationLevel.CODESIZE:
+        ReduceLiteralsCodesize(ac, fn).run_pass()
+
     DFTPass(ac, fn).run_pass()
 
 

vyper/venom/passes/__init__.py

@@ -1,6 +1,7 @@
 from .algebraic_optimization import AlgebraicOptimizationPass
 from .branch_optimization import BranchOptimizationPass
 from .dft import DFTPass
+from .literals_codesize import ReduceLiteralsCodesize
 from .make_ssa import MakeSSA
 from .mem2var import Mem2Var
 from .normalization import NormalizationPass

vyper/venom/passes/literals_codesize.py

@@ -0,0 +1,49 @@
+from vyper.utils import evm_not
+from vyper.venom.basicblock import IRLiteral
+from vyper.venom.passes.base_pass import IRPass
+
+# not takes 1 byte1, so it makes sense to use it when we can save at least
+# 1 byte
+NOT_THRESHOLD = 1
+
+# shl takes 3 bytes, so it makes sense to use it when we can save at least
+# 3 bytes
+SHL_THRESHOLD = 3
+
+
+class ReduceLiteralsCodesize(IRPass):
+    def run_pass(self):
+        for bb in self.function.get_basic_blocks():
+            self._process_bb(bb)
+
+    def _process_bb(self, bb):
+        i = 0
+        while i < len(bb.instructions):
+            inst = bb.instructions[i]
+            i += 1
+            if inst.opcode != "store":
+                continue
+
+            op = inst.operands[0]
+            if not isinstance(op, IRLiteral):
+                continue
+
+            val = op.value % (2**256)
+
+            # transform things like 0xffff...01 to (not 0xfe)
+            if len(hex(val)) // 2 - len(hex(evm_not(val))) // 2 > NOT_THRESHOLD:
+                inst.opcode = "not"
+                op.value = evm_not(val)
+                continue
+
+            # transform things like 0x123400....000 to 0x1234 << ...
+            binz = bin(val)[2:]
+            if (ix := len(binz) - binz.rfind("1")) > SHL_THRESHOLD * 8:
+                ix -= 1
+                # sanity check
+                assert (val >> ix) << ix == val, val
+                assert (val >> ix) & 1 == 1, val
+
+                inst.opcode = "shl"
+                inst.operands = [IRLiteral(val >> ix), IRLiteral(ix)]
+                continue

vyper/venom/passes/sccp/eval.py

@@ -5,6 +5,7 @@
     SizeLimits,
     evm_div,
     evm_mod,
+    evm_not,
     evm_pow,
     signed_to_unsigned,
     unsigned_to_signed,
@@ -95,11 +96,6 @@ def _evm_sar(shift_len: int, value: int) -> int:
     return value >> shift_len
 
 
-def _evm_not(value: int) -> int:
-    assert 0 <= value <= SizeLimits.MAX_UINT256, "Value out of bounds"
-    return SizeLimits.MAX_UINT256 ^ value
-
-
 ARITHMETIC_OPS: dict[str, Callable[[list[IROperand]], int]] = {
     "add": _wrap_binop(operator.add),
     "sub": _wrap_binop(operator.sub),
@@ -122,7 +118,7 @@ def _evm_not(value: int) -> int:
     "or": _wrap_binop(operator.or_),
     "and": _wrap_binop(operator.and_),
     "xor": _wrap_binop(operator.xor),
-    "not": _wrap_unop(_evm_not),
+    "not": _wrap_unop(evm_not),
     "signextend": _wrap_binop(_evm_signextend),
     "iszero": _wrap_unop(_evm_iszero),
     "shr": _wrap_binop(_evm_shr),