perf[venom]: improve CSE elimination speed (#4607)

in the CSE elimination pass, avoid `O(n*k)` (where `n` == number of
instructions, and `k` == number of available expressions) copies
by replacing the dict with a HAMT. HAMT is an immutable dict with
`O(1)` copy (copy by reference), and `O(log N)` insert/delete/update
operations, each returning a fresh reference to a new HAMT.

we pull in the external package `immutables`, which is stable, shares
maintainership with CPython, and seems mature (receives bugfixes every
now and again, but mostly not adding new features). we pull in the
external library since HAMTs should be generally useful for the types
of lattice operations we do, and we can probably use them to optimize
other passes as well.

this commit also adds the use of `lru_cache` for some `Flag` operations
(which are bizarrely a perf hotspot)

Author: charles-cooper

.github/workflows/build.yml

@@ -47,9 +47,13 @@ jobs:
           cache: "pip"
 
       - name: Generate Binary
-        run: >-
-          pip install --no-binary pycryptodome --no-binary cbor2 . &&
-          pip install pyinstaller &&
+        run: |
+          pip install \
+            --no-binary pycryptodome \
+            --no-binary cbor2 \
+            --no-binary immutables \
+            . && \
+          pip install pyinstaller && \
           make freeze
 
 

setup.py

@@ -96,6 +96,7 @@ def _global_version(version):
         "packaging>=23.1",
         "lark>=1.0.0,<2",
         "wheel",
+        "immutables",
     ],
     setup_requires=["setuptools_scm>=7.1.0,<8.0.0"],
     extras_require=extras_require,

vyper/venom/analysis/available_expression.py

@@ -2,7 +2,9 @@
 
 from collections import deque
 from dataclasses import dataclass
-from functools import cached_property
+from functools import cached_property, lru_cache
+
+import immutables
 
 import vyper.venom.effects as effects
 from vyper.venom.analysis.analysis import IRAnalysesCache, IRAnalysis
@@ -35,6 +37,33 @@
     assert opcode in NONIDEMPOTENT_INSTRUCTIONS
 
 
+# flag bitwise operations are somehow a perf bottleneck, cache them
+@lru_cache
+def _get_read_effects(opcode, ignore_msize):
+    ret = effects.reads.get(opcode, effects.EMPTY)
+    if ignore_msize:
+        ret &= ~Effects.MSIZE
+    return ret
+
+
+@lru_cache
+def _get_write_effects(opcode, ignore_msize):
+    ret = effects.writes.get(opcode, effects.EMPTY)
+    if ignore_msize:
+        ret &= ~Effects.MSIZE
+    return ret
+
+
+@lru_cache
+def _get_overlap_effects(opcode, ignore_msize):
+    return _get_read_effects(opcode, ignore_msize) & _get_write_effects(opcode, ignore_msize)
+
+
+@lru_cache
+def _get_effects(opcode, ignore_msize):
+    return _get_read_effects(opcode, ignore_msize) | _get_write_effects(opcode, ignore_msize)
+
+
 @dataclass
 class _Expression:
     opcode: str
@@ -94,18 +123,6 @@ def depth(self) -> int:
                     max_depth = d
         return max_depth + 1
 
-    def get_reads(self, ignore_msize) -> Effects:
-        ret = effects.reads.get(self.opcode, effects.EMPTY)
-        if ignore_msize:
-            ret &= ~Effects.MSIZE
-        return ret
-
-    def get_writes(self, ignore_msize) -> Effects:
-        ret = effects.writes.get(self.opcode, effects.EMPTY)
-        if ignore_msize:
-            ret &= ~Effects.MSIZE
-        return ret
-
     @property
     def is_commutative(self) -> bool:
         return self.opcode in COMMUTATIVE_INSTRUCTIONS
@@ -130,10 +147,10 @@ class _AvailableExpressions:
     and provides API for handling them
     """
 
-    exprs: dict[_Expression, list[IRInstruction]]
+    exprs: immutables.Map[_Expression, list[IRInstruction]]
 
     def __init__(self):
-        self.exprs = dict()
+        self.exprs = immutables.Map()
 
     def __eq__(self, other) -> bool:
         if not isinstance(other, _AvailableExpressions):
@@ -148,23 +165,27 @@ def __repr__(self) -> str:
         return res
 
     def add(self, expr: _Expression, src_inst: IRInstruction):
-        if expr not in self.exprs:
-            self.exprs[expr] = []
-        self.exprs[expr].append(src_inst)
+        with self.exprs.mutate() as mt:
+            if expr not in mt:
+                mt[expr] = []
+            else:
+                mt[expr] = mt[expr].copy()
+            mt[expr].append(src_inst)
+            self.exprs = mt.finish()
 
     def remove_effect(self, effect: Effects, ignore_msize):
         if effect == effects.EMPTY:
             return
         to_remove = set()
         for expr in self.exprs.keys():
-            read_effs = expr.get_reads(ignore_msize)
-            write_effs = expr.get_writes(ignore_msize)
-            op_effect = read_effs | write_effs
+            op_effect = _get_effects(expr.opcode, ignore_msize)
             if op_effect & effect != effects.EMPTY:
                 to_remove.add(expr)
 
-        for expr in to_remove:
-            del self.exprs[expr]
+        with self.exprs.mutate() as mt:
+            for expr in to_remove:
+                del mt[expr]
+            self.exprs = mt.finish()
 
     def get_source_instruction(self, expr: _Expression) -> IRInstruction | None:
         """
@@ -178,18 +199,19 @@ def get_source_instruction(self, expr: _Expression) -> IRInstruction | None:
 
     def copy(self) -> _AvailableExpressions:
         res = _AvailableExpressions()
-        for k, v in self.exprs.items():
-            res.exprs[k] = v.copy()
+        res.exprs = self.exprs
         return res
 
     @staticmethod
     def lattice_meet(lattices: list[_AvailableExpressions]):
         if len(lattices) == 0:
             return _AvailableExpressions()
         res = lattices[0].copy()
+        # compute intersection
         for item in lattices[1:]:
             tmp = res
             res = _AvailableExpressions()
+            mt = res.exprs.mutate()
             for expr, insts in item.exprs.items():
                 if expr not in tmp.exprs:
                     continue
@@ -199,7 +221,8 @@ def lattice_meet(lattices: list[_AvailableExpressions]):
                         new_insts.append(i)
                 if len(new_insts) == 0:
                     continue
-                res.exprs[expr] = new_insts
+                mt[expr] = new_insts
+            res.exprs = mt.finish()
         return res
 
 
@@ -279,7 +302,7 @@ def _handle_bb(self, bb: IRBasicBlock) -> bool:
 
             self._update_expr(inst, expr)
 
-            write_effects = expr.get_writes(self.ignore_msize)
+            write_effects = _get_write_effects(expr.opcode, self.ignore_msize)
             available_exprs.remove_effect(write_effects, self.ignore_msize)
 
             # nonidempotent instructions affect other instructions,
@@ -288,7 +311,7 @@ def _handle_bb(self, bb: IRBasicBlock) -> bool:
             if inst.opcode in NONIDEMPOTENT_INSTRUCTIONS:
                 continue
 
-            expr_effects = expr.get_writes(self.ignore_msize) & expr.get_reads(self.ignore_msize)
+            expr_effects = _get_overlap_effects(expr.opcode, self.ignore_msize)
             if expr_effects == effects.EMPTY:
                 available_exprs.add(expr, inst)