[mypyc] Keep generator local state in C locals when possible (#21982)

Initially generate IR for generators and coroutines where locals and
temporaries are stored in registers. Only spill registers that are live
across a yield/await to the generator object. Previously most registers
were stored in the generator object as attributes. The new spill is
implemented as a new transform. We still use the old spill transform for
op results that are not stored in registers.

This makes microbenchmarks where most work happens between awaits
significantly faster, as C locals can be used for many operations
instead of attributes. This also reduces the size of generator objects,
as fewer struct fields are needed.

Remove old spill-related helpers from the IR builder as unnecessary.

I used coding agent assist.

Author: JukkaL

mypyc/analysis/dataflow.py

@@ -7,6 +7,7 @@
 from typing import Any, Generic, TypeVar
 
 from mypyc.ir.ops import (
+    ERR_NEVER,
     Assign,
     AssignMulti,
     BasicBlock,
@@ -518,6 +519,55 @@ def analyze_live_regs(blocks: list[BasicBlock], cfg: CFG) -> AnalysisResult[Valu
     )
 
 
+def analyze_live_regs_with_exception_edges(
+    blocks: list[BasicBlock],
+) -> dict[BasicBlock, set[Value]]:
+    """Calculate block-entry liveness before exception handling is inserted.
+
+    Unlike get_cfg(), this models an error edge at each operation that can
+    raise. This matters when an assignment later in the same block overwrites a
+    value that the error handler can still read. Returns inserted for a yield
+    are treated as edges to their continuation blocks.
+
+    TODO: Unlike run_analysis(), this supports edges from the middle of basic
+          blocks, so we reimplement data flow analysis here. Figure out a way
+          to unify this with run_analysis().
+    """
+    visitor = LivenessVisitor()
+    live_in: dict[BasicBlock, set[Value]] = {block: set() for block in blocks}
+
+    while True:
+        changed = False
+        for block in reversed(blocks):
+            terminator = block.terminator
+            if isinstance(terminator, Return) and terminator.yield_target is not None:
+                successors = (terminator.yield_target,)
+            else:
+                successors = terminator.targets()
+
+            live: set[Value] = set()
+            for successor in successors:
+                live.update(live_in[successor])
+
+            for op in reversed(block.ops):
+                if (
+                    block.error_handler is not None
+                    and isinstance(op, RegisterOp)
+                    and op.error_kind != ERR_NEVER
+                ):
+                    live.update(live_in[block.error_handler])
+                gen, kill = op.accept(visitor)
+                live.difference_update(kill)
+                live.update(gen)
+
+            if live != live_in[block]:
+                live_in[block] = live
+                changed = True
+
+        if not changed:
+            return live_in
+
+
 # Analysis kinds
 MUST_ANALYSIS = 0
 MAYBE_ANALYSIS = 1

mypyc/irbuild/builder.py

@@ -75,7 +75,6 @@
     KEEP_ALIVE_WHOLE_EXPRESSION,
     MODULE_PREFIX,
     SELF_NAME,
-    TEMP_ATTR_NAME,
     shared_lib_name,
 )
 from mypyc.crash import catch_errors
@@ -152,7 +151,7 @@
     AssignmentTargetRegister,
     AssignmentTargetTuple,
 )
-from mypyc.irbuild.util import bytes_from_str, is_constant
+from mypyc.irbuild.util import bytes_from_str, get_func_def, is_constant
 from mypyc.irbuild.vec import vec_set_item
 from mypyc.namegen import exported_name
 from mypyc.options import CompilerOptions
@@ -249,11 +248,8 @@ def __init__(
         self.callable_class_names: set[str] = set()
         self.options = options
 
-        # These variables keep track of the number of lambdas, implicit indices, and implicit
-        # iterators instantiated so we avoid name conflicts. The indices and iterators are
-        # instantiated from for-loops.
+        # Keep track of the number of lambdas instantiated so we avoid name conflicts.
         self.lambda_counter = 0
-        self.temp_counter = 0
 
         # These variables are populated from the first-pass PreBuildVisitor.
         self.free_variables = pbv.free_variables
@@ -263,6 +259,7 @@ def __init__(
         self.fdefs_to_decorators = pbv.funcs_to_decorators
         self.module_import_groups = pbv.module_import_groups
         self.comprehension_to_fitem = pbv.comprehension_to_fitem
+        self.deleted_vars = pbv.deleted_vars
 
         self.singledispatch_impls = singledispatch_impls
 
@@ -749,11 +746,11 @@ def get_assignment_target(
         if line == -1:
             line = lvalue.line
         if isinstance(lvalue, NameExpr):
-            # If we are visiting a decorator, then the SymbolNode we really want to be looking at
-            # is the function that is decorated, not the entire Decorator node itself.
+            # Use the concrete implementation as the symbol-table key for
+            # decorated and overloaded functions.
             symbol = lvalue.node
-            if isinstance(symbol, Decorator):
-                symbol = symbol.func
+            if isinstance(symbol, Decorator | OverloadedFuncDef):
+                symbol = get_func_def(symbol)
             if symbol is None:
                 # Semantic analyzer doesn't create ad-hoc Vars for special forms.
                 assert lvalue.is_special_form
@@ -766,12 +763,14 @@ def get_assignment_target(
                         reg_type = self.type_to_rtype(symbol.type)
                     else:
                         reg_type = self.node_type(lvalue)
-                    # If the function is a generator function, then first define a new variable
-                    # in the current function's environment class. Next, define a target that
-                    # refers to the newly defined variable in that environment class. Add the
-                    # target to the table containing class environment variables, as well as the
-                    # current environment.
-                    if self.fn_info.is_generator or self.fn_info.is_coroutine:
+                    # A deleted error-overlap value needs the environment's
+                    # definedness bitmap. Other generator locals start in
+                    # registers and are promoted later if they cross a yield.
+                    if (
+                        self.fn_info.is_generator
+                        and reg_type.error_overlap
+                        and symbol in self.deleted_vars
+                    ):
                         return self.add_var_to_env_class(
                             symbol,
                             reg_type,
@@ -780,7 +779,6 @@ def get_assignment_target(
                             prefix=GENERATOR_ATTRIBUTE_PREFIX,
                         )
 
-                    # Otherwise define a new local variable.
                     return self.add_local_reg(symbol, reg_type)
                 else:
                     # Assign to a previously defined variable.
@@ -856,11 +854,6 @@ def read(
 
         assert False, "Unsupported lvalue: %r" % target
 
-    def read_nullable_attr(self, obj: Value, attr: str, line: int = -1) -> Value:
-        """Read an attribute that might have an error value without raising AttributeError."""
-        assert isinstance(obj.type, RInstance) and obj.type.class_ir.is_ext_class
-        return self.add(GetAttr(obj, attr, line, allow_error_value=True))
-
     def assign(self, target: Register | AssignmentTarget, rvalue_reg: Value, line: int) -> None:
         if isinstance(target, Register):
             self.add(Assign(target, self.coerce_rvalue(rvalue_reg, target.type, line), line))
@@ -1045,47 +1038,8 @@ def push_loop_stack(self, continue_block: BasicBlock, break_block: BasicBlock) -
     def pop_loop_stack(self) -> None:
         self.nonlocal_control.pop()
 
-    def make_spill_target(self, type: RType) -> AssignmentTarget:
-        """Moves a given Value instance into the private generator frame."""
-        frame = self.fn_info.generator_class
-        # Generator classes for overriding methods can inherit from one another. Include the
-        # module-qualified owning class name so unrelated helper spills don't alias an inherited
-        # struct field.
-        name = f"{TEMP_ATTR_NAME}1_{exported_name(frame.ir.fullname)}_{self.temp_counter}"
-        self.temp_counter += 1
-        target = self.add_var_to_class(Var(name), type, frame.ir, frame.self_reg)
-        return target
-
-    def spill(self, value: Value) -> AssignmentTarget:
-        """Moves a given Value instance into the private generator frame."""
-        target = self.make_spill_target(value.type)
-        # Shouldn't be able to fail
-        self.assign(target, value, NO_TRACEBACK_LINE_NO)
-        return target
-
-    def maybe_spill(self, value: Value) -> Value | AssignmentTarget:
-        """
-        Moves a given Value instance into the private frame for generator functions. For
-        non-generator functions, leaves the Value instance as it is.
-
-        Returns an AssignmentTarget associated with the Value for generator functions and the
-        original Value itself for non-generator functions.
-        """
-        if self.fn_info.is_generator:
-            return self.spill(value)
-        return value
-
-    def maybe_spill_assignable(self, value: Value) -> Register | AssignmentTarget:
-        """
-        Moves a given Value instance into the private frame for generator functions. For
-        non-generator functions, allocate a temporary Register.
-
-        Returns an AssignmentTarget associated with the Value for generator functions and an
-        assignable Register for non-generator functions.
-        """
-        if self.fn_info.is_generator:
-            return self.spill(value)
-
+    def ensure_register(self, value: Value) -> Register:
+        """Return an assignable register containing a value."""
         if isinstance(value, Register):
             return value
 

mypyc/irbuild/expression.py

@@ -1415,7 +1415,7 @@ def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehe
 
 
 def _dict_comp_body(builder: IRBuilder, o: DictionaryComprehension) -> Value:
-    d = builder.maybe_spill(builder.call_c(dict_new_op, [], o.line))
+    d = builder.call_c(dict_new_op, [], o.line)
     loop_params = list(zip(o.indices, o.sequences, o.condlists, o.is_async))
 
     def gen_inner_stmts() -> None:

mypyc/irbuild/for_helpers.py

@@ -321,7 +321,7 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None:
     if val is not None:
         return val
 
-    list_ops = builder.maybe_spill(builder.new_list_op([], gen.line))
+    list_ops = builder.new_list_op([], gen.line)
 
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
 
@@ -357,7 +357,7 @@ def translate_set_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value
     if raise_error_if_contains_unreachable_names(builder, gen):
         return builder.none()
 
-    set_ops = builder.maybe_spill(builder.new_set_op([], gen.line))
+    set_ops = builder.new_set_op([], gen.line)
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
 
     def gen_inner_stmts() -> None:
@@ -687,7 +687,7 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         """Generate post-loop cleanup (if needed)."""
 
-    def load_len(self, expr: Value | AssignmentTarget) -> Value:
+    def load_len(self, expr: Value) -> Value:
         """A helper to get collection length, used by several subclasses."""
         return self.builder.builder.builtin_len(
             self.builder.read(expr, self.line), self.line, use_pyssize_t=True
@@ -702,13 +702,11 @@ def need_cleanup(self) -> bool:
         return True
 
     def init(self, expr_reg: Value, target_type: RType) -> None:
-        # Define targets to contain the expression, along with the iterator that will be used
-        # for the for-loop. If we are inside of a generator function, spill these into the
-        # private generator frame.
+        # Define a target containing the iterator used by the for-loop. The generator spill
+        # transform will promote it to the private generator frame if needed.
         builder = self.builder
         iter_reg = builder.primitive_op(iter_op, [expr_reg], self.line)
-        builder.maybe_spill(expr_reg)
-        self.iter_target = builder.maybe_spill(iter_reg)
+        self.iter_target = iter_reg
         self.target_type = target_type
 
     def gen_condition(self) -> None:
@@ -751,10 +749,9 @@ def need_cleanup(self) -> bool:
         return True
 
     def init(self, expr_reg: Value, target_type: RType) -> None:
-        # Define target to contains the generator expression. It's also the iterator.
-        # If we are inside a generator function, spill these into the private generator frame.
-        builder = self.builder
-        self.iter_target = builder.maybe_spill(expr_reg)
+        # The generator expression is also the iterator. The generator spill transform will
+        # promote it to the private generator frame if needed.
+        self.iter_target = expr_reg
         self.target_type = target_type
 
     def gen_condition(self) -> None:
@@ -807,14 +804,11 @@ class ForAsyncIterable(ForGenerator):
     """Generate IR for an async for loop."""
 
     def init(self, expr_reg: Value, target_type: RType) -> None:
-        # Define targets to contain the expression, along with the
-        # iterator that will be used for the for-loop. We are inside
-        # of a generator function, so we will spill these into
-        # the private generator frame.
+        # Define a target containing the iterator used by the for-loop. The generator spill
+        # transform will promote it to the private generator frame if needed.
         builder = self.builder
         iter_reg = builder.call_c(aiter_op, [expr_reg], self.line)
-        builder.maybe_spill(expr_reg)
-        self.iter_target = builder.maybe_spill(iter_reg)
+        self.iter_target = iter_reg
         self.target_type = target_type
         self.stop_reg = Register(bool_rprimitive)
 
@@ -894,7 +888,7 @@ class ForSequence(ForGenerator):
     Supports iterating in both forward and reverse.
     """
 
-    length_reg: Value | AssignmentTarget | None
+    length_reg: Value | None
 
     def init(
         self, expr_reg: Value, target_type: RType, reverse: bool, length: Value | None = None
@@ -907,13 +901,12 @@ def init(
         # Record a Value indicating the length of the sequence, if known at compile time.
         self.length = length
         self.reverse = reverse
-        # Define target to contain the expression, along with the index that will be used
-        # for the for-loop. If we are inside of a generator function, spill these into the
-        # private generator frame.
-        self.expr_target = builder.maybe_spill(expr_reg)
+        # The generator spill transform will promote loop state to the private generator frame
+        # if needed.
+        self.expr_target = expr_reg
         if is_immutable_rprimitive(expr_reg.type):
             # If the expression is an immutable type, we can load the length just once.
-            self.length_reg = builder.maybe_spill(self.length or self.load_len(self.expr_target))
+            self.length_reg = self.length or self.load_len(self.expr_target)
         else:
             # Otherwise, even if the length is known, we must recalculate the length
             # at every iteration for compatibility with python semantics.
@@ -926,7 +919,7 @@ def init(
             else:
                 len_val = self.load_len(self.expr_target)
             index_reg = builder.builder.int_sub(len_val, 1)
-        self.index_target = builder.maybe_spill_assignable(index_reg)
+        self.index_target = builder.ensure_register(index_reg)
         self.target_type = target_type
 
     def gen_condition(self) -> None:
@@ -1010,15 +1003,15 @@ def init(self, expr_reg: Value, target_type: RType) -> None:
         builder = self.builder
         self.target_type = target_type
 
-        # Spill some values so they can be read across yield.
-        self.expr_target = builder.maybe_spill(expr_reg)
+        # The generator spill transform will promote loop state that crosses a yield.
+        self.expr_target = expr_reg
         offset = Integer(0)
-        self.offset_target = builder.maybe_spill_assignable(offset)
-        self.size = builder.maybe_spill(self.load_len(self.expr_target))
+        self.offset_target = builder.ensure_register(offset)
+        self.size = self.load_len(self.expr_target)
 
         # For dict class (not a subclass) this is the dictionary itself.
         iter_reg = builder.call_c(self.dict_iter_op, [expr_reg], self.line)
-        self.iter_target = builder.maybe_spill(iter_reg)
+        self.iter_target = iter_reg
 
     def gen_condition(self) -> None:
         """Get next key/value pair, set new offset, and check if we should continue."""
@@ -1132,7 +1125,7 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         self.start_reg = start_reg
         self.end_reg = end_reg
         self.step = step
-        self.end_target = builder.maybe_spill(end_reg)
+        self.end_target = end_reg
         if is_short_int_rprimitive(start_reg.type) and is_short_int_rprimitive(end_reg.type):
             index_type: RType = short_int_rprimitive
         elif is_fixed_width_rtype(end_reg.type):
@@ -1141,7 +1134,7 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
             index_type = int_rprimitive
         index_reg = Register(index_type, line=self.line)
         builder.assign(index_reg, start_reg, self.line)
-        self.index_reg = builder.maybe_spill_assignable(index_reg)
+        self.index_reg = index_reg
         # Initialize loop index to 0. Assert that the index target is assignable.
         self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
         builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
@@ -1195,7 +1188,7 @@ def init(self) -> None:
         # Create a register to store the state of the loop index and
         # initialize this register along with the loop index to 0.
         zero = Integer(0)
-        self.index_reg = builder.maybe_spill_assignable(zero)
+        self.index_reg = builder.ensure_register(zero)
         self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
 
     def gen_step(self) -> None:

mypyc/irbuild/main.py

@@ -47,6 +47,7 @@ def f(x: int) -> int:
 from mypyc.irbuild.visitor import IRBuilderVisitor
 from mypyc.irbuild.vtable import compute_vtable
 from mypyc.options import CompilerOptions
+from mypyc.transform.generator_spills import promote_generator_registers
 
 # The stubs for callable contextmanagers are busted so cast it to the
 # right type...
@@ -123,6 +124,13 @@ def build_ir(
         result[module.fullname] = module_ir
         class_irs.extend(builder.classes)
 
+    # Generator helper calls are fresh C activations after every suspension.
+    # Move values that must outlive one activation to the private frame before
+    # attribute-definedness and the remaining IR transforms run.
+    for class_ir in class_irs:
+        if class_ir.env_user_function is not None:
+            promote_generator_registers(class_ir.env_user_function, class_ir)
+
     analyze_always_defined_attrs(class_irs)
 
     # Compute vtables.