[mypyc] Do not preemptively diff & union op level dataflow sets (#20897)

Mypyc is currently unable to compile [SQLGlot's
AST](https://github.com/tobymao/sqlglot/pull/7143/changes) (~950 classes
in a single file) as it dies with OOM even on a 64 gb machine.

Upon investigating, one chokehold seems to be the following line which
generates the dataflow for each BB `Op`; Given that most instructions
generate empty `kill` and `gen` sets, executing the union & difference
preemptively ends up creating identical `cur` copies.

I have verified that this fix locally unblocks compilation; Do note that
mypyc still consumes 6-7 gbs of RAM so I'm still looking out for more
improvements.

Author: VaggelisD

mypyc/analysis/dataflow.py

@@ -3,8 +3,8 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Iterable, Iterator
-from typing import Generic, TypeVar
+from collections.abc import Iterable, Iterator, Set as AbstractSet
+from typing import Any, Generic, TypeVar
 
 from mypyc.ir.ops import (
     Assign,
@@ -174,12 +174,14 @@ def __str__(self) -> str:
         return f"before: {self.before}\nafter: {self.after}\n"
 
 
-GenAndKill = tuple[set[T], set[T]]
+GenAndKill = tuple[AbstractSet[T], AbstractSet[T]]
+
+_EMPTY: tuple[frozenset[Any], frozenset[Any]] = (frozenset(), frozenset())
 
 
 class BaseAnalysisVisitor(OpVisitor[GenAndKill[T]]):
     def visit_goto(self, op: Goto) -> GenAndKill[T]:
-        return set(), set()
+        return _EMPTY
 
     @abstractmethod
     def visit_register_op(self, op: RegisterOp) -> GenAndKill[T]:
@@ -317,16 +319,16 @@ def __init__(self, strict_errors: bool = False) -> None:
         self.strict_errors = strict_errors
 
     def visit_branch(self, op: Branch) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_return(self, op: Return) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_unreachable(self, op: Unreachable) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_register_op(self, op: RegisterOp) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_assign(self, op: Assign) -> GenAndKill[Value]:
         # Loading an error value may undefine the register.
@@ -337,10 +339,10 @@ def visit_assign(self, op: Assign) -> GenAndKill[Value]:
 
     def visit_assign_multi(self, op: AssignMulti) -> GenAndKill[Value]:
         # Array registers are special and we don't track the definedness of them.
-        return set(), set()
+        return _EMPTY
 
     def visit_set_mem(self, op: SetMem) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
 
 def analyze_maybe_defined_regs(
@@ -392,27 +394,27 @@ def __init__(self, args: set[Value]) -> None:
         self.args = args
 
     def visit_branch(self, op: Branch) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_return(self, op: Return) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_unreachable(self, op: Unreachable) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_register_op(self, op: RegisterOp) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_assign(self, op: Assign) -> GenAndKill[Value]:
         if op.dest in self.args:
             return set(), {op.dest}
-        return set(), set()
+        return _EMPTY
 
     def visit_assign_multi(self, op: AssignMulti) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_set_mem(self, op: SetMem) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
 
 def analyze_borrowed_arguments(
@@ -435,13 +437,13 @@ def analyze_borrowed_arguments(
 
 class UndefinedVisitor(BaseAnalysisVisitor[Value]):
     def visit_branch(self, op: Branch) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_return(self, op: Return) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_unreachable(self, op: Unreachable) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_register_op(self, op: RegisterOp) -> GenAndKill[Value]:
         return set(), {op} if not op.is_void else set()
@@ -453,7 +455,7 @@ def visit_assign_multi(self, op: AssignMulti) -> GenAndKill[Value]:
         return set(), {op.dest}
 
     def visit_set_mem(self, op: SetMem) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
 
 def non_trivial_sources(op: Op) -> set[Value]:
@@ -472,10 +474,10 @@ def visit_return(self, op: Return) -> GenAndKill[Value]:
         if not isinstance(op.value, (Integer, Float)):
             return {op.value}, set()
         else:
-            return set(), set()
+            return _EMPTY
 
     def visit_unreachable(self, op: Unreachable) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_register_op(self, op: RegisterOp) -> GenAndKill[Value]:
         gen = non_trivial_sources(op)
@@ -494,10 +496,10 @@ def visit_set_mem(self, op: SetMem) -> GenAndKill[Value]:
         return non_trivial_sources(op), set()
 
     def visit_inc_ref(self, op: IncRef) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
     def visit_dec_ref(self, op: DecRef) -> GenAndKill[Value]:
-        return set(), set()
+        return _EMPTY
 
 
 def analyze_live_regs(blocks: list[BasicBlock], cfg: CFG) -> AnalysisResult[Value]:
@@ -559,8 +561,16 @@ def run_analysis(
             ops = list(reversed(ops))
         for op in ops:
             opgen, opkill = op.accept(gen_and_kill)
-            gen = (gen - opkill) | opgen
-            kill = (kill - opgen) | opkill
+            if opkill:
+                gen -= opkill
+
+            if opgen:
+                gen |= opgen
+                kill -= opgen
+
+            if opkill:
+                kill |= opkill
+
         block_gen[block] = gen
         block_kill[block] = kill
 
@@ -624,7 +634,10 @@ def run_analysis(
         for idx, op in ops_enum:
             op_before[label, idx] = cur
             opgen, opkill = op.accept(gen_and_kill)
-            cur = (cur - opkill) | opgen
+            if opkill:
+                cur = cur - opkill
+            if opgen:
+                cur = cur | opgen
             op_after[label, idx] = cur
     if backward:
         op_after, op_before = op_before, op_after

mypyc/analysis/selfleaks.py

@@ -1,6 +1,12 @@
 from __future__ import annotations
 
-from mypyc.analysis.dataflow import CFG, MAYBE_ANALYSIS, AnalysisResult, run_analysis
+from mypyc.analysis.dataflow import (
+    CFG,
+    MAYBE_ANALYSIS,
+    AnalysisResult,
+    GenAndKill as _DataflowGenAndKill,
+    run_analysis,
+)
 from mypyc.ir.ops import (
     Assign,
     AssignMulti,
@@ -49,7 +55,7 @@
 )
 from mypyc.ir.rtypes import RInstance
 
-GenAndKill = tuple[set[None], set[None]]
+GenAndKill = _DataflowGenAndKill[None]
 
 CLEAN: GenAndKill = (set(), set())
 DIRTY: GenAndKill = ({None}, {None})